From 93f1427aac3481d92702045a207e1dc104ab7204 Mon Sep 17 00:00:00 2001 From: Nico Amorisco <102979655+nicamo@users.noreply.github.com> Date: Sat, 1 Aug 2026 23:27:29 +0100 Subject: [PATCH 1/4] Speed up PCS waveform evaluation --- freegsnke/control_loop/shape_category.py | 24 ++++++--- freegsnke/control_loop/systems_category.py | 21 +++++--- freegsnke/control_loop/useful_functions.py | 50 +++++++++++++++++-- .../control_loop/virtual_circuits_category.py | 21 +++++--- 4 files changed, 90 insertions(+), 26 deletions(-) diff --git a/freegsnke/control_loop/shape_category.py b/freegsnke/control_loop/shape_category.py index 71d8cfa..9735eec 100644 --- a/freegsnke/control_loop/shape_category.py +++ b/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) * ( diff --git a/freegsnke/control_loop/systems_category.py b/freegsnke/control_loop/systems_category.py index f85ddec..0414709 100644 --- a/freegsnke/control_loop/systems_category.py +++ b/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 diff --git a/freegsnke/control_loop/useful_functions.py b/freegsnke/control_loop/useful_functions.py index 718ddd5..3022ec1 100644 --- a/freegsnke/control_loop/useful_functions.py +++ b/freegsnke/control_loop/useful_functions.py @@ -27,14 +27,42 @@ # 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. + """ + + def __init__(self, value: Any) -> None: + self.value = np.asarray(value) + + def __call__(self, t: Any) -> np.ndarray: + """Return the constant value with the same leading shape as ``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 this interpolant for order zero, otherwise a zero interpolant.""" + + 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 +79,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 +105,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 +122,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, diff --git a/freegsnke/control_loop/virtual_circuits_category.py b/freegsnke/control_loop/virtual_circuits_category.py index d880b33..323b6a6 100644 --- a/freegsnke/control_loop/virtual_circuits_category.py +++ b/freegsnke/control_loop/virtual_circuits_category.py @@ -68,6 +68,12 @@ class VirtualCircuitsController: Optional argument to specify how often, in seconds, new VCs are computed with vc_generator. If None provided, defaults to zero and new VC computed at every time step. + Attributes + ---------- + interpolant_derivatives : dict + Derivatives of the coil-reference spline interpolants, rebuilt whenever + `update_interpolants` is called and reused during control steps. + """ def __init__( @@ -252,19 +258,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]) @@ -460,12 +467,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]) From 9382782252658780250f13dd9a80f83d96c494fa Mon Sep 17 00:00:00 2001 From: Nico Amorisco <102979655+nicamo@users.noreply.github.com> Date: Sat, 1 Aug 2026 23:28:21 +0100 Subject: [PATCH 2/4] Add PCS interpolation regression tests --- .../tests/test_control_loop_interpolation.py | 83 +++++++++++++++++++ 1 file changed, 83 insertions(+) create mode 100644 freegsnke/tests/test_control_loop_interpolation.py diff --git a/freegsnke/tests/test_control_loop_interpolation.py b/freegsnke/tests/test_control_loop_interpolation.py new file mode 100644 index 0000000..d0221cf --- /dev/null +++ b/freegsnke/tests/test_control_loop_interpolation.py @@ -0,0 +1,83 @@ +"""Tests for inexpensive control-loop waveform evaluation.""" + +import numpy as np + +from freegsnke.control_loop.shape_category import ShapeController +from freegsnke.control_loop.systems_category import SystemsController +from freegsnke.control_loop.useful_functions import ( + ConstantInterpolant, + interpolate_spline, + interpolate_step, +) +from freegsnke.control_loop.virtual_circuits_category import ( + VirtualCircuitsController, +) + + +def waveform(values): + """Return a three-point waveform for controller test data.""" + + return {"times": np.array([0.0, 1.0, 2.0]), "vals": np.asarray(values)} + + +def test_constant_interpolants_preserve_scalar_and_array_shapes(): + scalar = interpolate_spline(waveform([2.0, 2.0, 2.0])) + matrix = interpolate_step(waveform([np.eye(2), np.eye(2), np.eye(2)])) + + assert isinstance(scalar, ConstantInterpolant) + assert isinstance(matrix, ConstantInterpolant) + np.testing.assert_array_equal(scalar(0.5), np.array(2.0)) + np.testing.assert_array_equal(scalar([0.5, 1.5]), np.array([2.0, 2.0])) + np.testing.assert_array_equal(matrix(0.5), np.eye(2)) + np.testing.assert_array_equal(matrix([0.5, 1.5]), np.stack([np.eye(2)] * 2)) + np.testing.assert_array_equal(scalar.derivative()(0.5), np.array(0.0)) + + +def test_shape_controller_reuses_spline_derivatives(): + data = { + "shape": { + "ff": waveform([0.0, 2.0, 2.0]), + "ref": waveform([1.0, 1.0, 1.0]), + "blend": waveform([1.0, 1.0, 1.0]), + "k_prop": waveform([1.0, 1.0, 1.0]), + "k_int": waveform([0.0, 0.0, 0.0]), + "k_deriv": waveform([0.0, 0.0, 0.0]), + "damping": waveform([1.0, 1.0, 1.0]), + } + } + controller = ShapeController(data=data, ctrl_targets=["shape"]) + derivative = controller.interpolant_derivatives["shape"]["ff"] + + np.testing.assert_allclose( + controller.extract_values(0.5, ["shape"], "ff", deriv=True), [2.0] + ) + assert controller.interpolant_derivatives["shape"]["ff"] is derivative + + +def test_coil_controllers_reuse_spline_derivatives(): + limits = waveform([[-10.0], [-10.0], [-10.0]]) + systems = SystemsController( + data={ + "coil_pert": waveform([0.0, 2.0, 2.0]), + "min_coil_curr_lims": limits, + "max_coil_curr_lims": waveform([[10.0], [10.0], [10.0]]), + "max_coil_curr_ramp_lims": waveform([[5.0], [5.0], [5.0]]), + }, + ctrl_coils=["coil"], + ) + virtual_circuits = VirtualCircuitsController( + data={ + "coil_order": ["coil"], + "coil_ref": waveform([0.0, 2.0, 2.0]), + "shape": waveform([[1.0], [1.0], [1.0]]), + "plasma": waveform([[1.0], [1.0], [1.0]]), + }, + ctrl_coils=["coil"], + ctrl_targets=["shape"], + plasma_target=["plasma"], + ) + + np.testing.assert_allclose(systems.extract_values(0.5, ["coil"], True), [2.0]) + np.testing.assert_allclose( + virtual_circuits.extract_values(0.5, ["coil_ref"], True), [2.0] + ) From 7a5a8e3efe880bf470b5b670b177b5de61660de5 Mon Sep 17 00:00:00 2001 From: kpentland Date: Fri, 14 Aug 2026 15:55:04 +0100 Subject: [PATCH 3/4] isort error --- freegsnke/tests/test_control_loop_interpolation.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/freegsnke/tests/test_control_loop_interpolation.py b/freegsnke/tests/test_control_loop_interpolation.py index d0221cf..08085df 100644 --- a/freegsnke/tests/test_control_loop_interpolation.py +++ b/freegsnke/tests/test_control_loop_interpolation.py @@ -9,9 +9,7 @@ interpolate_spline, interpolate_step, ) -from freegsnke.control_loop.virtual_circuits_category import ( - VirtualCircuitsController, -) +from freegsnke.control_loop.virtual_circuits_category import VirtualCircuitsController def waveform(values): From a0b1ce79e07cddf379cbb24bff41d721f3a37dd0 Mon Sep 17 00:00:00 2001 From: kpentland Date: Fri, 14 Aug 2026 16:01:24 +0100 Subject: [PATCH 4/4] adding missing docstrings --- freegsnke/control_loop/useful_functions.py | 58 +++++++++++++++++++++- freegsnke/control_loop/vc_provider.py | 16 ++++++ 2 files changed, 72 insertions(+), 2 deletions(-) diff --git a/freegsnke/control_loop/useful_functions.py b/freegsnke/control_loop/useful_functions.py index 3022ec1..8376990 100644 --- a/freegsnke/control_loop/useful_functions.py +++ b/freegsnke/control_loop/useful_functions.py @@ -34,19 +34,73 @@ class ConstantInterpolant: 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 with the same leading shape as ``t``.""" + """ + 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 this interpolant for order zero, otherwise a zero interpolant.""" + """ + 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.") diff --git a/freegsnke/control_loop/vc_provider.py b/freegsnke/control_loop/vc_provider.py index 891903b..5f6ba73 100644 --- a/freegsnke/control_loop/vc_provider.py +++ b/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] )