From 829cdfcd17b53f24af896314873e8a8199dedd5a Mon Sep 17 00:00:00 2001
From: MatejTomes <tomes@ipp.cas.cz>
Date: Tue, 16 Feb 2021 17:13:29 +0100
Subject: [PATCH] Adds the possibility to load solps from balance.nc
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Changes the required structure of the config dictionary:
        config['plasma']['edge']['mds_server']
        config['plasma']['edge']['mds_solps_reference']
        config['plasma']['edge']['solps_directory']
are no longer used.

The value of config['plasma']['edge']['mesh'] is not checked and the
mesh is always returned.

The config['plasma']['edge']['SOLPS_format'] now decided the way
config['plasma']['edge']["solps_reference"] will be interpreted:

1. if config['plasma']['edge']['SOLPS_format'] is 'MDSplus' then the
config['plasma']['edge']["solps_reference"] is expected to be a
dictionary with two keys "mds_server" and "reference":
        config["plasma"]["edge"]["solps_reference"]["mds_server"] is the
        reference to the mdsplus server.
        config["plasma"]["edge"]["solps_reference"]["reference"] is the
        reference to the SOLPS simulation.
2. if config['plasma']['edge']['SOLPS_format'] is 'RawFiles' then the
config['plasma']['edge']["solps_reference"] is expected to be the path
to the folder containing the solps raw files as fort.44.
3. if config['plasma']['edge']['SOLPS_format'] is 'balance.nc' then the
config['plasma']['edge']["solps_reference"] is expected to be the path
to the solps balance.nc file.
---
 cherab/omfit/load_plasma.py | 39 ++++++++++++++++---------------------
 demos/cherab.json           |  6 ++++--
 2 files changed, 21 insertions(+), 24 deletions(-)

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