Fix CER loading

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

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)