Speed up PCS waveform evaluation

freegsnke/control_loop/shape_category.py

@@ -73,6 +73,10 @@ class ShapeController:
         A nested dictionary storing interpolation functions of each input waveform for each
         shape target.
         Structure: {target: {spline/step key: interpolant_function}}
+
+    interpolant_derivatives : dict
+        A nested dictionary storing derivatives of the spline interpolants. These are
+        constructed with `interpolants` so control steps only evaluate them.
     """
 
     def __init__(
@@ -196,21 +200,25 @@ def update_interpolants(self) -> None:
         """
         Recompute all interpolant functions from the current `self.data`.
 
-        This method clears the existing `self.interpolants` dictionary and
-        rebuilds it by applying either `interpolate_spline` or `interpolate_step`
-        depending on whether each key belongs to `self.keys_to_spline` or
-        `self.keys_to_step`.
+        This method rebuilds `self.interpolants` by applying either
+        `interpolate_spline` or `interpolate_step`, and rebuilds the cached
+        derivatives of every spline interpolant.
 
         """
 
-        # create a dictionary to store the spline funcions
+        # create dictionaries to store the interpolants and spline derivatives
         self.interpolants = {}
+        self.interpolant_derivatives = {}
 
         # interpolate the input data
         for targ in self.ctrl_targets:
             self.interpolants[targ] = {}
+            self.interpolant_derivatives[targ] = {}
             for key in self.keys_to_spline:
                 self.interpolants[targ][key] = interpolate_spline(self.data[targ][key])
+                self.interpolant_derivatives[targ][key] = self.interpolants[targ][
+                    key
+                ].derivative()
             for key in self.keys_to_step:
                 self.interpolants[targ][key] = interpolate_step(self.data[targ][key])
 
@@ -511,12 +519,12 @@ def extract_values(
         Notes
         -----
         - Assumes that `self.interpolants[target][key]` is a valid `scipy.interpolate` object.
-        - If `deriv=True`, the method calls `.derivative()` on the interpolant before evaluation.
+        - Spline derivatives are constructed by `update_interpolants` and reused here.
         """
 
         if deriv:
             return np.array(
-                [self.interpolants[target][key].derivative()(t) for target in targets]
+                [self.interpolant_derivatives[target][key](t) for target in targets]
             )
         else:
             return np.array([self.interpolants[target][key](t) for target in targets])
@@ -560,7 +568,7 @@ def plot_data(
         # find out which control is ON and when
         FF_reference = self.interpolants[targ]["ff"](t)
         FF_mask = (self.interpolants[targ]["blend"](t) < 1) * (
-            np.abs(self.interpolants[targ]["ff"].derivative()(t)) > 0
+            np.abs(self.interpolant_derivatives[targ]["ff"](t)) > 0
         )
         FB_reference = self.interpolants[targ]["ref"](t)
         FB_mask = (self.interpolants[targ]["blend"](t) > 0) * (

freegsnke/control_loop/systems_category.py

@@ -69,6 +69,10 @@ class SystemsController:
         A nested dictionary storing interpolation functions of each input waveform.
         Structure: {spline/step key: interpolant_function}
 
+    interpolant_derivatives : dict
+        Derivatives of the coil-perturbation spline interpolants, rebuilt whenever
+        `update_interpolants` is called and reused during control steps.
+
     """
 
     def __init__(
@@ -148,19 +152,20 @@ def update_interpolants(self) -> None:
         """
         Recompute all interpolant functions from the current `self.data`.
 
-        This method clears the existing `self.interpolants` dictionary and
-        rebuilds it by applying either `interpolate_spline` or `interpolate_step`
-        depending on whether each key belongs to `self.keys_to_spline` or
-        `self.keys_to_step`.
+        This method rebuilds `self.interpolants` by applying either
+        `interpolate_spline` or `interpolate_step`, and rebuilds the cached
+        derivatives of every spline interpolant.
 
         """
 
-        # create a dictionary to store the spline functions
+        # create dictionaries to store the interpolants and spline derivatives
         self.interpolants = {}
+        self.interpolant_derivatives = {}
 
         # interpolate the input data
         for key in self.keys_to_spline:
             self.interpolants[key] = interpolate_spline(self.data[key])
+            self.interpolant_derivatives[key] = self.interpolants[key].derivative(n=1)
         for key in self.keys_to_step:
             self.interpolants[key] = interpolate_step(self.data[key])
 
@@ -271,13 +276,13 @@ def extract_values(
         Notes
         -----
         - Assumes that `self.interpolants[target]` is a valid `scipy.interpolate` object.
-        - If `deriv=True`, the method calls `.derivative()` on the interpolant before evaluation.
+        - Spline derivatives are constructed by `update_interpolants` and reused here.
         """
 
         if deriv:
             return np.array(
                 [
-                    self.interpolants[target + "_pert"].derivative(n=1)(t)
+                    self.interpolant_derivatives[target + "_pert"](t)
                     for target in targets
                 ]
             )
@@ -327,7 +332,7 @@ def plot_data(self, tmin: float = -1.0, tmax: float = 1.0, nt: int = 1001) -> No
 
             # find out which control is ON and when
             if key in self.keys_to_spline:
-                FF_reference = self.interpolants[key].derivative()(t)
+                FF_reference = self.interpolant_derivatives[key](t)
                 FF_mask = np.abs(FF_reference) > 0
 
                 # shade region of FF control

freegsnke/control_loop/useful_functions.py

@@ -27,14 +27,96 @@
 # a single time-series entry, e.g. {"times": [...], "vals": [...]}
 Waveform = dict[str, Any]
 
+
+class ConstantInterpolant:
+    """
+    Callable interpolant for a waveform whose value is constant in time.
+
+    It preserves the output shapes of the SciPy interpolants used for
+    non-constant waveforms and provides a compatible `derivative` method.
+
+    Parameters
+    ----------
+    value : array_like
+        The constant value returned for any input time. Stored as a
+        NumPy array.
+
+    Attributes
+    ----------
+    value : np.ndarray
+        The constant value returned by the interpolant.
+    """
+
+    def __init__(self, value: Any) -> None:
+        """
+        Initialize the interpolant with a constant value.
+
+        Parameters
+        ----------
+        value : array_like
+            The constant value to store, converted to a NumPy array
+            via `np.asarray`.
+        """
+        self.value = np.asarray(value)
+
+    def __call__(self, t: Any) -> np.ndarray:
+        """
+        Return the constant value broadcast to match the shape of ``t``.
+
+        Parameters
+        ----------
+        t : array_like
+            Time point(s) at which to evaluate the interpolant. Only the
+            shape of ``t`` is used; its values do not affect the output.
+
+        Returns
+        -------
+        np.ndarray
+            Array of shape ``np.shape(t) + self.value.shape`` containing
+            copies of ``self.value``, one for each element of ``t``.
+        """
+
+        result = np.broadcast_to(self.value, np.shape(t) + self.value.shape)
+        return np.array(result, copy=True)
+
+    def derivative(self, n: int = 1) -> "ConstantInterpolant":
+        """
+        Return the ``n``-th derivative of this constant interpolant.
+
+        Parameters
+        ----------
+        n : int, optional
+            Order of the derivative. Must be non-negative. Default is 1.
+
+        Returns
+        -------
+        ConstantInterpolant
+            ``self`` if ``n == 0`` (the value is unchanged), otherwise a
+            new `ConstantInterpolant` whose value is zero everywhere,
+            since the derivative of a constant is zero for any order
+            greater than zero.
+
+        Raises
+        ------
+        ValueError
+            If ``n`` is negative.
+        """
+
+        if n < 0:
+            raise ValueError("Derivative order must be non-negative.")
+        if n == 0:
+            return self
+        return ConstantInterpolant(np.zeros_like(self.value, dtype=float))
+
+
 # an interpolant produced by `interpolate_step`/`interpolate_spline`: callable at a
 # time `t`, and (for splines only) supports `.derivative()`
-Interpolant = Union[interp1d, UnivariateSpline]
+Interpolant = Union[ConstantInterpolant, interp1d, UnivariateSpline]
 
 
 def interpolate_step(
     data: Waveform,
-) -> interp1d:
+) -> Interpolant:
     """
     Creates a step-wise interpolator for time-series data using 'previous' value interpolation.
 
@@ -51,11 +133,19 @@ def interpolate_step(
         Callable function f(t) that returns the step-wise interpolated value at time t.
         For t < min(times), returns the first value.
         For t > max(times), returns the last value.
+
+    Notes
+    -----
+    Constant waveforms use `ConstantInterpolant` to avoid repeated SciPy
+    interpolation overhead during control-loop execution.
     """
 
     times = np.array(data["times"])
     vals = np.stack(data["vals"])
 
+    if np.all(vals == vals[0]):
+        return ConstantInterpolant(vals[0])
+
     # build interpolator
     f_interp = interp1d(
         times,
@@ -69,7 +159,7 @@ def interpolate_step(
     return f_interp
 
 
-def interpolate_spline(data: Waveform) -> UnivariateSpline:
+def interpolate_spline(data: Waveform) -> Interpolant:
     """
     Creates a spline interpolator for time-series data in 'data'.
 
@@ -86,11 +176,19 @@ def interpolate_spline(data: Waveform) -> UnivariateSpline:
         Callable function f(t) that returns the spline interpolated value at time t.
         For t < min(times), returns the first value.
         For t > max(times), returns the last value.
+
+    Notes
+    -----
+    Constant waveforms use `ConstantInterpolant` to avoid repeated SciPy
+    interpolation overhead during control-loop execution.
     """
 
     times = np.array(data["times"])
     vals = np.array(data["vals"])
 
+    if np.all(vals == vals[0]):
+        return ConstantInterpolant(vals[0])
+
     # build interpolator
     f_interp = UnivariateSpline(
         times,

freegsnke/control_loop/vc_provider.py

@@ -267,6 +267,22 @@ def _create_target_calculator(
         else:
             # reorder the target calculator outputs if targets are different order or a subset
             def array_func(eq):
+                """
+                Evaluate all requested target quantities for a given equilibrium.
+
+                Parameters
+                ----------
+                eq : object
+                    Equilibrium (or similar state) object passed to each target
+                    calculator function.
+
+                Returns
+                -------
+                np.ndarray
+                    1D array of computed target values, in the same order as
+                    `targets`, obtained by calling ``self.target_calculator_dict[targ](eq)``
+                    for each ``targ`` in `targets`.
+                """
                 return np.array(
                     [self.target_calculator_dict[targ](eq) for targ in targets]
                 )

freegsnke/control_loop/virtual_circuits_category.py

@@ -66,6 +66,12 @@ class VirtualCircuitsController:
         attribute (set when the generator was initialised) controls how often, in
         seconds, new VCs are computed.
 
+    Attributes
+    ----------
+    interpolant_derivatives : dict
+        Derivatives of the coil-reference spline interpolants, rebuilt whenever
+        `update_interpolants` is called and reused during control steps.
+
     """
 
     def __init__(
@@ -242,19 +248,20 @@ def update_interpolants(self) -> None:
         """
         Recompute all interpolant functions from the current `self.data`.
 
-        This method clears the existing `self.interpolants` dictionary and
-        rebuilds it by applying either `interpolate_spline` or `interpolate_step`
-        depending on whether each key belongs to `self.keys_to_spline` or
-        `self.keys_to_step`.
+        This method rebuilds `self.interpolants` by applying either
+        `interpolate_spline` or `interpolate_step`, and rebuilds the cached
+        derivatives of every spline interpolant.
 
         """
 
-        # create a dictionary to store the spline functions
+        # create dictionaries to store the interpolants and spline derivatives
         self.interpolants = {}
+        self.interpolant_derivatives = {}
 
         # interpolate the input data
         for key in self.keys_to_spline:
             self.interpolants[key] = interpolate_spline(self.data[key])
+            self.interpolant_derivatives[key] = self.interpolants[key].derivative(n=1)
         for key in self.keys_to_step:
             self.interpolants[key] = interpolate_step(self.data[key])
 
@@ -427,12 +434,12 @@ def extract_values(
         Notes
         -----
         - Assumes that `self.interpolants[target]` is a valid `scipy.interpolate` object.
-        - If `deriv=True`, the method calls `.derivative()` on the interpolant before evaluation.
+        - Spline derivatives are constructed by `update_interpolants` and reused here.
         """
 
         if deriv:
             return np.array(
-                [self.interpolants[target].derivative(n=1)(t) for target in targets]
+                [self.interpolant_derivatives[target](t) for target in targets]
             )
         else:
             return np.array([self.interpolants[target](t) for target in targets])