diff --git a/cherab/omfit/load_plasma.py b/cherab/omfit/load_plasma.py
index 1f7dd70..f818697 100644
--- a/cherab/omfit/load_plasma.py
+++ b/cherab/omfit/load_plasma.py
@@ -1,14 +1,15 @@
+from cherab.solps import load_solps_from_mdsplus, load_solps_from_raw_output, load_solps_from_balance
 
 def load_edge_simulation(config, world):
 
     # Only try to do something if edge plasma has been selected as True
     try:
-        if not config['plasma']['edge']['present']:
+        if not config["plasma"]["edge"]["present"]:
             raise ValueError("This config file does not specify an edge plasma.")
     except KeyError:
         raise ValueError("This config file does not specify an edge plasma.")
 
-    if config['plasma']['edge']['type'] == "SOLPS":
+    if config["plasma"]["edge"]["type"] == "SOLPS":
         return _load_solps_simulation(config, world)
 
     else:
@@ -17,33 +18,27 @@ def load_edge_simulation(config, world):
 
 def _load_solps_simulation(config, world):
 
-    if config['plasma']['edge']['SOLPS_format'] == 'MDSplus':
+    if config["plasma"]["edge"]["SOLPS_format"] == "MDSplus":
 
-        print()
-        print("loading MDPsplus SOLPS plasma")
-
-        from cherab.solps import load_solps_from_mdsplus
-
-        mds_server = config['plasma']['edge']['mds_server']
-        mds_solps_reference = config['plasma']['edge']['mds_solps_reference']
+        mds_server = config["plasma"]["edge"]["solps_reference"]["mds_server"]
+        mds_solps_reference = config["plasma"]["edge"]["solps_reference"]["reference"]
 
         sim = load_solps_from_mdsplus(mds_server, mds_solps_reference)
-        plasma = sim.create_plasma(parent=world)
-        if config['plasma']['edge']['mesh']:
-            mesh = sim.mesh
-        else:
-            mesh = None
             
-    elif config['plasma']['edge']['SOLPS_format'] == 'Files':
-
-        from cherab.solps import load_solps_from_raw_output
+    elif config["plasma"]["edge"]["SOLPS_format"] == "RawFiles":
 
-        solps_directory = config['plasma']['edge']['solps_directory']
+        solps_directory = config["plasma"]["edge"]["solps_reference"]
         sim = load_solps_from_raw_output(solps_directory, debug=True)
-        plasma = sim.create_plasma(parent=world)
+
+    elif config["plasma"]["edge"]["SOLPS_format"] == "balance.nc":
+
+        solps_directory = config["plasma"]["edge"]["solps_reference"]
+        sim = load_solps_from_balance(solps_directory)
 
     else:
-        raise ValueError("Unrecognised SOLPS format '{}'.".format(config['plasma']['edge']['SOLPS_format']))
+        raise ValueError("Unrecognised SOLPS format '{}'.".format(config["plasma"]["edge"]["SOLPS_format"]))
 
-    return plasma, mesh
+    plasma = sim.create_plasma(parent=world)
+    mesh = sim.mesh
 
+    return plasma, mesh
\ No newline at end of file
diff --git a/demos/cherab.json b/demos/cherab.json
index 8a02fa4..98c4cab 100755
--- a/demos/cherab.json
+++ b/demos/cherab.json
@@ -585,8 +585,10 @@
    "present": true,
    "type": "SOLPS",
    "SOLPS_format": "MDSplus",
-   "mds_server": "solps-mdsplus.aug.ipp.mpg.de:8001",
-   "mds_solps_reference": 121853,
+   "solps_reference": {
+    "mds_server": "solps-mdsplus.aug.ipp.mpg.de:8001",
+    "mds_solps_reference": 121853
+   },
    "Rrange": [
     0,
     2.0