Fix C-MOD TCI-related methods

disruption_py/machine/cmod/config.toml

@@ -144,6 +144,24 @@ description = "Vertical stability parameter, vacuum field index normalized to cr
 units = "dimensionless"
 validity = [-3, 3]
 
+[cmod.physics.attributes.ne_line_int_tci1]
+description = "TCI line-integrated electron density sampled at YAG 1 laser timestamps."
+imas = "/interferometer/channel(i1)/n_e_line"
+units = "m^-2"
+
+[cmod.physics.attributes.ne_line_int_tci2]
+description = "TCI line-integrated electron density sampled at YAG 2 laser timestamps."
+imas = "/interferometer/channel(i1)/n_e_line"
+units = "m^-2"
+
+[cmod.physics.attributes.ne_line_int_ts1]
+description = "Thomson scattering line-integrated electron density along TCI chord, sampled at YAG 1 laser timestamps."
+units = "m^-2"
+
+[cmod.physics.attributes.ne_line_int_ts2]
+description = "Thomson scattering line-integrated electron density along TCI chord, sampled at YAG 2 laser timestamps."
+units = "m^-2"
+
 [cmod.physics.attributes.ne_peaking]
 description = "Peaking factor of the electron density profile measured by the Thomson scattering diagnostic."
 units = "dimensionless"

disruption_py/machine/cmod/efit.py

@@ -116,7 +116,7 @@ def efit_check(params: PhysicsMethodParams):
             A tuple containing valid indices and corresponding times.
         """
         values = [
-            params.data_conn.get(expr, tree_name="analysis")
+            params.data_conn.get(expr, tree_name="_efit_tree")
             for expr in [
                 r"_lf=\efit_aeqdsk:lflag",
                 r"_l0=((sum(_lf,1) - _lf[*,20] - _lf[*,1])==0)",
@@ -126,6 +126,6 @@ def efit_check(params: PhysicsMethodParams):
         _n = values[2].data()
         valid_indices = np.nonzero(_n)
         (times,) = params.data_conn.get_dims(
-            r"\efit_aeqdsk:lflag", tree_name="analysis"
+            r"\efit_aeqdsk:lflag", tree_name="_efit_tree"
         )
         return valid_indices, times[valid_indices]

disruption_py/machine/cmod/physics.py

@@ -1472,6 +1472,53 @@ def get_peaking_factors(params: PhysicsMethodParams):
             params.times, ts_time, ts_te, ts_ne, ts_z, efit_time, bminor, z0
         )
 
+    @staticmethod
+    @physics_method(
+        columns=[
+            "ne_line_int_ts1",
+            "ne_line_int_ts2",
+            "ne_line_int_tci1",
+            "ne_line_int_tci2",
+        ],
+        tokamak=Tokamak.CMOD,
+    )
+    def get_ts_tci_comparison(params: PhysicsMethodParams):
+        """
+        Return Thomson scattering and TCI line-integrated density measurements
+        interpolated onto the standard timebase.
+
+        Returns
+        -------
+        dict
+            ne_line_int_ts1 : TS line-integrated density for YAG 1 timestamps [m^-2]
+            ne_line_int_ts2 : TS line-integrated density for YAG 2 timestamps [m^-2]
+            ne_line_int_tci1 : TCI line-integrated density at YAG 1 timestamps [m^-2]
+            ne_line_int_tci2 : TCI line-integrated density at YAG 2 timestamps [m^-2]
+        """
+        nl_ts1, nl_ts2, nl_tci1, nl_tci2, time1, time2 = (
+            CmodThomsonDensityMeasure.compare_ts_tci(params)
+        )
+
+        # If time1 or time2 is an int (-1) that means there's no valid data, replace with [np.nan]
+        if isinstance(time1, int):
+            ts1, tci1 = [np.nan], [np.nan]
+        else:
+            ts1 = interp1(time1, nl_ts1, params.times)
+            tci1 = interp1(time1, nl_tci1, params.times)
+
+        if isinstance(time2, int):
+            ts2, tci2 = [np.nan], [np.nan]
+        else:
+            ts2 = interp1(time2, nl_ts2, params.times)
+            tci2 = interp1(time2, nl_tci2, params.times)
+
+        return {
+            "ne_line_int_ts1": ts1,
+            "ne_line_int_ts2": ts2,
+            "ne_line_int_tci1": tci1,
+            "ne_line_int_tci2": tci2,
+        }
+
     @staticmethod
     def _get_te_profile_params_ece(
         times,

disruption_py/machine/cmod/thomson.py

@@ -62,8 +62,8 @@ def compare_ts_tci(params: PhysicsMethodParams, nlnum=4):
             ts_time2 = tci_time[indices2]
             (valid_indices,) = np.where((ts_time2 >= t0) & (ts_time2 <= t1))
             if valid_indices.size > 0:
-                nl_tci1 = interp1(tci_t, tci, ts_time2[valid_indices])
-                nl_ts1 = interp1(nlts_t, nlts, ts_time2[valid_indices])
+                nl_tci2 = interp1(tci_t, tci, ts_time2[valid_indices])
+                nl_ts2 = interp1(nlts_t, nlts, ts_time2[valid_indices])
                 time2 = ts_time2[valid_indices]
         return nl_ts1, nl_ts2, nl_tci1, nl_tci2, time1, time2
 
@@ -99,24 +99,34 @@ def _parse_yags(params: PhysicsMethodParams):
                 if nyag1 == nyag2:
                     indices1 = 2 * np.arange(nyag1)
                     indices2 = indices1 + 1
-                else:
-                    indices1 = 2 * np.arange(nyag1) + (nyag1 > nyag2)
+                elif nyag1 < nyag2:
+                    indices1 = 2 * np.arange(nyag1)
                     indices2 = np.concatenate(
                         (
-                            2 * np.arange(nyag2) + (nyag1 < nyag2),
-                            2 * nyag2 + np.arange(nyag1 - nyag2 - 1),
+                            2 * np.arange(nyag1) + 1,
+                            2 * nyag1 + np.arange(nyag2 - nyag1),
                         )
                     )
-        (v_ind1,) = np.where(indices1 < nt)
-        if nyag1 > 0 and v_ind1.size > 0:
-            indices1 = indices1[v_ind1]
-        else:
-            indices1 = -1
-        (v_ind2,) = np.where(indices2 < nt)
-        if nyag2 > 0 and v_ind2.size > 0:
-            indices2 = indices2[v_ind2]
-        else:
-            indices2 = -1
+                else:  # nyag1 > nyag2
+                    indices2 = 2 * np.arange(nyag2) + 1
+                    indices1 = np.concatenate(
+                        (
+                            2 * np.arange(nyag2),
+                            2 * nyag2 + np.arange(nyag1 - nyag2),
+                        )
+                    )
+        if isinstance(indices1, np.ndarray):
+            (v_ind1,) = np.where(indices1 < nt)
+            if nyag1 > 0 and v_ind1.size > 0:
+                indices1 = indices1[v_ind1]
+            else:
+                indices1 = -1
+        if isinstance(indices2, np.ndarray):
+            (v_ind2,) = np.where(indices2 < nt)
+            if nyag2 > 0 and v_ind2.size > 0:
+                indices2 = indices2[v_ind2]
+            else:
+                indices2 = -1
         return nyag1, nyag2, indices1, indices2
 
     @staticmethod
@@ -184,13 +194,10 @@ def _map_ts2tci(params: PhysicsMethodParams, nlnum):
         n_e = [1e32]
         n_e_sig = [1e32]
         flag = 1
-        valid_indices, efit_times = CmodEfitMethods.efit_check(params)
+        _, efit_times = CmodEfitMethods.efit_check(params)
         ip = params.data_conn.get_data(r"\ip", "cmod")
         if np.mean(ip) > 0:
             flag = 0
-        efit_times = params.data_conn.get_data(
-            r"\efit_aeqdsk:time", tree_name="_efit_tree"
-        )
         t1 = np.amin(efit_times)
         t2 = np.amax(efit_times)
         psia, psia_t = params.data_conn.get_data_with_dims(
@@ -207,17 +214,22 @@ def _map_ts2tci(params: PhysicsMethodParams, nlnum):
             ".YAG_NEW.RESULTS.PROFILES:Z_SORTED", tree_name="electrons"
         )
         mts_core = len(zts_core)
-        zts_edge = params.data_conn.get_data(r"\fiber_z")
+        zts_edge = params.data_conn.get_data(r"\fiber_z", tree_name="electrons")
         mts_edge = len(zts_edge)
         try:
-            nets_edge = params.data_conn.get_data(r"\ts_ne")
-            nets_edge_err = params.data_conn.get_data(r"\ts_ne_err")
+            nets_edge = params.data_conn.get_data(r"\ts_ne", tree_name="electrons")
+            nets_edge_err = params.data_conn.get_data(
+                r"\ts_ne_err", tree_name="electrons"
+            )
         except mdsExceptions.MdsException:
             nets_edge = np.zeros((len(nets_core[:, 1]), mts_edge))
             nets_edge_err = nets_edge + 1e20
         mts = mts_core + mts_edge
-        rts = params.data_conn.get(".YAG.RESULTS.PARAM:R") + np.zeros((1, mts))
-        rtci = params.data_conn.get_data(".tci.results:rad")
+        rts_float = params.data_conn.get_data(
+            ".YAG.RESULTS.PARAM:R", tree_name="electrons"
+        )
+        rts = np.full((1, mts), rts_float)  # (1, mts) array of chord location in R
+        rtci = params.data_conn.get_data(".tci.results:rad", tree_name="electrons")
         nts = len(nets_core_t)
         zts = np.zeros((1, mts))
         zts[:, :mts_core] = zts_core
@@ -237,7 +249,10 @@ def _map_ts2tci(params: PhysicsMethodParams, nlnum):
         psits = CmodThomsonDensityMeasure._efit_rz2psi(params, rts, zts, nets_core_t)
         mtci = 101
         ztci = -0.4 + 0.8 * np.arange(0, mtci) / (mtci - 1)
-        rtci = rtci[nlnum] + np.zeros((1, mtci))
+        # There are 10 TCI channels like NL_01, NL_02, ..., NL_10
+        # nlnum is 1-indexed to match the channel numbering
+        # but we need to convert it to 0-indexed for array indexing
+        rtci = rtci[nlnum - 1] + np.zeros((1, mtci))
         psitci = CmodThomsonDensityMeasure._efit_rz2psi(params, rtci, ztci, nets_core_t)
         psia = interp1(psia_t, psia, nets_core_t)
         psi_0 = interp1(psia_t, psi_0, nets_core_t)
@@ -308,9 +323,9 @@ def _efit_rz2psi(params: PhysicsMethodParams, r, z, t, tree="analysis"):
             # Find the index of the closest time
             time_idx = np.argmin(np.abs(times - time))
             # Extract the corresponding psirz slice and transpose it
-            psirz = np.transpose(psirz[time_idx, :, :])
+            psirz_t = np.transpose(psirz[time_idx, :, :])
             # Perform cubic interpolation on the psirz slice
-            values = psirz.flatten()
+            values = psirz_t.flatten()
             psi[:, i] = scipy.interpolate.griddata(
                 points, values, (r, z), method="cubic"
             )