Test new efit mappings for ida

omas/machine_mappings/d3d.json

@@ -299,7 +299,7 @@
   "PYTHON": "ec_launcher_active_hardware(ods, {pulse})"
  },
  "ece": {
-    "PYTHON": "electron_cyclotron_emission_hardware(ods, {pulse}, {fast_ece!r})"
+  "PYTHON": "electron_cyclotron_emission_hardware(ods, {pulse}, {fast_ece!r})"
  },
  "ece.channel.:": {
   "PYTHON": "electron_cyclotron_emission_hardware(ods, {pulse}, {fast_ece!r})"
@@ -310,6 +310,9 @@
  "ece.channel.:.identifier": {
   "PYTHON": "electron_cyclotron_emission_hardware(ods, {pulse}, {fast_ece!r})"
  },
+ "ece.channel.:.if_bandwidth": {
+  "PYTHON": "electron_cyclotron_emission_hardware(ods, {pulse}, {fast_ece!r})"
+ },
  "ece.channel.:.name": {
   "PYTHON": "electron_cyclotron_emission_hardware(ods, {pulse}, {fast_ece!r})"
  },
@@ -319,6 +322,9 @@
  "ece.channel.:.time": {
   "PYTHON": "electron_cyclotron_emission_hardware(ods, {pulse}, {fast_ece!r})"
  },
+ "ece.ids_properties.homogeneous_time": {
+  "PYTHON": "electron_cyclotron_emission_hardware(ods, {pulse}, {fast_ece!r})"
+ },
  "ece.line_of_sight.first_point.phi": {
   "COCOSIO": 11,
   "PYTHON": "electron_cyclotron_emission_hardware(ods, {pulse}, {fast_ece!r})"
@@ -697,5 +703,8 @@
    2
   ],
   "treename": "{EFIT_tree}"
+ },
+ "wall.ids_properties.homogeneous_time": {
+  "VALUE": 1
  }
 }

omas/machine_mappings/d3d.py

@@ -773,7 +773,7 @@ def electron_cyclotron_emission_data(ods, pulse=133221, fast_ece=False, _measure
     TECE = '\\ECE::TOP.TECE.TECE' + fast_ece
 
     query = {}
-    for node, quantities in zip([setup, cal], [['ECEPHI', 'ECETHETA', 'ECEZH', 'FREQ'], ['NUMCH']]):
+    for node, quantities in zip([setup, cal], [['ECEPHI', 'ECETHETA', 'ECEZH', 'FREQ', "FLTRWID"], ['NUMCH']]):
         for quantity in quantities:
             query[quantity] = node + quantity
     query['TIME'] = f"dim_of({TECE + '01'})"
@@ -786,7 +786,7 @@ def electron_cyclotron_emission_data(ods, pulse=133221, fast_ece=False, _measure
         for ich in range(1, N_ch + 1):
             query[f'T{ich}'] = TECE + '{0:02d}'.format(ich)
         ece_data = mdsvalue('d3d', treename='ELECTRONS', pulse=pulse, TDI=query).raw()
-
+    ods['ece.ids_properties.homogeneous_time'] = 0
     # Not in mds+
     if not _measurements:
         points = [{}, {}]
@@ -813,6 +813,7 @@ def electron_cyclotron_emission_data(ods, pulse=133221, fast_ece=False, _measure
             ch['time'] = ece_map['TIME'] * 1.0e-3
             f[:] = ece_map['FREQ'][ich]
             ch['frequency']['data'] = f * 1.0e9
+            ch['if_bandwidth'] = ece_map['FLTRWID'][ich] * 1.0e9
 
 
 @machine_mapping_function(__regression_arguments__, pulse=133221)

omas/omas_imas.py

@@ -80,7 +80,7 @@ def imas_open(user, machine, pulse, run, occurrence={}, new=False, imas_major_ve
     return IDS(DBentry, occurrence)
 
 
-def imas_set(ids, path, value, skip_missing_nodes=False, allocate=False):
+def imas_set(ids, path, value, skip_missing_nodes=False, allocate=False, ids_is_subtype=False, only_allocate=True):
     """
     assign a value to a path of an open IMAS ids
 
@@ -100,21 +100,25 @@ def imas_set(ids, path, value, skip_missing_nodes=False, allocate=False):
     :return: path if set was done, otherwise None
     """
     # handle uncertain data
+    if type(path) != list:
+        path = p2l(path)
     if is_uncertain(value):
         path = copy.deepcopy(path)
-        tmp = imas_set(ids, path, nominal_values(value), skip_missing_nodes=skip_missing_nodes, allocate=allocate)
+        tmp = imas_set(ids, path, nominal_values(value), skip_missing_nodes=skip_missing_nodes, allocate=allocate, only_allocate=only_allocate)
         path[-1] = path[-1] + '_error_upper'
-        imas_set(ids, path, std_devs(value), skip_missing_nodes=skip_missing_nodes, allocate=allocate)
+        imas_set(ids, path, std_devs(value), skip_missing_nodes=skip_missing_nodes, allocate=allocate, only_allocate=only_allocate)
         return tmp
-
     ds = path[0]
     path = path[1:]
 
     # identify data dictionary to use, from this point on `m` points to the IDS
     debug_path = ''
-    if hasattr(ids, ds):
+    if hasattr(ids, ds) or ids_is_subtype:
         debug_path += '%s' % ds
-        m = getattr(ids, ds)
+        if ids_is_subtype:
+            m = ids
+        else:
+            m = getattr(ids, ds)
         if hasattr(m, 'time') and not isinstance(m.time, float) and not m.time.size:
             m.time= numpy.resize(m.time, 1)
             m.time[0] = -1.0
@@ -130,13 +134,32 @@ def imas_set(ids, path, value, skip_missing_nodes=False, allocate=False):
 
     # traverse IMAS structure until reaching the leaf
     out = m
+    done = allocate
     for kp, p in enumerate(path):
         location = l2i([ds] + path[: kp + 1])
         if isinstance(p, str):
-            if hasattr(out, p):
+            if p == ":":
+                if allocate and len(out) != len(value):
+                    out.resize(len(value))
+                    done = True
+                if kp == len(path) - 1:
+                    break
+                else:
+                    if len(value) == 1:
+                        out = out[0]
+                        break
+                    else:
+                        for i in range(value.shape[0]):
+                            if len(path[kp + 1:]) == 1:
+                                setattr(out[i], path[-1], value[i])
+                            else:
+                                imas_set(out[i], path[kp + 1:], value[i], skip_missing_nodes=False, allocate=allocate,only_allocate=only_allocate)
+                    return [ds] + path
+            elif hasattr(out, p):
                 if kp < (len(path) - 1):
                     debug_path += '.' + p
                     out = getattr(out, p)
+
             elif skip_missing_nodes is not False:
                 if skip_missing_nodes is None:
                     printe('WARNING: %s is not part of IMAS' % location)
@@ -157,7 +180,7 @@ def imas_set(ids, path, value, skip_missing_nodes=False, allocate=False):
                 out = out[p]
 
     # if we are allocating data, simply stop here
-    if allocate:
+    if done and only_allocate:
         return [ds] + path
 
     # assign data to leaf node

omas/omas_physics.py

@@ -213,8 +213,72 @@ def equilibrium_ggd_to_rectangular(ods, time_index=None, resolution=None, method
         profiles_2d['grid.dim2'] = z
     return ods_n
 
+@add_to__ODS__
+@preprocess_ods('equilibrium')
+def add_rho_pol_norm_to_equilbrium_profiles_1d_ods(ods, time_index):
+    try:
+        if (len(ods["equilibrium"]["time_slice"][time_index]["profiles_1d"]["rho_pol-norm"]) ==
+            len(ods["equilibrium"]["time_slice"][time_index]["profiles_1d"]["psi"])):
+            return
+        else:
+            raise LookupError
+    except LookupError:
+        ods["equilibrium"]["time_slice"][time_index]["profiles_1d"]["rho_pol_norm"] = (
+            ods.map_pol_flux_to_flux_coordinate(ods, time_index, "rho_pol_norm", 
+            ods["equilibrium"]["time_slice"][time_index]["profiles_1d"]["psi"])
+        )
+
+def mask_SOL(ods, time_index, psi_values):
+    """
+        Returns a numpy array of `dtype=bool`/
+        The array is true for all values inside and on the LCFS
+        :param ods: input ods
+
+        :param time_index: time slices to process
+
+        :param values: Psi values that need to masked
+
+        :return: mask that is True for values inside and on the LCFS
+    """
+    if (ods["equilibrium"]["time_slice"][time_index]["global_quantities"]["psi_axis"] <
+        ods["equilibrium"]["time_slice"][time_index]["global_quantities"]["psi_boundary"]):
+        return psi_values <= ods["equilibrium"]["time_slice"][time_index]["global_quantities"]["psi_boundary"]
+    else:
+        return psi_values >= ods["equilibrium"]["time_slice"][time_index]["global_quantities"]["psi_boundary"]
+
+@add_to__ODS__
+@preprocess_ods('equilibrium')
+def add_phi_to_equilbrium_profiles_1d_ods(ods, time_index):
+    """
+        Adds `profiles_1d.phi` to an ODS using q
+        :param ods: input ods
+
+        :param time_index: time slices to process
+    """
+    try:
+        if (len(ods["equilibrium"]["time_slice"][time_index]["profiles_1d"]["phi"]) ==
+            len(ods["equilibrium"]["time_slice"][time_index]["profiles_1d"]["psi"])):
+            return
+        else:
+            raise LookupError
+    except (LookupError, ValueError):
+        from scipy.interpolate import InterpolatedUnivariateSpline
+        #TODO: 
+        # - Any cocos needed here?
+        psi = ods["equilibrium"]["time_slice"][time_index]["profiles_1d"]["psi"]
+        mask = mask_SOL(ods, time_index, psi)
+        q_spline = InterpolatedUnivariateSpline(
+                ods["equilibrium"]["time_slice"][time_index]["profiles_1d"]["psi"][mask],
+                ods["equilibrium"]["time_slice"][time_index]["profiles_1d"]["q"][mask])
+        phi_spline = q_spline.antiderivative(1)
+        ods["equilibrium"]["time_slice"][time_index]["profiles_1d"]["phi"] = numpy.zeros(psi.shape)
+        ods["equilibrium"]["time_slice"][time_index]["profiles_1d"]["phi"][mask] = (
+            phi_spline(ods["equilibrium"]["time_slice"][time_index]["profiles_1d"]["psi"][mask]))
+        ods["equilibrium"]["time_slice"][time_index]["profiles_1d"]["phi"][mask==False] = numpy.inf
+
 
 def map_flux_coordinate_to_pol_flux(ods, time_index, origin, values):
+    import numpy as np
     """
     Maps from one magnetic coordinate system to psi
     :param ods: input ods
@@ -227,15 +291,31 @@ def map_flux_coordinate_to_pol_flux(ods, time_index, origin, values):
 
     :return: Transformed values
     """
-    if origin == "rho_pol":
-        return (
-            values**2
-            * (
-                ods["equilibrium"]["time_slice"][time_index]["global_quantities"]["psi_sep"]
-                - ods["equilibrium"]["time_slice"][time_index]["global_quantities"]["psi_axis"]
-            )
-            + ods["equilibrium"]["time_slice"][time_index]["global_quantities"]["psi_axis"]
-        )
+    if origin == "rho_pol_norm":
+        return (values**2 * (ods["equilibrium"]["time_slice"][time_index]["global_quantities"]["psi_boundary"]
+                - ods["equilibrium"]["time_slice"][time_index]["global_quantities"]["psi_axis"])
+                + ods["equilibrium"]["time_slice"][time_index]["global_quantities"]["psi_axis"])
+    elif origin == "rho_tor_norm":
+        phi = values**2
+        phi *= map_pol_flux_to_flux_coordinate(ods, time_index, "phi", 
+                np.array([ods["equilibrium"]["time_slice"][time_index]["global_quantities"]["psi_boundary"]]))
+        return map_flux_coordinate_to_pol_flux(ods, time_index, "phi", phi)
+    elif origin == "phi":
+        from scipy.interpolate import InterpolatedUnivariateSpline
+        psi_grid = ods["equilibrium"]["time_slice"][time_index]["profiles_1d"]["psi"]
+        psi_mask = mask_SOL(ods, time_index, psi_grid)
+        phi_grid = ods["equilibrium"]["time_slice"][time_index]["profiles_1d"]["phi"][psi_mask]
+        if phi_grid[-1] < phi_grid[0]:
+            psi_spl = InterpolatedUnivariateSpline(phi_grid[psi_mask][::-1], psi_grid[psi_mask][::-1])
+        else:
+            psi_spl = InterpolatedUnivariateSpline(phi_grid[psi_mask], psi_grid[psi_mask])
+        phi_min = np.min(phi_grid)
+        phi_max = np.max(phi_grid)
+        values_mask = np.logical_and(values >= phi_min, values <= phi_max)
+        psi = np.zeros(values.shape)
+        psi[:] = np.nan
+        psi[values_mask] = psi_spl(values[values_mask])
+        return psi
     else:
         raise NotImplementedError(f"Conversion from {origin} not yet implemented.")
 
@@ -255,14 +335,42 @@ def map_pol_flux_to_flux_coordinate(ods, time_index, destination, values):
 
         :return: Transformed values
     """
-    if destination == "rho_pol":
-        return np.sqrt(
-            (values - ods["equilibrium"]["time_slice"][time_index]["global_quantities"]["psi_axis"])
-            / (
-                ods["equilibrium"]["time_slice"][time_index]["global_quantities"]["psi_boundary"]
-                - ods["equilibrium"]["time_slice"][time_index]["global_quantities"]["psi_axis"]
-            )
-        )
+    if destination == "rho_pol_norm":
+        return np.sqrt((values - ods["equilibrium"]["time_slice"][time_index]["global_quantities"]["psi_axis"]) /
+                       (ods["equilibrium"]["time_slice"][time_index]["global_quantities"]["psi_boundary"]
+                        - ods["equilibrium"]["time_slice"][time_index]["global_quantities"]["psi_axis"]))
+    elif destination == "rho_tor_norm":
+        mask = mask_SOL(ods, time_index, values)
+        phi = map_pol_flux_to_flux_coordinate(ods, time_index, "phi", values[mask])
+        rho_tor_norm = np.zeros(values.shape)
+        phi_boundary = map_pol_flux_to_flux_coordinate(ods, time_index, "phi", 
+                np.array([ods["equilibrium"]["time_slice"][time_index]["global_quantities"]["psi_boundary"]]))
+        rho_tor_norm[mask] = np.sqrt(phi / phi_boundary)
+        rho_tor_norm[mask==False] = np.inf
+        return rho_tor_norm
+    elif destination == "phi":
+        from scipy.interpolate import InterpolatedUnivariateSpline
+        psi_grid = ods["equilibrium"]["time_slice"][time_index]["profiles_1d"]["psi"]
+        psi_grid_mask = mask_SOL(ods, time_index, psi_grid)
+        try:
+            phi_spl = InterpolatedUnivariateSpline(psi_grid[psi_grid_mask],
+                    ods["equilibrium"]["time_slice"][time_index]["profiles_1d"]["phi"][psi_grid_mask])
+        except ValueError:
+            add_phi_to_equilbrium_profiles_1d_ods(ods, time_index)
+            phi_spl = InterpolatedUnivariateSpline(psi_grid[psi_grid_mask],
+                    ods["equilibrium"]["time_slice"][time_index]["profiles_1d"]["phi"][psi_grid_mask])
+        mask = mask_SOL(ods, time_index, values)
+        phi = np.zeros(values.shape)
+        phi[mask] = phi_spl(values[mask])
+        phi_bound = phi_spl(ods["equilibrium"]["time_slice"][time_index]["global_quantities"]["psi_boundary"])
+        phi[mask==False] = np.inf * np.sign(phi_bound)
+        wrong_sign_mask = phi_bound * phi < 0
+        if np.any(wrong_sign_mask):
+            if np.any(np.abs(phi[wrong_sign_mask]/phi_bound) > 1.e-4):
+                raise ValueError("Unphysical phi encountered when mapping to phi")
+            else:
+                phi[wrong_sign_mask] = 0.0
+        return phi
     else:
         raise NotImplementedError(f"Conversion to {destination} not yet implemented.")
 
@@ -350,17 +458,17 @@ def derive_equilibrium_profiles_2d_quantity(ods, time_index, grid_index, quantit
         ods[f'equilibrium.time_slice.{time_index}.profiles_2d.{grid_index}.psi'],
     )
     cocos = define_cocos(11)
-    if quantity == "b_r":
-        ods[f'equilibrium.time_slice.{time_index}.profiles_2d.{grid_index}.b_r'] = (
+    if quantity == "b_field_r":
+        ods[f'equilibrium.time_slice.{time_index}.profiles_2d.{grid_index}.b_field_r'] = (
             psi_spl(r, z, dy=1, grid=False) * cocos['sigma_RpZ'] * cocos['sigma_Bp'] / ((2.0 * numpy.pi) ** cocos['exp_Bp'] * r)
         )
         return ods
-    elif quantity == "b_z":
-        ods[f'equilibrium.time_slice.{time_index}.profiles_2d.{grid_index}.b_z'] = (
+    elif quantity == "b_field_z":
+        ods[f'equilibrium.time_slice.{time_index}.profiles_2d.{grid_index}.b_field_z'] = (
             -psi_spl(r, z, dx=1, grid=False) * cocos['sigma_RpZ'] * cocos['sigma_Bp'] / ((2.0 * numpy.pi) ** cocos['exp_Bp'] * r)
         )
         return ods
-    elif quantity == "b_tor":
+    elif quantity == "b_field_tor":
         mask = numpy.logical_and(
             ods[f'equilibrium.time_slice.{time_index}.profiles_2d.{grid_index}.psi']
             < numpy.max(ods[f'equilibrium.time_slice.{time_index}.profiles_1d.psi']),
@@ -370,11 +478,11 @@ def derive_equilibrium_profiles_2d_quantity(ods, time_index, grid_index, quantit
         f_spl = InterpolatedUnivariateSpline(
             ods[f'equilibrium.time_slice.{time_index}.profiles_1d.psi'], ods[f'equilibrium.time_slice.{time_index}.profiles_1d.f']
         )
-        ods[f'equilibrium.time_slice.{time_index}.profiles_2d.{grid_index}.b_tor'] = numpy.zeros(r.shape)
-        ods[f'equilibrium.time_slice.{time_index}.profiles_2d.{grid_index}.b_tor'][mask] = (
+        ods[f'equilibrium.time_slice.{time_index}.profiles_2d.{grid_index}.b_field_tor'] = numpy.zeros(r.shape)
+        ods[f'equilibrium.time_slice.{time_index}.profiles_2d.{grid_index}.b_field_tor'][mask] = (
             f_spl(psi_spl(r[mask], z[mask], grid=False)) / r[mask]
         )
-        ods[f'equilibrium.time_slice.{time_index}.profiles_2d.{grid_index}.b_tor'][mask == False] = (
+        ods[f'equilibrium.time_slice.{time_index}.profiles_2d.{grid_index}.b_field_tor'][mask == False] = (
             ods[f'equilibrium.time_slice.{time_index}.profiles_1d.f'][-1] / r[mask == False]
         )
         return ods