From 7ddd712e7be809533e413d13b11d7704f8f77daa Mon Sep 17 00:00:00 2001 From: Nico Amorisco Date: Sat, 1 Aug 2026 00:30:07 +0100 Subject: [PATCH 1/4] Speed up linearization construction Reuse accepted finite-difference current perturbations and build independent current and profile Jacobian columns in configurable worker processes. Reset each finite-difference solve to the unperturbed plasma state and deterministic NK random state. Keep the existing relinearization trigger and full-update policy unchanged. --- ...inear_evolution_with_relinearisation.ipynb | 8 + freegsnke/nonlinear_solve.py | 575 ++++++++++++++---- .../tests/test_linearisation_perturbations.py | 228 +++++++ requirements.txt | 3 +- 4 files changed, 695 insertions(+), 119 deletions(-) create mode 100644 freegsnke/tests/test_linearisation_perturbations.py diff --git a/examples/example05c - linear_evolution_with_relinearisation.ipynb b/examples/example05c - linear_evolution_with_relinearisation.ipynb index 054048e6..46a57caa 100644 --- a/examples/example05c - linear_evolution_with_relinearisation.ipynb +++ b/examples/example05c - linear_evolution_with_relinearisation.ipynb @@ -168,10 +168,18 @@ " full_timestep=5e-4, \n", " plasma_resistivity=1e-6,\n", " max_mode_frequency=10**2.5,\n", + " n_linearization_workers=1,\n", " plasma_descriptor_function=plasma_descriptors\n", ")" ] }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "`n_linearization_workers` controls how many independent worker processes build both the current-response and profile-parameter-response columns. It applies to the first linearisation constructed by `nl_solver` and to later relinearisations. The default value of `1` retains the serial calculation. Values larger than one use that many CPU threads in total for Jacobian construction, with one native numerical-library thread per worker to avoid oversubscription. This limit is scoped to Jacobian construction; it does not alter the inner threading used by the ordinary GS solves that advance the evolution." + ] + }, { "cell_type": "markdown", "metadata": {}, diff --git a/freegsnke/nonlinear_solve.py b/freegsnke/nonlinear_solve.py index 917b3aef..d26371ca 100644 --- a/freegsnke/nonlinear_solve.py +++ b/freegsnke/nonlinear_solve.py @@ -19,7 +19,9 @@ along with FreeGSNKE. If not, see . """ +import multiprocessing import warnings +from concurrent.futures import ProcessPoolExecutor from copy import deepcopy import matplotlib.pyplot as plt @@ -27,6 +29,7 @@ from freegs4e import bilinear_interpolation from freegs4e.gradshafranov import GreensBr, GreensdBrdz from scipy.signal import convolve2d +from threadpoolctl import threadpool_limits from . import nk_solver_H as nk_solver from .circuit_eq_metal import metal_currents @@ -36,6 +39,19 @@ from .simplified_solve import simplified_solver_J1 +_parallel_linearization_solver = None + + +def _build_dIydI_column_worker(arguments): + """Build one current-response column in an isolated worker process.""" + return _parallel_linearization_solver._build_dIydI_column(*arguments) + + +def _build_dIydtheta_column_worker(arguments): + """Build one profile-response column in an isolated worker process.""" + return _parallel_linearization_solver._build_dIydtheta_column(*arguments) + + class nl_solver: """ Nonlinear solver for time-evolution of plasma equilibria and circuit dynamics. @@ -55,6 +71,9 @@ class nl_solver: - Interfaces to Newton–Krylov solvers for plasma flux and circuit equations """ + _MAX_STARTING_DI_RATIO = np.sqrt(10.0) + _MAX_REUSED_STARTING_DI_RATIO = 4.0 / 3.0 + def __init__( self, profiles, @@ -83,6 +102,7 @@ def __init__( collinearity_reg=1e-6, verbose=False, plasma_descriptor_function=None, + n_linearization_workers=1, ): """ Initialize the nonlinear solver. @@ -144,6 +164,10 @@ def __init__( Additional penalty for collinear terms in nonlinear solver. verbose : bool, default=False Print diagnostic output during initialization. + n_linearization_workers : int, default=1 + Number of worker processes used to build independent ``dIydI`` and + ``dIydtheta`` columns during initial and later linearisations. A value + of 1 retains the serial calculation. """ print("-----") @@ -216,6 +240,13 @@ def __init__( self.max_internal_timestep = max_internal_timestep self.set_plasma_resistivity(plasma_resistivity) self.target_dIy = target_dIy + if ( + isinstance(n_linearization_workers, (bool, np.bool_)) + or not isinstance(n_linearization_workers, (int, np.integer)) + or n_linearization_workers < 1 + ): + raise ValueError("'n_linearization_workers' must be a positive integer.") + self.n_linearization_workers = int(n_linearization_workers) # prepare for mode selection if max_mode_frequency is None: @@ -390,6 +421,7 @@ def __init__( self.approved_target_dIy = np.concatenate( (self.approved_target_dIy, [target_dIy]) ) + self.initial_starting_dI = np.copy(self.starting_dI) # starting dtheta values for Jacobian calculation self.approved_target_dtheta = target_dIy * np.ones(self.n_profiles_parameters) @@ -569,6 +601,9 @@ def __init__( self.retained_modes_mask ] self.starting_dI = self.starting_dI[self.retained_modes_mask] + self.initial_starting_dI = self.initial_starting_dI[ + self.retained_modes_mask + ] self.remove_modes(eq, self.retained_modes_mask[:-1]) @@ -1412,6 +1447,13 @@ def build_dIydtheta( return dIydtheta, rel_ndIy, dvdtheta + def _reset_linearization_solve_state(self): + """Reset auxiliary plasma state before a finite-difference GS solve.""" + self.profiles2 = self.profiles1.copy() + self.eq2.plasma_psi = np.copy(self.eq1.plasma_psi) + if hasattr(self, "_linearization_rng_state"): + self.NK.rng.bit_generator.state = deepcopy(self._linearization_rng_state) + def prepare_build_dIydI_j( self, j, @@ -1453,8 +1495,7 @@ def prepare_build_dIydI_j( current_ = np.copy(self.currents_vec) current_[j] += starting_dI - # reset the auxiliary equilibrium - self.eq2.plasma_psi = np.copy(self.eq1.plasma_psi) + self._reset_linearization_solve_state() if GS: # solve self.assign_currents_solve_GS(current_, rtol_NK) @@ -1475,6 +1516,54 @@ def prepare_build_dIydI_j( self.final_dI_record[j] = final_dI return dIy_0 / starting_dI, rel_ndIy_0 + def update_starting_dI(self): + """Reuse perturbations accepted by the previous linearisation. + + The first linearisation retains the geometry-based perturbations + prepared during solver construction. Each rebuilt column subsequently + records its accepted finite-difference amplitude in + ``final_dI_record``. Reusing that value gives the next relinearisation a + state-informed initial guess without extrapolating from a stale + Jacobian. + + Missing, zero, or incompatible accepted amplitudes leave the existing + perturbations unchanged. + + Returns + ------- + ndarray of bool + Mask identifying perturbations updated from the stored Jacobian. + """ + + updated = np.zeros_like(self.starting_dI, dtype=bool) + if np.shape(self.final_dI_record) != np.shape(self.starting_dI): + return updated + + accepted = np.abs(np.asarray(self.final_dI_record)) + updated = np.isfinite(accepted) & (accepted > 0) + self.starting_dI[updated] = accepted[updated] + return updated + + @classmethod + def starting_dI_requires_rescaling( + cls, + starting_dI, + scaled_dI, + max_ratio=None, + ): + """Return whether a calibrated perturbation requires a second GS solve.""" + + if max_ratio is None: + max_ratio = cls._MAX_STARTING_DI_RATIO + if ( + not np.isfinite(starting_dI) + or not np.isfinite(scaled_dI) + or starting_dI == 0 + ): + return True + ratio = np.abs(scaled_dI / starting_dI) + return ratio < 1.0 / max_ratio or ratio > max_ratio + def build_dIydI_j(self, j, rtol_NK): """ Compute the finite-difference derivative d(Iy)/dI_j using the prepared perturbation. @@ -1505,8 +1594,7 @@ def build_dIydI_j(self, j, rtol_NK): current_ = np.copy(self.currents_vec) current_[j] += final_dI - # reset the auxiliary equilibrium - self.eq2.plasma_psi = np.copy(self.eq1.plasma_psi) + self._reset_linearization_solve_state() # solve self.assign_currents_solve_GS(current_, rtol_NK) @@ -1537,6 +1625,269 @@ def new_plasma_descriptors( + profile_contribution ) + def _core_mask_matches(self): + """Return whether the reference and perturbed plasma masks match.""" + return np.array_equal( + self.profiles1.diverted_core_mask, + self.profiles2.diverted_core_mask, + ) + + def _build_dIydI_column( + self, + j, + target_relative_tolerance_linearization, + force_core_mask_linearization, + reused_starting_dI, + ): + """Build and return one independent current-response column.""" + this_target_dIy = float(self.approved_target_dIy[j]) + dIydIj, ndIy = self.prepare_build_dIydI_j( + j, + target_relative_tolerance_linearization, + this_target_dIy, + self.starting_dI[j], + GS=True, + ) + + if force_core_mask_linearization: + while not self._core_mask_matches(): + self.starting_dI[j] /= 1.5 + this_target_dIy /= 1.5 + dIydIj, ndIy = self.prepare_build_dIydI_j( + j, + target_relative_tolerance_linearization, + this_target_dIy, + self.starting_dI[j], + ) + + if reused_starting_dI and self.starting_dI_requires_rescaling( + self.starting_dI[j], + self.final_dI_record[j], + max_ratio=self._MAX_REUSED_STARTING_DI_RATIO, + ): + # Discard a stale amplitude that is no longer predictive and use + # the original geometry-based path. + self.starting_dI[j] = self.initial_starting_dI[j] + dIydIj, ndIy = self.prepare_build_dIydI_j( + j, + target_relative_tolerance_linearization, + this_target_dIy, + self.starting_dI[j], + GS=True, + ) + reused_starting_dI = False + + rel_ndIy = ndIy + if self.starting_dI_requires_rescaling( + self.starting_dI[j], + self.final_dI_record[j], + max_ratio=( + self._MAX_REUSED_STARTING_DI_RATIO + if reused_starting_dI + else self._MAX_STARTING_DI_RATIO + ), + ): + dIydIj, rel_ndIy = self.build_dIydI_j( + j, + target_relative_tolerance_linearization, + ) + if force_core_mask_linearization: + while not self._core_mask_matches(): + self.final_dI_record[j] /= 1.2 + dIydIj, rel_ndIy = self.build_dIydI_j( + j, + target_relative_tolerance_linearization, + ) + else: + self.final_dI_record[j] = self.starting_dI[j] + + starting_dI = float(self.starting_dI[j]) + final_dI = float(self.final_dI_record[j]) + self.starting_dI[j] = final_dI + perturbed_psi = np.copy(self.eq2.psi()) + dRZdI = np.array( + ( + (self.eq2.Rcurrent() - self.R0) / final_dI, + (self.eq2.Zcurrent() - self.Z0) / final_dI, + ) + ) + dvdI = ( + np.asarray(self._column_plasma_descriptor_function(self.eq2)) + - self.initial_plasma_descriptors + ) / final_dI + + return ( + j, + np.copy(dIydIj), + perturbed_psi, + dRZdI, + dvdI, + starting_dI, + final_dI, + float(ndIy), + float(rel_ndIy), + float(self.NK.initial_rel_residual), + float(self.NK.relative_change), + float(self.current_at_last_linearization[j]), + ) + + def _build_dIydI_columns( + self, + target_relative_tolerance_linearization, + force_core_mask_linearization, + reused_starting_dI, + plasma_descriptor_function, + ): + """Build current-response columns serially or in isolated processes.""" + arguments = [ + ( + int(j), + target_relative_tolerance_linearization, + force_core_mask_linearization, + bool(reused_starting_dI[j]), + ) + for j in self.arange_currents + ] + self._column_plasma_descriptor_function = plasma_descriptor_function + self._linearization_rng_state = deepcopy(self.NK.rng.bit_generator.state) + try: + if self.n_linearization_workers == 1 or len(arguments) < 2: + return [self._build_dIydI_column(*argument) for argument in arguments] + + if "fork" not in multiprocessing.get_all_start_methods(): + raise RuntimeError( + "Parallel linearization requires multiprocessing support for " + "the 'fork' start method. Set n_linearization_workers=1." + ) + + global _parallel_linearization_solver + _parallel_linearization_solver = self + # Each worker gets one native BLAS/OpenMP thread so the requested + # worker count is also the total CPU-thread budget. + with threadpool_limits(limits=1): + with ProcessPoolExecutor( + max_workers=min(self.n_linearization_workers, len(arguments)), + mp_context=multiprocessing.get_context("fork"), + ) as executor: + return list( + executor.map( + _build_dIydI_column_worker, + arguments, + chunksize=1, + ) + ) + finally: + _parallel_linearization_solver = None + self.NK.rng.bit_generator.state = self._linearization_rng_state + del self._linearization_rng_state + del self._column_plasma_descriptor_function + + def _profile_parameters_for_column(self, profiles, j, delta): + """Return a complete profile-parameter set with column ``j`` perturbed.""" + if self.profiles_param is not None: + parameters = { + "alpha_m": profiles.alpha_m, + "alpha_n": profiles.alpha_n, + self.profiles_param: getattr(profiles, self.profiles_param), + } + parameter_name = ("alpha_m", "alpha_n", self.profiles_param)[j] + parameters[parameter_name] += delta + return parameters + + alpha = profiles.alpha.copy() + beta = profiles.beta.copy() + if j < self.n_profiles_parameters_alpha: + alpha[j] += delta + if profiles.alpha_logic: + alpha[-1] -= delta + else: + beta_index = j - self.n_profiles_parameters_alpha + beta[beta_index] += delta + if profiles.beta_logic: + beta[-1] -= delta + return {"alpha": alpha, "beta": beta} + + def _profile_parameter_name(self, j): + """Return the user-facing name of independent profile parameter ``j``.""" + if self.profiles_param is not None: + return ("alpha_m", "alpha_n", self.profiles_param)[j] + if j < self.n_profiles_parameters_alpha: + return f"alpha_{j}" + return f"beta_{j - self.n_profiles_parameters_alpha}" + + def _build_dIydtheta_column(self, j, delta, rtol_NK): + """Build and return one independent profile-response column.""" + self._reset_linearization_solve_state() + self.check_and_change_profiles( + self._profile_parameters_for_column(self._column_profiles, j, delta) + ) + self.assign_currents_solve_GS(np.copy(self.currents_vec), rtol_NK) + + dIy = self.limiter_handler.Iy_from_jtor(self.profiles2.jtor) - self.Iy + dv = ( + np.asarray(self._column_plasma_descriptor_function(self.eq2)) + - self.initial_plasma_descriptors + ) + return ( + j, + dIy / delta, + float(np.linalg.norm(dIy) / self.nIy), + dv / delta, + float(self.NK.initial_rel_residual), + float(self.NK.relative_change), + ) + + def _build_dIydtheta_columns( + self, + profiles, + rtol_NK, + perturbations, + plasma_descriptor_function, + ): + """Build profile-response columns serially or in isolated processes.""" + arguments = [ + (int(j), float(perturbations[j]), rtol_NK) + for j in range(self.n_profiles_parameters) + ] + self._column_profiles = profiles.copy() + self._column_plasma_descriptor_function = plasma_descriptor_function + self._linearization_rng_state = deepcopy(self.NK.rng.bit_generator.state) + try: + if self.n_linearization_workers == 1 or len(arguments) < 2: + return [ + self._build_dIydtheta_column(*argument) for argument in arguments + ] + + if "fork" not in multiprocessing.get_all_start_methods(): + raise RuntimeError( + "Parallel linearization requires multiprocessing support for " + "the 'fork' start method. Set n_linearization_workers=1." + ) + + global _parallel_linearization_solver + _parallel_linearization_solver = self + with threadpool_limits(limits=1): + with ProcessPoolExecutor( + max_workers=min(self.n_linearization_workers, len(arguments)), + mp_context=multiprocessing.get_context("fork"), + ) as executor: + return list( + executor.map( + _build_dIydtheta_column_worker, + arguments, + chunksize=1, + ) + ) + finally: + _parallel_linearization_solver = None + self.NK.rng.bit_generator.state = self._linearization_rng_state + del self._linearization_rng_state + self.check_and_change_profiles( + self._profile_parameters_for_column(self._column_profiles, 0, 0.0) + ) + del self._column_profiles + del self._column_plasma_descriptor_function + def build_linearization( self, eq, @@ -1610,6 +1961,14 @@ def build_linearization( # dIydI = 1 if dIydI is None: if self.dIydI_ICs is None: + if force_core_mask_linearization: + self.starting_dI = np.copy(self.initial_starting_dI) + reused_starting_dI = np.zeros_like( + self.starting_dI, + dtype=bool, + ) + else: + reused_starting_dI = self.update_starting_dI() print( f"Building the {self.plasma_domain_size} x {self.n_metal_modes + 1} Jacobian (dIy/dI)", "of plasma current density (inside the LCFS)", @@ -1626,113 +1985,47 @@ def build_linearization( ) self.initial_currents_plasma_descriptor = np.copy(self.currents_vec) - for j in self.arange_currents: - this_target_dIy = 1.0 * self.approved_target_dIy[j] - dIydIj, ndIy = self.prepare_build_dIydI_j( - j, - target_relative_tolerance_linearization, - this_target_dIy, - self.starting_dI[j], - GS=True, - ) - core_check = ( - np.sum( - np.abs( - self.profiles1.diverted_core_mask.astype(float) - - self.profiles2.diverted_core_mask.astype(float) - ) - ) - == 0 - ) - if force_core_mask_linearization: - while core_check == False: - self.starting_dI[j] /= 1.5 - this_target_dIy /= 1.5 - dIydIj, ndIy = self.prepare_build_dIydI_j( - j, - target_relative_tolerance_linearization, - this_target_dIy, - self.starting_dI[j], - ) - core_check = ( - np.sum( - np.abs( - self.profiles1.diverted_core_mask.astype(float) - - self.profiles2.diverted_core_mask.astype( - float - ) - ) - ) - == 0 - ) - - if ( - np.abs(np.log10(self.final_dI_record[j] / self.starting_dI[j])) - > 0.5 - ): - dIydIj, rel_ndIy = self.build_dIydI_j( - j, - target_relative_tolerance_linearization, - ) - core_check = ( - np.sum( - np.abs( - self.profiles1.diverted_core_mask.astype(float) - - self.profiles2.diverted_core_mask.astype(float) - ) - ) - == 0 - ) - if force_core_mask_linearization: - while core_check == False: - self.final_dI_record[j] /= 1.2 - dIydIj, rel_ndIy = self.build_dIydI_j( - j, - target_relative_tolerance_linearization, - ) - core_check = ( - np.sum( - np.abs( - self.profiles1.diverted_core_mask.astype( - float - ) - - self.profiles2.diverted_core_mask.astype( - float - ) - ) - ) - == 0 - ) - else: - self.final_dI_record[j] = 1.0 * self.starting_dI[j] - rel_ndIy = ndIy - + column_results = self._build_dIydI_columns( + target_relative_tolerance_linearization, + force_core_mask_linearization, + reused_starting_dI, + plasma_descriptor_function, + ) + for ( + j, + dIydIj, + perturbed_psi, + dRZdI, + dvdI, + starting_dI, + final_dI, + ndIy, + rel_ndIy, + initial_rel_residual, + relative_change, + current_at_last_linearization, + ) in column_results: if verbose: print("") print(f"Mode: {j}") - print(f" Initial delta_current = {self.starting_dI[j]}") + print(f" Initial delta_current = {starting_dI}") print(f" Initial relative Iy change = {ndIy}") - print(f" Final delta_current = {self.final_dI_record[j]}") + print(f" Final delta_current = {final_dI}") print("") - if "rel_ndIy" in locals(): - print(f" Final relative Iy change = {rel_ndIy}") - else: - print(f" Final relative Iy change = {ndIy}") + print(f" Final relative Iy change = {rel_ndIy}") print( - f" Initial vs. Final GS residual: {self.NK.initial_rel_residual} vs. {self.NK.relative_change}" + f" Initial vs. Final GS residual: {initial_rel_residual} vs. {relative_change}" ) - self.dIydI[:, j] = np.copy(dIydIj) - self.psideltaI[j] = np.copy(self.eq2.psi()) - R0 = self.eq2.Rcurrent() - Z0 = self.eq2.Zcurrent() - self.dRZdI[0, j] = (R0 - self.R0) / self.final_dI_record[j] - self.dRZdI[1, j] = (Z0 - self.Z0) / self.final_dI_record[j] - - v0 = plasma_descriptor_function(self.eq2) - self.dvdId[:, j] = ( - v0 - self.initial_plasma_descriptors - ) / self.final_dI_record[j] + self.dIydI[:, j] = dIydIj + self.psideltaI[j] = perturbed_psi + self.dRZdI[:, j] = dRZdI + self.dvdId[:, j] = dvdI + self.starting_dI[j] = final_dI + self.final_dI_record[j] = final_dI + self.current_at_last_linearization[j] = ( + current_at_last_linearization + ) self.dIydI_ICs = np.copy(self.dIydI) else: @@ -1765,33 +2058,79 @@ def build_linearization( profiles_copy = profiles.copy() - # prepare to build the Jacobian by finding appropriate step size - dIydtheta, ndIy, dvdtheta = self.prepare_build_dIydtheta( - profiles=profiles_copy, - rtol_NK=target_relative_tolerance_linearization, - target_dIy=self.approved_target_dtheta, - starting_dtheta=self.starting_dtheta, - plasma_descriptor_function=plasma_descriptor_function, - verbose=verbose, + if self.profiles_param is not None: + self.initial_profiles_plasma_descriptor = np.array( + [ + profiles.alpha_m, + profiles.alpha_n, + getattr(profiles, self.profiles_param), + ] + ) + else: + self.initial_profiles_plasma_descriptor = np.concatenate( + ( + profiles.alpha[: self.n_profiles_parameters_alpha], + profiles.beta[: self.n_profiles_parameters_beta], + ) + ) + + # First estimate perturbations that produce the requested Iy change. + column_results = self._build_dIydtheta_columns( + profiles_copy, + target_relative_tolerance_linearization, + self.starting_dtheta, + plasma_descriptor_function, ) + for j, column, ndIy, descriptor_column, _, _ in column_results: + self.dIydtheta[:, j] = column + self.dvdtheta[:, j] = descriptor_column + self.final_dtheta_record[j] = ( + self.starting_dtheta[j] * self.approved_target_dtheta[j] / ndIy + ) + if verbose: + print("") + print(f"Profile parameter: {self._profile_parameter_name(j)}:") + print(f" Initial delta parameter = {self.starting_dtheta[j]}") + print(f" Initial relative Iy change = {ndIy}") + print( + f" Final delta parameter = {self.final_dtheta_record[j]}" + ) if ( np.abs(np.log10(self.final_dtheta_record / self.starting_dtheta)) > 0.5 ).any(): - dIydtheta, rel_ndIy, dvdtheta = self.build_dIydtheta( - profiles=profiles_copy, - rtol_NK=target_relative_tolerance_linearization, - plasma_descriptor_function=plasma_descriptor_function, - verbose=verbose, + column_results = self._build_dIydtheta_columns( + profiles_copy, + target_relative_tolerance_linearization, + self.final_dtheta_record, + plasma_descriptor_function, ) + for ( + j, + column, + rel_ndIy, + descriptor_column, + initial_rel_residual, + relative_change, + ) in column_results: + self.dIydtheta[:, j] = column + self.dvdtheta[:, j] = descriptor_column + if verbose: + print("") + print( + f"Profile parameter: {self._profile_parameter_name(j)}:" + ) + print(f" Final relative Iy change = {rel_ndIy}") + print( + " Initial vs. Final GS residual: " + f"{initial_rel_residual} vs. {relative_change}" + ) else: self.final_dtheta_record = 1.0 * self.starting_dtheta - self.dIydtheta = np.copy(dIydtheta) self.dIydtheta_ICs = np.copy(self.dIydtheta) - self.dvdtheta = np.copy(dvdtheta) if plasma_descriptor_function is not None: print( diff --git a/freegsnke/tests/test_linearisation_perturbations.py b/freegsnke/tests/test_linearisation_perturbations.py new file mode 100644 index 00000000..6db944dd --- /dev/null +++ b/freegsnke/tests/test_linearisation_perturbations.py @@ -0,0 +1,228 @@ +import concurrent.futures.process +import multiprocessing +from types import MethodType, SimpleNamespace + +import numpy as np +import pytest + +from freegsnke.nonlinear_solve import nl_solver + + +def bare_solver(): + """Construct only the state needed by the perturbation helpers.""" + solver = nl_solver.__new__(nl_solver) + solver.NK = SimpleNamespace(rng=np.random.default_rng(seed=0)) + return solver + + +def test_update_starting_dI_uses_previously_accepted_amplitudes(): + solver = bare_solver() + solver.final_dI_record = np.array([1e-3, -2e-3]) + solver.starting_dI = np.array([100.0, 200.0]) + + updated = solver.update_starting_dI() + + np.testing.assert_array_equal(updated, [True, True]) + np.testing.assert_allclose(solver.starting_dI, [1e-3, 2e-3]) + + +def test_update_starting_dI_preserves_invalid_columns(): + solver = bare_solver() + solver.final_dI_record = np.array([1e-3, 0.0, np.nan]) + solver.starting_dI = np.array([100.0, 200.0, 300.0]) + + updated = solver.update_starting_dI() + + np.testing.assert_array_equal(updated, [True, False, False]) + np.testing.assert_allclose(solver.starting_dI, [1e-3, 200.0, 300.0]) + + +def test_update_starting_dI_ignores_incompatible_record(): + solver = bare_solver() + solver.final_dI_record = np.array([1e-3]) + solver.starting_dI = np.array([100.0, 200.0]) + + updated = solver.update_starting_dI() + + assert not np.any(updated) + np.testing.assert_allclose(solver.starting_dI, [100.0, 200.0]) + + +def test_starting_dI_rescaling_guard(): + limit = nl_solver._MAX_STARTING_DI_RATIO + + assert not nl_solver.starting_dI_requires_rescaling(10.0, 10.0 * limit * 0.99) + assert not nl_solver.starting_dI_requires_rescaling(10.0, 10.0 / limit / 0.99) + assert nl_solver.starting_dI_requires_rescaling(10.0, 10.0 * limit * 1.01) + assert nl_solver.starting_dI_requires_rescaling(10.0, 10.0 / limit / 1.01) + assert nl_solver.starting_dI_requires_rescaling(0.0, 1.0) + assert nl_solver.starting_dI_requires_rescaling(10.0, np.nan) + + +def test_reused_starting_dI_uses_tighter_rescaling_guard(): + legacy_limit = nl_solver._MAX_STARTING_DI_RATIO + reused_limit = nl_solver._MAX_REUSED_STARTING_DI_RATIO + scaled_dI = 10.0 * 0.5 * (legacy_limit + reused_limit) + + assert not nl_solver.starting_dI_requires_rescaling(10.0, scaled_dI) + assert nl_solver.starting_dI_requires_rescaling( + 10.0, + scaled_dI, + max_ratio=reused_limit, + ) + + +def fake_column_builder(self, column, *args): + """Return enough information to identify one dispatched column.""" + return int(column) + + +class FakeProfiles: + """Minimal conventional profile object for dispatcher tests.""" + + alpha_m = 1.0 + alpha_n = 2.0 + betap = 0.5 + + def copy(self): + return SimpleNamespace( + alpha_m=self.alpha_m, + alpha_n=self.alpha_n, + betap=self.betap, + ) + + +def configure_profile_dispatch(solver, workers): + """Populate only the state needed by the profile-column dispatcher.""" + solver.n_linearization_workers = workers + solver.n_profiles_parameters = 3 + solver.profiles_param = "betap" + solver._build_dIydtheta_column = MethodType(fake_column_builder, solver) + solver.check_and_change_profiles = lambda _: None + + +def test_linearization_solve_state_resets_flux_profiles_and_rng(): + solver = bare_solver() + solver.eq1 = SimpleNamespace(plasma_psi=np.arange(4.0).reshape(2, 2)) + solver.eq2 = SimpleNamespace(plasma_psi=np.zeros((2, 2))) + solver.profiles1 = FakeProfiles() + solver.profiles2 = None + solver._linearization_rng_state = solver.NK.rng.bit_generator.state + expected_random_value = np.random.default_rng(seed=0).random() + solver.NK.rng.random(5) + + solver._reset_linearization_solve_state() + + np.testing.assert_array_equal(solver.eq2.plasma_psi, solver.eq1.plasma_psi) + assert solver.profiles2 is not solver.profiles1 + assert solver.NK.rng.random() == expected_random_value + + +def test_dIydI_columns_use_serial_dispatch_by_default(): + solver = bare_solver() + solver.n_linearization_workers = 1 + solver.arange_currents = np.array([2, 4]) + solver._build_dIydI_column = MethodType(fake_column_builder, solver) + + results = solver._build_dIydI_columns( + 1e-8, + False, + np.zeros(5, dtype=bool), + lambda _: np.array([0.0]), + ) + + assert results == [2, 4] + + +@pytest.mark.skipif( + "fork" not in multiprocessing.get_all_start_methods(), + reason="Parallel linearization requires the fork start method.", +) +def test_dIydI_columns_preserve_order_with_multiple_workers(monkeypatch): + solver = bare_solver() + solver.n_linearization_workers = 2 + solver.arange_currents = np.array([4, 2, 3]) + solver._build_dIydI_column = MethodType(fake_column_builder, solver) + monkeypatch.setattr( + concurrent.futures.process, + "_check_system_limits", + lambda: None, + ) + + results = solver._build_dIydI_columns( + 1e-8, + False, + np.zeros(5, dtype=bool), + lambda _: np.array([0.0]), + ) + + assert results == [4, 2, 3] + + +def test_profile_parameter_shift_preserves_lao_constraints(): + solver = bare_solver() + solver.profiles_param = None + solver.n_profiles_parameters_alpha = 2 + profiles = SimpleNamespace( + alpha=np.array([1.0, 2.0, -3.0]), + beta=np.array([4.0, -4.0]), + alpha_logic=True, + beta_logic=True, + ) + + alpha_shift = solver._profile_parameters_for_column(profiles, 1, 0.25) + beta_shift = solver._profile_parameters_for_column(profiles, 2, 0.5) + + np.testing.assert_allclose(alpha_shift["alpha"], [1.0, 2.25, -3.25]) + np.testing.assert_allclose(alpha_shift["beta"], profiles.beta) + np.testing.assert_allclose(beta_shift["alpha"], profiles.alpha) + np.testing.assert_allclose(beta_shift["beta"], [4.5, -4.5]) + np.testing.assert_allclose(profiles.alpha, [1.0, 2.0, -3.0]) + np.testing.assert_allclose(profiles.beta, [4.0, -4.0]) + + +def test_profile_parameter_shift_for_conventional_profile(): + solver = bare_solver() + solver.profiles_param = "betap" + profiles = FakeProfiles() + + shifted = solver._profile_parameters_for_column(profiles, 2, 0.05) + + assert shifted == {"alpha_m": 1.0, "alpha_n": 2.0, "betap": 0.55} + + +def test_dIydtheta_columns_use_serial_dispatch_by_default(): + solver = bare_solver() + configure_profile_dispatch(solver, workers=1) + + results = solver._build_dIydtheta_columns( + FakeProfiles(), + 1e-8, + np.ones(3), + lambda _: np.array([0.0]), + ) + + assert results == [0, 1, 2] + + +@pytest.mark.skipif( + "fork" not in multiprocessing.get_all_start_methods(), + reason="Parallel linearization requires the fork start method.", +) +def test_dIydtheta_columns_preserve_order_with_multiple_workers(monkeypatch): + solver = bare_solver() + configure_profile_dispatch(solver, workers=2) + monkeypatch.setattr( + concurrent.futures.process, + "_check_system_limits", + lambda: None, + ) + + results = solver._build_dIydtheta_columns( + FakeProfiles(), + 1e-8, + np.ones(3), + lambda _: np.array([0.0]), + ) + + assert results == [0, 1, 2] diff --git a/requirements.txt b/requirements.txt index 8d28366c..c914c957 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,8 +1,9 @@ scipy~=1.15.3 numpy~=1.26.4 +threadpoolctl~=3.6.0 matplotlib~=3.10.3 h5py~=3.13.0 deepdiff~=7.0.1 scikit-image~=0.25.2 notebook~=7.4.2 -cvxpy~=1.7.5 \ No newline at end of file +cvxpy~=1.7.5 From 75c6861acf9249d9b5bda035ca97885d4faa45ff Mon Sep 17 00:00:00 2001 From: Nico Amorisco Date: Sat, 1 Aug 2026 00:33:35 +0100 Subject: [PATCH 2/4] Document linearization workers in examples --- ...example05a - nonlinear_and_linear_evolution_with_GS.ipynb | 5 ++++- examples/example05b - linear_evolution_without_GS.ipynb | 5 ++++- .../example05c - linear_evolution_with_relinearisation.ipynb | 4 ++-- examples/example11 - pulse_design_tool.ipynb | 5 ++++- 4 files changed, 14 insertions(+), 5 deletions(-) diff --git a/examples/example05a - nonlinear_and_linear_evolution_with_GS.ipynb b/examples/example05a - nonlinear_and_linear_evolution_with_GS.ipynb index 18a3657a..c7f51d15 100644 --- a/examples/example05a - nonlinear_and_linear_evolution_with_GS.ipynb +++ b/examples/example05a - nonlinear_and_linear_evolution_with_GS.ipynb @@ -216,6 +216,7 @@ "- `plasma_resistivity`: resistivity of the plasma (which here is assumed to be constant during the time evolution but can be made time-dependent if known).\n", "- `min_dIy_dI`: threshold value below which passive structure normal modes are dropped. Modes with norm(d(Iy)/dI)<`min_dIy_dI` are dropped, which filters out modes that do not actually couple with the plasma.\n", "- `max_mode_frequency`: threshold value for characteristic frequencies above which passive structure normal modes are dropped (i.e. the fast modes).\n", + "- **`n_linearization_workers`**: number of independent worker processes used to build the current-response (`dIydI`) and profile-parameter-response (`dIydtheta`) Jacobian columns. The default, `1`, retains serial construction. Values larger than one parallelise both the initial linearisation and later relinearisations without limiting the inner threading of ordinary time-step GS solves.\n", "\n", "Other customisable inputs are available, do see the documentation or the later notebook on \"Growth Rates\" for more details. For example, one may explicitly set your own resistance and inductance matrices for the tokamak, rather than the geometrical value calculated internally in FreeGSNKE.\n", "\n", @@ -239,6 +240,7 @@ " full_timestep=5e-4, \n", " plasma_resistivity=1e-6,\n", " max_mode_frequency=10**2.5,\n", + " n_linearization_workers=1, # set >1 to parallelise Jacobian construction\n", ")" ] }, @@ -651,7 +653,8 @@ " profiles=profiles, \n", " full_timestep=.5e-3, \n", " plasma_resistivity=1e-6, \n", - " GSStaticSolver=GSStaticSolver \n", + " GSStaticSolver=GSStaticSolver,\n", + " n_linearization_workers=1, # set >1 to parallelise Jacobian construction\n", ")" ] }, diff --git a/examples/example05b - linear_evolution_without_GS.ipynb b/examples/example05b - linear_evolution_without_GS.ipynb index ebd783d8..55497193 100644 --- a/examples/example05b - linear_evolution_without_GS.ipynb +++ b/examples/example05b - linear_evolution_without_GS.ipynb @@ -171,7 +171,9 @@ "source": [ "### Time evolution\n", "\n", - "Having defined the plasma descriptors, we can now instantiate the evolutive solver object. By including the `plasma_descriptors` function as an argument, the relevant Jacobians will be calculated to enable this evolution. " + "Having defined the plasma descriptors, we can now instantiate the evolutive solver object. By including the `plasma_descriptors` function as an argument, the relevant Jacobians will be calculated to enable this evolution.\n", + "\n", + "The new **`n_linearization_workers`** argument controls the number of independent worker processes used to construct both the current-response and profile-parameter-response Jacobian columns. `1` retains serial construction; a larger value parallelises the initial linearisation and any later relinearisations. It does not limit the inner threading of ordinary time-step GS solves." ] }, { @@ -189,6 +191,7 @@ " full_timestep=5e-4, \n", " plasma_resistivity=1e-6,\n", " max_mode_frequency=10**2.5,\n", + " n_linearization_workers=1, # set >1 to parallelise Jacobian construction\n", " plasma_descriptor_function=plasma_descriptors\n", ")" ] diff --git a/examples/example05c - linear_evolution_with_relinearisation.ipynb b/examples/example05c - linear_evolution_with_relinearisation.ipynb index 46a57caa..3fc3cf21 100644 --- a/examples/example05c - linear_evolution_with_relinearisation.ipynb +++ b/examples/example05c - linear_evolution_with_relinearisation.ipynb @@ -168,7 +168,7 @@ " full_timestep=5e-4, \n", " plasma_resistivity=1e-6,\n", " max_mode_frequency=10**2.5,\n", - " n_linearization_workers=1,\n", + " n_linearization_workers=1, # set >1 to parallelise Jacobian construction\n", " plasma_descriptor_function=plasma_descriptors\n", ")" ] @@ -177,7 +177,7 @@ "cell_type": "markdown", "metadata": {}, "source": [ - "`n_linearization_workers` controls how many independent worker processes build both the current-response and profile-parameter-response columns. It applies to the first linearisation constructed by `nl_solver` and to later relinearisations. The default value of `1` retains the serial calculation. Values larger than one use that many CPU threads in total for Jacobian construction, with one native numerical-library thread per worker to avoid oversubscription. This limit is scoped to Jacobian construction; it does not alter the inner threading used by the ordinary GS solves that advance the evolution." + "The new **`n_linearization_workers`** argument controls how many independent worker processes build both the current-response and profile-parameter-response columns. It applies to the first linearisation constructed by `nl_solver` and to later relinearisations. The default value of `1` retains the serial calculation. Values larger than one use that many CPU threads in total for Jacobian construction, with one native numerical-library thread per worker to avoid oversubscription. This limit is scoped to Jacobian construction; it does not alter the inner threading used by the ordinary GS solves that advance the evolution." ] }, { diff --git a/examples/example11 - pulse_design_tool.ipynb b/examples/example11 - pulse_design_tool.ipynb index 1a57468a..ed2ec3d2 100644 --- a/examples/example11 - pulse_design_tool.ipynb +++ b/examples/example11 - pulse_design_tool.ipynb @@ -893,7 +893,9 @@ "### Initialise the nonlinear solver object\n", "Now choose the desired settings for the nonlinear solver object - recall the prior example notebook on this. \n", "\n", - "The simulation timestep was chosen based on the vertical instability timescale in the output below (you need to always choose a value that is 5-10x smaller than this timescale) - this keeps the simulation numerically stable. " + "The simulation timestep was chosen based on the vertical instability timescale in the output below (you need to always choose a value that is 5-10x smaller than this timescale) - this keeps the simulation numerically stable.\n", + "\n", + "The new **`n_linearization_workers`** argument controls the number of independent worker processes used to build both the current-response (`dIydI`) and profile-parameter-response (`dIydtheta`) Jacobian columns. `1` retains serial construction; set it higher to parallelise both the initial linearisation and later relinearisations. This setting is scoped to Jacobian construction and does not limit the inner threading of the GS solves used during ordinary evolution steps." ] }, { @@ -911,6 +913,7 @@ " full_timestep=5e-4,\n", " plasma_resistivity=1e-7,\n", " fix_n_vessel_modes=30, \n", + " n_linearization_workers=1, # set >1 to parallelise Jacobian construction\n", " plasma_descriptor_function=plasma_descriptors,\n", " )" ] From 300c067390c5ffc257fa86f7336c28d98e91d104 Mon Sep 17 00:00:00 2001 From: Nico Amorisco Date: Wed, 5 Aug 2026 09:53:59 +0100 Subject: [PATCH 3/4] Record refreshed linearization column state --- freegsnke/nonlinear_solve.py | 1 + .../tests/test_linearisation_perturbations.py | 26 +++++++++++++++++++ 2 files changed, 27 insertions(+) diff --git a/freegsnke/nonlinear_solve.py b/freegsnke/nonlinear_solve.py index d26371ca..98b9ec54 100644 --- a/freegsnke/nonlinear_solve.py +++ b/freegsnke/nonlinear_solve.py @@ -1704,6 +1704,7 @@ def _build_dIydI_column( starting_dI = float(self.starting_dI[j]) final_dI = float(self.final_dI_record[j]) self.starting_dI[j] = final_dI + self.current_at_last_linearization[j] = self.currents_vec[j] perturbed_psi = np.copy(self.eq2.psi()) dRZdI = np.array( ( diff --git a/freegsnke/tests/test_linearisation_perturbations.py b/freegsnke/tests/test_linearisation_perturbations.py index 6db944dd..fb4b0e63 100644 --- a/freegsnke/tests/test_linearisation_perturbations.py +++ b/freegsnke/tests/test_linearisation_perturbations.py @@ -72,6 +72,32 @@ def test_reused_starting_dI_uses_tighter_rescaling_guard(): ) +def test_accepted_first_perturbation_records_current_linearization_point(): + solver = bare_solver() + solver.approved_target_dIy = np.array([0.01]) + solver.starting_dI = np.array([10.0]) + solver.final_dI_record = np.array([10.0]) + solver.current_at_last_linearization = np.array([-1.0]) + solver.currents_vec = np.array([42.0]) + solver.R0 = 1.0 + solver.Z0 = 0.0 + solver.initial_plasma_descriptors = np.array([0.0]) + solver.eq2 = SimpleNamespace( + psi=lambda: np.zeros((2, 2)), + Rcurrent=lambda: 1.0, + Zcurrent=lambda: 0.0, + ) + solver.NK.initial_rel_residual = 0.0 + solver.NK.relative_change = 0.0 + solver._column_plasma_descriptor_function = lambda _: np.array([0.0]) + solver.prepare_build_dIydI_j = lambda *args, **kwargs: (np.ones(2), 0.01) + + result = solver._build_dIydI_column(0, 1e-8, False, False) + + assert result[-1] == solver.currents_vec[0] + assert solver.current_at_last_linearization[0] == solver.currents_vec[0] + + def fake_column_builder(self, column, *args): """Return enough information to identify one dispatched column.""" return int(column) From 89f4b957f85773cbef075d714fd29d1c858ff926 Mon Sep 17 00:00:00 2001 From: Nico Amorisco Date: Wed, 5 Aug 2026 15:49:01 +0100 Subject: [PATCH 4/4] Fix Lao profile finite-difference inputs --- freegsnke/nonlinear_solve.py | 11 +++-------- freegsnke/tests/test_linearisation_perturbations.py | 11 ++++++----- 2 files changed, 9 insertions(+), 13 deletions(-) diff --git a/freegsnke/nonlinear_solve.py b/freegsnke/nonlinear_solve.py index 98b9ec54..b6e332a2 100644 --- a/freegsnke/nonlinear_solve.py +++ b/freegsnke/nonlinear_solve.py @@ -38,7 +38,6 @@ from .Myy_builder import Myy_handler from .simplified_solve import simplified_solver_J1 - _parallel_linearization_solver = None @@ -1784,7 +1783,7 @@ def _build_dIydI_columns( del self._column_plasma_descriptor_function def _profile_parameters_for_column(self, profiles, j, delta): - """Return a complete profile-parameter set with column ``j`` perturbed.""" + """Return independent profile parameters with column ``j`` perturbed.""" if self.profiles_param is not None: parameters = { "alpha_m": profiles.alpha_m, @@ -1795,17 +1794,13 @@ def _profile_parameters_for_column(self, profiles, j, delta): parameters[parameter_name] += delta return parameters - alpha = profiles.alpha.copy() - beta = profiles.beta.copy() + alpha = profiles.alpha[: self.n_profiles_parameters_alpha].copy() + beta = profiles.beta[: self.n_profiles_parameters_beta].copy() if j < self.n_profiles_parameters_alpha: alpha[j] += delta - if profiles.alpha_logic: - alpha[-1] -= delta else: beta_index = j - self.n_profiles_parameters_alpha beta[beta_index] += delta - if profiles.beta_logic: - beta[-1] -= delta return {"alpha": alpha, "beta": beta} def _profile_parameter_name(self, j): diff --git a/freegsnke/tests/test_linearisation_perturbations.py b/freegsnke/tests/test_linearisation_perturbations.py index fb4b0e63..22ef467e 100644 --- a/freegsnke/tests/test_linearisation_perturbations.py +++ b/freegsnke/tests/test_linearisation_perturbations.py @@ -185,10 +185,11 @@ def test_dIydI_columns_preserve_order_with_multiple_workers(monkeypatch): assert results == [4, 2, 3] -def test_profile_parameter_shift_preserves_lao_constraints(): +def test_profile_parameter_shift_supplies_independent_lao_coefficients(): solver = bare_solver() solver.profiles_param = None solver.n_profiles_parameters_alpha = 2 + solver.n_profiles_parameters_beta = 1 profiles = SimpleNamespace( alpha=np.array([1.0, 2.0, -3.0]), beta=np.array([4.0, -4.0]), @@ -199,10 +200,10 @@ def test_profile_parameter_shift_preserves_lao_constraints(): alpha_shift = solver._profile_parameters_for_column(profiles, 1, 0.25) beta_shift = solver._profile_parameters_for_column(profiles, 2, 0.5) - np.testing.assert_allclose(alpha_shift["alpha"], [1.0, 2.25, -3.25]) - np.testing.assert_allclose(alpha_shift["beta"], profiles.beta) - np.testing.assert_allclose(beta_shift["alpha"], profiles.alpha) - np.testing.assert_allclose(beta_shift["beta"], [4.5, -4.5]) + np.testing.assert_allclose(alpha_shift["alpha"], [1.0, 2.25]) + np.testing.assert_allclose(alpha_shift["beta"], [4.0]) + np.testing.assert_allclose(beta_shift["alpha"], [1.0, 2.0]) + np.testing.assert_allclose(beta_shift["beta"], [4.5]) np.testing.assert_allclose(profiles.alpha, [1.0, 2.0, -3.0]) np.testing.assert_allclose(profiles.beta, [4.0, -4.0])