Fill out more of the EFIT mapping to MDS+

omas/machine_mappings/_efit.json

@@ -15,10 +15,6 @@
   "TDI": "dim_of(\\{EFIT_tree}::TOP.RESULTS.GEQDSK.BCENTR)/1000.",
   "treename": "{EFIT_tree}"
  },
- "equilibrium.vacuum_toroidal_field.b0": {
-  "TDI": "data(\\{EFIT_tree}::TOP.RESULTS.GEQDSK.BCENTR)",
-  "treename": "{EFIT_tree}"
- },
  "equilibrium.time_slice.:.time": {
   "TDI": "dim_of(\\{EFIT_tree}::TOP.RESULTS.GEQDSK.BCENTR)/1000.",
   "treename": "{EFIT_tree}"
@@ -39,13 +35,13 @@
  },
  "equilibrium.time_slice.:.boundary.x_point.:.r": {
   "NANFILTER": true,
-  "eval2TDI": "py2tdi(nan_where,'\\{EFIT_tree}::TOP.RESULTS.AEQDSK.RSEPS','\\{EFIT_tree}::TOP.RESULTS.GEQDSK.RSEPS',0)",
+  "eval2TDI": "py2tdi(nan_where,'\\{EFIT_tree}::TOP.RESULTS.AEQDSK.RSEPS','\\{EFIT_tree}::TOP.RESULTS.AEQDSK.RSEPS',0)",
   "treename": "{EFIT_tree}",
   "TRANSPOSE": [1,0]
  },
- "equilibrium.time_slice.:.boundary.x_point.:.r": {
+ "equilibrium.time_slice.:.boundary.x_point.:.z": {
   "NANFILTER": true,
-  "eval2TDI": "py2tdi(nan_where,'\\{EFIT_tree}::TOP.RESULTS.AEQDSK.ZSEPS','\\{EFIT_tree}::TOP.RESULTS.GEQDSK.ZSEPS',0)",
+  "eval2TDI": "py2tdi(nan_where,'\\{EFIT_tree}::TOP.RESULTS.AEQDSK.ZSEPS','\\{EFIT_tree}::TOP.RESULTS.AEQDSK.ZSEPS',0)",
   "treename": "{EFIT_tree}",
   "TRANSPOSE": [1,0]
  },
@@ -61,13 +57,13 @@
  },
  "equilibrium.time_slice.:.boundary_separatrix.x_point.:.r": {
   "NANFILTER": true,
-  "eval2TDI": "py2tdi(nan_where,'\\{EFIT_tree}::TOP.RESULTS.AEQDSK.RSEPS','\\{EFIT_tree}::TOP.RESULTS.GEQDSK.RSEPS',0)",
+  "eval2TDI": "py2tdi(nan_where,'\\{EFIT_tree}::TOP.RESULTS.AEQDSK.RSEPS','\\{EFIT_tree}::TOP.RESULTS.AEQDSK.RSEPS',0)",
   "treename": "{EFIT_tree}",
   "TRANSPOSE": [1,0]
  },
- "equilibrium.time_slice.:.boundary_separatrix.x_point.:.r": {
+ "equilibrium.time_slice.:.boundary_separatrix.x_point.:.z": {
   "NANFILTER": true,
-  "eval2TDI": "py2tdi(nan_where,'\\{EFIT_tree}::TOP.RESULTS.AEQDSK.ZSEPS','\\{EFIT_tree}::TOP.RESULTS.GEQDSK.ZSEPS',0)",
+  "eval2TDI": "py2tdi(nan_where,'\\{EFIT_tree}::TOP.RESULTS.AEQDSK.ZSEPS','\\{EFIT_tree}::TOP.RESULTS.AEQDSK.ZSEPS',0)",
   "treename": "{EFIT_tree}",
   "TRANSPOSE": [1,0]
  },
@@ -84,30 +80,30 @@
   "treename": "{EFIT_tree}"
  },
  "equilibrium.time_slice.:.boundary_separatrix.gap.1.name": {
-  "EVAL": "inboard",
+  "EVAL": "inboard"
  },
  "equilibrium.time_slice.:.boundary_separatrix.gap.1.value": {
   "TDI": "data(\\{EFIT_tree}::TOP.RESULTS.AEQDSK.GAPIN)/100.",
   "treename": "{EFIT_tree}"
  },
- "equilibrium.time_slice.:.boundary_separatrix.gap.1.name": {
-  "EVAL": "outboard",
+ "equilibrium.time_slice.:.boundary_separatrix.gap.2.name": {
+  "EVAL": "outboard"
  },
- "equilibrium.time_slice.:.boundary_separatrix.gap.1.value": {
+ "equilibrium.time_slice.:.boundary_separatrix.gap.2.value": {
   "TDI": "data(\\{EFIT_tree}::TOP.RESULTS.AEQDSK.GAPOUT)/100.",
   "treename": "{EFIT_tree}"
  },
- "equilibrium.time_slice.:.boundary_separatrix.gap.1.name": {
-  "EVAL": "top",
+ "equilibrium.time_slice.:.boundary_separatrix.gap.3.name": {
+  "EVAL": "top"
  },
- "equilibrium.time_slice.:.boundary_separatrix.gap.1.value": {
+ "equilibrium.time_slice.:.boundary_separatrix.gap.3.value": {
   "TDI": "data(\\{EFIT_tree}::TOP.RESULTS.AEQDSK.GAPTOP)/100.",
   "treename": "{EFIT_tree}"
  },
- "equilibrium.time_slice.:.boundary_separatrix.gap.1.name": {
-  "EVAL": "bottom",
+ "equilibrium.time_slice.:.boundary_separatrix.gap.4.name": {
+  "EVAL": "bottom"
  },
- "equilibrium.time_slice.:.boundary_separatrix.gap.1.value": {
+ "equilibrium.time_slice.:.boundary_separatrix.gap.4.value": {
   "TDI": "data(\\{EFIT_tree}::TOP.RESULTS.AEQDSK.GAPBOT)/100.",
   "treename": "{EFIT_tree}"
  },
@@ -252,11 +248,11 @@
   "treename": "{EFIT_tree}"
  },
  "equilibrium.time_slice.:.constraints.mse_polarisation_angle.measured": {
-  "TDI": "data(\\{EFIT_tree}::TOP.MEASUREMENTS.TANGAM)",
+  "TDI": "ATAN(data(\\{EFIT_tree}::TOP.MEASUREMENTS.TANGAM))",
   "treename": "{EFIT_tree}"
  },
  "equilibrium.time_slice.:.constraints.mse_polarisation_angle.measured_error_upper": {
-  "TDI": "data(\\{EFIT_tree}::TOP.MEASUREMENTS.SIGGAM)",
+  "TDI": "ATAN(data(\\{EFIT_tree}::TOP.MEASUREMENTS.SIGGAM))",
   "treename": "{EFIT_tree}"
  },
  "equilibrium.time_slice.:.constraints.mse_polarisation_angle.weight": {
@@ -272,27 +268,27 @@
   "treename": "{EFIT_tree}"
  },
  "equilibrium.time_slice.:.constraints.pf_current.measured": {
-  "TDI": "data(\\{EFIT_tree}::TOP.MEASUREMENTS.ECCURT/FCCURT/ACCURT)",
+  "eval2TDI": "py2tdi(stack_outer,'\\{EFIT_tree}::TOP.MEASUREMENTS.ECCURT','\\{EFIT_tree}::TOP.MEASUREMENTS.FCCURT','\\{EFIT_tree}::TOP.MEASUREMENTS.ACCURT')",
   "treename": "{EFIT_tree}"
  },
  "equilibrium.time_slice.:.constraints.pf_current.measured_error_upper": {
-  "TDI": "data(\\{EFIT_tree}::TOP.MEASUREMENTS.SIGECC/SIGFCC)",
+  "eval2TDI": "py2tdi(stack_outer,'\\{EFIT_tree}::TOP.MEASUREMENTS.SIGECC','\\{EFIT_tree}::TOP.MEASUREMENTS.SIGFCC')",
   "treename": "{EFIT_tree}"
  },
  "equilibrium.time_slice.:.constraints.pf_current.weight": {
-  "TDI": "data(\\{EFIT_tree}::TOP.MEASUREMENTS.FWTEC/FWTFC)",
+  "eval2TDI": "py2tdi(stack_outer,'\\{EFIT_tree}::TOP.MEASUREMENTS.FWTEC','\\{EFIT_tree}::TOP.MEASUREMENTS.FWTFC')",
   "treename": "{EFIT_tree}"
  },
  "equilibrium.time_slice.:.constraints.pf_current.reconstructed": {
-  "TDI": "data(\\{EFIT_tree}::TOP.MEASUREMENTS.CECURR/CCBRSP)",
+  "eval2TDI": "py2tdi(stack_outer,'\\{EFIT_tree}::TOP.MEASUREMENTS.CECURR','\\{EFIT_tree}::TOP.MEASUREMENTS.CCBRSP')",
   "treename": "{EFIT_tree}"
  },
  "equilibrium.time_slice.:.constraints.pf_current.chi_squared": {
-  "TDI": "data(\\{EFIT_tree}::TOP.MEASUREMENTS.CHIECC/CHIFCC)",
+  "eval2TDI": "py2tdi(stack_outer,'\\{EFIT_tree}::TOP.MEASUREMENTS.CHIECC','\\{EFIT_tree}::TOP.MEASUREMENTS.CHIFCC')",
   "treename": "{EFIT_tree}"
  },
  "equilibrium.time_slice.:.constraints.pressure.:.position.psi": {
-  "TDI": "data(\\{EFIT_tree}::TOP.MEASUREMENTS.RPRESS)",
+  "eval2TDI": "py2tdi(efit_psi_to_psi,'\\{EFIT_tree}::TOP.MEASUREMENTS.RPRESS','\\{EFIT_tree}::TOP.RESULTS.GEQDSK.SSIMAG','\\{EFIT_tree}::TOP.RESULTS.GEQDSK.SSIBRY')",
   "treename": "{EFIT_tree}"
  },
  "equilibrium.time_slice.:.constraints.pressure.:.measured": {
@@ -316,11 +312,11 @@
   "treename": "{EFIT_tree}"
  },
  "equilibrium.time_slice.:.constraints.j_tor.:.position.psi": {
-  "TDI": "data(\\{EFIT_tree}::TOP.MEASUREMENTS.SIZEROJ)",
+   "eval2TDI": "py2tdi(efit_psi_to_psi,'\\{EFIT_tree}::TOP.MEASUREMENTS.SIZEROJ','\\{EFIT_tree}::TOP.RESULTS.GEQDSK.SSIMAG','\\{EFIT_tree}::TOP.RESULTS.GEQDSK.SSIBRY')",
   "treename": "{EFIT_tree}"
  },
  "equilibrium.time_slice.:.constraints.j_tor.:.measured": {
-  "TDI": "data(\\{EFIT_tree}::TOP.MEASUREMENTS.VZEROJ)",
+  "TDI": "data(\\{EFIT_tree}::TOP.MEASUREMENTS.VZEROJ)*data(\\{EFIT_tree}::TOP.RESULTS.GEQDSK.CPASMA)/(data(\\{EFIT_tree}::TOP.RESULTS.AEQDSK.AREA)/10000.)",
   "treename": "{EFIT_tree}"
  },
  "equilibrium.time_slice.:.constraints.chi_squared_total": {
@@ -454,16 +450,16 @@
  "equilibrium.vacuum_toroidal_field.r0": {
   "TDI": "data(\\{EFIT_tree}::TOP.RESULTS.GEQDSK.RZERO)[0]",
   "treename": "{EFIT_tree}"
- }
+ },
  "equilibrium.time_slice.:.convergence.iterations_n": {
     "eval2TDI": "py2tdi(tile,1,'size(\\{EFIT_tree}::TOP.MEASUREMENTS.CERROR)')",
-    "treename": "{EFIT_tree}",
+    "treename": "{EFIT_tree}"
  },
  "equilibrium.time_slice.:.convergence.grad_shafranov_deviation_value": {
   "TDI": "data(\\{EFIT_tree}::TOP.RESULTS.AEQDSK.TERROR)",
   "treename": "{EFIT_tree}"
  },
  "equilibrium.time_slice.:.convergence.grad_shafranov_deviation_expression.index": {
-  "EVAL": "3",
- },
+  "EVAL": "3"
+ }
 }

omas/machine_mappings/python_tdi.py

@@ -4,6 +4,14 @@ def tile(a, n):
     a = a.data()
     return np.array([a for k in range(n)])
 
+def stack_outer(*args):
+    if len(args) < 2:
+        raise ValueError("stack_outer only makes sense with 2 ore more arguments")
+    import numpy as np
+    a = np.array(args[0].data)
+    for array in args[1:]:
+        a = np.concatenate([a,array],axis=1)
+
 
 def nan_where(a, b, n):
     import numpy as np
@@ -35,6 +43,11 @@ def geqdsk_psi(a, b, c):
     M = a[:, None] + np.linspace(0, 1, n).T * (b[:, None] - a[:, None])
     return M
 
+def efit_psi_to_psi(a, b, c):
+    a = a.data()
+    b = b.data()
+    c = c.data()
+    return (a - b)/(c - b)
 
 def py2tdi(func, *args):
     import inspect

omas/omas_machine.py

@@ -62,6 +62,15 @@ def python_tdi_namespace(branch):
 
     return _python_tdi_namespace[branch]
 
+def remove_nans(x):
+    import numpy as np
+    if np.isscalar(x):
+        if np.isnan(x):
+            raise ValueError("Behavior of Nan filter undefined for scalar nan values")
+        else:
+            return x
+    else:
+        return x[~np.isnan(x)]
 
 def machine_to_omas(ods, machine, pulse, location, options={}, branch='', user_machine_mappings=None, cache=None):
     """
@@ -75,7 +84,7 @@ def machine_to_omas(ods, machine, pulse, location, options={}, branch='', user_m
 
     :param location: ODS location to be populated
 
-    :param options: dictionary with options to use when loadinig the data
+    :param options: dictionary with options to use when loading the data
 
     :param branch: load machine mappings and mapping functions from a specific GitHub branch
 
@@ -99,13 +108,51 @@ def machine_to_omas(ods, machine, pulse, location, options={}, branch='', user_m
         options_with_defaults.update(options)
         options_with_defaults.update({'machine': machine, 'pulse': pulse, 'location': location})
         try:
-            mapped = mappings[location]
+            if not location.endswith(".*"):
+                mapped = mappings[location]
             break
-        except KeyError:
+        except KeyError as e:
             if branch == 'master':
-                raise
+                raise e
+            else:
+                print(f"Failed to load {location} from head. Attempting to resolve using the master branch.")
+                print(f"Error was:")
+                print(e)
     idm = (machine, branch)
+    if location.endswith(".*"):
+        root = location.split(".*")[0]
+        for key in mappings:
+            if root in key:
+                resolve_mapped(ods, machine, pulse, mappings, key, idm, options_with_defaults, branch, cache=cache)
+        return ods
+    else:
+        return resolve_mapped(ods, machine, pulse,  mappings, location, idm, options_with_defaults, branch, cache=cache)
+
+def resolve_mapped(ods, machine, pulse,  mappings, location, idm, options_with_defaults, branch, cache=None):
+    """
+    Routine to resolve a mapping
+
+    :param ods: input ODS to populate
+
+    :param machine: machine name
+
+    :param pulse: pulse number
+
+    :param mappings: Dictionary of available mappings
+
+    :param location: ODS location to be resolved
 
+    :param idm: Tuple with machine and branch
+
+    :param options_with_defaults: dictionary with options to use when loading the data including default settings
+
+    :param branch: load machine mappings and mapping functions from a specific GitHub branch
+
+    :param cache: if cache is a dictionary, this will be used to establiish a cash
+
+    :return: updated ODS and data before being assigned to the ODS
+    """
+    mapped = mappings[location]
     # cocosio
     cocosio = None
     if 'COCOSIO' in mapped:
@@ -194,7 +241,8 @@ def machine_to_omas(ods, machine, pulse, location, options={}, branch='', user_m
     # transpose filter
     nanfilter = lambda x: x
     if mapped.get('NANFILTER', False):
-        nanfilter = lambda x: x[~numpy.isnan(x)]
+        #lambda x: x[~numpy.isnan(x)]
+        nanfilter = remove_nans
 
     # assign data to ODS
     if not hasattr(data, 'shape'):