From f0e5394fa4dc5cede1df0bb39e29c21eac620ec6 Mon Sep 17 00:00:00 2001 From: Severin Denk Date: Tue, 29 Nov 2022 18:44:59 -0800 Subject: [PATCH 1/8] Extended magnetic mapping - This my crude attempt to allow mapping between rho_pol_norm <-> rho_tor_norm - This also includes phi - I qualitatively tested rho_pol_norm -> rho_tor_norm - Might have some COCOS issues - Probably overly complicated - Needs dynamic fetching of `phi` and `rho_pol_norm` --- omas/omas_physics.py | 117 +++++++++++++++++++++++++++++++++++++++++-- 1 file changed, 114 insertions(+), 3 deletions(-) diff --git a/omas/omas_physics.py b/omas/omas_physics.py index 8d7895c91..6f1e8ff21 100644 --- a/omas/omas_physics.py +++ b/omas/omas_physics.py @@ -213,7 +213,73 @@ 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 @@ -226,10 +292,28 @@ 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"] + 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, "psi", + np.array([ods["equilibrium"]["time_slice"][time_index]["global_quantities"]["psi_sep"]])) / np.max(phi) + 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"]["psi"][psi_mask] + psi_spl = InterpolatedUnivariateSpline(phi_grid, psi_grid[psi_mask]) + phi_min = np.min(phi_grid) + phi_max = np.max(phi_grid) + values_mask = np.logical_ant(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.") @@ -247,10 +331,37 @@ def map_pol_flux_to_flux_coordinate(ods, time_index, destination, values): :return: Transformed values """ - if destination == "rho_pol": + 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] = phi**2 / phi_boundary**2 + rho_tor_norm /= np.max(rho_tor_norm[mask]) + 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]) + rho_tor_norm = np.zeros(values.shape) + mask = mask_SOL(ods, time_index, values) + phi = np.zeros(values.shape) + phi[mask] = phi_spl(values[mask]) + phi[mask==False] = np.inf + return phi else: raise NotImplementedError(f"Conversion to {destination} not yet implemented.") From 2f0e0b3a525d59f30156a58bf050688e7b1f8aab Mon Sep 17 00:00:00 2001 From: Severin Denk Date: Wed, 4 Jan 2023 21:02:50 -0800 Subject: [PATCH 2/8] Fixed some horrible math errors from last commit - Tested against mapping form OMFITprofiles - Results are OK by eye. Numerical verification still pending. --- omas/omas_physics.py | 13 +++++++++---- 1 file changed, 9 insertions(+), 4 deletions(-) diff --git a/omas/omas_physics.py b/omas/omas_physics.py index 6f1e8ff21..66bacb0f6 100644 --- a/omas/omas_physics.py +++ b/omas/omas_physics.py @@ -341,8 +341,7 @@ def map_pol_flux_to_flux_coordinate(ods, time_index, destination, values): 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] = phi**2 / phi_boundary**2 - rho_tor_norm /= np.max(rho_tor_norm[mask]) + rho_tor_norm[mask] = np.sqrt(phi / phi_boundary) rho_tor_norm[mask==False] = np.inf return rho_tor_norm elif destination == "phi": @@ -356,11 +355,17 @@ def map_pol_flux_to_flux_coordinate(ods, time_index, destination, values): 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]) - rho_tor_norm = np.zeros(values.shape) mask = mask_SOL(ods, time_index, values) phi = np.zeros(values.shape) phi[mask] = phi_spl(values[mask]) - phi[mask==False] = np.inf + 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.") From a6344447cc43bcbe1b367740d0f2af7f68e40dec Mon Sep 17 00:00:00 2001 From: AreWeDreaming Date: Thu, 16 Feb 2023 12:02:49 -0800 Subject: [PATCH 3/8] Fixed error in rho_tor to rho_pol conversion --- omas/omas_physics.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/omas/omas_physics.py b/omas/omas_physics.py index 66bacb0f6..6a8fb8918 100644 --- a/omas/omas_physics.py +++ b/omas/omas_physics.py @@ -298,8 +298,7 @@ def map_flux_coordinate_to_pol_flux(ods, time_index, origin, values): + 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, "psi", - np.array([ods["equilibrium"]["time_slice"][time_index]["global_quantities"]["psi_sep"]])) / np.max(phi) + phi *= np.array([ods["equilibrium"]["time_slice"][time_index]["global_quantities"]["psi_boundary"]]) / np.max(phi) return map_flux_coordinate_to_pol_flux(ods, time_index, "phi", phi) elif origin == "phi": from scipy.interpolate import InterpolatedUnivariateSpline From a35e79db1d1c2b1f908b29a70dc3dd4765f8c855 Mon Sep 17 00:00:00 2001 From: AreWeDreaming Date: Thu, 16 Feb 2023 12:39:01 -0800 Subject: [PATCH 4/8] Fixed some bugs, but the result is still wrong --- omas/omas_physics.py | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/omas/omas_physics.py b/omas/omas_physics.py index 5d8de0a33..3ea9f020d 100644 --- a/omas/omas_physics.py +++ b/omas/omas_physics.py @@ -297,7 +297,8 @@ def map_flux_coordinate_to_pol_flux(ods, time_index, origin, values): + ods["equilibrium"]["time_slice"][time_index]["global_quantities"]["psi_axis"]) elif origin == "rho_tor_norm": phi = values**2 - phi *= np.array([ods["equilibrium"]["time_slice"][time_index]["global_quantities"]["psi_boundary"]]) / np.max(phi) + 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 @@ -307,9 +308,9 @@ def map_flux_coordinate_to_pol_flux(ods, time_index, origin, values): psi_spl = InterpolatedUnivariateSpline(phi_grid, psi_grid[psi_mask]) phi_min = np.min(phi_grid) phi_max = np.max(phi_grid) - values_mask = np.logical_ant(values > phi_min, values < phi_max) + values_mask = np.logical_and(values > phi_min, values < phi_max) psi = np.zeros(values.shape) - psi = np.nan + psi[:] = np.nan psi[values_mask] = psi_spl(values[values_mask]) return psi else: From df2cf6b74ab5d2b7eec11c5cc692bccebd8e50bd Mon Sep 17 00:00:00 2001 From: AreWeDreaming Date: Thu, 16 Feb 2023 12:57:56 -0800 Subject: [PATCH 5/8] Fixed broken psi <-> conversion - Results seem good now --- omas/omas_physics.py | 9 ++++++--- 1 file changed, 6 insertions(+), 3 deletions(-) diff --git a/omas/omas_physics.py b/omas/omas_physics.py index 3ea9f020d..615b4ac0e 100644 --- a/omas/omas_physics.py +++ b/omas/omas_physics.py @@ -304,11 +304,14 @@ def map_flux_coordinate_to_pol_flux(ods, time_index, origin, values): 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"]["psi"][psi_mask] - psi_spl = InterpolatedUnivariateSpline(phi_grid, psi_grid[psi_mask]) + 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) + 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]) From fbdd940bc41e787e3ccd137c0c443c25494a9bac Mon Sep 17 00:00:00 2001 From: Severin Denk Date: Mon, 17 Jul 2023 13:15:40 -0700 Subject: [PATCH 6/8] Some changes that allow imas_set for ids objects --- omas/omas_imas.py | 36 ++++++++++++++++++++++++++++-------- 1 file changed, 28 insertions(+), 8 deletions(-) diff --git a/omas/omas_imas.py b/omas/omas_imas.py index 5069c178c..f24620604 100644 --- a/omas/omas_imas.py +++ b/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,29 @@ 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: + 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]): + 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 +177,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 From 6b0bf1660e2be9109e0ff8ddd5a1db9d11779bb9 Mon Sep 17 00:00:00 2001 From: Severin Denk Date: Wed, 19 Jul 2023 11:46:17 -0700 Subject: [PATCH 7/8] Fixed missing raised exception --- omas/utilities/omas_mds.py | 1 + 1 file changed, 1 insertion(+) diff --git a/omas/utilities/omas_mds.py b/omas/utilities/omas_mds.py index 65ed58ae4..07a4cdfd6 100644 --- a/omas/utilities/omas_mds.py +++ b/omas/utilities/omas_mds.py @@ -177,6 +177,7 @@ def mdsk(value): for item in ['server', 'treename', 'pulse']: txt += [f' - {item}: {getattr(self, item)}'] txt += [f' - TDI: {TDI}'] + raise _excp.__class__(str(_excp) + '\n' + '\n'.join(txt)) finally: if out_results is not None: From ad2dc068f885a99504001cbc8ba8af142a156846 Mon Sep 17 00:00:00 2001 From: Severin Denk Date: Thu, 27 Jul 2023 15:34:53 -0700 Subject: [PATCH 8/8] Changes needed to get ITPA IDA to work - OMAS can now store omas data in IDS in memory - Changed `b_r` and related to `b_field_r` - Added ece if_bandwidth for ECE --- omas/machine_mappings/d3d.json | 11 ++++++++++- omas/machine_mappings/d3d.py | 5 +++-- omas/omas_imas.py | 7 +++++-- omas/omas_physics.py | 16 ++++++++-------- 4 files changed, 26 insertions(+), 13 deletions(-) diff --git a/omas/machine_mappings/d3d.json b/omas/machine_mappings/d3d.json index 9ee7f1ccf..45a72add6 100644 --- a/omas/machine_mappings/d3d.json +++ b/omas/machine_mappings/d3d.json @@ -302,7 +302,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})" @@ -313,6 +313,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})" }, @@ -322,6 +325,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})" @@ -699,5 +705,8 @@ 2 ], "treename": "{EFIT_tree}" + }, + "wall.ids_properties.homogeneous_time": { + "VALUE": 1 } } \ No newline at end of file diff --git a/omas/machine_mappings/d3d.py b/omas/machine_mappings/d3d.py index b89177404..8f23436e7 100644 --- a/omas/machine_mappings/d3d.py +++ b/omas/machine_mappings/d3d.py @@ -764,7 +764,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'})" @@ -777,7 +777,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 = [{}, {}] @@ -804,6 +804,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) diff --git a/omas/omas_imas.py b/omas/omas_imas.py index f24620604..b4e6898d8 100644 --- a/omas/omas_imas.py +++ b/omas/omas_imas.py @@ -139,7 +139,7 @@ def imas_set(ids, path, value, skip_missing_nodes=False, allocate=False, ids_is_ location = l2i([ds] + path[: kp + 1]) if isinstance(p, str): if p == ":": - if allocate: + if allocate and len(out) != len(value): out.resize(len(value)) done = True if kp == len(path) - 1: @@ -150,7 +150,10 @@ def imas_set(ids, path, value, skip_missing_nodes=False, allocate=False, ids_is_ break else: for i in range(value.shape[0]): - imas_set(out[i], path[kp + 1:], value[i], skip_missing_nodes=False, allocate=allocate,only_allocate=only_allocate) + 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): diff --git a/omas/omas_physics.py b/omas/omas_physics.py index cc09f6288..ce30f8825 100644 --- a/omas/omas_physics.py +++ b/omas/omas_physics.py @@ -458,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']), @@ -478,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