Revive DIII-D get_n1_bradial_parameters

disruption_py/machine/d3d/physics.py

@@ -902,6 +902,105 @@ def get_z_parameters(params: PhysicsMethodParams):
             z_cur_norm = z_cur / nominal_flattop_radius
         return {"zcur": z_cur, "zcur_normalized": z_cur_norm}
 
+    @staticmethod
+    @physics_method(columns=["btor"], tokamak=Tokamak.D3D)
+    def get_btor(params: PhysicsMethodParams):
+        """
+        Calculate the toroidal magnetic field.
+
+        Parameters
+        ----------
+        params : PhysicsMethodParams
+            Parameters containing MDS connection and shot information
+
+        Returns
+        -------
+        dict
+            A dictionary containing the toroidal magnetic field (`btor`).
+
+        References
+        -------
+        - original sources: [get_n1_bradial_d3d.m](https://github.com/MIT-PSFC/disruption
+        -py/blob/matlab/DIII-D/get_n1_bradial_d3d.m), [get_n1rms_d3d.m](https://github.com
+        /MIT-PSFC/disruption-py/blob/matlab/DIII-D/get_n1rms_d3d.m)
+        """
+        b_tor, t_b_tor = params.data_conn.get_data_with_dims(
+            f"ptdata('bt', {params.shot_id})"
+        )  # [T], [ms]
+        t_b_tor /= 1e3  # [ms] -> [s]
+        b_tor = interp1(t_b_tor, b_tor, params.times)
+        return {"btor": b_tor}
+
+    @staticmethod
+    @physics_method(
+        columns=["n_equal_1_normalized", "n_equal_1_mode"],
+        tokamak=Tokamak.D3D,
+    )
+    def get_n1_bradial_parameters(params: PhysicsMethodParams):
+        """
+        This method obtains the n=1 radial magnetic field information (`n_equal_1_mode`) from
+        the ESLD coils (units of Tesla). It also calculates the normalized radial magnetic field
+        by dividing by the toroidal B-field (`n_equal_1_normalized`).
+
+        Parameters
+        ----------
+        params : PhysicsMethodParams
+            The parameters containing the MDSplus connection, shot id and more.
+
+        Returns
+        -------
+        dict
+            A dictionary containing `n_equal_1_mode` and `n_equal_1_normalized`.
+
+        References
+        -------
+        - original source: [get_n1_bradial_d3d.m](https://github.com/MIT-PSFC/disruption-py
+        /blob/matlab/DIII-D/get_n1_bradial_d3d.m)
+        - pull requests: #[500](https://github.com/MIT-PSFC/disruption-py/pull/500)
+        - issues: #[261](https://github.com/MIT-PSFC/disruption-py/issues/261)
+        """
+        # The following shots are missing bradial calculations in MDSplus and
+        # must be loaded from a separate datafile
+        if 156199 <= params.shot_id <= 172430:
+            # TODO: load data from custom tree structures
+            raise NotImplementedError
+        if 176030 <= params.shot_id <= 176912:
+            # TODO: load data from NetCDF file
+            raise NotImplementedError
+
+        try:
+            # Get data from the ONFR system
+            n_equal_1_mode, t_n1 = params.data_conn.get_data_with_dims(
+                f"ptdata('onsbradial', {params.shot_id})",
+            )  # [G], [ms]
+        except mdsExceptions.MdsException:
+            # Fallback: get data from the legacy DUD system
+            params.logger.verbose(
+                "get_n1_bradial_parameters: Failed to get ONSBRADIAL signal. "
+                "Retrieving from DUSBRADIAL instead."
+            )
+            n_equal_1_mode, t_n1 = params.data_conn.get_data_with_dims(
+                f"ptdata('dusbradial', {params.shot_id})",
+            )  # [G], [ms]
+        t_n1 /= 1e3  # [ms] -> [s]
+        n_equal_1_mode *= 1.0e-4  # [G] -> [T]
+        n_equal_1_mode = interp1(t_n1, n_equal_1_mode, params.times)
+
+        # Calculate n_equal_1_normalized
+        try:
+            b_tor = D3DPhysicsMethods.get_btor(params)["btor"]
+            n_equal_1_normalized = n_equal_1_mode / np.abs(b_tor)
+        except mdsExceptions.MdsException:
+            params.logger.warning(
+                "get_n1_bradial_parameters: Failed to get b_tor signal "
+                "to compute n_equal_1_normalized. Returning NaNs."
+            )
+            n_equal_1_normalized = np.array([np.nan])
+        return {
+            "n_equal_1_mode": n_equal_1_mode,
+            "n_equal_1_normalized": n_equal_1_normalized,
+        }
+
     @staticmethod
     @physics_method(columns=["n1rms", "n1rms_normalized"], tokamak=Tokamak.D3D)
     def get_n1rms_parameters(params: PhysicsMethodParams):
@@ -932,18 +1031,14 @@ def get_n1rms_parameters(params: PhysicsMethodParams):
         n1rms = interp1(t_n1rms, n1rms, params.times)
         # Calculate n1rms_norm
         try:
-            b_tor, t_b_tor = params.data_conn.get_data_with_dims(
-                f"ptdata('bt', {params.shot_id})"
-            )
-            t_b_tor /= 1e3  # [ms] -> [s]
-            b_tor = interp1(t_b_tor, b_tor, params.times)  # [T]
+            b_tor = D3DPhysicsMethods.get_btor(params)["btor"]
             n1rms_norm = n1rms / np.abs(b_tor)
-        except mdsExceptions.MdsException as e:
+        except mdsExceptions.MdsException:
             params.logger.warning(
-                "Failed to get b_tor signal to compute n1rms_normalized"
+                "get_n1rms_parameters: Failed to get b_tor signal "
+                "to compute n1rms_normalized. Returning NaNs."
             )
-            params.logger.opt(exception=True).debug(e)
-            n1rms_norm = [np.nan]
+            n1rms_norm = np.array([np.nan])
         return {"n1rms": n1rms, "n1rms_normalized": n1rms_norm}
 
     # TODO: Need to test and unblock recalculating peaking factors