From c23239765e0de10d74b37140723774d6991adf7c Mon Sep 17 00:00:00 2001
From: Joseph McClenaghan <jmcclena@gmail.com>
Date: Wed, 5 Jun 2024 20:43:54 +0100
Subject: [PATCH 1/3] fixes for MAST maching mappings. Add MAST thomson scatter
 machine mapping

---
 omas/machine_mappings/mast.py | 51 ++++++++++++++++++++++++++++++++++-
 omas/omas_machine.py          |  5 +---
 2 files changed, 51 insertions(+), 5 deletions(-)

diff --git a/omas/machine_mappings/mast.py b/omas/machine_mappings/mast.py
index 778a71d1..00147fa5 100644
--- a/omas/machine_mappings/mast.py
+++ b/omas/machine_mappings/mast.py
@@ -312,13 +312,62 @@ def thomson_scattering_hardware(ods, pulse):
 
 
 @machine_mapping_function(__regression_arguments__, pulse=44653)
-def thomson_scattering_data(ods, pulse):
+def thomson_scattering_data(ods, pulse, server=None, port=None):
     """
     Loads MAST Thomson measurement data
 
     :param pulse: int
     """
 
+
+    if int(pulse) < 23000:
+        trace_te = "atm_te"
+        trace_ne = "atm_ne"
+        trace_r = "atm_R"
+    elif int(pulse) > 43000:
+        trace_te = "ayc/T_e"
+        trace_ne = "ayc/n_e"
+        trace_r = "ayc/r"
+    else:
+        trace_te = "ayc_te"
+        trace_ne = "ayc_ne"
+        trace_r = "ayc_r"
+
+    client = get_pyuda_client(server=server, port=port)
+
+    udasignal = client.get(trace_te, pulse)
+
+    Te_data = udasignal.data
+    Te_err = udasignal.errors
+    Te_time = udasignal.dims[0].data
+    udasignal = client.get(trace_ne, pulse)
+    ne_data = udasignal.data
+    ne_err = udasignal.errors
+    ne_time = udasignal.dims[0].data
+
+    r = client.get(trace_r, pulse).data
+
+    r = r[0]
+
+    for i, R in enumerate(r):
+        
+        ch = ods['thomson_scattering']['channel'][i]
+        ch['name'] = 'ch_' + str(i)
+        ch['identifier'] = 'ch_' + str(i)
+        ch['position']['r'] = R 
+        ch['position']['z'] = 0.0
+
+        mask = ~np.isnan(ne_data[:,i])
+
+        ch['n_e.time'] = ne_time[mask]
+        ch['n_e.data_error_upper'] =  ne_err[mask,i]
+        ch['n_e.data'] =  ne_data[mask,i]
+
+        mask = ~np.isinf(Te_data[:,i])
+        ch['t_e.time'] = Te_time[mask]
+        ch['t_e.data_error_upper'] =  Te_err[mask,i]
+        ch['t_e.data'] =  Te_data[mask,i]
+
     return
 
 @machine_mapping_function(__regression_arguments__, pulse=44653)
diff --git a/omas/omas_machine.py b/omas/omas_machine.py
index 4502acda..941be81f 100644
--- a/omas/omas_machine.py
+++ b/omas/omas_machine.py
@@ -6,7 +6,7 @@
 from .omas_utils import *
 from .omas_core import ODS, dynamic_ODS, omas_environment, omas_info_node, imas_json_dir, omas_rcparams
 from .omas_physics import cocos_signals
-from omas.machine_mappings import d3d, nstx, nstxu, east
+from omas.machine_mappings import d3d, nstx, nstxu, east, mast
 from omas.machine_mappings.d3d import __regression_arguments__
 from omas.utilities.machine_mapping_decorator import machine_mapping_function
 from omas.utilities.omas_mds import mdsvalue
@@ -215,9 +215,6 @@ def resolve_mapped(ods, machine, pulse,  mappings, location, idm, options_with_d
 
     # PYTHON
     elif 'PYTHON' in mapped:
-        if 'mast' in machine:
-            printe(f"MAST is currently not supported because of UDA", topic='machine')
-            return ods
         call = mapped['PYTHON'].format(**options_with_defaults)
         # python functions tend to set multiple locations at once
         # it is thus very beneficial to cache that

From a4a9bff1d83199f529899d509258ae8134dbd7a9 Mon Sep 17 00:00:00 2001
From: Joseph McClenaghan <jmcclena@gmail.com>
Date: Wed, 5 Jun 2024 21:25:09 +0100
Subject: [PATCH 2/3] put mast machine mapping import into try-except

---
 omas/omas_machine.py | 8 +++++++-
 1 file changed, 7 insertions(+), 1 deletion(-)

diff --git a/omas/omas_machine.py b/omas/omas_machine.py
index 941be81f..a8f31de3 100644
--- a/omas/omas_machine.py
+++ b/omas/omas_machine.py
@@ -6,7 +6,7 @@
 from .omas_utils import *
 from .omas_core import ODS, dynamic_ODS, omas_environment, omas_info_node, imas_json_dir, omas_rcparams
 from .omas_physics import cocos_signals
-from omas.machine_mappings import d3d, nstx, nstxu, east, mast
+from omas.machine_mappings import d3d, nstx, nstxu, east
 from omas.machine_mappings.d3d import __regression_arguments__
 from omas.utilities.machine_mapping_decorator import machine_mapping_function
 from omas.utilities.omas_mds import mdsvalue
@@ -16,6 +16,12 @@
 except:
     pass
 
+try:
+    from omas.machine_mappings import mast
+except ImportError:
+    pass
+
+
 __all__ = [
     'machine_expression_types',
     'machines',

From 1b75aad33aa2755c02f26e291cecf47bde82e955 Mon Sep 17 00:00:00 2001
From: Joseph McClenaghan <jmcclena@gmail.com>
Date: Wed, 5 Jun 2024 21:40:22 +0100
Subject: [PATCH 3/3] add print statement to mast mapping try except

---
 omas/omas_machine.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/omas/omas_machine.py b/omas/omas_machine.py
index a8f31de3..9ec54316 100644
--- a/omas/omas_machine.py
+++ b/omas/omas_machine.py
@@ -19,6 +19,7 @@
 try:
     from omas.machine_mappings import mast
 except ImportError:
+    print('Could not import mast machine mappings. Check that pyuda is installed')
     pass