Implement generic current time computation method

disruption_py/machine/cmod/config.toml

@@ -37,6 +37,10 @@ description = "Magnetic chi^2 (EFIT)."
 units = "dimensionless"
 validity = [0.1, 40]
 
+[cmod.physics.attributes.current_quench_time]
+description = "Time of the current quench event. NaN for non-disruptive discharges."
+units = "s"
+
 [cmod.physics.attributes.dbetap_dt]
 description = "Time derivative of the poloidal beta (EFIT)."
 units = "s^-1"
@@ -312,6 +316,14 @@ imas = "/equilibrium/time_slice(itime)/global_quantities/current_centre/z"
 units = "m"
 validity = [-0.5, 0.5]
 
+[cmod.physics.current_quench_time_thresholds]
+abs_ip0 = 0.1e6
+abs_ip_final = 50e3
+abs_ip_max = 0.1e6
+ip0_over_ip_max = 0.33
+ip0_over_max_didt = 0.05
+shot_duration = 0.1
+
 [cmod.physics.time_domain_thresholds]
 dipprog_dt = 50e3
 ip_prog = 100e3

disruption_py/machine/d3d/config.toml

@@ -10,6 +10,14 @@ driver = "FreeTDS"
 host = "d3drdb"
 port = 8001
 
+[d3d.physics.current_quench_time_thresholds]
+abs_ip0 = 0.1e6
+abs_ip_final = 100e3
+abs_ip_max = 0.1e6
+ip0_over_ip_max = 0.33
+ip0_over_max_didt = 0.05
+shot_duration = 0.5
+
 [d3d.physics.time_domain_thresholds]
 dipprog_dt = 2e3
 ip_prog = 100e3

disruption_py/machine/east/config.toml

@@ -7,6 +7,14 @@ driver = "MySQL"
 host = "202.127.205.10"
 port = 3306
 
+[east.physics.current_quench_time_thresholds]
+abs_ip0 = 0.2e6
+abs_ip_final = 100e3
+abs_ip_max = 0.2e6
+ip0_over_ip_max = 0.33
+ip0_over_max_didt = 0.05
+shot_duration = 0.6
+
 [east.physics.time_domain_thresholds]
 dipprog_dt = 1e3
 

disruption_py/machine/generic/physics.py

@@ -9,9 +9,12 @@
 from disruption_py.config import config
 from disruption_py.core.physics_method.decorator import physics_method
 from disruption_py.core.physics_method.params import PhysicsMethodParams
+from disruption_py.inout.mds import mdsExceptions
 from disruption_py.machine.cmod import CmodPhysicsMethods
 from disruption_py.machine.d3d import D3DPhysicsMethods
 from disruption_py.machine.east import EastPhysicsMethods
+from disruption_py.machine.east.util import EastUtilMethods
+from disruption_py.machine.generic.util import GenericUtilMethods
 from disruption_py.machine.mast.physics import MastPhysicsMethods
 from disruption_py.machine.tokamak import Tokamak
 
@@ -115,3 +118,155 @@ def get_time_domain(params: PhysicsMethodParams):
             time_domain[flattop_end:] = 3
 
         return {"time_domain": time_domain}
+
+    @staticmethod
+    @physics_method(columns=["current_quench_time"])
+    def get_current_quench_time(params: PhysicsMethodParams):
+        """
+        Determine and compute the current quench time of a shot. If a shot is determined
+        to be non-disruptive, return the current quench time as NaN. The criteria for a
+        disruptive shot are as follows:
+
+        - 1. Shot duration > duration_min: reject very short shots.
+        - 2. abs(Ip_max) > ip_threshold: reject very low current shots.
+        - 3. Ip0 / Ip_max > rampdown_threshold: reject shots that disrupt late in current ramp down.
+        - 4. -Ip0 / max_dIdt < tau_CQ_max: reject shots with relatively slow current decay.
+        - 5. abs(Ip_final) < Ip_final_max: reject minor disruptions.
+        - 6. abs(Ip0) > Ip_threshold: reject very low current disruptions.
+
+        Parameters
+        ----------
+        params : PhysicsMethodParams
+            The parameters containing the data connection, shot id and more.
+
+        Returns
+        -------
+        dict
+            A dictionary containing the `current_quench_time`.
+
+        References
+        -------
+        - original source:
+            - cmod: [test_disrupt.pro]
+            - d3d: [test_for_disruption.m]
+            - east: [test_for_disruption.m]
+            - kstar: [test_for_disruption_kstar.m]
+        - pull requests: #[545](https://github.com/MIT-PSFC/disruption-py/pull/545)
+        - issues: #[223](https://github.com/MIT-PSFC/disruption-py/issues/223)
+        """
+        # Initialize test criteria
+        criteria = {
+            "shot_duration": lambda duration, duration_min: duration > duration_min,
+            "abs_ip_max": lambda ip_max, ip_threshold: abs(ip_max) > ip_threshold,
+            "ip0_over_ip_max": lambda ip0_over_ip_max, rampdown_threshold: ip0_over_ip_max
+            > rampdown_threshold,
+            "ip0_over_max_didt": lambda ip0_over_max_didt, tau_cq_max: -ip0_over_max_didt
+            < tau_cq_max,
+            "abs_ip_final": lambda ip_final, ip_final_max: abs(ip_final) < ip_final_max,
+            "abs_ip0": lambda ip0, ip_threshold: abs(ip0) > ip_threshold,
+        }
+
+        # Get ip and ip timebase
+        if params.tokamak == Tokamak.CMOD:
+            try:
+                ip, t_ip = params.data_conn.get_data_with_dims(
+                    r"\ip", tree_name="magnetics"
+                )  # [A], [s]
+            except mdsExceptions.MdsException:
+                params.logger.warning(
+                    "Failed to get measured plasma current parameters. "
+                    "Skip current quench time computation."
+                )
+        elif params.tokamak == Tokamak.D3D:
+            try:
+                ip, t_ip = params.data_conn.get_data_with_dims(
+                    f"ptdata('ip', {params.shot_id})"
+                )  # [A], [ms]
+                t_ip = t_ip / 1.0e3  # [ms] -> [s]
+            except mdsExceptions.MdsException:
+                params.logger.warning(
+                    "Failed to get measured plasma current parameters. "
+                    "Skip current quench time computation."
+                )
+            # Subtract baseline offset to ip
+            (baseline_indices,) = np.where(t_ip <= 0)
+            if len(baseline_indices) > 0:
+                ip_baseline = np.mean(ip[baseline_indices])
+                ip -= ip_baseline
+        elif params.tokamak == Tokamak.EAST:
+            ip, t_ip = EastUtilMethods.retrieve_ip(params, params.shot_id)
+        elif params.tokamak == Tokamak.HBTEP:
+            ip, t_ip = params.data_conn.get_data_with_dims(
+                r"\top.sensors.rogowskis:ip", tree_name="hbtep2"
+            )  # [A], [s]
+        elif params.tokamak == Tokamak.MAST:
+            ip = params.data_conn.get_data("summary/ip")
+            t_ip = params.data_conn.get_data("summary/time")
+        else:
+            raise NotImplementedError
+
+        # Get test thresholds
+        thresholds = config(params.tokamak).physics.current_quench_time_thresholds
+
+        # Compute parameters for the tests
+        # Get the duration and polarity
+        end_of_current_params = GenericUtilMethods.get_end_of_current(
+            ip=ip, ip_time=t_ip, threshold=thresholds["abs_ip_max"]
+        )
+        duration = end_of_current_params["duration"]
+        polarity = end_of_current_params["polarity"]
+
+        # Find the maximum plasma current excluding the current spike
+        ip_upright = ip * polarity
+        (time_indices,) = np.where((t_ip > 0) & (t_ip < duration - 0.050))
+        ip_max = max(ip_upright[time_indices]) * polarity
+
+        # Find the plasma current right before the end of the discharge
+        (time_indices,) = np.where((t_ip > duration - 0.06) & (t_ip < duration - 0.04))
+        if len(time_indices) == 0:
+            raise ValueError(
+                "Invalid timebase for shot {param.shot_id}. Cannot compute current quench time."
+            )
+        candidate_ip0 = np.mean(ip_upright[time_indices]) * polarity
+
+        # Compute dI/dt during the latter part of the discharge
+        (time_indices,) = np.where((t_ip > duration - 0.05) & (t_ip < duration + 0.05))
+        didt_upright = np.diff(ip_upright[time_indices]) / np.diff(t_ip[time_indices])
+        indx = np.argmin(didt_upright)
+        candidate_max_didt = didt_upright[indx] * polarity
+        candidate_t_disrupt = t_ip[time_indices[indx]]
+
+        # Compute ip_final
+        (time_indices,) = np.where(
+            (t_ip > candidate_t_disrupt) & (t_ip < candidate_t_disrupt + 0.15)
+        )
+        ip_final = min(ip_upright[time_indices])
+
+        # Collect the computed parameters for the tests
+        parameters = {
+            "shot_duration": duration,
+            "abs_ip_max": ip_max,
+            "ip0_over_ip_max": candidate_ip0 / ip_max if ip_max != 0 else None,
+            "ip0_over_max_didt": (
+                candidate_ip0 / candidate_max_didt if candidate_max_didt != 0 else None
+            ),
+            "abs_ip_final": ip_final,
+            "abs_ip0": candidate_ip0,
+        }
+
+        # Run all 6 tests
+        for test, criterion in criteria.items():
+            parameter, threshold = parameters[test], thresholds[test]
+            # Use threshold = None to indicate that we will skip a test.
+            if threshold is None:
+                continue
+            # Use parameter = None to indicate a failed computation.
+            if parameter is None:
+                # TODO: consider raising an error.
+                return {"current_quench_time": [np.nan]}
+            # Run the test criteriion.
+            if not criterion(parameter, threshold):
+                # Failed the test, mark the shot as non-disruptive.
+                return {"current_quench_time": [np.nan]}
+
+        return {"current_quench_time": np.full(len(params.times), candidate_t_disrupt)}

disruption_py/machine/generic/util.py

@@ -0,0 +1,55 @@
+#!/usr/bin/env python3
+
+"""
+Module for helper, not physics, methods.
+"""
+
+import numpy as np
+import scipy
+
+
+class GenericUtilMethods:
+    """
+    A class of helper methods that might fetch and compute data from MDSplus
+    but are not physics methods.
+    """
+
+    @staticmethod
+    def get_end_of_current(ip, ip_time, threshold=1e5):
+        """
+        Python translation of the original MATLAB implementation based on end_of_current_d3d.m
+
+        Returns a dictionary of `duration` and `polarity`.
+        """
+        # Determine if there was any finite plasma current on this shot.  If
+        # not, then the shot was a "no plasma" shot, and the end-of-shot is set
+        # to 0 s.
+        (finite_indices,) = np.where((ip_time >= 0) & (abs(ip) > threshold))
+        if len(finite_indices) == 0:
+            return {"duration": 0, "ip_max": 0, "polarity": 1}
+        dt = np.diff(ip_time)
+        duration = sum(dt[finite_indices[:-1]])
+        if duration < 0.1:
+            # Assume < 100 ms is not a bona fide plasma
+            return {"duration": 0, "ip_max": 0, "polarity": 1}
+
+        # Determine the polarity of the plasma current.
+        polarity = np.sign(
+            scipy.integrate.trapezoid(ip[finite_indices], ip_time[finite_indices])
+        )
+        ip_upright = ip * polarity
+
+        # Find all the times that Ip is greater than the threshold.  The largest
+        # time value is the end of current.  But also check to see if the plasma
+        # current signal has been digitized for long enough to capture the end of
+        # the discharge.  If not, negate the end of shot value to indicate that
+        # that actual value cannot be determined, but it is longer than
+        # abs(value).
+        (indices,) = np.where((ip_upright >= threshold) & (ip_time > 0))
+        # Get the last index
+        max_idx = indices[-1] if len(indices) > 0 else None
+        duration = ip_time[max_idx]
+        if max_idx == len(ip_time) - 1:
+            duration = -duration  # TODO: what is this for?
+
+        return {"duration": duration, "polarity": polarity}

disruption_py/machine/hbtep/config.toml

@@ -4,6 +4,14 @@ mdsplus_connection_string = "maxwell.ap.columbia.edu:8003"
 [hbtep.inout.sql]
 host = ""
 
+[hbtep.physics.current_quench_time_thresholds]
+abs_ip0 = 0.1e6
+abs_ip_final = 50e3
+abs_ip_max = 0.1e6
+ip0_over_ip_max = 0.33
+ip0_over_max_didt = 0.05
+shot_duration = 0.1
+
 [hbtep.tests]
 expected_failure_columns = []
 test_columns = []

disruption_py/machine/mast/config.toml

@@ -149,6 +149,14 @@ units = "V"
 description = "Static LCFS loop voltage from EFIT. Defined as V_loop = -2 * pi * d(psi at LCFS) / dt"
 units = "V"
 
+[mast.physics.current_quench_time_thresholds]
+abs_ip0 = 0.1e6
+abs_ip_final = 50e3
+abs_ip_max = 0.1e6
+ip0_over_ip_max = 0.33
+ip0_over_max_didt = 0.05
+shot_duration = 0.1
+
 [mast.physics.time_domain_thresholds]
 dipprog_dt = 50e3
 ip_prog = 100e3