From 80ba33574b29245d3925370ea08cdcca4f15cfae Mon Sep 17 00:00:00 2001 From: AreWeDreaming Date: Tue, 19 May 2026 18:06:40 -0700 Subject: [PATCH 1/5] Follow IMPDENS data to its source - This prevents MDSplus from closing the IONS tree on you - Some light refactoring of omas_mds to allow `getnci` to be executed outside of mdsvalue --- omas/machine_mappings/d3d.py | 46 ++++++++++++++---- omas/utilities/omas_mds.py | 92 ++++++++++++++++++++++-------------- 2 files changed, 92 insertions(+), 46 deletions(-) diff --git a/omas/machine_mappings/d3d.py b/omas/machine_mappings/d3d.py index c59b18cb8..46ec7cafd 100644 --- a/omas/machine_mappings/d3d.py +++ b/omas/machine_mappings/d3d.py @@ -2,13 +2,14 @@ from inspect import unwrap from scipy.signal import medfilt from collections import OrderedDict +import re from omas import * from omas.omas_utils import printd, printe from omas.machine_mappings._common import * from uncertainties import unumpy from omas.utilities.machine_mapping_decorator import machine_mapping_function -from omas.utilities.omas_mds import mdsvalue +from omas.utilities.omas_mds import mdsvalue, exec_tdi from omas.omas_core import ODS from omas.omas_structure import add_extra_structures from omas.omas_physics import omas_environment @@ -1574,28 +1575,53 @@ def charge_exchange_data(ods, pulse, analysis_type='CERQUICK', _measurements=Tru # fetch TDIs = {} + + # look up reference + look_up = {} + # Number of channels in each system + n_ch = {} for sub in subsystems: - for channel in range(1,100): + n_ch[sub] = len(exec_tdi('d3d', 'IONS', pulse, f'getnci("CER.{analysis_type}.{sub}.CHANNEL*:TIME","LENGTH")')) + for channel in range(1, n_ch[sub]+1): for pos in ['TIME', 'R', 'Z', 'VIEW_PHI']: - TDIs[f'{sub}_{channel}_{pos}'] = f"\\IONS::TOP.CER.{analysis_type}.{sub}.CHANNEL{channel:02d}.{pos}" + TDIs[f'{sub}_{channel}_{pos}'] = f"CER.{analysis_type}.{sub}.CHANNEL{channel:02d}.{pos}" if _measurements: for pos in ['TEMP', 'TEMP_ERR', 'ROT', 'ROT_ERR']: if sub == 'TANGENTIAL' and pos == 'ROT': pos1 = 'ROTC' else: pos1 = pos - TDIs[f'{sub}_{channel}_{pos}__data'] = f"\\IONS::TOP.CER.{analysis_type}.{sub}.CHANNEL{channel:02d}.{pos1}" - TDIs[f'{sub}_{channel}_{pos}__time'] = f"dim_of(\\IONS::TOP.CER.{analysis_type}.{sub}.CHANNEL{channel:02d}.{pos1}, 0)/1000" + TDIs[f'{sub}_{channel}_{pos}__data'] = f"CER.{analysis_type}.{sub}.CHANNEL{channel:02d}.{pos1}" + TDIs[f'{sub}_{channel}_{pos}__time'] = f"dim_of(CER.{analysis_type}.{sub}.CHANNEL{channel:02d}.{pos1}, 0)/1000" for pos in ['FZ', 'ZEFF']: - TDIs[f'{sub}_{channel}_{pos}__data'] = f"\\IONS::TOP.IMPDENS.{analysis_type}.{pos}{sub[0]}{channel}" - TDIs[f'{sub}_{channel}_{pos}__time'] = f"dim_of(\\IONS::TOP.IMPDENS.{analysis_type}.{pos}{sub[0]}{channel}, 0)/1000" - + look_up[f'{sub}_{channel}_{pos}__data'] = f"TCL('decomp IMPDENS.{analysis_type}.{pos}{sub[0]}{channel}', _output), _output" + + references = mdsvalue('d3d', treename='IONS', pulse=pulse, TDI=look_up).raw() + impcon_TDIs = {} + impcon_tree_name = None + SIGNAL_PATTERN = re.compile(r'::TOP\.([A-Z0-9_.:]+?)[\s",].*?"([A-Z0-9_]+)"') + for key, path in references.items(): + if "error" in path: + continue + match = SIGNAL_PATTERN.search(path) + if not match: + print(f"Failed to resolve {key}'s true location from {path}") + continue + new_path = match.group(1) + tree_name = match.group(2) + if impcon_tree_name is None: + impcon_tree_name = tree_name + else: + assert impcon_tree_name==tree_name, "References to multiple IMCPON trees in one IMPDENS analysis type are not supported." + impcon_TDIs[key] = new_path + impcon_TDIs[key.replace("_data", "_time")] = f"dim_of({new_path},0)/1000" # fetch data = mdsvalue('d3d', treename='IONS', pulse=pulse, TDI=TDIs).raw() - + data = data | mdsvalue('d3d', treename=impcon_tree_name, pulse=pulse, TDI=impcon_TDIs).raw() + # assign for sub in subsystems: - for channel in range(1,100): + for channel in range(1, n_ch[sub]+1): postime = data[f'{sub}_{channel}_TIME'] if isinstance(postime, Exception): continue diff --git a/omas/utilities/omas_mds.py b/omas/utilities/omas_mds.py index 9992d4a55..d2be1fdbd 100644 --- a/omas/utilities/omas_mds.py +++ b/omas/utilities/omas_mds.py @@ -2,6 +2,12 @@ import os from omas.omas_utils import printd import numpy as np +try: + import MDSplus + +except ImportError: + print("Warning no MDSplus! No machine mappings available.") + MDSplus = None __all__ = [ 'mdstree', @@ -79,36 +85,57 @@ def tunnel_mds(server, treename): return server.format(**os.environ) +def get_cached_connection(server, pulse, treename): + for fallback in [0, 1]: + if server not in _mds_connection_cache: + _mds_connection_cache[server] = MDSplus.Connection(server) + _last_open_tree[server] = None + try: + conn = _mds_connection_cache[server] + if treename is not None: + open_tree_cached(conn, server, pulse, treename) + break + except Exception as _excp: + if server in _mds_connection_cache: + _last_open_tree[server] = None + _mds_connection_cache[server].reconnect() + if fallback: + raise + return conn - - +def resolve_server(machine, treename): + if 'nstx' in machine: + old_MDS_server = True + try: + # handle the case that server is just the machine name + machine_mappings_path = os.path.join(os.path.dirname(__file__), "../", "machine_mappings") + machine_mappings_path = os.path.join(machine_mappings_path, machine + ".json") + with open(machine_mappings_path, "r") as machine_file: + server = json.load(machine_file)["__mdsserver__"] + except Exception: + # handle case where server is actually a URL + server = tunnel_mds(server, treename) + if '.' not in server: + raise + + old_servers = ['skylark.pppl.gov:8500', 'skylark.pppl.gov:8501', 'skylark.pppl.gov:8000'] + if server in old_servers or machine in old_servers: + return server, True + return server, False class mdsvalue(dict): """ Execute MDSplus TDI functions """ - def __init__(self, server, treename, pulse, TDI, old_MDS_server=False): + def __init__(self, machine, treename, pulse, TDI, old_MDS_server=False): self.treename = treename self.pulse = pulse self.TDI = TDI - if 'nstx' in server: - old_MDS_server = True - try: - # handle the case that server is just the machine name - machine_mappings_path = os.path.join(os.path.dirname(__file__), "../", "machine_mappings") - machine_mappings_path = os.path.join(machine_mappings_path, server + ".json") - with open(machine_mappings_path, "r") as machine_file: - server = json.load(machine_file)["__mdsserver__"] - except Exception: - # hanlde case where server is actually a URL - if '.' not in server: - raise - self.server = tunnel_mds(server, self.treename) - old_servers = ['skylark.pppl.gov:8500', 'skylark.pppl.gov:8501', 'skylark.pppl.gov:8000'] - if server in old_servers or self.server in old_servers: - old_MDS_server = True - self.old_MDS_server = old_MDS_server + self.server, self.old_MDS_server = resolve_server(machine, treename) + if old_MDS_server: + self.old_MDS_server = old_MDS_server + def data(self): return self.raw(f'data({self.TDI})') @@ -144,7 +171,6 @@ def raw(self, TDI=None): import time t0 = time.time() - import MDSplus def mdsk(value): """ @@ -159,21 +185,7 @@ def mdsk(value): out_results = None # try connecting and re-try on fail - for fallback in [0, 1]: - if self.server not in _mds_connection_cache: - _mds_connection_cache[self.server] = MDSplus.Connection(self.server) - _last_open_tree[self.server] = None - try: - conn = _mds_connection_cache[self.server] - if self.treename is not None: - open_tree_cached(conn, self.server, self.pulse, self.treename) - break - except Exception as _excp: - if self.server in _mds_connection_cache: - _last_open_tree[self.server] = None - _mds_connection_cache[self.server].reconnect() - if fallback: - raise + conn = get_cached_connection(self.server, self.pulse, self.treename) # list of TDI expressions if isinstance(TDI, (list, tuple)): @@ -276,3 +288,11 @@ def __init__(self, server, treename, pulse): h[path[-1]] = mdsvalue(server, treename, pulse, TDI) else: h[path[-1]].TDI = TDI + +def exec_tdi(machine, treename, pulse, tdi_command): + """ + Simple helper function to extract number of active channels for things like CER + """ + server, _ = resolve_server(machine, treename) + conn = get_cached_connection(server, pulse, treename) + return conn.get(tdi_command) \ No newline at end of file From 9e21020fc07f30d69cdb64d2937e6c23ab18cf9d Mon Sep 17 00:00:00 2001 From: AreWeDreaming Date: Wed, 20 May 2026 12:10:04 -0700 Subject: [PATCH 2/5] Use TIME length as filter for active --- omas/machine_mappings/d3d.py | 8 +++++++- 1 file changed, 7 insertions(+), 1 deletion(-) diff --git a/omas/machine_mappings/d3d.py b/omas/machine_mappings/d3d.py index 92e385cce..f66f23758 100644 --- a/omas/machine_mappings/d3d.py +++ b/omas/machine_mappings/d3d.py @@ -1580,9 +1580,13 @@ def charge_exchange_data(ods, pulse, analysis_type='CERQUICK', _measurements=Tru look_up = {} # Number of channels in each system n_ch = {} + active_channels = {} for sub in subsystems: - n_ch[sub] = len(exec_tdi('d3d', 'IONS', pulse, f'getnci("CER.{analysis_type}.{sub}.CHANNEL*:TIME","LENGTH")')) + active_channels[sub] = np.asarray(exec_tdi('d3d', 'IONS', pulse, f'getnci("CER.{analysis_type}.{sub}.CHANNEL*:TIME","LENGTH")')) > 0 + n_ch[sub] = len(active_channels[sub]) for channel in range(1, n_ch[sub]+1): + if not active_channels[sub][channel - 1]: + continue for pos in ['TIME', 'R', 'Z', 'VIEW_PHI']: TDIs[f'{sub}_{channel}_{pos}'] = f"CER.{analysis_type}.{sub}.CHANNEL{channel:02d}.{pos}" if _measurements: @@ -1622,6 +1626,8 @@ def charge_exchange_data(ods, pulse, analysis_type='CERQUICK', _measurements=Tru # assign for sub in subsystems: for channel in range(1, n_ch[sub]+1): + if not active_channels[sub][channel - 1]: + continue postime = data[f'{sub}_{channel}_TIME'] if isinstance(postime, Exception): continue From 5bc43f6021f4786798abe46551e297e2c3fdd094 Mon Sep 17 00:00:00 2001 From: AreWeDreaming Date: Wed, 20 May 2026 12:19:02 -0700 Subject: [PATCH 3/5] Handle missing MDSplus more gracefully --- omas/utilities/omas_mds.py | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) diff --git a/omas/utilities/omas_mds.py b/omas/utilities/omas_mds.py index d2be1fdbd..212fdd5a2 100644 --- a/omas/utilities/omas_mds.py +++ b/omas/utilities/omas_mds.py @@ -5,9 +5,9 @@ try: import MDSplus -except ImportError: +except ImportError as e: print("Warning no MDSplus! No machine mappings available.") - MDSplus = None + MDSplus = e __all__ = [ 'mdstree', @@ -86,6 +86,8 @@ def tunnel_mds(server, treename): return server.format(**os.environ) def get_cached_connection(server, pulse, treename): + if type(MDSplus) == ModuleNotFoundError: + raise MDSplus for fallback in [0, 1]: if server not in _mds_connection_cache: _mds_connection_cache[server] = MDSplus.Connection(server) @@ -129,6 +131,8 @@ class mdsvalue(dict): """ def __init__(self, machine, treename, pulse, TDI, old_MDS_server=False): + if type(MDSplus) == ModuleNotFoundError: + raise MDSplus self.treename = treename self.pulse = pulse self.TDI = TDI From 42a807751b12c36cee660b0b2839b6be1b0ad0ca Mon Sep 17 00:00:00 2001 From: AreWeDreaming Date: Wed, 20 May 2026 13:41:21 -0700 Subject: [PATCH 4/5] Python <3.,9 compatibility --- omas/machine_mappings/d3d.py | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/omas/machine_mappings/d3d.py b/omas/machine_mappings/d3d.py index 92e385cce..b200768a5 100644 --- a/omas/machine_mappings/d3d.py +++ b/omas/machine_mappings/d3d.py @@ -3,6 +3,7 @@ from scipy.signal import medfilt from collections import OrderedDict import re +import sys from omas import * from omas.omas_utils import printd, printe @@ -1617,9 +1618,11 @@ def charge_exchange_data(ods, pulse, analysis_type='CERQUICK', _measurements=Tru impcon_TDIs[key.replace("_data", "_time")] = f"dim_of({new_path},0)/1000" # fetch data = mdsvalue('d3d', treename='IONS', pulse=pulse, TDI=TDIs).raw() - data = data | mdsvalue('d3d', treename=impcon_tree_name, pulse=pulse, TDI=impcon_TDIs).raw() + if sys.version_info >= (3, 9): + data = data | mdsvalue('d3d', treename=impcon_tree_name, pulse=pulse, TDI=impcon_TDIs).raw() + else: + data = {**data, **mdsvalue('d3d', treename=impcon_tree_name, pulse=pulse, TDI=impcon_TDIs).raw()} - # assign for sub in subsystems: for channel in range(1, n_ch[sub]+1): postime = data[f'{sub}_{channel}_TIME'] From 8c457d540ce5914c290e125322c46e8cf126b5db Mon Sep 17 00:00:00 2001 From: AreWeDreaming Date: Wed, 20 May 2026 19:08:20 -0700 Subject: [PATCH 5/5] Add ROT/TEMP_ERR_PS for unc --- omas/machine_mappings/d3d.py | 10 +++++++--- 1 file changed, 7 insertions(+), 3 deletions(-) diff --git a/omas/machine_mappings/d3d.py b/omas/machine_mappings/d3d.py index 47b355fed..d8e157852 100644 --- a/omas/machine_mappings/d3d.py +++ b/omas/machine_mappings/d3d.py @@ -1591,7 +1591,7 @@ def charge_exchange_data(ods, pulse, analysis_type='CERQUICK', _measurements=Tru for pos in ['TIME', 'R', 'Z', 'VIEW_PHI']: TDIs[f'{sub}_{channel}_{pos}'] = f"CER.{analysis_type}.{sub}.CHANNEL{channel:02d}.{pos}" if _measurements: - for pos in ['TEMP', 'TEMP_ERR', 'ROT', 'ROT_ERR']: + for pos in ['TEMP', 'TEMP_ERR', 'TEMP_ERR_PS', 'ROT', 'ROT_ERR', 'ROT_ERR_PS']: if sub == 'TANGENTIAL' and pos == 'ROT': pos1 = 'ROTC' else: @@ -1646,10 +1646,14 @@ def charge_exchange_data(ods, pulse, analysis_type='CERQUICK', _measurements=Tru if _measurements: if not isinstance(data[f'{sub}_{channel}_TEMP__data'], Exception): ch['ion.0.t_i.time'] = data[f'{sub}_{channel}_TEMP__time'] - ch['ion.0.t_i.data'] = unumpy.uarray(data[f'{sub}_{channel}_TEMP__data'], data[f'{sub}_{channel}_TEMP_ERR__data']) + ch['ion.0.t_i.data'] = unumpy.uarray(data[f'{sub}_{channel}_TEMP__data'], + data[f'{sub}_{channel}_TEMP_ERR_PS__data'] + + data[f'{sub}_{channel}_TEMP_ERR__data']) if not isinstance(data[f'{sub}_{channel}_ROT__data'], Exception): ch['ion.0.velocity_tor.time'] = data[f'{sub}_{channel}_ROT__time'] - ch['ion.0.velocity_tor.data'] = unumpy.uarray(data[f'{sub}_{channel}_ROT__data'] * 1000.0, data[f'{sub}_{channel}_ROT_ERR__data'] * 1000.0) # from Km/s to m/s + ch['ion.0.velocity_tor.data'] = unumpy.uarray(data[f'{sub}_{channel}_ROT__data'] * 1000.0, + data[f'{sub}_{channel}_ROT_ERR_PS__data'] * 1000.0 + + data[f'{sub}_{channel}_ROT_ERR__data'] * 1000.0) # from Km/s to m/s if not isinstance(data[f'{sub}_{channel}_FZ__data'], Exception): ch['ion.0.n_i_over_n_e.time'] = data[f'{sub}_{channel}_FZ__time'] ch['ion.0.n_i_over_n_e.data'] = data[f'{sub}_{channel}_FZ__data'] * 0.01