Make supplied linearisation reuse deterministic

examples/README.md

@@ -13,6 +13,7 @@ These example Jupyter notebooks are intended to be the **first port of call for
 | Example 05a | Learn how to use the nonlinear and linear (with GS) evolutive solver to simulate time-dependent equilibria. | Anyone |
 | Example 05b | Learn how to use the linear evolutive solver without solving GS at each timestep. | Anyone |
 | Example 05c | Learn how to use the linear evolutive solver (with or without solving GS) with automatic relinearisation enabled. | Anyone |
+| Example 05d | Learn how to reuse a calculated evolutive-solver linearisation with deterministic timescale-based mode truncation. | Anyone |
 | Example 06a/b | Simulate (static) MAST-U equilibria over an entire shot using inputs from EFIT++ (requires internal UKAEA MAST-U database). | UKAEA employees + collaborators |
 | Example 07a | Static inverse solve in the "Anamak" toy tokamak. | Anyone |
 | Example 07b | Static inverse solve in a SPARC-like tokamak. | Anyone |

freegsnke/nonlinear_solve.py

@@ -73,6 +73,30 @@ class nl_solver:
     _MAX_STARTING_DI_RATIO = np.sqrt(10.0)
     _MAX_REUSED_STARTING_DI_RATIO = 4.0 / 3.0
 
+    @staticmethod
+    def _validate_supplied_dIydI(dIydI, mode_selection):
+        """Reject bare current Jacobians without a deterministic mode basis."""
+        if dIydI is not None and mode_selection != "timescale":
+            raise ValueError(
+                "Supplying 'dIydI' is supported only with "
+                "'mode_selection=\"timescale\"'. Coupling-based selection "
+                "depends on the initial equilibrium and may change the current-"
+                "mode basis. Omit 'dIydI' to rebuild the linearisation, or use "
+                "timescale selection with the same settings as the source solver."
+            )
+
+    @staticmethod
+    def _copy_supplied_jacobian(name, value, expected_shape, compatibility):
+        """Validate and detach a user-supplied Jacobian array."""
+        if value is None:
+            return None
+        if np.shape(value) != expected_shape:
+            raise ValueError(
+                f"Supplied '{name}' has shape {np.shape(value)}, but "
+                f"{expected_shape} is required by {compatibility}."
+            )
+        return np.array(value, copy=True)
+
     def __init__(
         self,
         profiles,
@@ -172,7 +196,10 @@ def __init__(
         linearize : bool, default=True
             Whether to set up the linearised problem.
         dIydI : ndarray, optional
-            Plasma current Jacobian wrt coil and plasma currents.
+            Plasma current Jacobian wrt coil and plasma currents. Reuse is
+            supported only with ``mode_selection="timescale"``. The new solver
+            must use the same machine, grid, profile parameterisation, and
+            timescale-selection settings as the solver that built the matrix.
         dIydtheta : ndarray, optional
             Plasma current Jacobian wrt profile parameters.
         target_relative_tolerance_linearization : float, default=1e-8
@@ -230,6 +257,8 @@ def __init__(
         if mode_selection == "timescale":
             mode_removal = False
 
+        self._validate_supplied_dIydI(dIydI, mode_selection)
+
         # check threshold values
         if mode_selection == "coupling" and fix_n_vessel_modes < 0:
             if min_dIy_dI > threshold_dIy_dI:
@@ -593,13 +622,31 @@ def __init__(
 
         # self.dIydI is the Jacobian of the plasma current distribution
         # with respect to the independent currents (as in self.currents_vec)
-        self.dIydI_ICs = dIydI
-        self.dIydI = dIydI
+        current_jacobian_shape = (self.plasma_domain_size, self.n_metal_modes + 1)
+        self.dIydI_ICs = self._copy_supplied_jacobian(
+            "dIydI",
+            dIydI,
+            current_jacobian_shape,
+            "the selected timescale mode basis; reuse the same machine, grid, "
+            "and timescale-selection settings as the source solver",
+        )
+        self.dIydI = None if self.dIydI_ICs is None else np.copy(self.dIydI_ICs)
 
         # self.dIydtheta is the Jacobian of the plasma current distribution
         # with respect to the plasma current density profile parameters
-        self.dIydtheta_ICs = dIydtheta
-        self.dIydtheta = dIydtheta
+        profile_jacobian_shape = (
+            self.plasma_domain_size,
+            self.n_profiles_parameters,
+        )
+        self.dIydtheta_ICs = self._copy_supplied_jacobian(
+            "dIydtheta",
+            dIydtheta,
+            profile_jacobian_shape,
+            "the current grid and profile parameterisation",
+        )
+        self.dIydtheta = (
+            None if self.dIydtheta_ICs is None else np.copy(self.dIydtheta_ICs)
+        )
 
         # initialize and set up the linearization
         # input value for dIydI is used when available
@@ -1980,7 +2027,9 @@ def build_linearization(
             Plasma profiles associated with the equilibrium.
         dIydI : np.ndarray or None
             Optional input Jacobian of plasma current density with respect to metal currents.
-            If None, it will be computed internally.
+            If None, it will be computed internally. Supplying it is supported
+            only when this solver uses ``mode_selection="timescale"`` with the
+            same mode-selection settings as its source solver.
         dIydtheta : np.ndarray or None
             Optional input Jacobian of plasma current density with respect to plasma profile parameters.
             If None, it will be computed internally.
@@ -2003,6 +2052,8 @@ def build_linearization(
         in self.dRZdI.
         """
 
+        self._validate_supplied_dIydI(dIydI, self.mode_selection)
+
         # if (dIydI is None) and (self.dIydI is None):
         self.build_current_vec(eq, profiles)
         self.Iy = self.limiter_handler.Iy_from_jtor(profiles.jtor).copy()
@@ -2092,7 +2143,14 @@ def build_linearization(
             else:
                 self.dIydI = np.copy(self.dIydI_ICs)
         else:
-            self.dIydI = dIydI
+            expected_shape = (self.plasma_domain_size, self.n_metal_modes + 1)
+            self.dIydI = self._copy_supplied_jacobian(
+                "dIydI",
+                dIydI,
+                expected_shape,
+                "the selected timescale mode basis; reuse the same machine, grid, "
+                "and timescale-selection settings as the source solver",
+            )
             self.dIydI_ICs = np.copy(self.dIydI)
 
         # compose the vector of initial delta_theta (profile parameters) to be used for the finite difference calculation
@@ -2208,7 +2266,16 @@ def build_linearization(
             else:
                 self.dIydtheta = np.copy(self.dIydtheta_ICs)
         else:
-            self.dIydtheta = dIydtheta
+            expected_shape = (
+                self.plasma_domain_size,
+                self.n_profiles_parameters,
+            )
+            self.dIydtheta = self._copy_supplied_jacobian(
+                "dIydtheta",
+                dIydtheta,
+                expected_shape,
+                "the current grid and profile parameterisation",
+            )
             self.dIydtheta_ICs = np.copy(self.dIydtheta)
 
     def set_plasma_resistivity(self, plasma_resistivity):
@@ -2533,7 +2600,9 @@ def initialize_from_ICs(
         dIydI : np.array, optional
             Jacobian of plasma current distribution with respect to metal currents and total plasma current.
             Shape: (np.sum(plasma_domain_mask), n_metal_modes+1).
-            If not provided, it is computed from the given equilibrium.
+            If not provided, it is computed from the given equilibrium. A supplied
+            matrix requires ``mode_selection="timescale"`` and the same machine,
+            grid, normalization, and mode-selection settings used to build it.
         dIydtheta : np.array, optional
             Jacobian of plasma current distribution with respect to plasma profile parameters.
             If not provided, it is computed from the given equilibrium.

freegsnke/tests/test_mode_selection.py

@@ -3,11 +3,82 @@
 from types import SimpleNamespace
 
 import numpy as np
+import pytest
 
 from freegsnke.circuit_eq_metal import metal_currents
 from freegsnke.nonlinear_solve import nl_solver
 
 
+def test_supplied_dIydI_rejects_coupling_selection():
+    """A supplied current Jacobian requires deterministic timescale selection."""
+
+    eq = SimpleNamespace(
+        nx=2,
+        ny=2,
+        R=np.zeros((2, 2)),
+        Z=np.zeros((2, 2)),
+        dR=1.0,
+        dZ=1.0,
+        tokamak=SimpleNamespace(
+            n_active_coils=0,
+            n_coils=0,
+            coils_dict={},
+        ),
+        limiter_handler=SimpleNamespace(mask_inside_limiter=np.ones((2, 2), bool)),
+    )
+
+    with pytest.raises(ValueError, match="supported only.*timescale"):
+        nl_solver(
+            profiles=SimpleNamespace(),
+            eq=eq,
+            GSStaticSolver=SimpleNamespace(),
+            dIydI=np.zeros((4, 1)),
+            mode_selection="coupling",
+            mode_removal=False,
+        )
+
+
+def test_supplied_dIydI_validation_also_applies_after_construction():
+    """Later initialization cannot bypass the supplied-Jacobian contract."""
+
+    solver = nl_solver.__new__(nl_solver)
+    solver.mode_selection = "coupling"
+
+    with pytest.raises(ValueError, match="supported only.*timescale"):
+        solver.build_linearization(
+            eq=None,
+            profiles=None,
+            dIydI=np.zeros((4, 1)),
+            dIydtheta=None,
+            target_relative_tolerance_linearization=1e-8,
+            force_core_mask_linearization=False,
+            verbose=False,
+            plasma_descriptor_function=None,
+        )
+
+
+def test_supplied_jacobian_is_shape_checked_and_copied():
+    """The solver must not retain a mutable reference to a caller's cache."""
+
+    supplied = np.zeros((4, 2))
+    stored = nl_solver._copy_supplied_jacobian(
+        "dIydI",
+        supplied,
+        expected_shape=(4, 2),
+        compatibility="the selected mode basis",
+    )
+    stored[0, 0] = 1.0
+    assert supplied[0, 0] == 0.0
+
+    with pytest.raises(ValueError, match=r"shape \(4, 2\).+\(4, 3\)"):
+        nl_solver._copy_supplied_jacobian(
+            "dIydI",
+            supplied,
+            expected_shape=(4, 3),
+            compatibility="the selected mode basis",
+        )
+
+
 def test_no_gs_coupling_norm_uses_evaluated_perturbation():
     """The coupling norm must match the returned finite-difference column."""