freegsnke.nonlinear_solve module

Implements the core non-linear solver for the evolutive GS problem. Also handles the linearised evolution capabilites.

Copyright 2025 UKAEA, UKRI-STFC, and The Authors, as per the COPYRIGHT and README files.

This file is part of FreeGSNKE.

FreeGSNKE is distributed in the hope that it will be useful, but WITHOUT ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU Lesser General Public License for more details.

FreeGSNKE is free software: you can redistribute it and/or modify it under the terms of the GNU Lesser General Public License as published by the Free Software Foundation, either version 3 of the License, or (at your option) any later version.

You should have received a copy of the GNU Lesser General Public License along with FreeGSNKE. If not, see <http://www.gnu.org/licenses/>.

class freegsnke.nonlinear_solve.nl_solver(profiles, eq, GSStaticSolver, custom_coil_resist=None, custom_self_ind=None, full_timestep=0.0001, max_internal_timestep=0.0001, automatic_timestep=False, plasma_resistivity=1e-06, plasma_norm_factor=1000.0, blend_hatJ=0, max_mode_frequency=100.0, fix_n_vessel_modes=-1, threshold_dIy_dI=0.025, min_dIy_dI=0.01, mode_removal=True, linearize=True, dIydI=None, dIydtheta=None, target_relative_tolerance_linearization=1e-08, target_dIy=0.001, force_core_mask_linearization=False, l2_reg=1e-06, collinearity_reg=1e-06, verbose=False)

Bases: object

Nonlinear solver for time-evolution of plasma equilibria and circuit dynamics.

This class provides an interface to evolve both the linearised and nonlinear dynamic problems. It sets up plasma and metal circuit equations, handles vessel mode decomposition and selection, and enables coupled nonlinear simulation of plasma and machine dynamics.

Main features

• Linear and nonlinear timestepping of equilibrium dynamics

• Automatic timestep control based on growth rates

• Passive vessel mode decomposition and mode selection

• Coupling to FreeGSNKE profiles and equilibria

• Support for regularization in nonlinear solves

• Interfaces to Newton–Krylov solvers for plasma flux and circuit equations

F_function_curr(trial_currents, active_voltage_vec)

Evaluates the residual of the full non-linear plasma + circuit system for a given guess of currents at time t+dt.

This function casts the coupled plasma and circuit dynamics as a root-finding problem in the space of current values. The residual is defined as the difference between the currents predicted by the simplified circuit equations (given a plasma current distribution) and the input trial currents. A zero residual indicates that the trial currents and the corresponding plasma flux form a self-consistent solution of the full non-linear system.

The evaluation proceeds in two steps: 1. Compute the normalized plasma current distribution hatIy1 corresponding to the input trial_currents and the current plasma flux self.trial_plasma_psi. 2. Compute the iterated currents using the simplified circuit equations and the plasma distribution hatIy1. The residual is then: residual = iterated_currents - trial_currents.

Parameters:

• trial_currents (np.ndarray) – Vector of current values at time t+dt. The format matches self.currents_vec, typically including active coil currents, vessel mode currents, and the total plasma current (normalized by plasma_norm_factor).

• active_voltage_vec (np.ndarray) – Vector of voltages applied to the active coils over the timestep.

Returns:

Residual vector of currents, same format as self.currents_vec. Zero residual indicates a self-consistent solution of the plasma + circuit system for the timestep.

Return type:

np.ndarray

F_function_curr_GS(trial_currents, active_voltage_vec, rtol_NK)

Full non-linear system of circuit eqs written as root problem in the vector of current values at time t+dt. Note that, differently from self.F_function_curr, here the plasma flux is not imposed, but self-consistently solved for based on the input trial_currents. Iteration consists of: trial_currents -> plasma flux, through static GS [trial_currents, plasma_flux] -> hatIy1, through calculating plasma distribution hatIy1 -> iterated_currents, through ‘simplified’ circuit eqs Residual: iterated_currents - trial_currents Residual is zero if trial_currents solve the full non-linear problem.

Parameters:

• trial_currents (np.array) – Vector of current values. Same format as self.currents_vec.

• active_voltage_vec (np.array) – Vector of active voltages for the active coils, applied between t and t+dt.

• rtol_NK (float) – Relative tolerance to be used in the static GS problem.

Returns:

Residual in current values. Same format as self.currents_vec.

Return type:

np.array

F_function_psi(trial_plasma_psi, active_voltage_vec, rtol_NK)

Evaluates the residual of the full non-linear plasma + circuit system for a given guess of currents at time t+dt, solving the plasma flux self-consistently via the static Grad-Shafranov (GS) problem.

Unlike F_function_curr, the plasma flux is not imposed but computed from the trial currents. The residual is defined as the difference between the currents predicted by the simplified circuit equations (given the resulting plasma distribution) and the input trial currents. A zero residual indicates that the trial currents form a self-consistent solution of the full non-linear system, including the plasma response.

Iteration steps: 1. Compute the plasma flux corresponding to trial_currents by solving the static GS problem. 2. Compute the normalized plasma current distribution hatIy1 from the combination of trial currents and computed plasma flux. 3. Compute iterated currents from the simplified circuit equations using hatIy1. 4. Residual: residual = iterated_currents - trial_currents.

Parameters:

• trial_currents (np.ndarray) – Vector of current values at time t+dt. Format matches self.currents_vec, typically including active coil currents, vessel mode currents, and total plasma current (normalized by plasma_norm_factor).

• active_voltage_vec (np.ndarray) – Vector of voltages applied to the active coils over the timestep.

• rtol_NK (float) – Relative tolerance to be used in the static GS solver for the plasma flux.

Returns:

Residual vector of currents, same format as self.currents_vec. Zero residual indicates a self-consistent solution of the plasma + circuit system.

Return type:

np.ndarray

M_coils_plasma(eq, greens)

Calculates the mutual inductance matrix between all tokamak coils and plasma grid points.

The matrix represents either the first or second derivative of the mutual inductance with respect to the vertical coordinate z, depending on the Greens function provided. This is used in stability and force calculations involving coil-plasma interactions.

Parameters:

• eq (FreeGSNKE equilibrium Object) – The equilibrium containing plasma and tokamak coil information.

• greens (function) – The Greens function used for the calculation: - “GreensBr” for the first z-derivative of the magnetic field. - “GreensdBrdz” for the second z-derivative of the magnetic field.

Returns:

M – Mutual inductance matrix between each coil (rows) and plasma grid point (columns), including coil polarity, multipliers, and R-coordinate weighting. The returned matrix is multiplied by -2π as per the standard formulation.

Return type:

np.ndarray, shape (n_coils, n_plasma_pts)

Notes

• plasma_pts are taken from eq.limiter_handler.plasma_pts, i.e., the reduced plasma domain.

• Coil contributions are summed over all filaments in each coil.

__init__(profiles, eq, GSStaticSolver, custom_coil_resist=None, custom_self_ind=None, full_timestep=0.0001, max_internal_timestep=0.0001, automatic_timestep=False, plasma_resistivity=1e-06, plasma_norm_factor=1000.0, blend_hatJ=0, max_mode_frequency=100.0, fix_n_vessel_modes=-1, threshold_dIy_dI=0.025, min_dIy_dI=0.01, mode_removal=True, linearize=True, dIydI=None, dIydtheta=None, target_relative_tolerance_linearization=1e-08, target_dIy=0.001, force_core_mask_linearization=False, l2_reg=1e-06, collinearity_reg=1e-06, verbose=False)

Initialize the nonlinear solver.

Sets up the equilibrium, profiles, circuit equations, vessel modes, Jacobians, and linearization for the coupled plasma–machine system.

Parameters:

• profiles (FreeGSNKE profiles) – Initial plasma profiles used to set up the linearisation.

• eq (FreeGSNKE equilibrium) – Equilibrium object providing grid, machine geometry, and limiter info.

• GSStaticSolver (FreeGSNKE static solver) – Static Grad–Shafranov solver for equilibrium solves.

• custom_coil_resist (ndarray, optional) – Resistances for all coils (active + passive). Defaults to tokamak values.

• custom_self_ind (ndarray, optional) – Mutual inductance matrix for coils. Defaults to tokamak values.

• full_timestep (float, default=1e-4) – Time increment for full steps (dt).

• max_internal_timestep (float, default=1e-4) – Maximum sub-step size for advancing circuit equations.

• automatic_timestep (tuple(float, float) or False, default=False) – If set, determines timestep size from growth rates.

• plasma_resistivity (float, default=1e-6) – Plasma resistivity value.

• plasma_norm_factor (float, default=1e3) – Normalization factor for plasma current.

• blend_hatJ (float, default=0) – Coefficient for blending plasma current distributions at t and t+dt.

• max_mode_frequency (float, optional) – Threshold frequency for retaining vessel modes.

• fix_n_vessel_modes (int, default=-1) – Fix number of passive vessel modes; -1 = auto-selection.

• threshold_dIy_dI (float, default=0.025) – Relative coupling threshold for including vessel modes (must be a number in [0,1]).

• min_dIy_dI (float, default=0.01) – Minimum coupling threshold for excluding vessel modes (must be a number in [0,1]).

• mode_removal (bool, default=True) – If True, remove weakly coupled vessel modes after Jacobian calculation.

• linearize (bool, default=True) – Whether to set up the linearised problem.

• dIydI (ndarray, optional) – Plasma current Jacobian wrt coil and plasma currents.

• dIydtheta (ndarray, optional) – Plasma current Jacobian wrt profile parameters.

• target_relative_tolerance_linearization (float, default=1e-8) – Relative tolerance for GS solve during Jacobian computation.

• target_dIy (float, default=1e-3) – Target perturbation size for Jacobian finite differences.

• force_core_mask_linearization (bool, default=False) – Enforce same core mask during finite-difference Jacobian evaluation.

• l2_reg (float, default=1e-6) – L2 Tikhonov regularization for nonlinear solver.

• collinearity_reg (float, default=1e-6) – Additional penalty for collinear terms in nonlinear solver.

• verbose (bool, default=False) – Print diagnostic output during initialization.

assign_currents(currents_vec, eq, profiles)

Assigns the input currents to the equilibrium and plasma profiles.

The function updates: - The total plasma current in eq and profiles. - The metal (vessel + active coil) currents in the tokamak, converting

from normal mode representation to physical coil currents.

Parameters:

• currents_vec (np.array) –

  Vector of current values to assign. Format:

  (active coil currents, vessel normal mode currents, total plasma current / plasma_norm_factor)

• eq (FreeGSNKE equilibrium Object) – Equilibrium object to be modified. Its _current attribute and tokamak.current_vec are updated.

• profiles (FreeGSNKE profiles Object) – Profiles object to be modified. Its total plasma current Ip is updated.

• Effects (Side)

• ------------

• currents. (- Updates self.vessel_currents_vec with the reconstructed physical)

• in-place. (- Modifies the input eq and profiles)

assign_currents_solve_GS(currents_vec, rtol_NK)

Assigns the input currents to the auxiliary equilibrium (self.eq2) and profiles (self.profiles2), then solves the static Grad-Shafranov (GS) problem to find the resulting plasma flux and current distribution.

Parameters:

• currents_vec (np.array) –

  Vector of current values to assign. Format matches self.assign_currents:

  (active coil currents, vessel normal mode currents, total plasma current / plasma_norm_factor)

• rtol_NK (float) – Relative tolerance to be used in the static GS solver.

• Effects (Side)

• ------------

• self.profiles2.Ip. (- Updates self.eq2._current and)

• currents. (- Updates self.eq2.tokamak.current_vec with the reconstructed physical)

• equation (- Solves the GS)

• self.profiles2.jtor. (updating self.eq2.plasma_psi and)

build_current_vec(eq, profiles)

Constructs the vector of currents for the dynamics solver.

The vector self.currents_vec includes, in order: - Active coil currents - Selected vessel normal mode currents - Total plasma current normalized by plasma_norm_factor

Parameters:

• eq (FreeGSNKE equilibrium object) – The equilibrium used to extract all metal currents before any mode truncation.

• profiles (FreeGS4E profiles object) – Profiles of the initial equilibrium. Used to extract the total plasma current.

build_dIydI_j(j, rtol_NK)

Compute the finite-difference derivative d(Iy)/dI_j using the prepared perturbation.

Uses the perturbation ΔI_j inferred by prepare_build_dIydI_j to compute the actual finite-difference derivative of the poloidal current vector Iy with respect to coil current I_j. Solves the Grad–Shafranov equilibrium at the perturbed current.

Parameters:

• j (int) – Index of the coil current to be varied, corresponding to self.currents_vec.

• rtol_NK (float) – Relative tolerance for the static Grad–Shafranov solver.

Returns:

• dIydIj (ndarray) – Finite-difference derivative d(Iy)/dI_j, a 1D array over all grid points in the reduced plasma domain.

• rel_ndIy (float) – Relative norm of the induced change in Iy, normalized by self.nIy.

build_dIydI_noGS(force_core_mask_linearization, starting_dI, core_mask, verbose)

Compute a first estimate of the Jacobian norm dIy/dI without solving GS.

This routine evaluates the plasma current response to perturbations in each coil or mode current using only the modified tokamak Green’s functions (no Grad–Shafranov solves). The resulting Jacobian norm is used for an initial sifting of passive vessel modes before a full linearisation.

If force_core_mask_linearization is True, the perturbation size for each mode is adjusted to ensure that the diverted core mask of the perturbed equilibrium matches the reference core mask. In that case, self.starting_dI and self.approved_target_dIy are updated accordingly.

Parameters:

• force_core_mask_linearization (bool) – Whether to enforce identical core masks between perturbed and reference equilibria when computing finite-difference derivatives. If True, perturbation amplitudes are reduced until the core mask matches.

• starting_dI (ndarray) – Initial perturbation amplitudes for each independent coil/mode current. This array is updated in place depending on the masking strategy.

• core_mask (ndarray of bool) – Boolean mask indicating the core region of the reference equilibrium.

• verbose (bool) – If True, print diagnostic information about mode perturbations and accepted perturbation sizes.

• Updates

• -------

• self.dIydI_noGS (ndarray, shape (n_plasma_points, n_coils)) – Approximate Jacobian of plasma current distribution wrt coil currents, computed without GS solves.

• self.ndIydI_no_GS (ndarray, shape (n_coils,)) – Norm of dIy/dI for each coil/mode, used in mode selection.

• self.rel_ndIy (ndarray, shape (n_coils,)) – Relative plasma current perturbation norms for each mode.

• self.starting_dI (ndarray) – Adjusted perturbation amplitudes after enforcing masking strategy.

• self.approved_target_dIy (ndarray) – Updated target norms for plasma current perturbations (if core mask enforcement applied).

build_dIydtheta(profiles, rtol_NK, verbose=False)

Compute the finite-difference Jacobian d(Iy)/dθ using pre-scaled perturbations.

This function calculates the derivative of the poloidal current vector Iy with respect to the plasma profile parameters θ. Perturbations Δθ are taken from self.final_dtheta_record, which are inferred previously by prepare_build_dIydtheta to produce controlled changes in Iy. For each perturbed profile, the Grad–Shafranov equilibrium is solved and the resulting Iy change is measured. The final Jacobian is the normalized change per unit Δθ.

Parameters:

• profiles (FreeGS4E profiles object) – Profiles object of the linearization-point equilibrium.

• rtol_NK (float) – Relative tolerance for the Newton–Krylov Grad–Shafranov solver.

• verbose (bool, optional) – If True, print intermediate diagnostic information (default: False).

Returns:

• dIydtheta (ndarray, shape (len(Iy), n_profiles_parameters)) – Finite-difference derivative of Iy with respect to each profile parameter.

• rel_ndIy (ndarray, shape (n_profiles_parameters,)) – Relative norm of ΔIy induced by each perturbation, normalized by self.nIy.

Notes

• This function modifies the plasma profiles temporarily during the finite-difference calculation, then resets them to the original state.

• Supports both the standard (alpha_m, alpha_n, paxis/betap/Beta0) parameters and Lao-style profile coefficients (multiple alpha and beta terms).

build_linearization(eq, profiles, dIydI, dIydtheta, target_relative_tolerance_linearization, force_core_mask_linearization, verbose)

Builds the Jacobians d(Iy)/dI and d(Iy)/dtheta for linearizing the plasma-current response around a given equilibrium. These Jacobians are used to set up the solver for the linearised problem, providing initial slopes for both coil currents and plasma profile parameters.

This function optionally computes dIydI and dIydtheta if they are not provided. Perturbations are applied to coil currents and profile parameters, and the corresponding changes in the poloidal current vector Iy are used to form finite-difference derivatives. The perturbations are adjusted to ensure stability and, if requested, to maintain the plasma core mask.

Parameters:

• eq (FreeGS4E equilibrium object) – Equilibrium around which to linearize.

• profiles (FreeGS4E profiles object) – 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.

• 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.

• target_relative_tolerance_linearization (float) – Relative tolerance used in the static Grad–Shafranov solver during finite-difference calculations.

• force_core_mask_linearization (bool) – If True, adjusts the perturbations so that the perturbed plasma core mask remains identical to the original mask.

• verbose (bool) – If True, prints intermediate results including initial and final perturbations, relative changes in Iy, and Grad–Shafranov residuals.

Notes

• This function updates self.dIydI and self.dIydtheta (or copies from initial conditions

if already computed) and stores final perturbation magnitudes in self.final_dI_record and self.final_dtheta_record. - The core mask consistency is checked when force_core_mask_linearization is True. - The function also updates the derivatives of coil positions with respect to currents in self.dRZdI.

calc_lumped_plasma_resistance(norm_red_Iy0, norm_red_Iy1)

Compute the lumped plasma resistance using the plasma resistance matrix R_yy.

The lumped resistance is calculated by contracting the 1D plasma resistance vector (self.plasma_resistance_1d) with two normalized plasma current distribution vectors. This provides an effective scalar resistance for the given current profiles.

Parameters:

• norm_red_Iy0 (np.array) – Normalized plasma current distribution vector at time t_0. The entries should be normalized so that their sum equals 1.

• norm_red_Iy1 (np.array) – Normalized plasma current distribution vector at time t_1. The entries should be normalized so that their sum equals 1.

Returns:

Lumped plasma resistance corresponding to the given current distributions.

Return type:

float

calculate_Leuer_parameter()

Calculates the Leuer stability parameter for the plasma, which quantifies the passive vertical stability provided by surrounding metals and active coils.

The Leuer parameter, as defined in Leuer (1989, “Passive Vertical Stability in the Next Generation Tokamaks”, Eq. 6), is the ratio of stabilising force gradients induced by metals to the de-stabilising force gradients caused by the plasma and coil currents.

The calculation involves: - Mutual inductance derivatives between coils and plasma points (first and second z-derivatives). - Mutual inductances between metal coils themselves. - Plasma current distribution and coil currents. - Stabilising forces (from active and passive coils) and de-stabilising forces. - Computation of Leuer parameters in different configurations.

Attributes Updated

actives_stab_forcefloat

Stabilising force due to active coils.

passives_stab_forcefloat

Stabilising force due to passive coils.

all_coils_stab_forcefloat

Stabilising force due to all metal coils (active + passive).

actives_destab_forcefloat

De-stabilising force due to active coils.

passives_destab_forcefloat

De-stabilising force due to passive coils.

all_coils_destab_forcefloat

De-stabilising force due to all metal coils.

Leuer_passive_stab_over_active_destabfloat

Ratio of passive stabilising force to active coil de-stabilising force.

Leuer_metals_stab_over_active_destabfloat

Ratio of total metal stabilising force to active coil de-stabilising force.

Leuer_metals_stab_over_metals_destabfloat

Ratio of total metal stabilising force to total metal de-stabilising force.

References

Leuer, J. A., “Passive Vertical Stability in the Next Generation Tokamaks,” 1989, 10.13182/FST89-A39747.

calculate_hatIy(trial_currents, plasma_psi)

Computes the normalized plasma current distribution (hatIy) corresponding to a given set of currents and plasma flux.

Parameters:

• trial_currents (np.array) – Full vector of currents (active coils, vessel modes, plasma current).

• plasma_psi (np.array) – Plasma flux on the full grid (2D array).

Returns:

Normalized plasma current distribution on the reduced domain.

Return type:

np.array

calculate_hatIy_GS(trial_currents, rtol_NK, record_for_updates=False)

Computes the normalized plasma current distribution (hatIy) corresponding to a given set of currents by solving the static Grad-Shafranov problem.

Parameters:

• trial_currents (np.array) – Full vector of currents (active coils, vessel modes, plasma current).

• rtol_NK (float) – Relative tolerance for the static GS solver.

• record_for_updates (bool, optional) – If True, records auxiliary updates (not always needed).

Returns:

Normalized plasma current distribution on the reduced plasma domain.

Return type:

np.array

calculate_rel_tolerance_GS(trial_plasma_psi, a_res_GS=None)

Computes the relative residual of the plasma flux for the static Grad-Shafranov (GS) problem, comparing the GS residual to the actual change in plasma flux due to dynamics. This metric quantifies the convergence of the timestepper for the plasma flux update.

The relative residual is defined as:

r_res_GS = max(|GS_residual|) / max(|Δpsi|)

where Δpsi = psi(t+dt) - psi(t). If the GS residual a_res_GS is not provided, it is computed internally using the static GS solver.

Parameters:

• trial_plasma_psi (np.ndarray) – Plasma flux at the current timestep, psi(t+dt), shape (nx, ny).

• a_res_GS (np.ndarray, optional) – Residual of the static GS problem at t+dt. If None, it will be calculated internally. Shape should match the flattened plasma flux.

Returns:

Relative plasma flux residual. Values close to 0 indicate good convergence of the plasma flux update.

Return type:

float

calculate_rel_tolerance_currents(current_residual, curr_eps)

Computes the relative residual of the current update compared to the actual step taken in the currents. This quantifies the convergence of the timestepper by comparing the residual to the magnitude of the current change.

The relative residual is defined as:

relative_residual_i = |current_residual_i / max(|ΔI_i|, curr_eps)|

where ΔI_i = trial_currents_i - currents_vec_m1_i, and curr_eps prevents division by very small steps.

Parameters:

• current_residual (np.ndarray) – Residual of the current values at the current timestep. Same format as self.currents_vec.

• curr_eps (float) – Minimum allowable step size in the currents. Used to avoid artificially large relative residuals when the step is very small.

Returns:

Relative residual of the current update. Same format as self.currents_vec. Values close to 0 indicate good convergence of the timestep.

Return type:

np.ndarray

check_and_change_active_coil_resistances(active_coil_resistances)

Checks if new active coil resistances are provided and updates them if needed.

This method compares the input array of active coil resistances with the current resistances stored in self.evol_metal_curr. If the input is different, it resets the resistances, updates the solvers accordingly, and prints the new coil resistances.

Parameters:

active_coil_resistances (np.array or None) – Array of new resistances for the active coils. If None, no changes are made.

Notes

• If the input resistances are identical to the current ones, the method does nothing.

• When resistances are updated, self.set_solvers() is called to ensure the

solvers reflect the new electrical properties. - The updated resistances are printed for confirmation.

check_and_change_plasma_resistivity(plasma_resistivity, relative_threshold_difference=0.01)

Check if the plasma resistivity differs from the current value and update it if necessary.

Compares the new plasma_resistivity with the current self.plasma_resistivity. If the relative difference exceeds relative_threshold_difference, the resistivity is reset using reset_plasma_resistivity, which also updates the associated solver objects.

Parameters:

• plasma_resistivity (float) –

  New scalar value of the plasma resistivity [Ohm·m]. The resistance at each grid point in the reduced domain is computed as

  R_yy = 2 * π * plasma_resistivity * R / (dR * dZ)

  where R is the major radius of the grid point and dR*dZ is the area of the corresponding domain element.

• relative_threshold_difference (float, optional) – Relative threshold for updating the resistivity. Default is 0.01 (1%). If the relative change in resistivity exceeds this value, the resistivity is reset.

check_and_change_profiles(profiles_parameters=None)

Updates the plasma current profile parameters at time t+dt if new values are provided.

This method checks whether a dictionary of new profile parameters is supplied. If so, it updates both the evolving equilibrium profiles (self.profiles1) and the auxiliary profiles (self.profiles2) accordingly. For profiles of type “Lao85”, the internal profile initialization routine is called after updating the parameters. A flag is set to indicate that a change occurred.

Parameters:

profiles_parameters (dict or None, optional) – Dictionary of profile parameters to update. Keys and values should match the attributes of the profile object (see get_profiles_values for structure). If None, no changes are made and the profiles remain unchanged.

Notes

• Sets self.profiles_change_flag = 1 if parameters are updated, otherwise 0.

• Both the main (profiles1) and auxiliary (profiles2) profiles are updated

to ensure consistency during timestep calculations.

currents_from_hatIy(hatIy1, active_voltage_vec)

Computes the full set of currents at time t+dt from a guess of the normalized plasma current distribution, using the simplified circuit solver.

Parameters:

• hatIy1 (np.array) – Guess for the normalized plasma current distribution at time t+dt (sum=1, reduced plasma domain).

• active_voltage_vec (np.array) – Voltages applied to the active coils between t and t+dt.

Returns:

• np.array – Full current vector at t+dt, matching the format of self.currents_vec.

• Workflow

• ——–

• 1. Computes a blended plasma distribution using make_blended_hatIy.

• 2. Computes Myy_hatIy_left = self.handleMyy.dot(blended_hatIy).

• 3. Calls self.simplified_solver_J1.stepper to solve for all currents, given the blended distribution

• and applied coil voltages.

get_profiles_values(profiles)

Extracts and stores relevant properties from a FreeGS4E profiles object.

Sets internal attributes describing the profile type, number of independent parameters, and their values, which are used in linearisation and stepper calculations.

Parameters:

profiles (FreeGS4E profiles object) – The profiles object of the initial equilibrium. This provides the parameters defining the plasma current density profile, e.g., alpha_m, alpha_n, paxis, betap, Beta0, or Lao coefficients.

get_vessel_currents(eq)

Extracts and stores all metal currents from a given equilibrium.

Retrieves currents for both active coils and passive vessel structures from the input equilibrium’s tokamak object. The values are stored in self.vessel_currents_vec.

Parameters:

eq (FreeGSNKE equilibrium object) – The equilibrium from which to extract metal currents. Uses eq.tokamak to access current values.

hatIy1_iterative_cycle(hatIy1, active_voltage_vec, rtol_NK)

Performs one iteration of the cycle: 1. Uses a guessed plasma current distribution at t+dt (hatIy1) to compute all currents. 2. Solves the static Grad-Shafranov (GS) problem to find the resulting plasma flux and updated plasma current distribution.

Parameters:

• hatIy1 (np.array) – Guess for the normalized plasma current distribution at t+dt. Must sum to 1 (reduced plasma domain only).

• active_voltage_vec (np.array) – Voltages applied to the active coils between t and t+dt.

• rtol_NK (float) – Relative tolerance for the static GS solver.

initialize_from_ICs(eq, profiles, target_relative_tolerance_linearization=1e-07, dIydI=None, dIydtheta=None, force_core_mask_linearization=False, verbose=False)

Initialize the dynamics solver from a given equilibrium and plasma profiles.

This method prepares all internal structures required for advancing the coupled plasma–vessel dynamics, including:

• Building the vector of currents (self.currents_vec) comprising

active coil currents, selected vessel normal modes, and normalized plasma current. - Setting up auxiliary equilibrium (eq2) and profile (profiles2) objects for intermediate calculations. - Computing the linearized Jacobians d(Iy)/dI and d(Iy)/dtheta if not provided, and transferring them to the linear solver.

Parameters:

• eq (FreeGS4E equilibrium object) – The initial equilibrium, used to assign all metal currents.

• profiles (FreeGS4E profiles object) – Profiles of the initial equilibrium, used to assign the total plasma current.

• target_relative_tolerance_linearization (float, default=1e-7) – Relative tolerance for static Grad–Shafranov solves during initialization and Jacobian computation. This does not affect the solver’s runtime tolerance.

• 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.

• 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.

• force_core_mask_linearization (bool, default=False) – If True, enforces that the perturbed plasma core mask remains identical when computing d(Iy)/dI.

• verbose (bool, default=False) – Enables progress printouts during initialization.

• Effects (Side)

• ------------

• self.eq1 (- Sets up)

• self.profiles1

• self.eq2

• self.profiles2.

• self.currents_vec (- Initializes)

• self.Iy

• self.hatIy

• self.blended_hatIy.

• (self.linearised_sol). (- Builds the linearization and transfers it to the linear solver)

• self.handleMyy. (- Updates the Myy matrix through)

make_blended_hatIy_(hatIy1, blend)

Produces a weighted average of the current plasma distribution at time t (self.hatIy) and a guess for the distribution at time t+dt (hatIy1).

Parameters:

• hatIy1 (np.array) – Guess for the normalized plasma current distribution at time t+dt. Must sum to 1. Only covers the reduced plasma domain.

• blend (float) – Weighting factor between 0 and 1. - blend=0 → only uses hatIy1 - blend=1 → only uses hatIy (current time)

• Effects (Side)

• ------------

• combination (- Sets self.blended_hatIy to the weighted) – blended_hatIy = (1 - blend) * hatIy1 + blend * self.hatIy

nlstepper(active_voltage_vec, profiles_parameters=None, plasma_resistivity=None, target_relative_tol_currents=0.005, target_relative_tol_GS=0.003, working_relative_tol_GS=0.001, target_relative_unexplained_residual=0.5, max_n_directions=3, step_size_psi=2.0, step_size_curr=0.8, scaling_with_n=0, blend_GS=0.5, curr_eps=1e-05, max_no_NK_psi=5.0, clip=5, verbose=0, linear_only=False, max_solving_iterations=50, custom_active_coil_resistances=None)

Advance the system by one timestep using a nonlinear Newton-Krylov (NK) stepper.

If linear_only=True, only the linearised dynamic problem is advanced. Otherwise, a full nonlinear solution is sought using an iterative NK-based algorithm. On convergence, the timestep is advanced by self.dt_step and the updated currents, equilibrium, and profile objects are assigned to self.currents_vec, self.eq1, and self.profiles1.

Algorithm overview

The solver proceeds as follows:

1. Solve the linearised problem to obtain an initial guess for the currents and solve the associated static Grad–Shafranov (GS) problem, yielding trial_plasma_psi and trial_currents (including tokamak_psi). 2. If the pair [trial_plasma_psi, tokamak_psi] fails the GS tolerance check, update trial_plasma_psi toward the GS solution. 3. At fixed currents, update trial_plasma_psi via NK iterations on the root problem in plasma flux. 4. At fixed plasma flux, update currents via NK iterations on the root problem in currents. 5. If either the current residuals or the GS tolerance check fail, return to step 2. 6. On convergence, record the solution into self.currents_vec, self.eq1, and self.profiles1.

param active_voltage_vec:

Vector of applied voltages on the active coils between t and t+dt.

type active_voltage_vec:

np.ndarray

param profiles_parameters:

If None, profile parameters remain unchanged. Otherwise, dictionary specifying updated parameters for the profiles object. See get_profiles_values for dictionary structure. This enables time-dependent profile parameters.

type profiles_parameters:

dict or None, optional

param plasma_resistivity:

Updated plasma resistivity. If None, resistivity is left unchanged. Enables time- dependent resistivity.

type plasma_resistivity:

float or array-like, optional

param target_relative_tol_currents:

Required relative tolerance on currents for convergence of the dynamic problem. Computed as residual / (currents(t+dt) - currents(t)).

type target_relative_tol_currents:

float, optional, default=0.005

param target_relative_tol_GS:

Required relative tolerance on plasma flux for convergence of the static GS problem. Computed as residual / Δψ where Δψ is the flux change between timesteps.

type target_relative_tol_GS:

float, optional, default=0.003

param working_relative_tol_GS:

Tolerance used when solving intermediate GS problems during the step. Must be stricter than target_relative_tol_GS.

type working_relative_tol_GS:

float, optional, default=0.001

param target_relative_unexplained_residual:

NK solver stopping criterion: inclusion of additional Krylov basis vectors stops once more than 1 - target_relative_unexplained_residual of the residual is canceled.

type target_relative_unexplained_residual:

float, optional, default=0.5

param max_n_directions:

Maximum number of Krylov basis vectors used in NK updates.

type max_n_directions:

int, optional, default=3

param step_size_psi:

Step size for finite difference calculations in the NK solver applied to ψ, measured in units of the residual norm.

type step_size_psi:

float, optional, default=2.0

param step_size_curr:

Step size for finite difference calculations in the NK solver applied to currents, measured in units of the residual norm.

type step_size_curr:

float, optional, default=0.8

param scaling_with_n:

Exponent controlling step scaling in NK updates: candidate step is scaled by (1 + n_iterations)**scaling_with_n.

type scaling_with_n:

int, optional, default=0

param blend_GS:

Blending coefficient used when updating trial_plasma_psi toward the GS solution. Must be in [0, 1].

type blend_GS:

float, optional, default=0.5

param curr_eps:

Regularisation parameter for relative current convergence checks, preventing division by small current changes.

type curr_eps:

float, optional, default=1e-5

param max_no_NK_psi:

Threshold for triggering NK updates on ψ. Activated if relative_psi_residual > max_no_NK_psi * target_relative_tol_GS.

type max_no_NK_psi:

float, optional, default=5.0

param clip:

Maximum allowed step size for each accepted Krylov basis vector, in units of the exploratory step.

type clip:

float, optional, default=5

param verbose:

Verbosity level. * 0: silent * 1: report convergence progress per NK cycle * 2: include detailed intermediate output

type verbose:

int, optional, default=0

param linear_only:

If True, only the linearised solution is used (skipping nonlinear solves).

type linear_only:

bool, optional, default=False

param max_solving_iterations:

Maximum number of nonlinear NK cycles before the solve is terminated.

type max_solving_iterations:

int, optional, default=50

param custom_active_coil_resistances:

If provided, overrides default active coil resistances with those specifed. Enables time-dependent coil resistances (can be used for switching coils “on” and “off”).

type custom_active_coil_resistances:

array-like or None, optional

Notes

On convergence, the method updates internal state: - self.currents_vec stores the evolved currents. - self.eq1 stores the new Grad–Shafranov equilibrium. - self.profiles1 stores the updated profile object.

raises RuntimeError:

If the nonlinear solve does not converge within max_solving_iterations.

prepare_build_dIydI_j(j, rtol_NK, target_dIy, starting_dI, GS=True)

Prepare the finite-difference derivative d(Iy)/dI_j by estimating the perturbation ΔI_j.

This function determines the size of a current perturbation ΔI_j that produces a target change in the poloidal current vector Iy with norm ||ΔIy|| = target_dIy. Optionally, the full Grad–Shafranov (GS) problem can be solved to update the equilibrium, or a simplified approach using the modified tokamak flux can be used.

Parameters:

• j (int) – Index of the coil current to be varied, corresponding to self.currents_vec.

• rtol_NK (float) – Relative tolerance for the static Grad–Shafranov solver.

• target_dIy (float) – Target norm of the induced change in Iy used to scale the perturbation.

• starting_dI (float) – Initial guess for the coil current perturbation ΔI_j.

• GS (bool, optional) – If True, solve the full Grad–Shafranov problem; if False, use modified tokamak flux without solving GS (default: True).

Returns:

• dIy_scaled (ndarray) – Initial finite-difference estimate of ΔIy / ΔI_j.

• rel_ndIy (float) – Relative norm of the induced change in Iy, normalized by self.nIy.

prepare_build_dIydtheta(profiles, rtol_NK, target_dIy, starting_dtheta, verbose=False)

Prepare finite-difference evaluation of d(Iy)/dθ, where θ parameterises the plasma current density profile.

Perturbs the profile parameters by small trial shifts starting_dtheta, solves the Grad–Shafranov equilibrium for each perturbed profile, and measures the corresponding change in Iy (poloidal current distribution). The trial perturbations are then rescaled so that the induced change in Iy has prescribed norm target_dIy. This sets up consistent perturbation sizes for later derivative calculations.

Parameters:

• profiles (FreeGS4E profiles object) – Profiles object of the linearisation-point equilibrium.

• rtol_NK (float) – Relative tolerance for the Newton–Krylov GS solves.

• target_dIy (ndarray) – Desired norms of ΔIy for each perturbed direction.

• starting_dtheta (ndarray) – Initial perturbations of the profile parameters, used to infer the scaling between Δθ and ΔIy.

• verbose (bool, optional) – If True, print intermediate diagnostic information (default: False).

Returns:

• dIy_0 (ndarray, shape (len(Iy), n_profiles_parameters)) – Raw changes in Iy from the initial perturbations.

• rel_ndIy_0 (ndarray, shape (n_profiles_parameters,)) – Relative norms of ΔIy per unit perturbation, normalised by self.nIy.

• Updates

• ——-

• self.final_dtheta_record (ndarray) – Adjusted perturbations Δθ that will yield ΔIy with the target norms.

remove_modes(eq, selected_modes_mask)

Remove unselected normal modes and update circuit equations.

Given an equilibrium and a mask over the current mode set, this method reduces the dimensionality of the system by discarding modes marked as inactive in selected_modes_mask. The circuit equations, current vector, and nonlinear solver are reinitialized consistently with the reduced system size.

Parameters:

• eq (FreeGSNKE equilibrium) – Equilibrium object containing plasma state information.

• selected_modes_mask (ndarray of bool, shape (n_modes,)) – Boolean mask indicating which modes to keep. Entries corresponding to True are retained, and those corresponding to False are removed. Must have the same shape as self.currents_vec at the time of call.

• Updates

• -------

• self.n_metal_modes (int) – Number of remaining independent modes after reduction.

• self.extensive_currents_dim (int) – Dimension of the reduced current vector (n_metal_modes + 1).

• self.arange_currents (ndarray) – Index array for the reduced system, ranging from 0 to self.extensive_currents_dim - 1.

• self.currents_vec (ndarray) – Zero-initialized current vector of reduced dimensionality.

• self.circuit_eq_residual (ndarray) – Residual vector for the reduced circuit equations.

• self.currents_nk_solver (nk_solver.nksolver) – Nonlinear solver instance reinitialized for the reduced system size.

• self.simplified_solver_J1 (simplified_solver_J1) – Rebuilt solver instance consistent with reduced system.

• self.linearised_sol (linear_solver) – Rebuilt linearised solver instance consistent with reduced system.

reset_plasma_resistivity(plasma_resistivity)

Reset the plasma resistivity and update all relevant solver objects.

This updates self.plasma_resistance_1d, the diagonal of the plasma resistance matrix restricted to the reduced domain (inside the limiter), and also propagates the updated resistivity to the linearized and simplified solvers.

Parameters:

plasma_resistivity (float) –

Scalar value of the plasma resistivity [Ohm·m]. The resistance at each grid point in the reduced domain is computed as

R_yy = 2 * π * plasma_resistivity * R / (dR * dZ)

where R is the major radius of the grid point and dR*dZ is the area of the corresponding domain element.

reset_timestep(full_timestep, max_internal_timestep)

Reset the timestep parameters for the simulation.

Updates the full timestep used to advance the dynamics and the maximum allowed internal sub-timestep used for circuit equation integration. Resets the corresponding timesteps in all relevant solver objects.

Parameters:

• full_timestep (float) – The time interval dt over which the stepper advances the full system dynamics. Applies to both linear and nonlinear steppers. A Grad-Shafranov (GS) equilibrium is recalculated at every full_timestep.

• max_internal_timestep (float) – Maximum size of sub-steps when advancing a single full_timestep. These sub-steps are used to integrate the circuit equations.

set_linear_solution(active_voltage_vec, dtheta_dt)

Compute an initial nonlinear solve guess using the linearised dynamics.

Advances the system one timestep using the linearised solver to generate a trial current vector at t + Δt, starting from self.currents_vec at time t. This trial solution is then used as the initial guess for the nonlinear solver, which updates the Grad–Shafranov (GS) equilibrium at t + Δt.

Parameters:

• active_voltage_vec (ndarray) – External voltage applied to the active coils during the timestep.

• dtheta_dt (ndarray) – Time derivatives of the plasma current density profile parameters.

• Updates

• -------

• self.trial_currents (ndarray) – Trial coil/mode currents at t + Δt obtained from the linearised solver.

• self.eq2.plasma_psi (ndarray) – Plasma poloidal flux after solving GS with the trial currents.

• self.trial_plasma_psi (ndarray) – Copy of the GS solution for the plasma flux surface configuration corresponding to the trial currents.

set_plasma_resistivity(plasma_resistivity)

Set the resistivity of the plasma and update the corresponding diagonal plasma resistance vector used in circuit calculations.

The vector self.plasma_resistance_1d represents the diagonal of the plasma resistance matrix R_yy, restricted to the reduced domain defined by plasma_domain_mask (grid points inside the limiter).

Parameters:

plasma_resistivity (float) –

Scalar value of the plasma resistivity [Ohm·m]. The resistance at each grid point in the reduced domain is computed as

R_yy = 2 * π * plasma_resistivity * R / (dR * dZ)

where R is the major radius of the grid point and dR*dZ is the area of the corresponding domain element.

set_solvers()

Initialize and configure the time-integration solvers.

Creates solver instances for both the simplified nonlinear system (simplified_solver_J1) and the linearised system (linear_solver). Both solvers are constructed using machine inductance/resistance matrices and plasma parameters, and are configured to advance currents consistently with the timestep settings.

After creation, the linearised solver is set to operate around the current linearisation point, using the Jacobians and reference plasma state.

Updates

self.simplified_solver_J1simplified_solver_J1

Nonlinear solver instance for reduced-order plasma–circuit dynamics.

self.linearised_sollinear_solver

Linear solver instance for coupled plasma–circuit dynamics.

step_complete_assign(working_relative_tol_GS, from_linear=False)

Finalize the timestep advancement and update the equilibrium and current state.

This function completes the evolution over a timestep dt_step by: - Recording the time-evolved currents (self.trial_currents) in self.currents_vec. - Updating the equilibrium (self.eq1) and plasma profiles (self.profiles1) with the corresponding plasma flux (self.trial_plasma_psi) and derived current distribution. - Updating the normalized plasma current distribution (self.hatIy).

Parameters:

• working_relative_tol_GS (float) – Fractional tolerance for the Grad–Shafranov solver during this timestep. The effective GS tolerance is set relative to the maximum change in plasma flux.

• from_linear (bool, default=False) – If True, only the linearised solution is applied. Reduces the number of full GS solves by copying auxiliary equilibrium and profiles to the main state.

• Effects (Side)

• ------------

• self.step_no. (- Advances self.time and increments)

• values. (- Updates self.currents_vec_m1 and self.Iy_m1 to store previous step)

• self.currents_vec (- Updates)

• self.Iy

• self.hatIy

• self.rtol_NK. (and)

• state. (- Modifies self.eq1 and self.profiles1 to reflect the timestep-evolved)

unstable_mode_deformations(starting_dI=50, rtol_NK=1e-07, target_dIy=0.002)

Applies the first unstable mode to evaluate plasma centroid deformations and the corresponding plasma current distribution response.

This method calculates the derivatives of the current-averaged plasma coordinates (R, Z) with respect to the magnitude of the current in the unstable mode (Im), i.e., dR/dIm and dZ/dIm. It also records the plasma current distribution after applying Im and constructs a “rigidly displaced” version of the original current map for comparison.

Parameters:

• starting_dI (float, optional) – Initial perturbation amplitude in the unstable mode used to estimate the slope of ||delta(Iy)|| / delta(Im). Default is 50.

• rtol_NK (float, optional) – Relative tolerance to be used in the static Grad-Shafranov (GS) problem. Default is 1e-7.

• target_dIy (float, optional) – Target norm of the plasma current change vector (delta(Iy)) used to scale the mode perturbation. Default is 2e-3.

• Updated (Attributes)

• ------------------

• dRZd_unstable_mode (np.ndarray) – Array [dR/dIm, dZ/dIm] representing the sensitivity of the plasma centroid to the applied unstable mode current.

• deformable_vs_rigid_jtor (tuple of np.ndarray) – Tuple containing: - The plasma current distribution with the unstable mode applied. - The “rigidly displaced” plasma current distribution obtained by shifting the original distribution to match the new centroid (R, Z) positions.