diff --git a/freegsnke/__init__.py b/freegsnke/__init__.py
index 11b4640..019a277 100644
--- a/freegsnke/__init__.py
+++ b/freegsnke/__init__.py
@@ -1,3 +1,10 @@
+"""
+FreeGSNKE package.
+
+Provides tools for solving Grad-Shafranov equilibria and related
+numerical methods for plasma physics modelling.
+"""
+
import importlib.metadata
__version__ = importlib.metadata.version("freegsnke")
diff --git a/freegsnke/build_machine.py b/freegsnke/build_machine.py
index 33841f1..370cff3 100644
--- a/freegsnke/build_machine.py
+++ b/freegsnke/build_machine.py
@@ -589,6 +589,40 @@ def build_active_coil_dict(active_coils):
def copy_tokamak(tokamak: Machine):
+ """
+ Create a copy of a tokamak Machine object with controlled deep/shallow copying.
+
+ This function duplicates the core Machine object while carefully controlling
+ which internal structures are shared and which are copied. It is intended to
+ produce a new independent tokamak instance suitable for simulation branching
+ or parameter variation.
+
+ Copy behaviour:
+ - Uses `tokamak.copy()` as the base object copy.
+ - Performs shallow copies of coil container structures.
+ - Copies selected numerical arrays using `copy_into`.
+ - Shares some reference-type attributes (e.g. probes).
+
+ Parameters
+ ----------
+ tokamak : Machine
+ Original tokamak machine object to be copied.
+
+ Returns
+ -------
+ Machine
+ New tokamak instance with copied geometry and selected attributes.
+
+ Notes
+ -----
+ - `coils_dict` is shallow-copied (dictionary structure copied, values shared).
+ - `coils_list` is shallow-copied (list container copied, objects shared).
+ - `coil_resist`, `coil_self_ind`, and `current_vec` are copied using
+ `copy_into` (typically deep or array-level copy depending on implementation).
+ - `probes` is shared by reference (not copied).
+ - This function does not guarantee full independence of all mutable
+ substructures unless `copy_into` enforces deep copying.
+ """
new_tokamak = tokamak.copy()
new_tokamak.coils_dict = tokamak.coils_dict.copy()
diff --git a/freegsnke/circuit_eq_metal.py b/freegsnke/circuit_eq_metal.py
index f1e41ba..3cbf4a6 100644
--- a/freegsnke/circuit_eq_metal.py
+++ b/freegsnke/circuit_eq_metal.py
@@ -1,5 +1,5 @@
"""
-Defines the metal_currents Object, which handles the circuit equations of
+Defines the metal_currents object, which handles the circuit equations of
all metal structures in the tokamak - both active PF coils and passive structures.
Copyright 2025 UKAEA, UKRI-STFC, and The Authors, as per the COPYRIGHT and README files.
@@ -28,6 +28,28 @@
class metal_currents:
+ """
+ Time evolution model for electrical currents in conducting structures
+ within a tokamak (active coils, passive vessel, and optional plasma coupling).
+
+ This class solves the dynamical circuit equations for metallic components
+ using an implicit-Euler time integrator. It can operate in multiple modes:
+
+ - Coil-only (vacuum vessel response)
+ - Vessel eigenmode representation of passive structures
+ - Fully coupled plasma–metal system (optional)
+
+ The system evolves currents in a reduced basis determined by:
+ - physical active coil currents
+ - vessel eigenmodes (if enabled)
+ - optional plasma current coupling
+
+ Key features:
+ - Construction of resistance and inductance operators
+ - Optional diagonalisation of passive vessel dynamics
+ - Optional coupling to plasma current evolution
+ - Support for multi-step implicit time integration
+ """
def __init__(
self,
@@ -42,45 +64,42 @@ def __init__(
coil_self_ind=None,
verbose=True,
):
- """Sets up framework to solve the dynamical evolution of all metal currents.
- Can be used by itself to solve metal circuit equations for vacuum shots,
- i.e. when in the absence of the plasma.
+ """
+ Initialise the dynamical evolution model for metallic currents.
+
+ This class solves the time evolution of currents in conducting structures
+ (coils, vessel, and passive structures). It can also be used independently
+ to solve vacuum vessel circuit equations in the absence of plasma.
Parameters
----------
-
- eq : FreeGSNKE equilibrium Object
- Initial equilibrium. This is used to set the domain/grid properties
- as well as the machine properties.
- Furthermore, eq will be used to set up the linearization used by the linear evolutive solver.
- It can be changed later by initializing a new set of initial conditions.
- Note however that, to change either the machine or limiter properties
- it will be necessary to instantiate a new nl_solver object.
+ eq : FreeGSNKE equilibrium object
+ Initial equilibrium used to define grid geometry, machine configuration,
+ and linearisation for the evolution solver. Changing machine geometry
+ requires re-instantiation.
flag_vessel_eig : bool
- Flag re whether vessel eigenmodes are used or not.
+ If True, include vessel eigenmodes in the circuit model.
flag_plasma : bool
- Whether to include plasma in circuit equation. If True, plasma_pts
- must be provided.
- plasma_pts : freegsnke.limiter_handler.plasma_pts
- Domain points in the domain that are included in the evolutive calculations.
- A typical choice would be all domain points inside the limiter. Defaults to None.
+ If True, include plasma coupling in the circuit equations.
+ In this case, `plasma_pts` must be provided.
max_mode_frequency : float
- Maximum frequency of vessel eigenmodes to include in circuit equation.
- Defaults to 1. Unit is s^-1.
+ Maximum vessel eigenmode frequency to include (s⁻¹).
max_internal_timestep : float
- Maximum value of the 'internal' timestep for implicit euler solver. Defaults to .0001.
- The 'internal' timestep is the one actually used by the solver.
+ Internal timestep used by the implicit Euler solver.
full_timestep : float
- Timestep by which the equations are advanced. If full_timestep>max_internal_timestep
- multiple 'internal' steps are executed. Defaults to .0001.
- coil_resist : np.array
- 1d array of resistance values for all conducting elements in the machine,
- including both active coils and passive structures.
- Defaults to None, meaning the values calculated by default in tokamak will be sourced and used.
- coil_self_ind : np.array
- 2d matrix of mutual inductances between all pairs of machine conducting elements,
- including both active coils and passive structures
- Defaults to None, meaning the values calculated by default in tokamak will be sourced and used.
+ Physical timestep for advancing the system. May be split into
+ multiple internal steps if larger than `max_internal_timestep`.
+ plasma_pts : array-like, optional
+ Grid points included in plasma evolution coupling (typically points
+ inside the limiter).
+ coil_resist : ndarray, optional
+ Resistances of all conducting elements (coils + passive structures).
+ If None, default machine values are used.
+ coil_self_ind : ndarray, optional
+ Mutual inductance matrix for all conducting elements.
+ If None, default machine values are used.
+ verbose : bool
+ If True, enable diagnostic output during setup.
"""
self.n_coils = eq.tokamak.n_coils
@@ -144,12 +163,52 @@ def __init__(
def build_rm1l(
self,
):
- # active + passive
+ """
+ Construct the R⁻¹L coupling matrix for the circuit model.
+
+ This method builds the product of the inverse resistance matrix with the
+ self-inductance (mutual inductance) matrix for all conducting elements
+ (active coils and passive structures).
+
+ The resulting matrix is used in the formulation of the linear circuit
+ evolution equations.
+
+ Returns
+ -------
+ None
+ The result is stored in `self.rm1l_non_symm`.
+ """
+
self.rm1l_non_symm = np.diag(self.coil_resist**-1.0) @ self.coil_self_ind
def make_selected_mode_mask(self, mode_coupling_masks, verbose):
- """Creates a mask for the vessel normal modes to include in the circuit
- equations, based on the maximum frequency of the selected modes.
+ """
+ Build selection mask for vessel normal modes used in circuit equations.
+
+ This method selects which passive structure eigenmodes are included in the
+ reduced-order circuit model based on a frequency cutoff and optional
+ plasma-coupling criteria.
+
+ Active coil variables are always included, while passive vessel modes are
+ filtered according to:
+ - maximum mode frequency (`self.max_mode_frequency`)
+ - optional coupling thresholds provided in `mode_coupling_masks`
+
+ Parameters
+ ----------
+ mode_coupling_masks : tuple of ndarray or None
+ Optional pair of boolean masks used to:
+ - reintroduce strongly coupled modes
+ - remove weakly coupled modes
+ verbose : bool
+ If True, print diagnostic information about mode selection.
+
+ Returns
+ -------
+ None
+ Updates internal attributes:
+ - `self.selected_modes_mask`
+ - `self.n_independent_vars`
"""
# this is for passives alone
selected_modes_mask = self.normal_modes.w_passive < self.max_mode_frequency
@@ -205,13 +264,42 @@ def make_selected_mode_mask(self, mode_coupling_masks, verbose):
def initialize_for_eig(
self, selected_modes_mask=None, mode_coupling_masks=None, verbose=True
):
- """Initializes the metal_currents object for the case where vessel
- eigenmodes are used.
+ """
+ Initialise the metal current system in eigenmode representation.
+
+ This method prepares the reduced-order circuit model when vessel
+ eigenmodes are included. It constructs the mode transformation matrices,
+ applies optional mode selection/reduction, and builds the system matrices
+ required for the implicit time integration solver.
+
+ The system is transformed between physical currents and eigenmodes via:
+ - P: transformation from eigenmodes to physical currents
+ - Pm1: inverse transformation (restricted to selected modes)
Parameters
----------
- mode_coupling_metric_masks : (np.ndarray of booles, np.ndarray of booles)
+ selected_modes_mask : ndarray of bool, optional
+ Explicit mask selecting which modes to retain. If None, the mask is
+ constructed using `mode_coupling_masks`. If False, all modes are used.
+ mode_coupling_masks : tuple of ndarray of bool, optional
+ Pair of masks used to:
+ - include strongly coupled modes
+ - exclude weakly coupled modes
+ Only used when `selected_modes_mask is None`.
+ verbose : bool
+ If True, print diagnostic information about mode reduction.
+
+ Returns
+ -------
+ None
+ Updates internal solver state including:
+ - mode selection masks
+ - transformation matrices (P, Pm1)
+ - system matrix (Lambdam1)
+ - time integrator solver
+ - forcing term selection
"""
+
if selected_modes_mask is None:
# this is the case when mode_coupling_masks are used to build self.selected_modes_mask
self.make_selected_mode_mask(mode_coupling_masks, verbose)
@@ -270,6 +358,24 @@ def initialize_for_eig(
self.forcing_term = self.forcing_term_eig_no_plasma
def reset_active_coil_resistances(self, active_coil_resistances):
+ """
+ Update the resistances of the active coils and rebuild derived system matrices.
+
+ This method replaces the active-coil portion of the full resistance vector
+ while keeping passive/vessel resistances unchanged. It then updates all
+ dependent operators used in the circuit evolution model.
+
+ The update triggers a rebuild of:
+ - the full resistance vector
+ - the inverse resistance matrix
+ - the resistance–inductance coupling operator
+ - the reduced system matrix used in vessel mode dynamics
+
+ Parameters
+ ----------
+ active_coil_resistances : ndarray
+ Updated resistances for the active coils only (length = n_active_coils).
+ """
self.coil_resist = np.concatenate(
(active_coil_resistances, self.coil_resist[self.n_active_coils :])
)
@@ -279,8 +385,24 @@ def reset_active_coil_resistances(self, active_coil_resistances):
self.Lambdam1 = self.Pm1 @ (self.rm1l_non_symm @ self.P)
def initialize_for_no_eig(self):
- """Initializes the metal currents object for the case where vessel
- eigenmodes are not used."""
+ """
+ Initialise the metal current system without eigenmode decomposition.
+
+ This method constructs and solves the full circuit equations in the
+ physical current basis, without projecting onto vessel eigenmodes.
+
+ The governing equation is:
+ Mmatrix · dI/dt + Rmatrix · I = F
+
+ where:
+ - Mmatrix is the full mutual inductance matrix
+ - Rmatrix is the diagonal resistance matrix
+
+ Returns
+ -------
+ None
+ Updates internal solver and forcing term configuration.
+ """
# Equation is Mmatrix Idot + Rmatrix I = F
self.solver = implicit_euler_solver(
@@ -296,88 +418,47 @@ def initialize_for_no_eig(self):
self.forcing_term = self.forcing_term_no_eig_no_plasma
def reset_timesteps(self, max_internal_timestep, full_timestep):
- """Resets the timesteps
+ """
+ Update solver time-stepping parameters.
+
+ This method resets both the internal solver timestep and the external
+ evolution timestep used to advance the circuit equations. If the full
+ timestep exceeds the internal timestep, multiple substeps are performed
+ automatically by the solver.
Parameters
----------
max_internal_timestep : float
- Maximum value of the 'internal' timestep for implicit euler solver.
- The 'internal' timestep is the one actually used by the solver.
+ Maximum timestep used internally by the implicit Euler solver.
full_timestep : float
- Timestep by which the equations are advanced. If full_timestep>max_internal_timestep
- multiple 'internal' steps are executed.
+ External timestep used to advance the system in time.
"""
self.solver.set_timesteps(
full_timestep=full_timestep, max_internal_timestep=max_internal_timestep
)
- # def reset_mode(
- # self,
- # flag_vessel_eig,
- # flag_plasma,
- # plasma_pts=None,
- # max_mode_frequency=1,
- # max_internal_timestep=0.0001,
- # full_timestep=0.0001,
- # ):
- # """Resets init inputs.
-
- # flag_vessel_eig : bool
- # Flag re whether vessel eigenmodes are used or not.
- # flag_plasma : bool
- # Whether to include plasma in circuit equation. If True, plasma_pts
- # must be provided.
- # plasma_pts : freegsnke.limiter_handler.plasma_pts
- # Domain points in the domain that are included in the evolutive calculations.
- # A typical choice would be all domain points inside the limiter. Defaults to None.
- # max_mode_frequency : float
- # Maximum frequency of vessel eigenmodes to include in circuit equation.
- # Defaults to 1. Unit is s^-1.
- # max_internal_timestep : float
- # Maximum value of the 'internal' timestep for implicit euler solver. Defaults to .0001.
- # The 'internal' timestep is the one actually used by the solver.
- # full_timestep : float
- # Timestep by which the equations are advanced. If full_timestep>max_internal_timestep
- # multiple 'internal' steps are executed. Defaults to .0001.
- # """
- # control = self.max_internal_timestep != max_internal_timestep
- # self.max_internal_timestep = max_internal_timestep
-
- # control += self.full_timestep != full_timestep
- # self.full_timestep = full_timestep
-
- # control += flag_plasma != self.flag_plasma
- # self.flag_plasma = flag_plasma
-
- # if control * flag_plasma:
- # self.plasma_pts = plasma_pts
- # self.Mey_matrix = self.Mey(eq)
-
- # control += flag_vessel_eig != self.flag_vessel_eig
- # self.flag_vessel_eig = flag_vessel_eig
-
- # if flag_vessel_eig:
- # control += max_mode_frequency != self.max_mode_frequency
- # self.max_mode_frequency = max_mode_frequency
- # if control * flag_vessel_eig:
- # self.initialize_for_eig(self.selected_modes_mask)
- # else:
- # self.initialize_for_no_eig()
-
def forcing_term_eig_plasma(self, active_voltage_vec, Iydot):
- """Right-hand-side of circuit equation in eigenmode basis with plasma.
+ """
+ Compute forcing term in eigenmode basis including plasma coupling.
+
+ This method constructs the effective right-hand side of the circuit
+ equations in eigenmode coordinates when plasma dynamics are included.
+
+ The forcing is built from:
+ - applied coil voltages
+ - inductive coupling to plasma current evolution
Parameters
----------
- active_voltage_vec : np.ndarray
- Vector of active coil voltages.
- Iydot : np.ndarray
- Vector of rate of change of plasma currents.
+ active_voltage_vec : ndarray
+ Voltages applied to the active coils.
+ Iydot : ndarray
+ Time derivative of plasma current degrees of freedom.
Returns
-------
- all_Us : np.ndarray
- Effective voltages in eigenmode basis.
+ all_Us : ndarray
+ Forcing term expressed in the eigenmode basis.
"""
all_Us = np.zeros_like(self.empty_U)
all_Us[: self.n_active_coils] = active_voltage_vec
@@ -386,20 +467,24 @@ def forcing_term_eig_plasma(self, active_voltage_vec, Iydot):
return all_Us
def forcing_term_eig_no_plasma(self, active_voltage_vec, Iydot=0):
- """Right-hand-side of circuit equation in eigenmode basis without
- plasma.
+ """
+ Compute forcing term in eigenmode basis without plasma coupling.
+
+ This method constructs the right-hand side of the circuit equations in
+ eigenmode coordinates when plasma effects are neglected. Only coil
+ voltages are included.
Parameters
----------
- active_voltage_vec : np.ndarray
- Vector of active coil voltages.
- Iydot : np.ndarray, optional
- This is not used.
+ active_voltage_vec : ndarray
+ Voltages applied to the active coils.
+ Iydot : ndarray, optional
+ Unused placeholder for interface compatibility.
Returns
-------
- all_Us : np.ndarray
- Effective voltages in eigenmode basis.
+ all_Us : ndarray
+ Forcing term expressed in the eigenmode basis.
"""
all_Us = self.empty_U.copy()
all_Us[: self.n_active_coils] = active_voltage_vec
@@ -407,19 +492,24 @@ def forcing_term_eig_no_plasma(self, active_voltage_vec, Iydot=0):
return all_Us
def forcing_term_no_eig_plasma(self, active_voltage_vec, Iydot):
- """Right-hand-side of circuit equation in normal mode basis with plasma.
+ """
+ Compute forcing term in coil basis including plasma coupling.
+
+ This method builds the right-hand side of the circuit equations in the
+ physical (non-eigenmode) coil basis, including inductive coupling to the
+ plasma evolution.
Parameters
----------
- active_voltage_vec : np.ndarray
- Vector of active coil voltages.
- Iydot : np.ndarray
- Vector of rate of change of plasma currents.
+ active_voltage_vec : ndarray
+ Voltages applied to the active coils.
+ Iydot : ndarray
+ Time derivative of plasma current degrees of freedom.
Returns
-------
- all_Us : np.ndarray
- Effective voltages in metals basis.
+ all_Us : ndarray
+ Forcing term in the physical coil basis.
"""
all_Us = self.empty_U.copy()
all_Us[: self.n_active_coils] = active_voltage_vec
@@ -427,20 +517,24 @@ def forcing_term_no_eig_plasma(self, active_voltage_vec, Iydot):
return all_Us
def forcing_term_no_eig_no_plasma(self, active_voltage_vec, Iydot=0):
- """Right-hand-side of circuit equation in normal mode basis without
- plasma.
+ """
+ Compute forcing term in coil basis without plasma coupling.
+
+ This method constructs the right-hand side of the circuit equations in the
+ physical (non-eigenmode) coil basis when plasma effects are neglected.
+ Only applied coil voltages are included.
Parameters
----------
- active_voltage_vec : np.ndarray
- Vector of active coil voltages.
- Iydot : np.ndarray, optional
- This is not used.
+ active_voltage_vec : ndarray
+ Voltages applied to the active coils.
+ Iydot : ndarray, optional
+ Unused placeholder for interface consistency.
Returns
-------
- all_Us : np.ndarray
- Effective voltages in metals basis.
+ all_Us : ndarray
+ Forcing term in the physical coil basis.
"""
all_Us = self.empty_U.copy()
all_Us[: self.n_active_coils] = active_voltage_vec
@@ -534,37 +628,3 @@ def Mey(
greenm *= coils_dict[labelj]["multiplier"][np.newaxis, :]
mey[j] = np.sum(greenm, axis=-1)
return 2 * np.pi * mey
-
- # def Mey_f(
- # self,
- # eq,
- # green_f
- # ):
- # """Calculates values of the function green_f for the matrix of
- # plasma_pts x all vessel coils. For clarity, the function Mey is Mey_f(green_f = Greens)
-
- # Parameters
- # -------
- # eq : class
- # FreeGSNKE equilibrium Object
- # green_f : function
- # with same structure as Greens, i.e. Greens(R1,Z1, R2,Z2)
-
- # Returns
- # -------
- # Mey : np.ndarray
- # Array of 'inductance values' between plasma grid points and all vessel coils
- # """
- # coils_dict = eq.tokamak.coils_dict
- # mey = np.zeros((eq.tokamak.n_coils, len(self.plasma_pts)))
- # for j, labelj in enumerate(eq.tokamak.coils_list):
- # greenm = green_f(
- # coils_dict[labelj]["coords"][0][np.newaxis, :],
- # coils_dict[labelj]["coords"][1][np.newaxis, :],
- # self.plasma_pts[:, 0, np.newaxis],
- # self.plasma_pts[:, 1, np.newaxis],
- # )
- # greenm *= coils_dict[labelj]["polarity"][np.newaxis, :]
- # greenm *= coils_dict[labelj]["multiplier"][np.newaxis, :]
- # mey[j] = np.sum(greenm, axis=-1)
- # return 2 * np.pi * mey
diff --git a/freegsnke/circuit_eq_plasma.py b/freegsnke/circuit_eq_plasma.py
index 172cfb9..0b367b3 100644
--- a/freegsnke/circuit_eq_plasma.py
+++ b/freegsnke/circuit_eq_plasma.py
@@ -1,6 +1,6 @@
"""
-Defines the plasma_current Object, which handles the lumped parameter model
-used as an effective circuit equation for the plasma.
+Defines a few properties of the plasma current equation (i.e. the lumped parameter model
+used as an effective circuit equation for the plasma).
Copyright 2025 UKAEA, UKRI-STFC, and The Authors, as per the COPYRIGHT and README files.
@@ -23,69 +23,37 @@
import numpy as np
from freegs4e.gradshafranov import Greens
-# class plasma_current:
-# """Implements the plasma circuit equation in projection on $I_{y}^T$:
-
-# $$I_{y}^T/I_p (M_{yy} \dot{I_y} + M_{ye} \dot{I_e} + R_p I_y) = 0$$
-# """
-
-# def __init__(self, plasma_pts, Rm1, P, plasma_resistance_1d, Mye):
-# """Implements the object dealing with the plasma circuit equation in projection on $I_y$,
-# I_y being the plasma toroidal current density distribution:
-
-# $$I_{y}^T/I_p (M_{yy} \dot{I_y} + M_{ye} \dot{I_e} + R_p I_y) = 0$$
-
-# Parameters
-# ----------
-# plasma_pts : freegsnke.limiter_handler.plasma_pts
-# Domain points in the domain that are included in the evolutive calculations.
-# A typical choice would be all domain points inside the limiter. Defaults to None.
-# Rm1 : np.ndarray
-# The diagonal matrix of all metal vessel resistances to the power of -1 ($R^{-1}$).
-# P : np.ndarray
-# Matrix used to change basis from normal mode currents to vessel metal currents.
-# plasma_resistance_1d : np.ndarray
-# Vector of plasma resistance values for all grid points in the reduced plasma domain.
-# plasma_resistance_1d = 2pi resistivity R/dA for all plasma_pts
-# Mye : np.ndarray
-# Matrix of mutual inductances between plasma grid points and all vessel coils.
-
-# """
-
-# self.plasma_pts = plasma_pts
-# self.Rm1 = Rm1
-# self.P = P
-# self.Mye = Mye
-# self.Ryy = plasma_resistance_1d
-# self.Myy_matrix = self.Myy()
-
-# def reset_modes(self, P):
-# """Allows a reset of the attributes set up at initialization time following a change
-# in the properties of the selected normal modes for the passive structures.
-
-# Parameters
-# ----------
-# P : np.ndarray
-# New change of basis matrix.
-# """
-# self.P = P
-
-
-def Myy(
- plasma_pts,
-):
- """Calculates the matrix of mutual inductances between all plasma grid points
+
+def Myy(plasma_pts):
+ """
+ Compute the mutual inductance matrix between plasma grid points.
+
+ The matrix is constructed using the Green's function evaluated between
+ all pairs of plasma points in cylindrical (R, Z) coordinates.
Parameters
----------
- plasma_pts : np.ndarray
- Array with R and Z coordinates of all the points inside the limiter
+ plasma_pts : ndarray, shape (N, 2)
+ Array of plasma grid point coordinates in cylindrical geometry:
+ - plasma_pts[:, 0] = R coordinates
+ - plasma_pts[:, 1] = Z coordinates
Returns
-------
- Myy : np.ndarray
- Array of mutual inductances between plasma grid points
+ Myy : ndarray, shape (N, N)
+ Mutual inductance matrix between all plasma grid points.
+ The matrix is symmetric.
+
+ Notes
+ -----
+ The matrix is computed as:
+
+ Myy_ij = 2π * G(R_i, Z_i, R_j, Z_j)
+
+ where G is the Green's function returned by `Greens`.
+
"""
+
greenm = Greens(
plasma_pts[:, np.newaxis, 0],
plasma_pts[:, np.newaxis, 1],
@@ -96,6 +64,43 @@ def Myy(
def grid_greens(R, Z):
+ """
+ Compute the Green's function matrix for a structured (R, Z) grid.
+
+ This function evaluates the Green's function between a set of radial
+ grid points and an axially discretised Z-grid, using broadcasting to
+ construct all pairwise interactions efficiently.
+
+ Parameters
+ ----------
+ R : ndarray, shape (N, M)
+ Radial grid values. Only the first column is used (R[:, 0]).
+ Z : ndarray, shape (N, M)
+ Axial grid values defining a uniform Z spacing.
+
+ Returns
+ -------
+ ggreens : ndarray
+ Green's function evaluated on the expanded grid, scaled by 2π.
+ Shape depends on broadcasting of the underlying `Greens` function.
+
+ Notes
+ -----
+ - The Z-grid is assumed to be uniformly spaced in the second dimension.
+ - The spacing is computed as:
+ dz = Z[0, 1] - Z[0, 0]
+ - The number of Z points is inferred as:
+ nZ = Z.shape[1]
+ - The Green's function is evaluated as:
+ G(R_i, z_k, R_j, z=0)
+ where the axial coordinate for the second argument is constructed as:
+ z_k = dz * arange(nZ)
+
+ Warning
+ -------
+ - Only `R[:, 0]` is used; any variation in R across the second dimension is ignored.
+ - Assumes uniform spacing in Z; non-uniform grids will produce incorrect results.
+ """
dz = Z[0, 1] - Z[0, 0]
nZ = np.shape(Z)[1]
diff --git a/freegsnke/copying.py b/freegsnke/copying.py
index 07d84c1..1e3e2fd 100644
--- a/freegsnke/copying.py
+++ b/freegsnke/copying.py
@@ -1,3 +1,24 @@
+"""
+Useful functions for copying attributed between FreeGSNKE objects.
+
+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 .
+"""
+
import copy
import logging
@@ -9,28 +30,45 @@
def copy_into(
obj, new_obj, attr: str, *, mutable=False, strict=True, allow_deepcopy=False
):
- """Copies an attribute from one object to another.
+ """
+ Copy an attribute from one object into another.
+
+ This function transfers a named attribute from `obj` to `new_obj`,
+ optionally performing a deep copy for mutable objects.
Parameters
- ==========
+ ----------
obj : object
- The object to copy the attribute from
+ Source object to copy the attribute from.
new_obj : object
- The object to copy the attribute to
+ Destination object to copy the attribute into.
attr : str
- The attribute name to copy (e.g. copy `obj.attr` into `new_obj`)
-
- mutable : bool
- If an attribute is mutable it needs to be copied explicitly between objects,
- a reference is not sufficient.
-
- strict : bool
- Raise an error if `True` and `attr` is not an attribute of `obj`
-
- allow_deepcopy : bool
- Raise an error if `True` and a deepcopy is required to copy the attribute across.
+ Name of the attribute to copy (equivalent to obj.attr).
+
+ mutable : bool, optional
+ If True, treat the attribute as mutable and ensure it is copied
+ rather than referenced. For NumPy arrays, a shallow copy is used
+ when possible.
+
+ strict : bool, optional
+ If True, raise an AttributeError when `obj` does not have `attr`.
+ If False, silently return without copying.
+
+ allow_deepcopy : bool, optional
+ If True, allow deepcopying of non-NumPy mutable objects when required.
+ If False, raise a TypeError when deepcopy would be needed.
+
+ Raises
+ ------
+ TypeError
+ If `mutable=True`, the attribute is not a safe NumPy array copy,
+ and `allow_deepcopy=False`.
+
+ Returns
+ -------
+ None
"""
if not hasattr(obj, attr) and not strict:
diff --git a/freegsnke/inverse.py b/freegsnke/inverse.py
index 2866920..7c77c5e 100644
--- a/freegsnke/inverse.py
+++ b/freegsnke/inverse.py
@@ -254,7 +254,32 @@ def __init__(
@staticmethod
def _extract_isoflux_constraints_weights(isoflux_set: np.ndarray):
+ """
+ Extract isoflux constraint locations and associated weights.
+
+ The input array encodes a set of isoflux constraints in one of two formats:
+ - (3, N): first two rows are constraint coordinates, third row contains weights
+ - (2, N): constraint coordinates only, with unit weights assumed
+
+ Parameters
+ ----------
+ isoflux_set : ndarray
+ Array of isoflux constraints with shape (2, N) or (3, N), where N is
+ the number of constraint points.
+ Returns
+ -------
+ constraints : ndarray
+ Constraint coordinates with shape (2, N).
+ weights : ndarray
+ Weights associated with each constraint point, shape (N,).
+ Defaults to an array of ones if not provided in input.
+
+ Raises
+ ------
+ ValueError
+ If input does not have shape (2, N) or (3, N).
+ """
if isoflux_set.shape[0] == 3:
return isoflux_set[0:2, :], isoflux_set[2, :]
elif isoflux_set.shape[0] == 2:
diff --git a/freegsnke/jtor_refinement.py b/freegsnke/jtor_refinement.py
index 33120d7..dc8d732 100644
--- a/freegsnke/jtor_refinement.py
+++ b/freegsnke/jtor_refinement.py
@@ -1,25 +1,62 @@
+"""
+Contains various functions for refining the plasma current density map.
+
+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 .
+"""
+
import numpy as np
from .copying import copy_into
class Jtor_refiner:
- """Class to allow for the refinement of the toroidal plasma current Jtor.
- Currently applied to the Lao85 profile family when 'refine_flag=True'.
- Only grid cells that are crossed by the separatrix are refined.
+ """
+ Refines the toroidal plasma current (Jtor) on a structured grid.
+
+ This class implements a local sub-grid refinement strategy for Jtor,
+ currently used primarily for Lao85-type profiles when refinement is enabled.
+
+ Refinement is applied selectively to grid cells, typically those:
+ - intersected by the separatrix (LCFS),
+ - exhibiting large Jtor magnitude, or
+ - showing strong spatial gradients.
+
+ The goal is to improve resolution of sharp features without globally
+ increasing grid resolution.
"""
def __init__(self, eq, nnx, nny):
- """Instantiates the object and prepares necessary quantities.
+ """
+ Initialise the Jtor refiner and precompute geometric and indexing data.
Parameters
----------
- eq : freegs4e Equilibrium object
- Specifies the domain properties
- nnx : even integer
- refinement factor in the R direction
- nny : even integer
- refinement factor in the Z direction
+ eq : freegs4e.Equilibrium
+ Equilibrium object defining the computational grid and limiter geometry.
+ nnx : int (even)
+ Refinement factor in the R-direction (number of subcells per cell).
+ nny : int (even)
+ Refinement factor in the Z-direction (number of subcells per cell).
+
+ Notes
+ -----
+ The grid is assumed to be uniform in R and Z, and refinement is performed
+ by subdividing each selected cell into an nnx × nny subgrid.
"""
self.eqR = eq.R
@@ -46,6 +83,23 @@ def __init__(self, eq, nnx, nny):
self.edges_mask[:, -1] = 0
def copy(self):
+ """
+ Create a deep copy of the Jtor_refiner object.
+
+ Returns
+ -------
+ Jtor_refiner
+ A new instance with the same configuration and precomputed refinement
+ structures.
+
+ Notes
+ -----
+ - Array-like attributes such as grid geometry and masks are copied using
+ `copy_into`.
+ - Derived refinement structures are recomputed via `prepare_for_refinement()`.
+ - Optional attributes (e.g. LCFS and refinement masks) are copied if present,
+ otherwise they are ignored (`strict=False`).
+ """
obj = type(self).__new__(type(self))
copy_into(self, obj, "eqR", mutable=True)
@@ -76,7 +130,26 @@ def copy(self):
def prepare_for_refinement(
self,
):
- """Prepares necessary quantities to operate refinement."""
+ """
+ Precompute geometric, interpolation, and masking structures used in Jtor refinement.
+
+ This method builds all static quantities required for sub-grid refinement,
+ including:
+ - coarse-grid index maps (Ridx, Zidx),
+ - sub-cell coordinate systems,
+ - bilinear interpolation weights,
+ - limiter-based inside/outside masks at refined resolution,
+ - quadrant decomposition masks for vectorised operations.
+
+ These arrays are reused across refinement calls to avoid recomputation.
+
+ Notes
+ -----
+ - Assumes a uniform structured grid in (R, Z).
+ - Sub-cell structure is defined by (nnx × nny) refinement per cell.
+ - Limiter geometry is accessed via `self.path.contains_points`.
+ """
+
self.Ridx = np.tile(np.arange(self.nx), (self.ny, 1)).T
self.Zidx = np.tile(np.arange(self.ny), (self.nx, 1))
@@ -140,19 +213,30 @@ def prepare_for_refinement(
self.quartermasks = quartermasks
def get_indexes_for_refinement(self, mask_to_refine):
- """Generates the indexes of psi values to be used for bilinear interpolation.
+ """
+ Construct index arrays for bilinear interpolation on refined cells.
+
+ For each selected coarse grid cell, this function returns the indices of
+ the 2×2 stencil (four surrounding vertices) required to perform bilinear
+ interpolation of ψ on the refined subgrid.
Parameters
----------
- mask_to_refine : np.array
- Mask of all domain cells to be refined
+ mask_to_refine : np.ndarray (bool)
+ Boolean mask of coarse grid cells selected for refinement.
Returns
-------
- np.array
- indexes of psi values to be used for bilinear interpolation
- 4 points per cell to refine, already set in 2-by-2 matrix for vectorised interpolation
- dimensions = (no of cells to refine, 2, 2)
+ RRidxs : np.ndarray
+ R-index stencil for each refined cell, shape
+ (n_cells, 2, 2, 2) depending on internal packing.
+ ZZidxs : np.ndarray
+ Z-index stencil matching RRidxs structure.
+
+ Notes
+ -----
+ Each cell contributes a structured 2×2 vertex stencil used for vectorised
+ bilinear interpolation of ψ on the nnx × nny subgrid.
"""
RRidxs = np.concatenate(
(
@@ -321,16 +405,27 @@ def get_indexes_for_refinement(self, mask_to_refine):
return RRidxs, ZZidxs
def build_jtor_value_mask(self, unrefined_jtor, threshold, quantiles=(0.5, 0.9)):
- """Selects the cells that need to be refined based on their value of jtor.
- Selection is such that it includes cells where jtor exceeds the value calculated
- based on the quantiles and threshold[0].
+ """
+ Construct a refinement mask based on the magnitude of Jtor.
+
+ Cells are selected for refinement if their Jtor value exceeds a scaled
+ threshold defined relative to two quantiles of the Jtor distribution.
Parameters
----------
- unrefined_jtor : np.array
- The jtor distribution
- thresholds : float
- the relevant value (in the tuple) used to identify where to apply refinement
+ unrefined_jtor : np.ndarray
+ Coarse-grid toroidal current density field.
+ threshold : float
+ Scaling factor applied to the inter-quantile range to determine
+ refinement sensitivity.
+ quantiles : tuple of float, optional
+ Two quantiles (q_low, q_high) used to define a reference range.
+ Default is (0.5, 0.9).
+
+ Returns
+ -------
+ np.ndarray (bool)
+ Boolean mask indicating cells selected for refinement.
"""
jtor_quantiles = np.quantile(unrefined_jtor.reshape(-1), quantiles)
@@ -340,16 +435,32 @@ def build_jtor_value_mask(self, unrefined_jtor, threshold, quantiles=(0.5, 0.9))
return mask
def build_jtor_gradient_mask(self, unrefined_jtor, threshold, quantiles=(0.5, 0.9)):
- """Selects the cells that need to be refined based on their value of the gradient of jtor.
- Selection is such that it includes cells where the norm of the gradient exceeds the value calculated
- based on the quantiles and threshold[1].
+ """
+ Construct a refinement mask based on local finite-difference variations of Jtor.
+
+ Cells are selected when the magnitude of local Jtor differences between
+ neighbouring grid points exceeds a threshold derived from the distribution
+ of these differences.
+
+ The method approximates spatial variation using forward finite differences
+ in both R and Z directions, and applies a quantile-based thresholding
+ strategy (via `build_jtor_value_mask`) to identify large-gradient regions.
Parameters
----------
- unrefined_jtor : np.array
- The jtor distribution
- thresholds : float
- the relevant value (in the tuple) used to identify where to apply refinement
+ unrefined_jtor : np.ndarray
+ Coarse-grid toroidal current density field.
+ threshold : float
+ Scaling factor applied to the inter-quantile range of gradient
+ magnitudes used for refinement.
+ quantiles : tuple of float, optional
+ (q_low, q_high) quantiles used to normalise gradient magnitude
+ thresholds. Default is (0.5, 0.9).
+
+ Returns
+ -------
+ np.ndarray (bool)
+ Boolean mask indicating cells selected for refinement.
"""
gradient_mask = np.zeros_like(unrefined_jtor)
@@ -372,13 +483,32 @@ def build_jtor_gradient_mask(self, unrefined_jtor, threshold, quantiles=(0.5, 0.
return gradient_mask > 0
def build_LCFS_mask(self, core_mask):
- """Builds a mask composed of all gridpoints connected to edges crossed by the LCFS.
- These trigger refinement.
+ """
+ Construct a refinement mask identifying grid points adjacent to the LCFS.
+
+ The mask is built by detecting interfaces between plasma and vacuum
+ cells in the binary `core_mask`. Any grid cell that shares an edge
+ with a cell of opposite classification (inside/outside plasma core)
+ is marked for refinement.
+
+ The resulting mask is additive: cells adjacent to multiple LCFS edges
+ may accumulate values greater than 1. This is primarily used as a
+ selection indicator, not a strict boolean mask.
Parameters
----------
- core_mask : np.array
- Plasma core mask on the standard domain (self.nx, self.ny)
+ core_mask : np.ndarray of shape (nx, ny)
+ Binary plasma core mask on the structured grid.
+ Expected values:
+ - 1 (or True): inside plasma core
+ - 0 (or False): outside plasma core
+
+ Returns
+ -------
+ np.ndarray of shape (nx, ny)
+ LCFS refinement indicator array. Entries are non-zero where grid
+ cells are adjacent to a change in `core_mask` across a grid edge.
+ Higher values indicate multiple adjacent LCFS crossings.
"""
core_mask = core_mask.astype(float)
@@ -402,16 +532,44 @@ def build_LCFS_mask(self, core_mask):
return lcfs_mask
def build_mask_to_refine(self, unrefined_jtor, core_mask, thresholds):
- """Selects the cells that need to be refined, using the user-defined thresholds
+ """
+ Construct the global refinement mask combining LCFS location,
+ Jtor magnitude, and Jtor gradient criteria.
+
+ This method aggregates multiple refinement indicators into a single
+ cell-wise mask. A cell is marked for refinement if it satisfies any
+ of the following conditions:
+
+ 1. It lies adjacent to the LCFS (core–vacuum interface)
+ 2. Its Jtor value exceeds a threshold based on distribution quantiles
+ 3. Its Jtor gradient exceeds a threshold based on distribution quantiles
+
+ Boundary cells are excluded from refinement.
+
+ The intermediate masks are also stored as attributes for diagnostics:
+ - self.lcfs_mask
+ - self.value_mask
+ - self.gradient_mask
Parameters
----------
- unrefined_jtor : np.array
- The jtor distribution
- core_mask : np.array
- Plasma core mask on the standard domain (self.nx, self.ny)
- thresholds : tuple (threshold for jtor criterion, threshold for gradient criterion)
- tuple of values used to identify where to apply refinement, by default None
+ unrefined_jtor : np.ndarray of shape (nx, ny)
+ Toroidal current density on the coarse grid.
+
+ core_mask : np.ndarray of shape (nx, ny)
+ Binary mask identifying plasma core cells.
+ Typically 1 inside plasma, 0 outside.
+
+ thresholds : tuple of float
+ (jtor_threshold, gradient_threshold)
+ Scaling factors applied to inter-quantile ranges used to define
+ refinement sensitivity.
+
+ Returns
+ -------
+ None
+ The result is stored in:
+ self.mask_to_refine : np.ndarray (bool)
"""
mask_to_refine = np.zeros_like(unrefined_jtor)
@@ -437,20 +595,49 @@ def build_mask_to_refine(self, unrefined_jtor, core_mask, thresholds):
self.mask_to_refine = mask_to_refine.astype(bool)
def build_bilinear_psi_interp(self, psi, core_mask, unrefined_jtor, thresholds):
- """Builds the mask of cells on which to operate refinement.
- Cells that are crossed by the separatrix and cells with large gradient on jtor are considered.
- Refines psi in the same cells.
+ """
+ Construct a refined representation of the poloidal flux `psi` on a
+ sub-grid using bilinear interpolation in selected refinement cells.
+
+ Cells are selected for refinement based on a combined criterion:
+ - proximity to the LCFS (core–vacuum interface)
+ - large values of Jtor (based on quantile thresholding)
+ - large gradients in Jtor (based on quantile thresholding)
+
+ For each selected coarse-grid cell, a higher-resolution sub-grid is
+ generated and psi is reconstructed via bilinear interpolation.
Parameters
----------
- psi : np.array
- Psi on the standard domain (self.nx, self.ny)
- core_mask : np.array
- Plasma core mask on the standard domain (self.nx, self.ny)
- unrefined_jtor : np.array
- The jtor distribution
- thresholds : tuple (threshold for jtor criterion, threshold for gradient criterion)
- tuple of values used to identify where to apply refinement, by default None
+ psi : np.ndarray of shape (nx, ny)
+ Poloidal flux on the coarse grid.
+
+ core_mask : np.ndarray of shape (nx, ny)
+ Binary plasma core mask defining LCFS location.
+
+ unrefined_jtor : np.ndarray of shape (nx, ny)
+ Toroidal current density on the coarse grid.
+
+ thresholds : tuple of float
+ (jtor_threshold, gradient_threshold)
+ Scaling factors used in refinement criteria via inter-quantile ranges.
+
+ Returns
+ -------
+ format_bilinear_psi : np.ndarray of shape (n_refined, nnx, nny)
+ Bilinearly interpolated psi values on refined sub-grids for each
+ selected coarse cell.
+
+ refined_R : np.ndarray of shape (n_refined, nnx, nny)
+ R-coordinate values corresponding to each refined sub-grid point.
+
+ Notes
+ -----
+ - The refinement mask is computed internally via `build_mask_to_refine`.
+ - Each selected coarse cell is subdivided into an `nnx × nny` sub-grid.
+ - Bilinear interpolation is performed using precomputed vertex weights
+ (`self.xxxx`, `self.yyyy`).
+ - Output is structured per refined cell, not a full fine global grid.
"""
self.build_mask_to_refine(unrefined_jtor, core_mask, thresholds)
@@ -491,19 +678,36 @@ def build_bilinear_psi_interp(self, psi, core_mask, unrefined_jtor, thresholds):
return format_bilinear_psi, refined_R
def build_from_refined_jtor(self, unrefined_jtor, refined_jtor):
- """Averages the refined maps to the (nx, ny) domain grid.
+ """
+ Reconstruct a coarse-grid Jtor field by averaging refined sub-grid values
+ back onto the original (nx, ny) mesh.
+
+ For each selected coarse cell, the corresponding refined sub-grid
+ (nnx × nny) is first masked to remove points outside the limiter.
+ The remaining values are then spatially averaged and used to replace
+ the coarse-grid value.
Parameters
----------
- unrefined_jtor : np.array
- (nx, ny) jtor map from unresolved method
- refined_jtor : np.array
- maps of the refined jtor, dimension = (no cells to refine, nnx, nny)
+ unrefined_jtor : np.ndarray of shape (nx, ny)
+ Original coarse-grid toroidal current density.
+ refined_jtor : np.ndarray of shape (n_refined, nnx, nny)
+ Refined Jtor values on sub-grids for each selected coarse cell.
Returns
-------
- Refined jtor on the (nx, ny) domain grid
+ np.ndarray of shape (nx, ny)
+ Updated Jtor field where selected cells have been replaced by
+ averaged refined values and all other cells remain unchanged.
+
+ Notes
+ -----
+ - Refinement contributions are masked using `self.full_masks` to
+ exclude points outside the limiter.
+ - Each coarse cell is updated independently; no smoothing is applied
+ between neighboring refined regions.
+ - The output preserves the original grid structure.
"""
# mask out refinement points that are outside the limiter
masked_refined_jtor = (
diff --git a/freegsnke/jtor_update.py b/freegsnke/jtor_update.py
index 8ec0234..1daae7a 100644
--- a/freegsnke/jtor_update.py
+++ b/freegsnke/jtor_update.py
@@ -1,5 +1,5 @@
"""
-Defines the FreeGSNKE profile Object, which inherits from the FreeGS4E profile object.
+Defines the FreeGSNKE profile object, which inherits from the FreeGS4E profile object.
Copyright 2025 UKAEA, UKRI-STFC, and The Authors, as per the COPYRIGHT and README files.
@@ -32,22 +32,86 @@
class Jtor_universal:
+ """
+ Wrapper class providing a unified interface for toroidal current density (Jtor)
+ evaluation, with optional refinement.
+
+ This class selects between two implementations of the toroidal current density
+ model depending on whether refinement is enabled:
+
+ - Unrefined Jtor: fast, standard evaluation
+ - Refined Jtor: higher-resolution or corrected evaluation using additional
+ numerical processing
+
+ The interface ensures that downstream code can call `Jtor()` without needing
+ to know which implementation is being used.
+ """
+
def __init__(self, refine_jtor=False):
"""Sets default unrefined Jtor."""
self._refine_jtor = refine_jtor
def Jtor(self, *args, **kwargs):
+ """
+ Evaluate toroidal current density (Jtor), dispatching to either the
+ refined or unrefined implementation.
+
+ This method acts as a unified interface:
+ - If `_refine_jtor` is True, it calls `Jtor_refined`
+ - Otherwise, it calls `Jtor_unrefined`
+
+ Parameters
+ ----------
+ *args, **kwargs
+ Arguments passed directly to the selected Jtor implementation.
+
+ Returns
+ -------
+ ndarray
+ Toroidal current density evaluated on the plasma grid.
+ """
if self._refine_jtor:
return self.Jtor_refined(*args, **kwargs)
else:
return self.Jtor_unrefined(*args, **kwargs)
def copy(self, obj=None):
- """Creates a copy the object.
+ """
+ Create a copy of the current Jtor_universal instance.
- obj : Jtor_universal
- An instance of self that the attributes are copied into instead of
- creating a new object
+ This method performs a selective copy of internal attributes,
+ combining shallow copies, deep copies, and shared references
+ depending on the nature of each attribute.
+
+ Parameters
+ ----------
+ obj : Jtor_universal, optional
+ Existing instance to copy attributes into. If None, a new
+ instance of the same class is created.
+
+ Returns
+ -------
+ Jtor_universal
+ A copied instance of the current object.
+
+ Notes
+ -----
+ Copy semantics:
+ - Immutable / scalar attributes are copied directly
+ - NumPy arrays and simple containers are shallow-copied where safe
+ - Selected complex objects are deep-copied (e.g. via copy.deepcopy)
+ - Some large system objects are shared by reference (e.g. limiter_handler)
+
+ Optional attributes:
+ The method safely ignores missing attributes using strict=False
+ for certain fields.
+
+ Warning
+ -------
+ This is not a full deep copy. Some internal objects are shared
+ between the original and the copy, particularly:
+ - limiter_handler
+ - any attributes not explicitly copied via `copy_into`
"""
obj = type(self).__new__(type(self)) if obj is None else obj
@@ -110,12 +174,18 @@ def copy(self, obj=None):
return obj
def set_masks(self, eq):
- """Universal function to set all masks related to the limiter.
+ """
+ Initialise grid geometry and limiter-related masks from an equilibrium object.
+
+ This method constructs and stores all spatial grid metadata, index mappings,
+ and limiter masks required for subsequent calculations. It modifies the
+ object in-place.
Parameters
----------
eq : FreeGSNKE Equilibrium object
- Specifies the domain properties
+ Equilibrium defining the computational domain, including 1D and 2D
+ coordinate grids and limiter geometry.
"""
self.dR = eq.R_1D[1] - eq.R_1D[0]
self.dZ = eq.Z_1D[1] - eq.Z_1D[0]
@@ -158,18 +228,28 @@ def set_masks(self, eq):
) = 1
def select_refinement(self, eq, refine_jtor, nnx, nny):
- """Initializes the object that handles the subgrid refinement of jtor
+ """
+ Initialise optional subgrid refinement for toroidal current density (jtor).
+
+ This method enables and configures subgrid refinement of the plasma
+ current density if requested, constructing the refinement handler and
+ associated parameters.
Parameters
----------
- eq : freegs4e Equilibrium object
- Specifies the domain properties
+ eq : FreeGS4E Equilibrium object
+ Equilibrium object defining the computational domain and geometry.
refine_jtor : bool
- Flag to select whether to apply sug-grid refinement of plasma current distribution jtor
- nnx : even integer
- refinement factor in the R direction
- nny : even integer
- refinement factor in the Z direction
+ If True, enable subgrid refinement of the toroidal current density.
+ If False, refinement is disabled.
+ nnx : int (even)
+ Refinement factor in the R-direction.
+ nny : int (even)
+ Refinement factor in the Z-direction.
+
+ Returns
+ -------
+ None
"""
self._refine_jtor = refine_jtor
if refine_jtor:
@@ -177,30 +257,24 @@ def select_refinement(self, eq, refine_jtor, nnx, nny):
self.set_refinement_thresholds()
def set_refinement_thresholds(self, thresholds=(1.0, 1.0)):
- """Sets the default criteria for refinement -- used when not directly set.
+ """
+ Set the criteria used to control jtor subgrid refinement.
+
+ These thresholds determine where refinement is applied based on the
+ current density and its gradient.
Parameters
----------
- thresholds : tuple (threshold for jtor criterion, threshold for gradient criterion)
- tuple of values used to identify where to apply refinement
+ thresholds : tuple of float, optional
+ (jtor_threshold, gradient_threshold), where each value controls
+ the activation of refinement criteria.
+
+ Returns
+ -------
+ None
"""
self.refinement_thresholds = thresholds
- # def all_open(self, contours):
- # checks = []
- # for contour in contours:
- # checks.append(
- # np.any(
- # [
- # np.any(contour[:, 0] <= 1),
- # np.any(contour[:, 0] >= self.nx - 2),
- # np.any(contour[:, 1] <= 1),
- # np.any(contour[:, 1] >= self.ny - 2),
- # ]
- # )
- # )
- # return np.all(checks), checks
-
def diverted_critical(
self,
R,
@@ -212,37 +286,38 @@ def diverted_critical(
starting_dx=0.05,
):
"""
- Replaces Jtor_part1 when that fails. Implements a new algorithm to define the LCFS.
- This is considerably more time consuming, but essential when the default routines in
- critical fail, as for example when the Xpt is not correctly identified.
+ Compute LCFS, O-point, X-point, and core mask using a contour-based fallback algorithm.
+ This method replaces the standard LCFS/X-point detection routine when it fails,
+ providing a more robust (but more expensive) contour-tracking approach.
Parameters
----------
- R : np.ndarray
- Radial coordinates of the grid points.
- Z : np.ndarray
- Vertical coordinates of the grid points.
- psi : np.ndarray
- Total poloidal field flux at each grid point [Webers/2pi].
+ R : ndarray
+ Radial grid coordinates.
+ Z : ndarray
+ Vertical grid coordinates.
+ psi : ndarray
+ Poloidal flux on the computational grid.
psi_bndry : float, optional
- Value of the poloidal field flux at the boundary of the plasma (last closed
- flux surface).
- mask_outside_limiter : np.ndarray
- Mask of points outside the limiter, if any.
+ Prescribed boundary flux value. If None, it is computed internally.
+ mask_outside_limiter : ndarray, optional
+ Boolean mask identifying points outside the limiter.
+ rel_tolerance_xpt : float, optional
+ Relative tolerance controlling convergence of X-point search.
+ starting_dx : float, optional
+ Initial step size in normalized flux space for contour search.
Returns
-------
- np.array
- Each row represents an O-point of the form [R, Z, ψ(R,Z)] [m, m, Webers/2pi].
- np.array
- Each row represents an X-point of the form [R, Z, ψ(R,Z)] [m, m, Webers/2pi].
- np.bool
- An array, the same shape as the computational grid, indicating the locations
- at which the core plasma resides (True) and where it does not (False).
- float
- Value of the poloidal field flux at the boundary of the plasma (last closed
- flux surface).
+ opt : ndarray, shape (1, 3)
+ O-point coordinates and flux value [R, Z, psi].
+ xpt : ndarray, shape (1, 3)
+ X-point coordinates and flux value [R, Z, psi].
+ diverted_core_mask : ndarray, shape (nx, ny)
+ Boolean mask identifying plasma core region.
+ psi_bndry : float
+ Flux value at the last closed flux surface (LCFS).
"""
# prepare psi_map to use
@@ -346,6 +421,43 @@ def diverted_critical_complete(
rel_tolerance_xpt=1e-4,
starting_dx=0.05,
):
+ """
+ Robust LCFS, O-point, X-point, and core mask detection with fallback logic.
+
+ This method attempts to compute plasma boundary information using the
+ primary routine `Jtor_part1`. If this fails (raises an exception), it
+ falls back to a more robust but computationally expensive contour-based
+ method implemented in `diverted_critical`.
+
+ Parameters
+ ----------
+ R : ndarray
+ Radial grid coordinates.
+ Z : ndarray
+ Vertical grid coordinates.
+ psi : ndarray
+ Poloidal flux on the computational grid.
+ psi_bndry : float, optional
+ Prescribed boundary flux value. If None, it is computed internally.
+ mask_outside_limiter : ndarray, optional
+ Boolean mask identifying points outside the limiter.
+ rel_tolerance_xpt : float, optional
+ Convergence tolerance for X-point search in fallback method.
+ starting_dx : float, optional
+ Initial step size for contour-based X-point search.
+
+ Returns
+ -------
+ opt : ndarray, shape (1, 3)
+ O-point coordinates and flux value [R, Z, psi].
+ xpt : ndarray, shape (1, 3)
+ X-point coordinates and flux value [R, Z, psi].
+ diverted_core_mask : ndarray
+ Boolean mask of plasma core region.
+ psi_bndry : float
+ Flux value at the last closed flux surface (LCFS).
+ """
+
try:
opt, xpt, diverted_core_mask, psi_bndry = self.Jtor_part1(
R, Z, psi, psi_bndry, mask_outside_limiter
@@ -363,97 +475,6 @@ def diverted_critical_complete(
return opt, xpt, diverted_core_mask, psi_bndry
- # def diverted_critical_old(self, R, Z, psi, psi_bndry=None, mask_outside_limiter=None, xpt_tol=1e-4):
- # # this
-
- # # find O- and X-points of equilibrium
- # opt, xpt = critical.fastcrit(
- # R, Z, psi, self.mask_inside_limiter, #self.Ip
- # )
- # len_xpt = len(xpt)
- # len_opt = len(opt)
-
- # # find core plasma mask (using user-defined psi_bndry)
- # if psi_bndry is not None:
- # diverted_core_mask = critical.inside_mask(
- # R, Z, psi, opt, xpt, mask_outside_limiter, psi_bndry
- # )
-
- # elif len_xpt:
- # del_psi = np.max(psi)-np.min(psi)
- # # order xpt according to psi
- # xpt = xpt[np.argsort(xpt[:,2])]
- # i = -1
- # xpt_found = False
- # while xpt_found==False and i psi_bndry)
-
- # else:
- # # no useful xpt found
- # psi_bndry = psi[0, 0]
- # diverted_core_mask = None
-
- # else:
- # # No X-points
- # psi_bndry = psi[0, 0]
- # diverted_core_mask = None
-
- # return opt, xpt, diverted_core_mask, psi_bndry
-
def Jtor_build(
self,
Jtor_part1,
@@ -466,33 +487,56 @@ def Jtor_build(
mask_outside_limiter,
limiter_mask_out,
):
- """Universal function that calculates the plasma current distribution,
- common to all of the different types of profile parametrizations used in FreeGSNKE.
+ """
+ Construct the toroidal current density (Jtor) using a modular profile pipeline.
+
+ This function is the main assembly routine for the plasma current density.
+ It combines:
+ - geometric reconstruction of plasma boundaries (Jtor_part1),
+ - limiter-aware core masking (core_mask_limiter),
+ - and evaluation of the current profile itself (Jtor_part2).
+
+ The implementation is designed to support multiple profile parametrisations
+ in a unified framework.
Parameters
----------
- Jtor_part1 : method
- method from the freegs4e Profile class
- returns opt, xpt, diverted_core_mask
- Jtor_part2 : method
- method from each individual profile class
- returns jtor itself
- core_mask_limiter : method
- method of the limiter_handler class
- returns the refined core_mask where jtor>0 accounting for the limiter
- R : np.ndarray
- R coordinates of the domain grid points
- Z : np.ndarray
- Z coordinates of the domain grid points
- psi : np.ndarray
- Poloidal field flux / 2*pi at each grid points (for example as returned by Equilibrium.psi())
- psi_bndry : float, optional
- Value of the poloidal field flux at the boundary of the plasma (last closed flux surface), by default None
- mask_outside_limiter : np.ndarray
- Mask of points outside the limiter, if any, optional
- limiter_mask_out : np.ndarray
- The mask identifying the border of the limiter, including points just inside it, the 'last' accessible to the plasma.
- Same size as psi.
+ Jtor_part1 : callable
+ Function that computes geometric plasma descriptors:
+ returns (opt, xpt, diverted_core_mask).
+ Jtor_part2 : callable
+ Function that evaluates the toroidal current density jtor.
+ core_mask_limiter : callable
+ Function that refines the core mask using limiter geometry.
+ R : ndarray
+ Radial grid coordinates.
+ Z : ndarray
+ Vertical grid coordinates.
+ psi : ndarray
+ Poloidal flux on the grid.
+ psi_bndry : float
+ Flux value at the last closed flux surface (LCFS).
+ mask_outside_limiter : ndarray
+ Boolean mask for points outside the limiter.
+ limiter_mask_out : ndarray
+ Limiter boundary mask (including edge-adjacent region).
+
+ Returns
+ -------
+ jtor : ndarray
+ Toroidal current density on the grid.
+ opt : ndarray
+ O-point coordinates and flux value.
+ xpt : ndarray
+ X-point coordinates and flux value.
+ psi_bndry : float
+ Updated LCFS flux value.
+ diverted_core_mask : ndarray
+ Core plasma mask from geometric reconstruction.
+ limiter_core_mask : ndarray
+ Core mask refined with limiter constraints.
+ flag_limiter : bool or int
+ Indicator of whether limiter correction was applied successfully.
"""
opt, xpt, diverted_core_mask, self.diverted_psi_bndry = Jtor_part1(
@@ -534,23 +578,29 @@ def Jtor_build(
)
def Jtor_unrefined(self, R, Z, psi, psi_bndry=None):
- """Replaces the FreeGS4E call, while maintaining the same input structure.
+ """
+ Compute the toroidal current density without subgrid refinement.
+
+ This method provides a direct replacement for the FreeGS4E current
+ computation interface, using the internal Jtor pipeline but disabling
+ any refined current reconstruction.
Parameters
----------
- R : np.ndarray
- R coordinates of the domain grid points
- Z : np.ndarray
- Z coordinates of the domain grid points
- psi : np.ndarray
- Poloidal field flux / 2*pi at each grid points (for example as returned by Equilibrium.psi())
+ R : ndarray
+ Radial grid coordinates.
+ Z : ndarray
+ Vertical grid coordinates.
+ psi : ndarray
+ Poloidal flux on the grid (Webers / 2π).
psi_bndry : float, optional
- Value of the poloidal field flux at the boundary of the plasma (last closed flux surface), by default None
+ Flux value at the last closed flux surface (LCFS). If None,
+ it is determined internally.
Returns
-------
- ndarray
- 2d map of toroidal current values
+ jtor : ndarray
+ Toroidal current density on the computational grid.
"""
(
self.jtor,
@@ -562,10 +612,8 @@ def Jtor_unrefined(self, R, Z, psi, psi_bndry=None):
self.flag_limiter,
) = self.Jtor_build(
self.diverted_critical_complete,
- # self.Jtor_part1,
self.Jtor_part2,
self.limiter_handler.core_mask_limiter,
- # self.core_mask_limiter,
R,
Z,
psi,
@@ -576,26 +624,33 @@ def Jtor_unrefined(self, R, Z, psi, psi_bndry=None):
return self.jtor
def Jtor_refined(self, R, Z, psi, psi_bndry=None, thresholds=None):
- """Implements the call to the Jtor method for the case in which the subgrid refinement is used.
+ """
+ Compute toroidal current density using subgrid refinement.
- Parameters
+ This method evaluates the unrefined current density first and then applies
+ a subgrid refinement procedure in regions where higher resolution is required.
+ The refinement is controlled by a threshold-based criterion acting on the
+ current density and its gradients.
+
+ Parameters
----------
- R : np.ndarray
- R coordinates of the domain grid points
- Z : np.ndarray
- Z coordinates of the domain grid points
- psi : np.ndarray
- Poloidal field flux / 2*pi at each grid points (for example as returned by Equilibrium.psi())
+ R : ndarray
+ Radial grid coordinates.
+ Z : ndarray
+ Vertical grid coordinates.
+ psi : ndarray
+ Poloidal flux on the grid (typically ψ / 2π).
psi_bndry : float, optional
- Value of the poloidal field flux at the boundary of the plasma (last closed flux surface), by default None
- thresholds : tuple (threshold for jtor criterion, threshold for gradient criterion)
- tuple of values used to identify where to apply refinement
- when None, the default refinement_thresholds are used
+ Flux value at the last closed flux surface (LCFS). If None,
+ it is determined internally.
+ thresholds : tuple of float, optional
+ (jtor_threshold, gradient_threshold) controlling where refinement is
+ applied. If None, the default `self.refinement_thresholds` is used.
Returns
-------
- ndarray
- 2d map of toroidal current values
+ jtor : ndarray
+ Refined toroidal current density on the computational grid.
"""
unrefined_jtor = self.Jtor_unrefined(R, Z, psi, psi_bndry)
@@ -643,20 +698,19 @@ def Jtor_refined(self, R, Z, psi, psi_bndry=None, thresholds=None):
class ConstrainBetapIp(freegs4e.jtor.ConstrainBetapIp, Jtor_universal):
- """FreeGSNKE profile class adapting the original FreeGS object with the same name,
- with a few modifications, to:
- - retain memory of critical point calculation;
- - deal with limiter plasma configurations
+ """
+ Betap–Ip constrained toroidal current profile with FreeGSNKE extensions.
"""
def __init__(self, eq, *args, **kwargs):
- """Instantiates the object.
+ """
+ Initialise the constrained profile.
Parameters
----------
eq : FreeGSNKE Equilibrium object
- Specifies the domain properties
+ Equilibrium object defining grid geometry and limiter structure.
"""
freegs4e.jtor.ConstrainBetapIp.__init__(self, *args, **kwargs)
Jtor_universal.__init__(self)
@@ -668,6 +722,15 @@ def __init__(self, eq, *args, **kwargs):
self.set_masks(eq=eq)
def copy(self):
+ """
+ Create a deep-ish copy of the profile object.
+
+ Returns
+ -------
+ ConstrainBetapIp
+ A new instance with copied scalar parameters and shared/replicated
+ internal state depending on attribute type.
+ """
obj = super().copy()
copy_into(self, obj, "profile_parameter")
@@ -686,11 +749,34 @@ def copy(self):
def Lao_parameters(
self, n_alpha, n_beta, alpha_logic=True, beta_logic=True, Ip_logic=True, nn=100
):
- """Finds best fitting alpha, beta parameters for a Lao85 profile,
- to reproduce the input pprime_ and ffprime_
- n_alpha and n_beta represent the number of free parameters
+ """
+ Fit Lao85 profile parameters to the current pprime and ffprime profiles.
+
+ This method constructs a discrete sampling of the normalized flux coordinate
+ and evaluates the current profile derivatives. It then fits a Lao85-type
+ polynomial representation to obtain optimal alpha and beta parameters.
+
+ Parameters
+ ----------
+ n_alpha : int
+ Number of free parameters used in the pprime (alpha) expansion.
+ n_beta : int
+ Number of free parameters used in the ffprime (beta) expansion.
+ alpha_logic : bool, optional
+ If True, enforce boundary-consistent modification of the alpha basis.
+ beta_logic : bool, optional
+ If True, enforce boundary-consistent modification of the beta basis.
+ Ip_logic : bool, optional
+ If True, apply total current normalisation during fitting.
+ nn : int, optional
+ Number of sampling points in the normalized flux coordinate.
- See Lao_parameters_finder.
+ Returns
+ -------
+ alpha : ndarray
+ Fitted alpha coefficients for the pprime expansion.
+ beta : ndarray
+ Fitted beta coefficients for the ffprime expansion.
"""
pn_ = np.linspace(0, 1, nn)
@@ -712,20 +798,19 @@ def Lao_parameters(
class ConstrainPaxisIp(freegs4e.jtor.ConstrainPaxisIp, Jtor_universal):
- """FreeGSNKE profile class adapting the original FreeGS object with the same name,
- with a few modifications, to:
- - retain memory of critical point calculation;
- - deal with limiter plasma configurations
+ """
+ Paxis–Ip constrained toroidal current profile with FreeGSNKE extensions.
"""
def __init__(self, eq, *args, **kwargs):
- """Instantiates the object.
+ """
+ Initialise the constrained profile.
Parameters
----------
eq : FreeGSNKE Equilibrium object
- Specifies the domain properties
+ Equilibrium object defining grid geometry and limiter structure.
"""
freegs4e.jtor.ConstrainPaxisIp.__init__(self, *args, **kwargs)
Jtor_universal.__init__(self)
@@ -737,6 +822,16 @@ def __init__(self, eq, *args, **kwargs):
self.set_masks(eq=eq)
def copy(self):
+ """
+ Create a copy of the current profile instance.
+
+ Returns
+ -------
+ ConstrainPaxisIp
+ A copied instance with duplicated scalar parameters and appropriately
+ handled internal state (deep or shallow depending on attribute type).
+ """
+
obj = super().copy()
copy_into(self, obj, "profile_parameter")
@@ -755,11 +850,35 @@ def copy(self):
def Lao_parameters(
self, n_alpha, n_beta, alpha_logic=True, beta_logic=True, Ip_logic=True, nn=100
):
- """Finds best fitting alpha, beta parameters for a Lao85 profile,
- to reproduce the input pprime_ and ffprime_
- n_alpha and n_beta represent the number of free parameters
+ """
+ Fit Lao85 profile coefficients from pprime and ffprime evaluations.
+
+ This method samples the normalized flux coordinate and evaluates the
+ pressure derivative (pprime) and flux function derivative (ffprime).
+ It then computes best-fit Lao85 polynomial coefficients using a linear
+ fitting procedure.
+
+ Parameters
+ ----------
+ n_alpha : int
+ Number of coefficients in the alpha (pprime) expansion.
+ n_beta : int
+ Number of coefficients in the beta (ffprime) expansion.
+ alpha_logic : bool, optional
+ If True, enforces boundary-consistent modification of the alpha basis.
+ beta_logic : bool, optional
+ If True, enforces boundary-consistent modification of the beta basis.
+ Ip_logic : bool, optional
+ If True, applies total current normalisation in the fitting procedure.
+ nn : int, optional
+ Number of points used to sample the normalized flux coordinate.
- See Lao_parameters_finder.
+ Returns
+ -------
+ alpha : ndarray
+ Fitted coefficients for the pprime expansion.
+ beta : ndarray
+ Fitted coefficients for the ffprime expansion.
"""
pn_ = np.linspace(0, 1, nn)
@@ -781,20 +900,19 @@ def Lao_parameters(
class Fiesta_Topeol(freegs4e.jtor.Fiesta_Topeol, Jtor_universal):
- """FreeGSNKE profile class adapting the FreeGS4E object with the same name,
- with a few modifications, to:
- - retain memory of critical point calculation;
- - deal with limiter plasma configurations
+ """
+ Fiesta Topeol constrained toroidal current profile with FreeGSNKE extensions.
"""
def __init__(self, eq, *args, **kwargs):
- """Instantiates the object.
+ """
+ Initialise the Fiesta-Topeol constrained current profile.
Parameters
----------
eq : FreeGSNKE Equilibrium object
- Specifies the domain properties
+ Equilibrium object defining grid geometry and limiter structure.
"""
freegs4e.jtor.Fiesta_Topeol.__init__(self, *args, **kwargs)
Jtor_universal.__init__(self)
@@ -806,6 +924,15 @@ def __init__(self, eq, *args, **kwargs):
self.set_masks(eq=eq)
def copy(self):
+ """
+ Create a copy of the Fiesta-Topeol profile instance.
+
+ Returns
+ -------
+ Fiesta_Topeol
+ A copied instance with duplicated scalar parameters and appropriately
+ handled internal state (deep or shallow depending on attribute type).
+ """
obj = super().copy()
copy_into(self, obj, "profile_parameter")
@@ -823,11 +950,35 @@ def copy(self):
def Lao_parameters(
self, n_alpha, n_beta, alpha_logic=True, beta_logic=True, Ip_logic=True, nn=100
):
- """Finds best fitting alpha, beta parameters for a Lao85 profile,
- to reproduce the input pprime_ and ffprime_
- n_alpha and n_beta represent the number of free parameters
+ """
+ Fit Lao85 profile coefficients from sampled pprime and ffprime data.
+
+ This method evaluates the pressure derivative (pprime) and flux function
+ derivative (ffprime) on a uniform grid in the normalized flux coordinate,
+ then fits Lao85 polynomial coefficients using a linear least-squares
+ procedure.
- See Lao_parameters_finder.
+ Parameters
+ ----------
+ n_alpha : int
+ Number of coefficients in the alpha (pprime) expansion.
+ n_beta : int
+ Number of coefficients in the beta (ffprime) expansion.
+ alpha_logic : bool, optional
+ If True, enforces boundary-consistent modification of the alpha basis.
+ beta_logic : bool, optional
+ If True, enforces boundary-consistent modification of the beta basis.
+ Ip_logic : bool, optional
+ If True, applies total current normalisation in the fitting procedure.
+ nn : int, optional
+ Number of sample points in the normalized flux coordinate.
+
+ Returns
+ -------
+ alpha : ndarray
+ Fitted coefficients for the pprime expansion.
+ beta : ndarray
+ Fitted coefficients for the ffprime expansion.
"""
pn_ = np.linspace(0, 1, nn)
@@ -849,32 +1000,40 @@ def Lao_parameters(
class Lao85(freegs4e.jtor.Lao85, Jtor_universal):
- """FreeGSNKE profile class adapting the FreeGS4E object with the same name,
- with a few modifications, to:
- - retain memory of critical point calculation;
- - deal with limiter plasma configurations
+ """
+ Lao 1985 constrained toroidal current profile with FreeGSNKE extensions.
"""
def __init__(self, eq, *args, refine_jtor=False, nnx=None, nny=None, **kwargs):
- """Instantiates the object.
+ """
+ Initialise the Lao85 current profile.
Parameters
----------
- eq : freegs4e Equilibrium object
- Specifies the domain properties
+ eq : FreeGSNKE Equilibrium object
+ Equilibrium object defining grid geometry and limiter structure.
refine_jtor : bool
- Flag to select whether to apply sug-grid refinement of plasma current distribution jtor
- nnx : even integer
- refinement factor in the R direction
- nny : even integer
- refinement factor in the Z direction
+ If True, enable subgrid refinement of the toroidal current density.
+ nnx : int
+ Refinement factor in the R-direction (must be even if used).
+ nny : int
+ Refinement factor in the Z-direction (must be even if used).
"""
freegs4e.jtor.Lao85.__init__(self, *args, **kwargs)
self.set_masks(eq=eq)
self.select_refinement(eq, refine_jtor, nnx, nny)
def copy(self):
+ """
+ Create a copy of the Lao85 profile instance.
+
+ Returns
+ -------
+ Lao85
+ A copied instance with duplicated profile parameters and internal
+ state (with controlled shallow/deep copying depending on attribute type).
+ """
obj = super().copy()
copy_into(self, obj, "Ip")
@@ -896,18 +1055,27 @@ def copy(self):
return obj
def Topeol_parameters(self, nn=100, max_it=100, tol=1e-5):
- """Fids best combination of
- (alpha_m, alpha_n, beta_0)
- to instantiate a Topeol profile object as similar as possible to self
+ """
+ Fit optimal Topeol profile parameters from target pprime and ffprime data.
+
+ This method determines the best-fitting parameters
+ (alpha_m, alpha_n, beta_0) for a Topeol current profile by minimising
+ a mismatch between the model and the target profile derivatives
+ evaluated on a sampled normalized flux grid.
Parameters
----------
nn : int, optional
- number of points to sample 0,1 interval in the normalised psi, by default 100
- max_it : int,
- maximum number of iterations in the optimization
- tol : float
- iterations stop when change in the optimised parameters in smaller than tol
+ Number of sampling points in the normalized flux interval (0, 1).
+ max_it : int, optional
+ Maximum number of optimisation iterations.
+ tol : float, optional
+ Convergence tolerance on parameter updates.
+
+ Returns
+ -------
+ pars : ndarray, shape (3,)
+ Optimised parameters (alpha_m, alpha_n, beta_0).
"""
x = np.linspace(1 / (100 * nn), 1 - 1 / (100 * nn), nn)
@@ -926,19 +1094,19 @@ def Topeol_parameters(self, nn=100, max_it=100, tol=1e-5):
class TensionSpline(freegs4e.jtor.TensionSpline, Jtor_universal):
- """FreeGSNKE profile class adapting the FreeGS4E object with the same name,
- with a few modifications, to:
- - retain memory of critical point calculation;
- - deal with limiter plasma configurations
+ """
+ Tension spline constrained toroidal current profile with FreeGSNKE extensions.
+
"""
def __init__(self, eq, *args, **kwargs):
- """Instantiates the object.
+ """
+ Initialise the tension spline current profile.
Parameters
----------
eq : FreeGSNKE Equilibrium object
- Specifies the domain properties
+ Equilibrium object defining grid geometry and limiter structure.
"""
freegs4e.jtor.TensionSpline.__init__(self, *args, **kwargs)
@@ -958,6 +1126,15 @@ def __init__(self, eq, *args, **kwargs):
self.set_masks(eq=eq)
def copy(self):
+ """
+ Create a copy of the TensionSpline profile instance.
+
+ Returns
+ -------
+ TensionSpline
+ A copied instance with all spline parameters duplicated and internal
+ state consistently updated.
+ """
obj = super().copy()
copy_into(self, obj, "Ip")
@@ -999,7 +1176,28 @@ def assign_profile_parameter(
ffp_values_2,
ffp_sigma,
):
- """Assigns to the profile object new values for the profile parameters"""
+ """
+ Assign new spline parameters to the profile object.
+
+ Parameters
+ ----------
+ pp_knots : ndarray
+ Knot locations for pprime spline.
+ pp_values : ndarray
+ Spline values for pprime.
+ pp_values_2 : ndarray
+ Second derivative (or auxiliary) values for pprime spline.
+ pp_sigma : ndarray
+ Regularisation / smoothing parameters for pprime spline.
+ ffp_knots : ndarray
+ Knot locations for ffprime spline.
+ ffp_values : ndarray
+ Spline values for ffprime.
+ ffp_values_2 : ndarray
+ Second derivative (or auxiliary) values for ffprime spline.
+ ffp_sigma : ndarray
+ Regularisation / smoothing parameters for ffprime spline.
+ """
self.pp_knots = pp_knots
self.pp_values = pp_values
self.pp_values_2 = pp_values_2
@@ -1022,19 +1220,19 @@ def assign_profile_parameter(
class GeneralPprimeFFprime(freegs4e.jtor.GeneralPprimeFFprime, Jtor_universal):
- """FreeGSNKE profile class adapting the FreeGS4E object with the same name,
- with a few modifications, to:
- - retain memory of critical point calculation;
- - deal with limiter plasma configurations
+ """
+ General unconstrained toroidal current profile with FreeGSNKE extensions.
+
"""
def __init__(self, eq, *args, **kwargs):
- """Instantiates the object.
+ """
+ Initialise the general pprime/ffprime current profile.
Parameters
----------
eq : FreeGSNKE Equilibrium object
- Specifies the domain properties
+ Equilibrium object defining grid geometry and limiter structure.
"""
freegs4e.jtor.GeneralPprimeFFprime.__init__(self, *args, **kwargs)
@@ -1044,6 +1242,15 @@ def __init__(self, eq, *args, **kwargs):
self.set_masks(eq=eq)
def copy(self):
+ """
+ Create a copy of the GeneralPprimeFFprime profile instance.
+
+ Returns
+ -------
+ GeneralPprimeFFprime
+ A copied instance with all profile data and grid-dependent state
+ duplicated and reinitialised appropriately.
+ """
obj = super().copy()
copy_into(self, obj, "profile_parameter")
@@ -1066,6 +1273,11 @@ def copy(self):
def assign_profile_parameter(
self,
):
- """Assigns to the profile object new values for the profile parameters"""
+ """
+ Reset profile parameter container.
+
+ This profile is fully non-parametric, so no scalar or vector parameter
+ set is required; the parameter container is explicitly cleared.
+ """
self.profile_parameter = []
diff --git a/freegsnke/limiter_func.py b/freegsnke/limiter_func.py
index 7b0417b..d8f2c69 100644
--- a/freegsnke/limiter_func.py
+++ b/freegsnke/limiter_func.py
@@ -24,6 +24,20 @@
class Limiter_handler:
+ """
+ Handles limiter-related geometric operations for FreeGSNKE profile evaluation.
+
+ This class provides functionality for enforcing and querying the plasma
+ boundary defined by the limiter geometry. It is primarily used by profile
+ objects to determine whether grid points lie inside the allowable plasma
+ region and to support limiter-dependent calculations.
+
+ Each profile object typically instantiates its own Limiter_handler.
+
+ Notes
+ -----
+ The limiter is treated as a closed boundary defining the valid plasma domain.
+ """
def __init__(self, eq, limiter):
"""Object to handle additional calculations due to the limiter.
diff --git a/freegsnke/linear_solve.py b/freegsnke/linear_solve.py
index 0841dcb..49736dd 100644
--- a/freegsnke/linear_solve.py
+++ b/freegsnke/linear_solve.py
@@ -26,10 +26,23 @@
class linear_solver:
- """Interface between the linearised system of circuit equations and the implicit-Euler
- time stepper. Calculates the linear growth rate and solves the linearised dynamical problem.
- It needs the Jacobian of the plasma current distribution with respect to all of the
- independent currents, dIy/dI.
+ """
+ Linearised plasma–metal circuit solver.
+
+ This class provides the interface between the linearised system of
+ coupled plasma and metal circuit equations and the implicit-Euler
+ time integrator. It constructs the matrices required for time
+ integration, computes linear growth rates, and advances the
+ linearised dynamical system in time.
+
+ The metal circuit equations are expressed in a reduced basis defined
+ by the mode decomposition matrix ``P``, in which the active coils
+ remain in the physical current basis and the passive structures are
+ represented by vessel eigenmodes.
+
+ The solver requires the Jacobian of the plasma current distribution
+ with respect to the independent current variables, ``dIy/dI``, which
+ determines the linear plasma response.
"""
def __init__(
@@ -45,38 +58,45 @@ def __init__(
max_internal_timestep=0.0001,
full_timestep=0.0001,
):
- """Instantiates the linear_solver object, with inputs computed mostly
- within the circuit_equation_metals object.
- Based on the input plasma properties and coupling matrices, it prepares:
- - an instance of the implicit Euler solver implicit_euler_solver()
- - internal time-stepper for the implicit-Euler
+ """
+ Initialise the linearised circuit solver.
+
+ Precomputes matrices describing the coupling between plasma currents
+ and metal circuits in the reduced modal basis, and creates the
+ implicit-Euler integrator used to advance the linearised system.
Parameters
----------
- eq : class
- FreeGSNKE equilibrium object.
- Lambdam1: np.array
- State matrix of the circuit equations for the metal in normal mode form:
- P is the identity on the active coils and diagonalises the isolated dynamics
- of the passive coils, R^{-1}L_{passive}
- P: np.array
- change of basis matrix, as defined above, with modes appropriately removed
- Pm1: np.array
- Inverse of the change of basis matrix, as defined above, with modes appropriately removed
- Rm1: np.array
- matrix of all metal resitances to the power of -1. Diagonal.
- Mey: np.array
- matrix of inductance values between grid points in the reduced plasma domain and all metal coils
- (active coils and passive-structure filaments)
- Calculated by the metal_currents object
- plasma_norm_factor: float
- an overall factor to work with a rescaled plasma current, so that
- it's within a comparable range with metal currents
- max_internal_timestep: float
- internal integration timestep of the implicit-Euler solver, to be used
- as substeps over the <> interval
- full_timestep: float
- full timestep requested to the implicit-Euler solver
+ coil_numbers : tuple[int, int]
+ Tuple ``(n_active_coils, n_coils)`` giving the number of active
+ coils and the total number of conducting elements.
+ Lambdam1 : ndarray
+ State matrix of the metal circuit equations in the reduced basis.
+ Active coils remain in the physical current basis, while passive
+ structures are represented by vessel eigenmodes.
+ P : ndarray
+ Change-of-basis matrix mapping currents from the reduced modal
+ basis to the physical conductor basis.
+ Pm1 : ndarray
+ Inverse of ``P``, mapping currents from the physical conductor
+ basis to the reduced modal basis.
+ Rm1 : ndarray
+ Inverse resistance matrix of all conducting elements. This is
+ typically diagonal.
+ Mey : ndarray
+ Mutual inductance matrix between plasma current elements and
+ metal conductors (active coils and passive structures).
+ plasma_norm_factor : float
+ Scaling factor applied to the plasma current variable to improve
+ conditioning and keep its magnitude comparable to metal currents.
+ plasma_resistance_1d : ndarray
+ Effective plasma resistance associated with each plasma current
+ element in the reduced plasma domain.
+ max_internal_timestep : float, optional
+ Maximum timestep used internally by the implicit-Euler solver.
+ Longer requested timesteps are subdivided into smaller steps.
+ full_timestep : float, optional
+ External timestep associated with a single solver advance.
"""
self.max_internal_timestep = max_internal_timestep
@@ -127,26 +147,36 @@ def __init__(
self.profiles_forcing = np.zeros(self.n_independent_vars + 1)
def reset_plasma_resistivity(self, plasma_resistance_1d):
- """Resets the value of the plasma resistivity,
- throught the vector of 'geometric restistances' in the plasma domain
+ """
+ Update the plasma resistivity profile used by the linearised solver.
+
+ Replaces the vector of effective plasma resistances associated with
+ the reduced plasma domain. This automatically invalidates the current
+ linearisation point so that system matrices will be rebuilt using the
+ updated resistivity.
Parameters
----------
plasma_resistance_1d : ndarray
- Vector of 2pi resistivity R/dA for all domain grid points in the reduced plasma domain
+ Vector of effective plasma resistance coefficients for all grid
+ points in the reduced plasma domain. Each entry represents the
+ geometric resistance factor (2π * resistivity divided by cell area).
"""
self.plasma_resistance_1d = plasma_resistance_1d
self.set_linearization_point(None, None, None, None)
def reset_timesteps(self, max_internal_timestep, full_timestep):
- """Resets the integration timesteps, calling self.solver.set_timesteps
+ """
+ Update the timesteps used by the implicit-Euler integrator.
Parameters
----------
max_internal_timestep : float
- integration substep of the ciruit equation, calling an implicit-Euler solver
+ Maximum internal timestep used by the implicit-Euler solver.
+ Larger timesteps are subdivided into steps no larger than this value.
full_timestep : float
- integration timestep of the circuit equation
+ External timestep associated with a single advance of the
+ circuit equations.
"""
self.max_internal_timestep = max_internal_timestep
self.full_timestep = full_timestep
@@ -155,24 +185,38 @@ def reset_timesteps(self, max_internal_timestep, full_timestep):
)
def set_linearization_point(self, dIydI, dIydtheta, hatIy0, Myy_hatIy0):
- """Initialises an implicit-Euler solver with the appropriate matrices for the linearised dynamic problem.
+ """
+ Set or update the linearisation point for the coupled plasma–metal system.
+
+ This method defines the equilibrium state around which the system is
+ linearised and rebuilds the system matrices used by the implicit-Euler
+ solver.
Parameters
----------
- dIydI : np.array
- partial derivatives of plasma-cell currents on the reduced plasma domain with respect to all intependent metal currents
- (active coil currents, vessel normal modes, total plasma current divided by plasma_norm_factor).
- These would typically come from having solved the forward Grad-Shafranov problem. Finite difference Jacobian.
- Calculated by the nl_solver object
- dIydtheta : np.array
- partial derivatives of plasma-cell currents on the reduced plasma domain with respect to all plasma current density
- profile parameters
- hatIy0 : np.array
- Plasma current distribution on the reduced plasma domain (1d) of the equilibrium around which the dynamics is linearised.
- This is normalised by the total plasma current, so that this vector sums to 1.
- Myy_hatIy0 : np.array
- The matrix product np.dot(Myy, hatIy0) in the same reduced domain as hatIy0
- This is provided by Myy_handler
+ dIydI : ndarray or None
+ Jacobian of plasma cell currents with respect to independent metal
+ currents (active coils, vessel modes, and normalised total plasma
+ current). Typically computed from a Grad–Shafranov solve using finite
+ differences.
+ dIydtheta : ndarray or None
+ Jacobian of plasma cell currents with respect to plasma profile
+ parameterisation variables.
+ hatIy0 : ndarray or None
+ Normalised equilibrium plasma current distribution over the reduced
+ plasma domain. This vector is scaled to sum to one.
+ Myy_hatIy0 : ndarray or None
+ Precomputed product of the plasma–plasma coupling matrix with the
+ equilibrium current distribution, provided by the plasma coupling
+ handler.
+
+ Notes
+ -----
+ Any argument set to None is left unchanged from the previous linearisation
+ point.
+
+ This call rebuilds the system matrix and reinitialises the implicit-Euler
+ solver using the updated linearised dynamics.
"""
if dIydI is not None:
self.dIydI = dIydI
@@ -278,6 +322,11 @@ def stepper(
voltages applied to the active coils
dtheta_dt : np.array
Vector of plasma current density profile parameters derivateives with respect to t.
+
+ Returns
+ -------
+ Itpdt : ndarray
+ Updated state vector after one full implicit-Euler timestep.
"""
# baseline forcing term (from the active coil voltages)
diff --git a/freegsnke/magnetic_probes.py b/freegsnke/magnetic_probes.py
index 4e4009f..ba7d3a1 100644
--- a/freegsnke/magnetic_probes.py
+++ b/freegsnke/magnetic_probes.py
@@ -31,41 +31,58 @@
class Probes:
"""
- Class to implement magnetic probes:
- - flux loops: compute psi(R,Z)
- - pickup coils: compute B(R,phi,Z).nhat (nhat is unit vector orientation of the probe)
-
- Inputs:
- - equilibrium object - contains grid, plasma profile, plasma and coil currents, coil positions.
- N.B:- in init/setup the eq object only provides machine /domain/position information
- - in methods the equilibrium provides currents and other aspects that evolve in solve().
-
- Attributes:
- - floops,pickups = dictionaries with name, position, orientation of the probes
- - floops_positions etc. = extract individual arrays of positions, orientations etc.
- - floops_order / pickups_order = list of fluxloop/pickups names, if individual probe value is required
- - greens_psi_coils_floops = greens functions for psi, evaluated at flux loop positions
- - greens_br/bz_coils_pickups = greens functions for Br/Bz, evaluated at pickup coil positions
- - greens_psi_plasma_floops = dictionary of greens functions for psi from plasma current, evaluated at flux loop positions
- - greens_BrBz_plasma_pickups = dictionary of greens functions for Br and Bz from plasma, evaluated at pickup coil positions
-
- - more greens function attributes would be added if new probes area added.
-
- Methods:
- - get_coil_currents(eq): returns current values in all the coils from equilibrium object.
- - get_plasma_current(eq): returns toroidal current values at each plasma grid point, taken from equilibrium input.
- - create_greens_all_coils(): returns array with greens function for each coil and probe combination.
- - psi_all_coils(eq): returns list of psi values at each flux loop position, summed over all coils.
- - psi_from_plasma(eq): returns list of psi values at each flux loop position from plasma itself.
- - create_greens_B_oriented_plasma(eq) : creates oriented greens functions for pickup coils.
- - calculate_fluxloop_value(eq): returns total flux at each probe position (sum of previous two outputs).
- - calculate_pickup_value(eq): returns pickup values at each probe position.
-
-
- - Br(eq)/ Bz(eq) : computes radial/z component of magnetic field, sum of coil and plasma contributions
- - Btor(eq) : extracts toroidal magnetic field (outside of plasma), evaluated at
-
- Methods currently have floop or pickup positions as default, but these can be changed with optional argument.
+ Representation of magnetic diagnostics for a tokamak equilibrium.
+
+ This class implements synthetic magnetic probe signals, including:
+
+ - Flux loops: measure poloidal flux ψ(R, Z)
+ - Pickup coils: measure magnetic field components projected onto a
+ probe orientation vector (B · n̂)
+
+ The signals are computed from contributions of both coil currents and
+ plasma current distributions via precomputed Green's functions.
+
+ Parameters
+ ----------
+ coils_dict : dict
+ Dictionary describing active (and possibly passive) coils, including
+ geometry and identifiers used for Green's function evaluation.
+
+ magnetic_probe_data : dict, optional
+ Dictionary containing pre-defined probe configurations. Must include
+ keys:
+ - "flux_loops"
+ - "pickups"
+
+ magnetic_probe_path : str, optional
+ Path to a pickled file containing `magnetic_probe_data`.
+
+ Notes
+ -----
+ Exactly one of `magnetic_probe_data` or `magnetic_probe_path` must be
+ provided. If both are given, a ValueError is raised. If neither is
+ provided, no probe geometry is initialised.
+
+ Attributes
+ ----------
+ floops : dict
+ Flux loop definitions (positions, identifiers, etc.).
+ pickups : dict
+ Pickup coil definitions (positions, orientations, etc.).
+ coil_names : list of str
+ Ordered list of coil identifiers.
+ coils_dict : dict
+ Copy of input coil dictionary.
+
+ Other Attributes
+ -----------------
+ The class also stores multiple Green's function caches used to compute:
+
+ - Flux loop signals from coils and plasma
+ - Pickup coil signals (Br, Bz projections)
+ - Combined diagnostic outputs
+
+ These are populated during later setup routines.
"""
def __init__(
@@ -111,10 +128,56 @@ def __init__(
def initialise_setup(self, eq):
"""
- Setup attributes: positions, orientations and greens functions
- - set probe positions and orientations
- - set coil positions from equilibrium object
- - create arrays/dictionaries of greens functions with positions of coils/plasma currents and probes
+ Initialise probe geometry and precompute Green's functions for a given
+ equilibrium configuration.
+
+ This method validates compatibility between the probe set and the supplied
+ equilibrium, then constructs all probe-related quantities (positions,
+ orientations, and Green's function caches) required for synthetic
+ diagnostic evaluation.
+
+ Parameters
+ ----------
+ eq : Equilibrium
+ Equilibrium object providing tokamak geometry, coil configuration,
+ and (optionally) plasma state information.
+
+ Raises
+ ------
+ AssertionError
+ If the coil configuration in `eq.tokamak` does not match the coil
+ configuration used to define the probes.
+
+ Notes
+ -----
+ This routine performs expensive precomputation of Green's functions and
+ should typically be called once per equilibrium setup.
+
+ Flux Loops
+ ----------
+ - Extracts flux loop positions and ordering from `self.floops`
+ - Computes coil contribution Green's functions:
+ `greens_psi_coils_floops`
+ - Computes plasma contribution Green's functions:
+ `greens_psi_plasma_floops` (cached per equilibrium key)
+
+ Pickup Coils
+ -------------
+ - Extracts pickup positions and orientation vectors
+ - Computes coil contributions:
+ `greens_br_coils_pickup`, `greens_bz_coils_pickup`
+ - Computes oriented magnetic field responses:
+ `greens_B_coils_oriented`, `greens_B_plasma_oriented`
+
+ Caching
+ -------
+ All plasma-dependent Green's functions are stored in dictionaries keyed by
+ an equilibrium identifier (`eq_key`) to avoid recomputation when possible.
+
+ Returns
+ -------
+ None
+ All results are stored in-place as object attributes.
"""
check = DeepDiff(eq.tokamak.coils_dict, self.coils_dict) == {}
@@ -125,9 +188,6 @@ def initialise_setup(self, eq):
eq_key = self.create_eq_key(eq)
- # self.limiter_handler = {}
- # self.limiter_handler[eq_key] = eq.limiter_handler
-
# FLUX LOOPS
# positions, number of probes, ordering
self.floop_pos = np.array([probe["position"] for probe in self.floops])
@@ -153,11 +213,6 @@ def initialise_setup(self, eq):
self.greens_br_coils_pickup, self.greens_bz_coils_pickup = (
self.greens_BrBz_all_coils(eq, "pickups")
)
- # not initialised unless needed to save memory
- # (
- # self.greens_br_plasma_pickup[eq_key],
- # self.greens_bz_plasma_pickup[eq_key],
- # ) = self.create_greens_BrBz_plasma(eq, "pickups")
self.greens_B_plasma_oriented = {}
self.greens_B_plasma_oriented[eq_key] = self.create_greens_B_oriented_plasma(
@@ -167,18 +222,32 @@ def initialise_setup(self, eq):
eq, "pickups"
)
- # Other probes - to add in future...
-
- """
- Things for all probes
- - coil current array
- - plasma current array
- - eq grid key
- """
-
def get_coil_currents(self, eq):
"""
- create list of coil currents from the equilibrium
+ Extract coil current values from an equilibrium object in a fixed ordering.
+
+ This method builds a vector of coil currents corresponding to
+ `self.coil_names`, ensuring a consistent ordering between the probe
+ model and the equilibrium representation.
+
+ Parameters
+ ----------
+ eq : Equilibrium
+ Equilibrium object containing a `tokamak` attribute with coil objects
+ indexed by name.
+
+ Returns
+ -------
+ np.ndarray
+ Array of coil currents ordered according to `self.coil_names`.
+
+ Notes
+ -----
+ - The ordering of currents is determined by `self.coil_names`, which is
+ fixed during probe initialisation.
+ - Each coil is accessed via `eq.tokamak[label].current`.
+ - This could be replaced by a vectorised call such as
+ `eq.tokamak.getcurrents()` if ordering consistency is guaranteed.
"""
array_of_coil_currents = np.zeros(len(self.coil_names))
for i, label in enumerate(self.coil_names):
@@ -189,55 +258,133 @@ def get_coil_currents(self, eq):
def get_plasma_current(self, eq):
"""
- From equilibirium object contains toroidal current distribution on the grid, reduced to the limiter region only.
+ Extract the toroidal plasma current distribution from an equilibrium object.
+
+ The returned quantity is the grid-based plasma current density mapped to
+ the physically valid domain (typically restricted by the limiter region).
+
+ Parameters
+ ----------
+ eq : Equilibrium
+ Equilibrium object containing the plasma current density profile and
+ limiter geometry.
+
+ Returns
+ -------
+ np.ndarray
+ Toroidal current distribution mapped onto the limiter-constrained grid.
+
+ Notes
+ -----
+ - The current density is taken from `eq._profiles.jtor`.
+ - The mapping to circuit-relevant currents is performed via:
+ `eq.limiter_handler.Iy_from_jtor(...)`
+ - This ensures only physically valid (in-limiter) contributions are used.
"""
return eq.limiter_handler.Iy_from_jtor(eq._profiles.jtor)
def create_eq_key(self, eq):
"""
- Produces tuple (Rmin,Rmax,Zmin,Zmax,nx,ny) from equilibrium to access correct greens function.
+ Generate a hashable identifier for an equilibrium grid configuration.
+
+ The key uniquely encodes the spatial domain and resolution of the
+ equilibrium, and is used to cache and retrieve precomputed Green's
+ functions associated with a specific grid.
+
+ Parameters
+ ----------
+ eq : Equilibrium
+ Equilibrium object containing grid metadata (domain bounds and
+ discretisation).
+
+ Returns
+ -------
+ tuple
+ A hashable key of the form:
+ (R_min, R_max, Z_min, Z_max, nx, ny)
+
+ Notes
+ -----
+ - nx and ny are inferred from `eq.R_1D` and `eq.Z_1D`.
+ - This key assumes that Green's functions are valid only when both
+ spatial bounds and resolution match exactly.
+ - Small floating-point differences in bounds will produce distinct keys.
"""
nx, ny = len(eq.R_1D), len(eq.Z_1D)
eq_key = (eq.Rmin, eq.Rmax, eq.Zmin, eq.Zmax, nx, ny)
return eq_key
- """
- Things for flux loops
- """
-
def create_greens_psi_single_coil(self, eq, coil_key, probe="floops"):
"""
- Create array of greens functions for given coil evaluate at all probe positions
- - pos_R and pos_Z are arrays of R,Z coordinates of the probes
- - defines array of greens for each filament at each probe.
- - multiplies by polarity and multiplier
- - then sums over filaments to return greens function for probes from a given coil
- - flux loops by default, can apply to other probes too with minor modification
+ Compute the Green's function contribution of a single coil to probe signals.
+
+ This method evaluates the poloidal flux response ψ at all probe locations
+ due to a unit current in a specified coil.
+
+ Parameters
+ ----------
+ eq : Equilibrium
+ Equilibrium object containing the tokamak model and coil definitions.
+ coil_key : str
+ Identifier of the coil in `eq.tokamak` for which the Green's function
+ is computed.
+ probe : {"floops"}, optional
+ Type of diagnostic probe to evaluate. Currently only flux loops
+ ("floops") are supported.
+
+ Returns
+ -------
+ np.ndarray
+ Array of ψ values at each probe location due to a unit current in the
+ specified coil.
+
+ Notes
+ -----
+ - Probe positions are taken from `self.floop_pos` when `probe="floops"`.
+ - The computation is delegated to:
+ `eq.tokamak[coil_key].controlPsi(R, Z)`
+ - The result corresponds to a linear response (unit current assumption).
"""
+
if probe == "floops":
pos_R = self.floop_pos[:, 0]
pos_Z = self.floop_pos[:, 1]
- # pol = self.coils_dict[coil_key]["polarity"][np.newaxis, :]
- # mul = self.coils_dict[coil_key]["multiplier"][np.newaxis, :]
-
- # greens_filaments = Greens(
- # self.coils_dict[coil_key]["coords"][0][np.newaxis, :],
- # self.coils_dict[coil_key]["coords"][1][np.newaxis, :],
- # pos_R[:, np.newaxis],
- # pos_Z[:, np.newaxis],
- # )
- # greens_filaments *= pol
- # greens_filaments *= mul
- # greens_psi_coil = np.sum(greens_filaments, axis=1)
greens_psi_coil = eq.tokamak[coil_key].controlPsi(pos_R, pos_Z)
return greens_psi_coil
def create_greens_psi_all_coils(self, eq, probe="floops"):
"""
- Create 2d array of greens functions for all coils and at all probe positions
- - array[i][j] is greens function for coil i evaluated at probe position j
+ Compute the Green's function matrix relating all coils to all probe
+ locations for poloidal flux measurements.
+
+ This builds a full response matrix where each entry represents the
+ contribution of a unit current in a coil to the flux measured at a
+ specific diagnostic probe.
+
+ Parameters
+ ----------
+ eq : Equilibrium
+ Equilibrium object containing coil geometry and Green's function
+ evaluation methods.
+ probe : {"floops"}, optional
+ Type of diagnostic probe. Currently only flux loops ("floops") are
+ supported.
+
+ Returns
+ -------
+ np.ndarray
+ 2D array of shape (n_coils, n_probes), where entry [i, j] gives the
+ poloidal flux ψ at probe j due to a unit current in coil i.
+
+ Notes
+ -----
+ - Each row corresponds to a coil in `self.coils_dict`.
+ - Each column corresponds to a flux loop in `self.floop_pos`.
+ - The matrix assumes linear superposition of coil contributions.
+ - Computation is performed via repeated calls to
+ `create_greens_psi_single_coil`.
"""
array = np.array([]).reshape(0, self.number_floops)
@@ -249,9 +396,31 @@ def create_greens_psi_all_coils(self, eq, probe="floops"):
def psi_floop_all_coils(self, eq, probe="floops"):
"""
- compute flux function summed over all coils.
- returns array of flux values at the positions of the floop probes by default.
- new probes can be used instead (just change which greens function is used)
+ Compute the total poloidal flux at all flux loop locations due to all coils.
+
+ This method forms the linear superposition of coil contributions using the
+ precomputed Green's function matrix and the current coil state from the
+ equilibrium.
+
+ Parameters
+ ----------
+ eq : Equilibrium
+ Equilibrium object providing the current coil state via `get_coil_currents`.
+ probe : {"floops"}, optional
+ Diagnostic type. Currently only flux loops ("floops") are supported.
+
+ Returns
+ -------
+ np.ndarray
+ Total poloidal flux ψ at each flux loop location due to all coils.
+
+ Notes
+ -----
+ - Uses the Green's function matrix:
+ `self.greens_psi_coils_floops`
+ - Coil contributions are combined via linear superposition:
+ ψ_total = Σ_i I_i G_i
+ - The ordering of coils is defined by `self.coil_names`.
"""
array_of_coil_currents = self.get_coil_currents(eq)
if probe == "floops":
@@ -265,9 +434,35 @@ def psi_floop_all_coils(self, eq, probe="floops"):
def create_green_psi_plasma(self, eq, probe="floops"):
"""
- Compute greens function at probes from the plasma currents .
- - plasma current source in grid from solve. grid points contained in eq object
- - evaluated on flux loops by default, can apply to other probes too with minor modification
+ Compute the Green's function mapping plasma current density to probe
+ measurements of poloidal flux ψ.
+
+ This constructs the response of each diagnostic probe to unit plasma
+ current sources distributed over the equilibrium grid (restricted to the
+ limiter-allowed region).
+
+ Parameters
+ ----------
+ eq : Equilibrium
+ Equilibrium object containing the plasma current grid and limiter
+ mapping.
+ probe : {"floops"}, optional
+ Diagnostic type. Currently only flux loops ("floops") are supported.
+
+ Returns
+ -------
+ np.ndarray
+ Green's function array mapping plasma current elements to flux loop
+ measurements. Shape is typically (n_plasma_cells, n_probes).
+
+ Notes
+ -----
+ - Plasma source points are taken from:
+ `eq.limiter_handler.plasma_pts`
+ - Probe positions are taken from `self.floop_pos`.
+ - The returned object represents a linear operator for:
+ ψ_probe = G_plasma → probe · I_plasma
+ - Only in-limiter plasma points are included in the computation.
"""
if probe == "floops":
@@ -286,9 +481,36 @@ def create_green_psi_plasma(self, eq, probe="floops"):
def psi_from_plasma(self, eq, probe="floops"):
"""
- Calculate flux function contribution from the plasma
- - returns array of flux values from plasma at position of floop probes
- - evaluated on flux loops by default, can apply to other probes too with minor modification
+ Compute the contribution of plasma current to poloidal flux measurements
+ at diagnostic probes.
+
+ This evaluates the linear mapping from distributed plasma current density
+ to flux loop signals using precomputed or lazily generated Green's
+ functions.
+
+ Parameters
+ ----------
+ eq : Equilibrium
+ Equilibrium object containing plasma current distribution and limiter
+ mapping information.
+ probe : {"floops"}, optional
+ Diagnostic type. Currently only flux loops ("floops") are supported.
+
+ Returns
+ -------
+ np.ndarray
+ Contribution to poloidal flux ψ at each probe due to plasma current.
+
+ Notes
+ -----
+ - Plasma current distribution is obtained via:
+ `get_plasma_current(eq)`
+ - Green's functions are cached per equilibrium grid using:
+ `create_eq_key(eq)`
+ - If no cached plasma Green's function exists for the given equilibrium,
+ it is computed on demand via `create_green_psi_plasma`.
+ - The final signal is computed by linear superposition:
+ ψ = Σ_i G_i I_i
"""
eq_key = self.create_eq_key(eq)
plasma_current_distribution = self.get_plasma_current(eq)
@@ -312,49 +534,67 @@ def psi_from_plasma(self, eq, probe="floops"):
def calculate_fluxloop_value(self, eq):
"""
- total flux for all floop probes
+ Compute the total flux loop signals by combining coil and plasma
+ contributions.
+
+ This method returns the full synthetic flux loop measurement at all
+ diagnostic locations by summing the magnetic response from coil currents
+ and plasma current distribution.
+
+ Parameters
+ ----------
+ eq : Equilibrium
+ Equilibrium object containing both coil currents and plasma current
+ density.
+
+ Returns
+ -------
+ np.ndarray
+ Total poloidal flux ψ at each flux loop location.
+
+ Notes
+ -----
+ - Coil contribution:
+ `psi_floop_all_coils(eq)`
+ - Plasma contribution:
+ `psi_from_plasma(eq)`
+ - Assumes linear superposition of magnetic fields.
"""
return self.psi_floop_all_coils(eq) + self.psi_from_plasma(eq)
- """
- Things for pickup coils
- """
-
def create_greens_BrBz_single_coil(self, eq, coil_key, probe="pickups"):
"""
- Create array of greens functions for given coil evaluate at all pickup positions
- - defines array of greens for each filament at each probe.
- - multiplies by polarity and multiplier
- - then sums over filaments to return greens function for probes from a given coil
- - evaluated on pickups by default, can apply to other probes too with minor modification
+ Compute Green's functions for the magnetic field (Br, Bz) from a single coil,
+ evaluated at a set of probe locations.
+
+ This function:
+ - Evaluates the radial (Br) and vertical (Bz) magnetic field contributions
+ from all filaments in the specified coil.
+ - Computes these contributions at the chosen probe positions (default: pickup coils).
+ - Returns the Green's function matrices mapping coil filament currents
+ to magnetic field components at the probe locations.
+
+ Parameters
+ ----------
+ eq : object
+ Equilibrium object containing the tokamak configuration.
+ coil_key : str
+ Key identifying the coil within `eq.tokamak`.
+ probe : str, optional
+ Location set where the field is evaluated.
+ Currently only "pickups" is implemented (default).
+
+ Returns
+ -------
+ greens_br_coil : ndarray
+ Green's function matrix for Br at probe locations.
+ greens_bz_coil : ndarray
+ Green's function matrix for Bz at probe locations.
"""
if probe == "pickups":
pos_R = self.pickup_pos[:, 0]
pos_Z = self.pickup_pos[:, 2]
- # pol = self.coils_dict[coil_key]["polarity"][np.newaxis, :]
- # mul = self.coils_dict[coil_key]["multiplier"][np.newaxis, :]
- # greens_br_filaments = GreensBr(
- # self.coils_dict[coil_key]["coords"][0][np.newaxis, :],
- # self.coils_dict[coil_key]["coords"][1][np.newaxis, :],
- # pos_R[:, np.newaxis],
- # pos_Z[:, np.newaxis],
- # )
- # greens_br_filaments *= pol
- # greens_br_filaments *= mul
- # greens_br_coil = np.sum(greens_br_filaments, axis=1)
-
- # greens_bz_filaments = GreensBz(
- # self.coils_dict[coil_key]["coords"][0][np.newaxis, :],
- # self.coils_dict[coil_key]["coords"][1][np.newaxis, :],
- # pos_R[:, np.newaxis],
- # pos_Z[:, np.newaxis],
- # )
-
- # greens_bz_filaments *= pol
- # greens_bz_filaments *= mul
- # greens_bz_coil = np.sum(greens_bz_filaments, axis=1)
-
greens_br_coil = eq.tokamak[coil_key].controlBr(pos_R, pos_Z)
greens_bz_coil = eq.tokamak[coil_key].controlBz(pos_R, pos_Z)
@@ -362,9 +602,30 @@ def create_greens_BrBz_single_coil(self, eq, coil_key, probe="pickups"):
def greens_BrBz_all_coils(self, eq, probe="pickups"):
"""
- Create 2d array of greens functions for all coils and at all probe positions
- - array[i][j] is greens function for coil i evaluated at probe position j
- - evaluated on pickups by default, can apply to other probes too with minor modification
+ Compute Green's function matrices for (Br, Bz) contributions from all coils
+ evaluated at a set of probe locations.
+
+ This function assembles the coil-wise Green's functions into global matrices:
+ - Each row corresponds to a coil in `self.coils_dict`
+ - Each column corresponds to a probe location
+ - Entries represent the magnetic field response (Br or Bz) at a probe due to a given coil
+
+ Parameters
+ ----------
+ eq : object
+ Equilibrium object containing the tokamak configuration.
+ probe : str, optional
+ Set of probe locations where fields are evaluated.
+ Currently only "pickups" is implemented (default).
+
+ Returns
+ -------
+ array_r : ndarray
+ Global Green's function matrix for radial field component Br,
+ shape (n_coils, n_probes).
+ array_z : ndarray
+ Global Green's function matrix for vertical field component Bz,
+ shape (n_coils, n_probes).
"""
if probe == "pickups":
array_r = np.array([]).reshape(0, self.number_pickups)
@@ -379,7 +640,30 @@ def greens_BrBz_all_coils(self, eq, probe="pickups"):
def create_greens_B_oriented_coils(self, eq, probe="pickups"):
"""
- perform dot product of greens function vector with pickup coil orientation
+ Compute the directional Green's function for coils projected onto probe orientations.
+
+ This function evaluates the magnetic field Green's functions (Br, Bz) from all coils
+ and projects them onto the local orientation of each probe. This yields the component
+ of the magnetic field aligned with the probe direction.
+
+ Mathematically, for each coil i and probe j:
+ G_oriented[i, j] = Br[i, j] * or_R[j] + Bz[i, j] * or_Z[j]
+
+ where (or_R, or_Z) defines the unit orientation vector of the probe in (R, Z).
+
+ Parameters
+ ----------
+ eq : object
+ Equilibrium object containing the tokamak configuration.
+ probe : str, optional
+ Probe set at which the field is evaluated.
+ Currently only "pickups" is implemented (default).
+
+ Returns
+ -------
+ prod : ndarray
+ Oriented Green's function matrix of shape (n_coils, n_probes),
+ representing the projection of (Br, Bz) onto probe orientations.
"""
if probe == "pickups":
or_R = self.pickup_or[:, 0]
@@ -392,8 +676,34 @@ def create_greens_B_oriented_coils(self, eq, probe="pickups"):
def BrBz_coils(self, eq, probe="pickups"):
"""
- Magnetic fields from coils, radial and z components only.
- evaluated on pickup positions by default.
+ Compute magnetic field components (Br, Bz) produced by all coils
+ at a set of probe locations.
+
+ The fields are obtained by contracting precomputed Green's functions
+ with the coil current vector:
+
+ Br[j] = sum_i G_br[i, j] * I_i
+ Bz[j] = sum_i G_bz[i, j] * I_i
+
+ where:
+ - i indexes coils
+ - j indexes probe locations
+ - I_i is the current in coil i
+
+ Parameters
+ ----------
+ eq : object
+ Equilibrium object containing the tokamak configuration.
+ probe : str, optional
+ Probe set at which fields are evaluated.
+ Currently only "pickups" is implemented (default).
+
+ Returns
+ -------
+ br_coil : ndarray
+ Radial magnetic field at probe locations, shape (n_probes,).
+ bz_coil : ndarray
+ Vertical magnetic field at probe locations, shape (n_probes,).
"""
coil_currents = self.get_coil_currents(eq)[:, np.newaxis]
if probe == "pickups":
@@ -403,9 +713,40 @@ def BrBz_coils(self, eq, probe="pickups"):
def create_greens_BrBz_plasma(self, eq, probe="pickups"):
"""
- Compute greens function at probes from the plasma currents .
- - plasma current source in grid from solve. grid points contained in eq object
- - evaluated on pickups by default, can apply to other probes too with minor modification
+ Compute Green's functions for magnetic field components (Br, Bz)
+ produced by plasma current elements and evaluated at probe locations.
+
+ This constructs the plasma-to-probe coupling matrices by evaluating
+ the Biot–Savart kernel between discrete plasma grid points and probes.
+
+ The resulting Green's functions map plasma current density (at grid points)
+ to magnetic field at probes:
+
+ Br[j] = sum_i G_br[i, j] * J_i
+ Bz[j] = sum_i G_bz[i, j] * J_i
+
+ where:
+ - i indexes plasma grid / limiter-handler points
+ - j indexes probe locations
+ - J_i represents plasma current contribution at grid point i
+
+ Parameters
+ ----------
+ eq : object
+ Equilibrium object containing:
+ - plasma grid points in `eq.limiter_handler.plasma_pts`
+ probe : str, optional
+ Probe set where fields are evaluated.
+ Currently only "pickups" is implemented (default).
+
+ Returns
+ -------
+ greens_br : object / ndarray
+ Green's function for radial field component Br
+ mapping plasma grid points to probes.
+ greens_bz : object / ndarray
+ Green's function for vertical field component Bz
+ mapping plasma grid points to probes.
"""
if probe == "pickups":
pos_R = self.pickup_pos[:, 0]
@@ -432,7 +773,36 @@ def create_greens_BrBz_plasma(self, eq, probe="pickups"):
def create_greens_B_oriented_plasma(self, eq, probe="pickups"):
"""
- perform dot product of greens function vector with orientation
+ Compute the oriented Green's function for plasma current contributions
+ projected onto probe directions.
+
+ This function first computes the plasma-to-probe Green's functions
+ for the magnetic field components (Br, Bz), and then projects them
+ onto the local orientation of each probe.
+
+ The resulting kernel represents the magnetic field component aligned
+ with the probe orientation:
+
+ G_oriented[i, j] = Br[i, j] * or_R[j] + Bz[i, j] * or_Z[j]
+
+ where:
+ - i indexes plasma grid points
+ - j indexes probe locations
+ - (or_R, or_Z) is the unit orientation vector at each probe
+
+ Parameters
+ ----------
+ eq : object
+ Equilibrium object containing plasma grid information.
+ probe : str, optional
+ Probe set at which the field is evaluated.
+ Currently only "pickups" is implemented (default).
+
+ Returns
+ -------
+ prod : ndarray
+ Oriented Green's function matrix mapping plasma currents
+ to field components aligned with probe orientation.
"""
br, bz = self.create_greens_BrBz_plasma(eq)
@@ -444,7 +814,40 @@ def create_greens_B_oriented_plasma(self, eq, probe="pickups"):
def BrBz_plasma(self, eq, probe="pickups"):
"""
- Magnetic fields from plasma
+ Compute magnetic field components (Br, Bz) generated by plasma currents
+ at a set of probe locations.
+
+ This function uses precomputed (or lazily generated) Green's functions
+ mapping plasma grid current elements to probe measurements, then contracts
+ them with the plasma current distribution.
+
+ The field is computed as:
+
+ Br[j] = sum_i G_br[i, j] * I_i
+ Bz[j] = sum_i G_bz[i, j] * I_i
+
+ where:
+ - i indexes plasma grid points
+ - j indexes probe locations
+ - I_i is the plasma current at grid point i
+
+ The Green's functions are cached per equilibrium via `eq_key` to avoid
+ recomputation when the plasma grid is unchanged.
+
+ Parameters
+ ----------
+ eq : object
+ Equilibrium object containing plasma current and grid information.
+ probe : str, optional
+ Probe set at which fields are evaluated.
+ Currently only "pickups" is implemented (default).
+
+ Returns
+ -------
+ br_plasma : ndarray
+ Radial magnetic field at probe locations, shape (n_probes,).
+ bz_plasma : ndarray
+ Vertical magnetic field at probe locations, shape (n_probes,).
"""
eq_key = self.create_eq_key(eq)
plasma_current = self.get_plasma_current(eq)[:, np.newaxis]
@@ -465,9 +868,36 @@ def BrBz_plasma(self, eq, probe="pickups"):
def Br(self, eq, probe="pickups"):
"""
- Method to compute total radial magnetic field from coil and plasma
- returns array with Br at each pickup coil probe
- - evaluated on pickups by default, can apply to other probes too with minor modification
+ Compute the total radial magnetic field (Br) from both coils and plasma
+ at a set of probe locations.
+
+ The total field is the sum of contributions from discrete coils and the
+ distributed plasma current:
+
+ Br_total[j] = sum_i G_coil[i, j] * I_coil[i]
+ + sum_k G_plasma[k, j] * I_plasma[k]
+
+ where:
+ - i indexes coils
+ - k indexes plasma grid points
+ - j indexes probe locations
+
+ Green's functions for the plasma contribution are cached per equilibrium
+ using `eq_key` and computed lazily if not already available.
+
+ Parameters
+ ----------
+ eq : object
+ Equilibrium object containing coil and plasma current information,
+ as well as grid definitions.
+ probe : str, optional
+ Probe set at which the field is evaluated.
+ Currently only "pickups" is implemented (default).
+
+ Returns
+ -------
+ br_total : ndarray
+ Total radial magnetic field at probe locations, shape (n_probes,).
"""
coil_currents = self.get_coil_currents(eq)[:, np.newaxis]
plasma_current = self.get_plasma_current(eq)[:, np.newaxis]
@@ -488,9 +918,36 @@ def Br(self, eq, probe="pickups"):
def Bz(self, eq, probe="pickups"):
"""
- Method to compute total z component of magnetic field from coil and plasma
- returns array with Bz at each pickup coil probe
- - evaluated on pickups by default, can apply to other probes too with minor modification
+ Compute the total vertical magnetic field (Bz) from both coils and plasma
+ at a set of probe locations.
+
+ The total field is the sum of contributions from discrete coils and the
+ distributed plasma current:
+
+ Bz_total[j] = sum_i G_coil[i, j] * I_coil[i]
+ + sum_k G_plasma[k, j] * I_plasma[k]
+
+ where:
+ - i indexes coils
+ - k indexes plasma grid points
+ - j indexes probe locations
+
+ Plasma Green's functions are cached per equilibrium using `eq_key` and
+ are computed lazily if not already available.
+
+ Parameters
+ ----------
+ eq : object
+ Equilibrium object containing coil and plasma current information,
+ as well as grid definitions.
+ probe : str, optional
+ Probe set at which the field is evaluated.
+ Currently only "pickups" is implemented (default).
+
+ Returns
+ -------
+ bz_total : ndarray
+ Total vertical magnetic field at probe locations, shape (n_probes,).
"""
coil_currents = self.get_coil_currents(eq)[:, np.newaxis]
plasma_current = self.get_plasma_current(eq)[:, np.newaxis]
@@ -512,9 +969,32 @@ def Bz(self, eq, probe="pickups"):
def Btor(self, eq, probe="pickups"):
"""
- Probes outside of plasma therefore Btor = fvac/R
- returns array of btor for each probe position
- - evaluated on pickups by default, can apply to other probes too with minor modification
+ Compute the toroidal magnetic field (Btor) at probe locations.
+
+ The toroidal field is assumed to follow the vacuum scaling law:
+
+ Btor(R) = f_vac / R
+
+ where:
+ - f_vac is the vacuum toroidal field function from the equilibrium
+ - R is the major radius of the probe location
+
+ This approximation assumes probes are located outside the plasma,
+ where the toroidal field is purely vacuum-like.
+
+ Parameters
+ ----------
+ eq : object
+ Equilibrium object providing the vacuum toroidal field profile
+ via `eq._profiles.fvac()`.
+ probe : str, optional
+ Probe set at which the field is evaluated.
+ Currently only "pickups" is implemented (default).
+
+ Returns
+ -------
+ btor : ndarray
+ Toroidal magnetic field at probe locations, shape (n_probes,).
"""
if probe == "pickups":
pos_R = self.pickup_pos[:, 0]
@@ -522,21 +1002,38 @@ def Btor(self, eq, probe="pickups"):
btor = eq._profiles.fvac() / pos_R
return btor
- # def calculate_pickup_value_v1(self,eq,probe = 'pickups'):
- # """
- # OLD VERSION
- # Method to compute and return B.n, for pickup coils
- # """
- # orientation = self.pickup_or.transpose()
- # Btotal = np.vstack((self.Br(eq),self.Btor(eq,probe),self.Bz(eq)))
+ def calculate_pickup_value(self, eq, probe="pickups"):
+ """
+ Compute the magnetic field projection (B · n) at pickup probes.
- # BdotN = np.sum(orientation*Btotal, axis = 0)
+ This function evaluates the total magnetic field at pickup locations
+ and projects it onto the local probe orientation vector n.
- # return BdotN
+ The total pickup signal is composed of three contributions:
- def calculate_pickup_value(self, eq, probe="pickups"):
- """
- Compute B.n at pickup probes, using oriented greens functions.
+ (B · n)_j = (B_coil · n)_j + (B_plasma · n)_j + (B_tor · n)_j
+
+ where:
+ - j indexes pickup probes
+ - coil and plasma terms are computed using precomputed oriented
+ Green's functions
+ - the toroidal field contribution is projected via the pickup
+ orientation component in the toroidal direction
+
+ Parameters
+ ----------
+ eq : object
+ Equilibrium object containing coil currents, plasma currents,
+ and magnetic field profiles.
+ probe : str, optional
+ Probe set at which the signal is evaluated.
+ Currently only "pickups" is implemented (default).
+
+ Returns
+ -------
+ signal : ndarray
+ Scalar pickup signal (B · n) at each probe location,
+ shape (n_probes,).
"""
coil_current = self.get_coil_currents(eq)[:, np.newaxis]
plasma_current = self.get_plasma_current(eq)[:, np.newaxis]
@@ -563,18 +1060,27 @@ def calculate_pickup_value(self, eq, probe="pickups"):
def plot(self, axis=None, show=True, floops=True, pickups=True, pickups_scale=0.05):
"""
- Plot the magnetic probes.
+ Plot magnetic diagnostic probes (fluxloops and pickup coils).
+
+ This is a convenience wrapper around `freegs4e.plotting.plotProbes`.
- axis - Specify the axis on which to plot
- show - Call matplotlib.pyplot.show() before returning
- floops - Plot the fluxloops
- pickups - Plot the pickup coils
+ Parameters
+ ----------
+ axis : matplotlib.axes.Axes, optional
+ Matplotlib axis to draw on. If None, a new figure/axis is created.
+ show : bool, optional
+ If True, calls `matplotlib.pyplot.show()` before returning.
+ floops : bool, optional
+ If True, plots fluxloop diagnostics.
+ pickups : bool, optional
+ If True, plots pickup coil diagnostics.
+ pickups_scale : float, optional
+ Scaling factor for visual representation of pickup coil orientation.
Returns
-------
-
- axis object from Matplotlib
-
+ axis : matplotlib.axes.Axes
+ Matplotlib axis containing the plotted probes.
"""
from freegs4e.plotting import plotProbes
diff --git a/freegsnke/mastu_tools.py b/freegsnke/mastu_tools.py
index a100eb5..559f4e4 100644
--- a/freegsnke/mastu_tools.py
+++ b/freegsnke/mastu_tools.py
@@ -1641,6 +1641,39 @@ def load_currents_voltages_and_TS_signals(
coil_list = np.sort(list(att_dict.keys()))
def get_voltages(shotn):
+ """
+ Retrieve coil voltage signals for a given shot number.
+
+ This function queries an external data source for voltage signals
+ associated with a set of actuators defined in `att_dict`. Depending
+ on the actuator type, it fetches either PS coil voltages or generic
+ output voltages.
+
+ The returned data is organised into a dictionary keyed by actuator name.
+
+ Parameters
+ ----------
+ shotn : int
+ Shot number identifying the discharge for which voltage data is
+ requested.
+
+ Returns
+ -------
+ outdict : dict
+ Dictionary of voltage signals. Each key corresponds to an actuator
+ name and maps to a dictionary with:
+
+ - data : ndarray
+ Voltage time series
+ - times : ndarray
+ Time base associated with the signal
+ - units : str
+ Physical units of the voltage signal
+
+ Notes
+ -----
+ Missing signals are silently ignored.
+ """
outdict = {}
for attk in att_dict:
tinner = att_dict[attk]
@@ -1659,6 +1692,39 @@ def get_voltages(shotn):
return outdict
def get_req_voltages(shotn):
+ """
+ Retrieve requested (reference/commanded) coil voltages for a given shot.
+
+ This function queries an external data source for the prescribed voltage
+ waveforms associated with each actuator defined in `att_dict`.
+
+ If a signal cannot be retrieved, a default placeholder signal is returned
+ and a warning message is printed.
+
+ Parameters
+ ----------
+ shotn : int
+ Shot number identifying the discharge for which voltage data is
+ requested.
+
+ Returns
+ -------
+ outdict : dict
+ Dictionary of requested voltage signals. Each key corresponds to a
+ coil name and maps to a dictionary with:
+
+ - data : ndarray
+ Requested voltage time series (or placeholder if missing)
+ - times : ndarray
+ Time base associated with the signal (or placeholder if missing)
+ - units : str
+ Physical units of the signal
+
+ Notes
+ -----
+ Missing signals are replaced with a default placeholder array and a
+ warning is printed to standard output.
+ """
outdict = {}
for coil in att_dict:
tinner = att_dict[coil]
@@ -1679,6 +1745,36 @@ def get_req_voltages(shotn):
return outdict
def get_coilcurrs_AMC(shotn):
+ """
+ Retrieve AMC Rogowski coil current measurements for a given shot.
+
+ This function queries an external data source for measured coil currents
+ from AMC Rogowski sensors. Data is organised hierarchically by actuator
+ group and individual coil channel.
+
+ Parameters
+ ----------
+ shotn : int
+ Shot number identifying the discharge for which current data is
+ requested.
+
+ Returns
+ -------
+ outdict : dict
+ Nested dictionary of coil current signals structured as:
+
+ outdict[actuator_group][coil_name] -> dict with:
+ - data : ndarray
+ Measured current time series
+ - times : ndarray
+ Time base associated with the signal
+ - units : str
+ Physical units of the measurement
+
+ Notes
+ -----
+ Missing coil signals are silently ignored.
+ """
outdict = {}
for attk in att_dict:
outdict[attk] = {}
@@ -1696,6 +1792,39 @@ def get_coilcurrs_AMC(shotn):
return outdict
def get_coilcurrs(shotn):
+ """
+ Retrieve measured coil currents for a given shot.
+
+ This function queries an external data source for coil current measurements,
+ selecting the appropriate signal path depending on actuator type.
+
+ For PS coils, currents are retrieved from the standard CURRENT signal.
+ For non-PS coils, the signal path depends on the actuator configuration,
+ with a special case for coil 'P6'.
+
+ Parameters
+ ----------
+ shotn : int
+ Shot number identifying the discharge for which coil current data is
+ requested.
+
+ Returns
+ -------
+ outdict : dict
+ Dictionary of coil current signals. Each key corresponds to a coil
+ name and maps to a dictionary containing:
+
+ - data : ndarray
+ Current time series
+ - times : ndarray
+ Time base associated with the signal
+ - units : str
+ Physical units of the signal
+
+ Notes
+ -----
+ Missing signals are silently ignored.
+ """
outdict = {}
for attk in att_dict:
tinner = att_dict[attk]
@@ -1718,6 +1847,39 @@ def get_coilcurrs(shotn):
return outdict
def get_rogs(shotn):
+ """
+ Retrieve Rogowski coil measurements (external and internal) for a given shot.
+
+ This function queries AMC Rogowski diagnostics and returns both external
+ (ROGEXT) and internal (ROGINT) signals where available. Data is organised
+ hierarchically by actuator group and Rogowski channel.
+
+ Parameters
+ ----------
+ shotn : int
+ Shot number identifying the discharge for which Rogowski data is
+ requested.
+
+ Returns
+ -------
+ outdict : dict
+ Nested dictionary containing Rogowski coil data:
+
+ outdict[actuator_group]["rogext"][channel] -> dict
+ External Rogowski signals:
+ - data : ndarray
+ - times : ndarray
+
+ outdict[actuator_group]["rogint"][channel] -> dict (if available)
+ Internal Rogowski signals:
+ - data : ndarray
+ - times : ndarray
+
+ Notes
+ -----
+ Missing channels are silently ignored. Internal Rogowski data is only
+ included if available for the given actuator.
+ """
outdict = {}
for attk in att_dict:
outdict[attk] = {}
@@ -2335,16 +2497,28 @@ def get_element_vertices(
return [rr, zz, dR, dZ]
-# ------------
def find_strikepoints(R, Z, psi, psi_boundary, wall):
"""
- This function can be used to find the strikepoints of an equilibrium using
- the:
- - R and Z grids (2D)
- - psi_total (2D) (i.e. the poloidal flux map)
- - psi_boundary (single value)
- - limiter/wall coordinates (N x 2)
+ Find the strikepoints of an equilibrium with the wall.
+
+ Parameters
+ ----------
+ R : np.ndarray (nx, nz)
+ 2D array of major radius coordinates [m]
+ Z : np.ndarray (nx, nz)
+ 2D array of vertical coordinates [m]
+ psi : np.ndarray (nx, nz)
+ 2D poloidal flux map [Wb]
+ psi_boundary : float
+ Value of psi at the plasma boundary [Wb]
+ wall : np.ndarray (N, 2)
+ Wall/limiter coordinates as (R, Z) pairs [m]
+ Returns
+ -------
+ np.ndarray or None
+ Array of strikepoint coordinates, shape (M, 2), or None if no
+ intersections are found.
"""
# find contour object for psi_boundary
@@ -2383,20 +2557,39 @@ def find_strikepoints(R, Z, psi, psi_boundary, wall):
def Separatrix(R, Z, psi, ntheta, psival=1.0, theta_grid=None, input_opoint=None):
- """Find the R, Z coordinates of the separatrix for equilbrium
- eq. Returns a tuple of (R, Z, R_X, Z_X), where R_X, Z_X are the
- coordinates of the X-point on the separatrix. Points are equally
- spaced in geometric poloidal angle.
-
- If opoint, xpoint or psi are not given, they are calculated from eq
-
- eq - Equilibrium object
- opoint - List of O-point tuples of (R, Z, psi)
- xpoint - List of X-point tuples of (R, Z, psi)
- ntheta - Number of points to find
- psi - Grid of psi on (R, Z)
- axis - A matplotlib axis object to plot points on
- input_opoint - a user-chosen magnetic axis from which to find separatrix
+ """
+ Compute separatrix coordinates for a given equilibrium flux map.
+
+ This function traces the ψ = const surface corresponding to the separatrix
+ (default psival = 1.0 in normalized coordinates) starting from the magnetic
+ axis and marching outwards in poloidal angle.
+
+ The separatrix is sampled at equally spaced geometric poloidal angles,
+ with care taken to avoid sampling exactly at the X-point locations.
+
+ Parameters
+ ----------
+ R : ndarray
+ Radial grid coordinates (2D array).
+ Z : ndarray
+ Vertical grid coordinates (2D array).
+ psi : ndarray
+ Poloidal flux on the (R, Z) grid.
+ ntheta : int
+ Number of poloidal angle samples along the separatrix.
+ psival : float, optional
+ Target normalized flux value defining the separatrix.
+ theta_grid : ndarray, optional
+ User-specified poloidal angle grid. If None, a uniform grid is used.
+ input_opoint : tuple, optional
+ User-specified magnetic axis (R0, Z0). If None, it is inferred.
+
+ Returns
+ -------
+ points : ndarray
+ Array of separatrix points (R, Z) sampled along θ.
+ theta_grid : ndarray
+ Poloidal angle grid used for the construction.
"""
opoint, xpoint = critical.find_critical(R, Z, psi)
@@ -2457,11 +2650,39 @@ def Separatrix(R, Z, psi, ntheta, psival=1.0, theta_grid=None, input_opoint=None
def find_psisurface(psifunc, R, Z, r0, z0, r1, z1, psival=1.0, n=100):
"""
- eq - Equilibrium object
- (r0,z0) - Start location inside separatrix
- (r1,z1) - Location outside separatrix
+ Find an intersection point of a ψ = const surface along a straight line.
+
+ This routine samples a straight line from an initial point inside the
+ separatrix to a point outside it, evaluates the flux along the line, and
+ interpolates to locate the position where ψ equals the target value.
+
+ Parameters
+ ----------
+ psifunc : callable
+ Interpolated flux function ψ(R, Z).
+ R : ndarray
+ Radial grid (used for domain clipping).
+ Z : ndarray
+ Vertical grid (used for domain clipping).
+ r0 : float
+ Starting radial coordinate (assumed inside target surface).
+ z0 : float
+ Starting vertical coordinate (assumed inside target surface).
+ r1 : float
+ Ending radial coordinate (outside target surface).
+ z1 : float
+ Ending vertical coordinate (outside target surface).
+ psival : float, optional
+ Target flux value defining the surface.
+ n : int, optional
+ Number of points sampled along the line.
- n - Number of starting points to use
+ Returns
+ -------
+ r : float
+ Radial coordinate of ψ = psival intersection.
+ z : float
+ Vertical coordinate of ψ = psival intersection.
"""
# Clip (r1,z1) to be inside domain
# Shorten the line so that the direction is unchanged
@@ -2500,8 +2721,22 @@ def find_psisurface(psifunc, R, Z, r0, z0, r1, z1, psival=1.0, n=100):
def max_euclidean_distance(points1, points2):
"""
- Calculate the maximum Euclidean distance between corresponding points in two sets.
- Exclude points with 'None' values.
+ Compute the maximum Euclidean distance between corresponding points in two sets.
+
+ Points containing NaN values are excluded before the distance calculation.
+
+ Parameters
+ ----------
+ points1 : ndarray
+ First set of (x, y) points.
+ points2 : ndarray
+ Second set of (x, y) points with the same shape as points1.
+
+ Returns
+ -------
+ float
+ Maximum Euclidean distance between valid corresponding points.
+ Returns NaN if no valid points remain.
"""
valid_indices = np.logical_not(
np.any(np.isnan(points1), axis=1) | np.any(np.isnan(points2), axis=1)
@@ -2515,8 +2750,22 @@ def max_euclidean_distance(points1, points2):
def median_euclidean_distance(points1, points2):
"""
- Calculate the maximum Euclidean distance between corresponding points in two sets.
- Exclude points with 'None' values.
+ Compute the median Euclidean distance between corresponding points in two sets.
+
+ Points containing NaN values are excluded before computing distances.
+
+ Parameters
+ ----------
+ points1 : ndarray
+ First set of (x, y) points.
+ points2 : ndarray
+ Second set of (x, y) points with the same shape as points1.
+
+ Returns
+ -------
+ float
+ Median Euclidean distance between valid corresponding points.
+ Returns NaN if no valid points remain.
"""
valid_indices = np.logical_not(
np.any(np.isnan(points1), axis=1) | np.any(np.isnan(points2), axis=1)
@@ -2530,11 +2779,28 @@ def median_euclidean_distance(points1, points2):
def separatrix_areas(separatrix_1, separatrix_2):
"""
- This function can be used to find the poloidal area of each separatrix given
- and the similiarity of the two as calculated using equation 8 in Bardsely et al
- 2024 (Nuclear Fusion) ("Decoupled magnetic control of spherical tokamak
- divertors via vacuum harmonic constraints"). Inputs:
- - separatrix_1 (and 2): np.array of (n x 2) (r,z) points
+ Compute a geometric similarity metric between two separatrix shapes.
+
+ This function constructs convex hull polygons from two sets of (R, Z)
+ points, then compares their overlap using union and intersection areas.
+ The resulting metric η is based on the non-overlapping area relative
+ to the total area, as defined in Bardsley et al. (2024, Nuclear Fusion).
+
+ Parameters
+ ----------
+ separatrix_1 : ndarray
+ Array of shape (N, 2) containing (R, Z) points of the first separatrix.
+ separatrix_2 : ndarray
+ Array of shape (M, 2) containing (R, Z) points of the second separatrix.
+
+ Returns
+ -------
+ eta : float
+ Normalised non-overlap metric between the two separatrices.
+ polygon1 : shapely.geometry.Polygon
+ Convex hull polygon of separatrix_1.
+ polygon2 : shapely.geometry.Polygon
+ Convex hull polygon of separatrix_2.
"""
# create Polygon objects from the points using Shapely package
@@ -2562,28 +2828,30 @@ def interpolate_data(
order=5,
):
"""
- This function will interpolate the 'data' and 'times' using an nth order polynomial,
- where n='order'. It can take optional arguments that will only use the times/data within
- a certain time window [t_start, t_final].
+ Fit a polynomial to time-series data and return an evaluatable interpolant.
+
+ This function selects a time window (if provided), fits a polynomial of
+ specified order to the data within that window, and returns the resulting
+ polynomial object. The output can be evaluated on scalar or array inputs.
Parameters
----------
- times : np.array
- Times at which 'data' are recorded.
- data : np.array
- Data to be interpolated at 'times'.
- t_start : float
- Start of time window to interpolate over.
- t_final : float
- End of time window to interpolate over.
- order : int
- Non-negative integer value for polynomial order to use in interpolation.
+ times : ndarray
+ Array of time values corresponding to the data samples.
+ data : ndarray
+ Data values sampled at the given times.
+ t_start : float, optional
+ Start of the time window. If None, the full range is used from the beginning.
+ t_final : float, optional
+ End of the time window. If None, the full range is used up to the end.
+ order : int, optional
+ Degree of the polynomial used for fitting.
Returns
-------
- function
- Returns a function that represents the interpolated polynomial. It can
- be called with a float or np.array input.
+ numpy.polynomial.Polynomial
+ Polynomial object representing the fitted interpolation, callable on
+ scalar or array inputs.
"""
# find closest time indices
diff --git a/freegsnke/nonlinear_solve.py b/freegsnke/nonlinear_solve.py
index ced7196..4540b93 100644
--- a/freegsnke/nonlinear_solve.py
+++ b/freegsnke/nonlinear_solve.py
@@ -1589,6 +1589,7 @@ def build_linearization(
)
else:
self.final_dI_record[j] = 1.0 * self.starting_dI[j]
+ rel_ndIy = ndIy
if verbose:
print("")
diff --git a/freegsnke/passive_structure.py b/freegsnke/passive_structure.py
index 0a3788f..8d9c559 100644
--- a/freegsnke/passive_structure.py
+++ b/freegsnke/passive_structure.py
@@ -84,6 +84,31 @@ def __init__(
self.greens = {}
def copy(self):
+ """
+ Create a shallow copy of the control object without reinitialising
+ geometry or recomputing Green's functions.
+
+ This method avoids calling the constructor for performance reasons and
+ instead manually copies attributes that define the control geometry and
+ cached electromagnetic response.
+
+ Returns
+ -------
+ object
+ A new instance of the same class with duplicated state.
+
+ Notes
+ -----
+ - Geometry-related arrays (e.g. `R`, `Z`, `vertices`, `filaments`) are
+ shared by reference and are assumed to be immutable.
+ - The `greens` dictionary is shallow-copied:
+ - Safe usage: replacing entries (e.g. `greens["psi"] = new_array`)
+ - Unsafe usage: in-place modification (e.g. `greens["psi"][:] = ...`)
+ - This copy is therefore *not fully independent* and should be used only
+ when immutability assumptions hold.
+ - The attribute `control` is currently assigned from `self.turns`
+ (this may be intentional or a likely typo depending on design intent).
+ """
# dont instantiate the new object, it will be slow
new_obj = type(self).__new__(type(self))
@@ -115,15 +140,25 @@ def copy(self):
def create_RZ_key(self, R, Z):
"""
- Produces tuple (Rmin,Rmax,Zmin,Zmax,nx,ny) to access correct dictionary entry of greens function.
+ Create a hashable key identifying a specific R–Z grid for caching Green's functions.
+
+ The key is based on the grid bounds and size, and is used to ensure that
+ Green's function evaluations are reused only when the underlying spatial
+ grid is identical.
Parameters
----------
- R : array
- eq.R, radial coordinate on the domain grid
- Z : array
- eq.Z, radial coordinate on the domain grid
-
+ R : np.ndarray
+ Radial coordinate grid (typically `eq.R`).
+ Z : np.ndarray
+ Vertical coordinate grid (typically `eq.Z`).
+
+ Returns
+ -------
+ tuple
+ A hashable identifier of the form:
+ (R_min, R_max, Z_min, Z_max, N),
+ where N = total number of grid points in `R`.
"""
RZ_key = (np.min(R), np.max(R), np.min(Z), np.max(Z), np.size(R))
return RZ_key
@@ -131,7 +166,28 @@ def create_RZ_key(self, R, Z):
def build_refining_filaments(
self,
):
- """Builds the grid used for the refinement"""
+ """
+ Construct the set of refining filaments used to discretise the control
+ region at higher resolution.
+
+ This routine generates a refined filament representation of the control
+ geometry, typically used to improve spatial resolution of control
+ response calculations.
+
+ Returns
+ -------
+ np.ndarray
+ Array of filament coordinates generated by the refinement procedure.
+
+ Notes
+ -----
+ - The filaments are generated from the control polygon defined by
+ `self.Rpolygon` and `self.Zpolygon`.
+ - The refinement density is controlled by `self.n_refine`.
+ - The refinement strategy is determined by `self.refine_mode`.
+ - The function `generate_refinement` also returns an area estimate,
+ which is currently ignored.
+ """
filaments, area = generate_refinement(
self.Rpolygon, self.Zpolygon, self.n_refine, self.refine_mode
@@ -139,14 +195,30 @@ def build_refining_filaments(
return filaments
def build_control_psi(self, R, Z):
- """Builds controlPsi for a new set of R, Z grids.
+ """
+ Compute and cache the Green's function for the poloidal flux (ψ)
+ induced by the control filaments on a specified R–Z grid.
+
+ The result is obtained by evaluating the contribution from each filament
+ (via `Greens`) and averaging across all filaments.
Parameters
----------
- R : array
- Grid on which to calculate the greens, i.e. eq.R
- Z : array
- Grid on which to calculate the greens, i.e. eq.Z
+ R : np.ndarray
+ 2D array defining the radial coordinate grid (e.g. equilibrium R grid).
+ Z : np.ndarray
+ 2D array defining the vertical coordinate grid (e.g. equilibrium Z grid).
+
+ Notes
+ -----
+ - Each filament is evaluated over the full grid using broadcasting.
+ - Filament contributions are averaged to produce the total response.
+ - The result is cached in `self.greens[(R, Z)]["psi"]`.
+
+ Returns
+ -------
+ None
+ The computed Green's function is stored internally.
"""
greens_psi = Greens(
@@ -164,14 +236,30 @@ def build_control_psi(self, R, Z):
self.greens[RZ_key] = {"psi": greens_psi}
def build_control_br(self, R, Z):
- """Builds controlBr for a new set of R, Z grids.
+ """
+ Compute and cache the Green's function for the radial magnetic field (Br)
+ induced by the control filaments on a specified R–Z grid.
+
+ The field is obtained by evaluating the Biot–Savart contribution from each
+ filament (via `GreensBr`) and averaging over all filaments.
Parameters
----------
- R : array
- Grid on which to calculate the greens, i.e. eq.R
- Z : array
- Grid on which to calculate the greens, i.e. eq.Z
+ R : np.ndarray
+ 2D array defining the radial coordinate grid (e.g. equilibrium R grid).
+ Z : np.ndarray
+ 2D array defining the vertical coordinate grid (e.g. equilibrium Z grid).
+
+ Notes
+ -----
+ - Each filament is evaluated on the full R–Z grid using broadcasting.
+ - Contributions are averaged across filaments to form the final field.
+ - The result is cached in `self.greens[(R, Z)]["Br"]` to avoid recomputation.
+
+ Returns
+ -------
+ None
+ The computed Green's function is stored in `self.greens`.
"""
greens_br = GreensBr(
@@ -189,14 +277,31 @@ def build_control_br(self, R, Z):
self.greens[RZ_key] = {"Br": greens_br}
def build_control_bz(self, R, Z):
- """Builds controlBz for a new set of R, Z grids.
+ """
+ Compute and cache the Green's function for the vertical magnetic field (Bz)
+ induced by the control filaments on a specified R–Z grid.
+
+ The result is computed by averaging the filament-wise Biot–Savart
+ contributions and stored in the internal `self.greens` cache using a
+ grid-dependent key.
Parameters
----------
- R : array
- Grid on which to calculate the greens, i.e. eq.R
- Z : array
- Grid on which to calculate the greens, i.e. eq.Z
+ R : np.ndarray
+ 2D array defining the radial coordinate grid (e.g. equilibrium R grid).
+ Z : np.ndarray
+ 2D array defining the vertical coordinate grid (e.g. equilibrium Z grid).
+
+ Notes
+ -----
+ - The computation is performed using `GreensBz` evaluated for each filament.
+ - Filament contributions are averaged to produce the final field.
+ - Results are cached in `self.greens[(R, Z)]["Bz"]` to avoid recomputation.
+
+ Returns
+ -------
+ None
+ The result is stored internally in `self.greens`.
"""
greens_bz = GreensBz(
@@ -215,8 +320,23 @@ def build_control_bz(self, R, Z):
def controlPsi(self, R, Z):
"""
- Retrieve poloidal flux at (R,Z) due to a unit current
- or calculate where necessary.
+ Return the poloidal flux ψ at a given observation point
+ due to a unit current in the control element.
+
+ The value is retrieved from a cached Green's function table if available;
+ otherwise it is computed and stored.
+
+ Parameters
+ ----------
+ R : float
+ Major radius coordinate of the evaluation point.
+ Z : float
+ Vertical coordinate of the evaluation point.
+
+ Returns
+ -------
+ np.ndarray or float
+ Green's function contribution to ψ at (R, Z) for unit current.
"""
RZ_key = self.create_RZ_key(R, Z)
@@ -229,8 +349,32 @@ def controlPsi(self, R, Z):
def controlBr(self, R, Z):
"""
- Retrieve Br at (R,Z) due to a unit current
- or calculate where necessary.
+ Retrieve the radial magnetic field component (Br) at a given
+ point (R, Z) due to a unit current source.
+
+ This method accesses a cached Green’s function if available.
+ If the requested value is not already cached, it is computed
+ on demand and stored for future reuse.
+
+ Parameters
+ ----------
+ R : float
+ Radial coordinate of the evaluation point.
+ Z : float
+ Vertical coordinate of the evaluation point.
+
+ Returns
+ -------
+ float or np.ndarray
+ Radial magnetic field component Br at (R, Z) due to a unit current.
+
+ Notes
+ -----
+ - Values are stored in `self.greens` using a key derived from (R, Z).
+ - If the entry is missing, `self.build_control_br(R, Z)` is called
+ to populate the cache.
+ - Repeated queries at the same location are retrieved from cache,
+ avoiding recomputation.
"""
RZ_key = self.create_RZ_key(R, Z)
@@ -243,8 +387,31 @@ def controlBr(self, R, Z):
def controlBz(self, R, Z):
"""
- Retrieve Bz at (R,Z) due to a unit current
- or calculate where necessary.
+ Retrieve the vertical magnetic field component (Bz) at a given
+ point (R, Z) due to a unit current source.
+
+ The method uses a cached Green’s function lookup when available.
+ If the requested value is not present in the cache, it is computed
+ on demand and stored for future use.
+
+ Parameters
+ ----------
+ R : float
+ Radial coordinate of evaluation point.
+ Z : float
+ Vertical coordinate of evaluation point.
+
+ Returns
+ -------
+ float or np.ndarray
+ Bz field contribution at (R, Z) due to a unit current source.
+
+ Notes
+ -----
+ - Results are cached in `self.greens` indexed by a key derived from (R, Z).
+ - If the entry is missing, `self.build_control_bz(R, Z)` is called,
+ which populates the cache.
+ - Subsequent calls with the same (R, Z) avoid recomputation.
"""
RZ_key = self.create_RZ_key(R, Z)
@@ -256,7 +423,32 @@ def controlBz(self, R, Z):
return greens_
def plot(self, axis=None, show=False):
- """Plot the passive structure polygon"""
+ """
+ Plot the passive structure polygon on a Matplotlib axis.
+
+ This method constructs a `matplotlib.patches.Polygon` from the
+ stored vertex coordinates and adds it to the provided axis. If no
+ axis is supplied, a new Matplotlib figure and axis are created.
+
+ Parameters
+ ----------
+ axis : matplotlib.axes.Axes, optional
+ Existing Matplotlib axis to plot onto. If None, a new figure
+ and axis are created.
+ show : bool, optional
+ Unused. Present for API compatibility; does not affect behaviour.
+
+ Returns
+ -------
+ matplotlib.axes.Axes
+ Axis containing the plotted polygon.
+
+ Notes
+ -----
+ - The polygon is also stored as `self.polygon` for later access.
+ - The polygon is rendered with fixed styling (grey fill, black edge).
+ - This method does not call `plt.show()` even if `show=True`.
+ """
if axis is None:
fig = plt.figure()
diff --git a/freegsnke/refine_passive.py b/freegsnke/refine_passive.py
index 78b00d5..3e8d646 100644
--- a/freegsnke/refine_passive.py
+++ b/freegsnke/refine_passive.py
@@ -1,5 +1,5 @@
"""
-Defines some of the functionality needed by the FreeGSNKE passive_structure object.
+Defines functionality needed for building passive conductive structures.
Copyright 2025 UKAEA, UKRI-STFC, and The Authors, as per the COPYRIGHT and README files.
@@ -19,43 +19,75 @@
along with FreeGSNKE. If not, see .
"""
+import numpy as np
from matplotlib.path import Path
from scipy.stats.qmc import LatinHypercube
-# fixing the seed for reproducibility purposes
+# Latin hypercube sampling engine (fixed seed for reproducibility)
engine = LatinHypercube(d=2, seed=42)
-import numpy as np
+def generate_refinement(R, Z, n_refine, refine_mode):
+ """
+ Generate a refined set of points in (R, Z) space using a selected strategy.
+ The refinement strategy is selected via `refine_mode` and delegates to
+ the corresponding refinement routine.
+
+ Parameters
+ ----------
+ R : ndarray
+ Radial coordinates of the input grid or points.
+ Z : ndarray
+ Vertical coordinates of the input grid or points.
+ n_refine : int
+ Number of refined points (or refinement resolution parameter,
+ depending on the selected method).
+ refine_mode : str
+ Refinement strategy selector:
+ - "G" : Gaussian-based refinement
+ - "LH" : Latin Hypercube refinement
+
+ Returns
+ -------
+ ndarray
+ Refined (R, Z) point set produced by the selected method.
+ """
-def generate_refinement(R, Z, n_refine, refine_mode):
if refine_mode == "G":
return generate_refinement_G(R, Z, n_refine)
elif refine_mode == "LH":
return generate_refinement_LH(R, Z, n_refine)
else:
- print("refinement mode not recognised!, please use G or LH.")
+ print("refine_mode not recognised!, please use G or LH.")
def generate_refinement_LH(R, Z, n_refine):
- """Uses a latine hypercube to fill the shape defined by the input vertices R, Z
- with exactly n_refine points.
+ """
+ Generate refinement points inside a polygon using Latin Hypercube sampling.
+
+ This method samples candidate points using a Latin Hypercube strategy,
+ maps them into the bounding box of the polygon defined by (R, Z), and
+ retains only points lying inside the polygon.
+
+ Sampling is repeated until at least `n_refine` valid interior points
+ are obtained (or a maximum iteration limit is reached).
Parameters
----------
R : array
- R coordinates of the vertices
+ R coordinates of the polygon vertices.
Z : array
- Z coordinates of the vertices
+ Z coordinates of the polygon vertices.
n_refine : int
- Number of refining points generated
+ Number of interior refinement points to generate.
Returns
-------
- array
- refining points
-
+ points : ndarray, shape (n_refine, 2)
+ Generated refinement points inside the polygon.
+ area : float
+ Estimated area of the polygon region returned by `find_area`.
"""
area, path, vmin, vmax, dv, meanR, meanZ = find_area(R, Z, n_refine)
@@ -73,22 +105,29 @@ def generate_refinement_LH(R, Z, n_refine):
def generate_refinement_G(R, Z, n_refine):
- """Generates a regular square grid refinement, so to include approximately
- n_refine points in the shape with vertices R,Z
+ """
+ Generate a structured grid refinement inside a polygon.
+
+ This method constructs a regular Cartesian grid over the bounding box
+ of the polygon defined by (R, Z), and retains only the points that lie
+ inside the polygon. The grid resolution is increased iteratively until
+ at least `n_refine` interior points are obtained.
Parameters
----------
R : array
- R coordinates of the vertices
+ R coordinates of the polygon vertices.
Z : array
- Z coordinates of the vertices
+ Z coordinates of the polygon vertices.
n_refine : int
- Number of desired refining points
+ Target number of refinement points.
Returns
-------
- array
- refining points
+ points : ndarray, shape (N, 2)
+ Refinement points inside the polygon (N ≥ n_refine).
+ area : float
+ Estimated area of the polygon region returned by `find_area`.
"""
area, path, vmin, vmax, dv, meanR, meanZ = find_area(R, Z, n_refine)
@@ -122,16 +161,39 @@ def generate_refinement_G(R, Z, n_refine):
def find_area(R, Z, n_refine):
- """Finds area inside polygon and builds the path.
+ """
+ Estimate polygon area and construct a point-in-polygon test path.
+
+ This function computes a bounding-box-based Monte Carlo estimate of the
+ area enclosed by a polygon defined by vertices (R, Z). It also constructs
+ a `matplotlib.path.Path` object for point-in-polygon queries and estimates
+ a centroid-like mean of accepted sample points.
Parameters
----------
R : array
- R coordinates of the vertices
+ R coordinates of the polygon vertices.
Z : array
- Z coordinates of the vertices
+ Z coordinates of the polygon vertices.
n_refine : int
- Number of desired refining points
+ Target number of refinement points used to control Monte Carlo sampling.
+
+ Returns
+ -------
+ area : float
+ Estimated area of the polygon.
+ path : matplotlib.path.Path
+ Path object for point-in-polygon evaluation.
+ vmin : ndarray
+ Minimum bounds of the vertex coordinates (R_min, Z_min).
+ vmax : ndarray
+ Maximum bounds of the vertex coordinates (R_max, Z_max).
+ dv : ndarray
+ Bounding box size (vmax - vmin).
+ meanR : float
+ Mean R coordinate of accepted Monte Carlo samples.
+ meanZ : float
+ Mean Z coordinate of accepted Monte Carlo samples.
"""
if n_refine is None:
n_refine = 100
diff --git a/freegsnke/switch_profile.py b/freegsnke/switch_profile.py
index f13e359..1cdff3b 100644
--- a/freegsnke/switch_profile.py
+++ b/freegsnke/switch_profile.py
@@ -1,6 +1,6 @@
"""
Implements some functionality needed by the FreeGSNKE profile object to find optimised coefficients
-when switching between profile parametrizations.
+when switching between different profile parametrisations.
Copyright 2025 UKAEA, UKRI-STFC, and The Authors, as per the COPYRIGHT and README files.
@@ -33,37 +33,51 @@ def Lao_parameters_finder(
beta_logic=True,
Ip_logic=True,
):
- """Finds best fitting alpha, beta parameters for a Lao85 profile,
- to reproduce the input pprime_ and ffprime_
- n_alpha and n_beta represent the number of free parameters
- Simple linear fit.
+ """
+ Fit Lao85 profile coefficients to prescribed pprime and ffprime profiles.
+
+ The profiles are represented as polynomial expansions in the normalised
+ flux coordinate ``pn_``. The coefficients are obtained via a linear
+ least-squares fit.
Parameters
----------
- pn_ : np.array
- Normalised plasma function, array of values btw 0 and 1.
- pprime_ : np.array
- pprime(pn_)
- ffprime_ : np.array
- ffprime(pn_)
+ pn_ : ndarray
+ Normalised flux coordinate values in the interval [0, 1].
+ pprime_ : ndarray
+ Values of pprime evaluated at ``pn_``.
+ ffprime_ : ndarray
+ Values of ffprime evaluated at ``pn_``.
n_alpha : int
- Number of free parameters for the pprime term
+ Number of free coefficients in the pprime expansion.
n_beta : int
- Number of free parameters for the ffprime term
+ Number of free coefficients in the ffprime expansion.
alpha_logic : bool, optional
- add polynomial term in Lao85 profile so that ppprime(1)=0, by default True
+ If True, include the Lao85 constraint term so that the fitted
+ pprime profile satisfies the boundary condition at ``pn = 1``.
+ Default is True.
beta_logic : bool, optional
- add polynomial term so that ffpprime(1)=0,, by default True
+ If True, include the Lao85 constraint term so that the fitted
+ ffprime profile satisfies the boundary condition at ``pn = 1``.
+ Default is True.
Ip_logic : bool, optional
- if False, scale coefficients, by default True
- nn : int, optional
- number of points to be used for fit, by default 100
+ If False, rescale the fitted coefficients using the original
+ profile amplitudes. Default is True.
Returns
-------
- np.array, np.array
- optimal alpha and beta parameters
- Note these need to be used with the same alpha and beta logic terms provided as inputs.
+ alpha : ndarray
+ Fitted coefficients for the pprime profile.
+ beta : ndarray
+ Fitted coefficients for the ffprime profile.
+
+ Notes
+ -----
+ The returned coefficients must be used with the same values of
+ ``alpha_logic`` and ``beta_logic`` that were used during fitting.
+
+ This routine normalises ``pprime_`` and ``ffprime_`` by their values
+ at ``pn_[0]`` before fitting.
"""
pprime0_ = pprime_[0]
@@ -92,15 +106,77 @@ def Lao_parameters_finder(
return alpha, beta
-# Functions for the Lao85.Topeol_parameters optimiser
-# Used to find best fitting set of parameters for a Topeol profile
-# to reproduce the input pprime_ and ffprime_ of a Lao85 profile
-# Non-linear optimization
def Topeol_std(x, alpha_m, alpha_n, beta_0):
+ """
+ Evaluate the standard Topeol profile.
+
+ The profile is defined as
+
+ T(x) = (1 - x^{alpha_m})^{alpha_n}.
+
+ Parameters
+ ----------
+ x : float or ndarray
+ Independent variable, typically a normalised coordinate in the
+ interval [0, 1].
+ alpha_m : float
+ Inner exponent controlling the profile shape.
+ alpha_n : float
+ Outer exponent controlling the profile peaking and edge behaviour.
+ beta_0 : float
+ Unused parameter. Included for consistency with related Topeol
+ profile functions and fitting routines.
+
+ Returns
+ -------
+ float or ndarray
+ Value of the standard Topeol profile evaluated at ``x``.
+
+ Notes
+ -----
+ This function does not currently depend on ``beta_0``.
+ """
return (1 - x**alpha_m) ** alpha_n
def d2Ldb2(t, x, alpha_m, alpha_n, beta_0, Tstd=None):
+ """
+ Evaluate the second derivative of the loss function with respect to
+ beta_0.
+
+ Parameters
+ ----------
+ t : float or ndarray
+ Target values used in the loss function.
+ x : float or ndarray
+ Independent variable, typically a normalised coordinate in the
+ interval [0, 1].
+ alpha_m : float
+ Inner exponent of the standard Topeol profile.
+ alpha_n : float
+ Outer exponent of the standard Topeol profile.
+ beta_0 : float
+ Amplitude parameter of the fitted profile.
+ Tstd : float or ndarray, optional
+ Precomputed value of ``Topeol_std(x, alpha_m, alpha_n, beta_0)``.
+ If not provided, it is evaluated internally.
+
+ Returns
+ -------
+ float or ndarray
+ Second derivative of the loss function with respect to beta_0.
+
+ Notes
+ -----
+ This function returns
+
+ d²L/dβ₀² = 2 Tstd²,
+
+ where Tstd is the standard Topeol profile. Providing ``Tstd`` avoids
+ recomputing the profile when multiple derivatives are evaluated at
+ the same point.
+ """
+
if Tstd is None:
Tstd = Topeol_std(x, alpha_m, alpha_n, beta_0)
res = 2 * Tstd**2
@@ -108,6 +184,49 @@ def d2Ldb2(t, x, alpha_m, alpha_n, beta_0, Tstd=None):
def d2Ldbdn(t, x, alpha_m, alpha_n, beta_0, Tstd=None):
+ """
+ Evaluate the mixed second derivative of the loss function with respect to
+ beta_0 and alpha_n.
+
+ Parameters
+ ----------
+ t : float or ndarray
+ Target values used in the loss function.
+ x : float or ndarray
+ Independent variable, typically a normalised coordinate in the interval [0, 1].
+ alpha_m : float
+ Inner exponent of the standard Topeol profile.
+ alpha_n : float
+ Outer exponent of the standard Topeol profile.
+ beta_0 : float
+ Amplitude parameter of the fitted profile.
+ Tstd : float or ndarray, optional
+ Precomputed value of Topeol_std(x, alpha_m, alpha_n, beta_0).
+ If not provided, it is evaluated internally.
+
+ Returns
+ -------
+ float or ndarray
+ Mixed second derivative of the loss function with respect to beta_0 and alpha_n.
+
+ Notes
+ -----
+ This function evaluates
+
+ d²L / (d beta_0 d alpha_n)
+ = (2 * beta_0 * Tstd - t) * (2 * Tstd * log(1 - x^alpha_m)),
+
+ where
+
+ Tstd = (1 - x^alpha_m)^alpha_n.
+
+ The logarithm term comes from differentiating the profile with respect to alpha_n.
+
+ Warning
+ -------
+ This expression becomes singular when x = 1, since log(0) is undefined.
+ """
+
if Tstd is None:
Tstd = Topeol_std(x, alpha_m, alpha_n, beta_0)
res = 2 * beta_0 * Tstd - t
@@ -116,6 +235,50 @@ def d2Ldbdn(t, x, alpha_m, alpha_n, beta_0, Tstd=None):
def d2Ldbdm(t, x, alpha_m, alpha_n, beta_0, Tstd=None):
+ """
+ Evaluate the mixed second derivative of the loss function with respect to
+ beta_0 and alpha_m.
+
+ Parameters
+ ----------
+ t : float or ndarray
+ Target values used in the loss function.
+ x : float or ndarray
+ Independent variable, typically a normalised coordinate in the interval [0, 1].
+ alpha_m : float
+ Inner exponent of the standard Topeol profile.
+ alpha_n : float
+ Outer exponent of the standard Topeol profile.
+ beta_0 : float
+ Amplitude parameter of the fitted profile.
+ Tstd : float or ndarray, optional
+ Precomputed value of Topeol_std(x, alpha_m, alpha_n, beta_0).
+ If not provided, it is evaluated internally.
+
+ Returns
+ -------
+ float or ndarray
+ Mixed second derivative of the loss function with respect to beta_0 and alpha_m.
+
+ Notes
+ -----
+ This function evaluates the mixed derivative
+
+ d²L / (d beta_0 d alpha_m),
+
+ using the standard Topeol profile
+
+ T(x) = (1 - x^alpha_m)^alpha_n.
+
+ The expression includes logarithmic and power-law sensitivity terms arising
+ from differentiation with respect to alpha_m.
+
+ Warning
+ -------
+ The term log(x) becomes singular at x = 0, so the expression is undefined
+ at that point.
+ """
+
if Tstd is None:
Tstd = Topeol_std(x, alpha_m, alpha_n, beta_0)
res = 2 * beta_0 * Tstd - t
@@ -125,6 +288,49 @@ def d2Ldbdm(t, x, alpha_m, alpha_n, beta_0, Tstd=None):
def d2Ldm2(t, x, alpha_m, alpha_n, beta_0, Tstd=None):
+ """
+ Evaluate the second derivative of the loss function with respect to
+ alpha_m.
+
+ Parameters
+ ----------
+ t : float or ndarray
+ Target values used in the loss function.
+ x : float or ndarray
+ Independent variable, typically a normalised coordinate in the interval [0, 1].
+ alpha_m : float
+ Inner exponent of the standard Topeol profile.
+ alpha_n : float
+ Outer exponent of the standard Topeol profile.
+ beta_0 : float
+ Amplitude parameter of the fitted profile.
+ Tstd : float or ndarray, optional
+ Precomputed value of Topeol_std(x, alpha_m, alpha_n, beta_0).
+ If not provided, it is evaluated internally.
+
+ Returns
+ -------
+ float or ndarray
+ Second derivative of the loss function with respect to alpha_m.
+
+ Notes
+ -----
+ This function evaluates
+
+ d²L / d alpha_m²,
+
+ using the standard Topeol profile
+
+ T(x) = (1 - x^alpha_m)^alpha_n.
+
+ The expression contains logarithmic terms (log(x)) and power-law factors
+ arising from repeated differentiation with respect to alpha_m.
+
+ Warning
+ -------
+ The function is undefined at x = 0 due to log(x), and may suffer
+ numerical instability near x = 0 or x = 1.
+ """
if Tstd is None:
Tstd = Topeol_std(x, alpha_m, alpha_n, beta_0)
res = beta_0 * Tstd
@@ -136,6 +342,49 @@ def d2Ldm2(t, x, alpha_m, alpha_n, beta_0, Tstd=None):
def d2Ldn2(t, x, alpha_m, alpha_n, beta_0, Tstd=None):
+ """
+ Evaluate the second derivative of the loss function with respect to
+ alpha_n.
+
+ Parameters
+ ----------
+ t : float or ndarray
+ Target values used in the loss function.
+ x : float or ndarray
+ Independent variable, typically a normalised coordinate in the interval [0, 1].
+ alpha_m : float
+ Inner exponent of the standard Topeol profile.
+ alpha_n : float
+ Outer exponent of the standard Topeol profile.
+ beta_0 : float
+ Amplitude parameter of the fitted profile.
+ Tstd : float or ndarray, optional
+ Precomputed value of Topeol_std(x, alpha_m, alpha_n, beta_0).
+ If not provided, it is evaluated internally.
+
+ Returns
+ -------
+ float or ndarray
+ Second derivative of the loss function with respect to alpha_n.
+
+ Notes
+ -----
+ This function evaluates
+
+ d²L / d alpha_n²,
+
+ using the standard Topeol profile
+
+ T(x) = (1 - x^alpha_m)^alpha_n.
+
+ The dependence on alpha_n enters through logarithmic terms of the form
+ log(1 - x^alpha_m), which are squared in this second derivative.
+
+ Warning
+ -------
+ The expression becomes singular when x = 1 since log(0) is undefined.
+ """
+
if Tstd is None:
Tstd = Topeol_std(x, alpha_m, alpha_n, beta_0)
res = 2 * beta_0 * Tstd - t
@@ -146,6 +395,51 @@ def d2Ldn2(t, x, alpha_m, alpha_n, beta_0, Tstd=None):
def d2Ldmdn(t, x, alpha_m, alpha_n, beta_0, Tstd=None):
+ """
+ Evaluate the mixed second derivative of the loss function with respect to
+ alpha_m and alpha_n.
+
+ Parameters
+ ----------
+ t : float or ndarray
+ Target values used in the loss function.
+ x : float or ndarray
+ Independent variable, typically a normalised coordinate in the interval [0, 1].
+ alpha_m : float
+ Inner exponent of the standard Topeol profile.
+ alpha_n : float
+ Outer exponent of the standard Topeol profile.
+ beta_0 : float
+ Amplitude parameter of the fitted profile.
+ Tstd : float or ndarray, optional
+ Precomputed value of Topeol_std(x, alpha_m, alpha_n, beta_0).
+ If not provided, it is evaluated internally.
+
+ Returns
+ -------
+ float or ndarray
+ Mixed second derivative of the loss function with respect to alpha_m and alpha_n.
+
+ Notes
+ -----
+ This function evaluates the mixed derivative
+
+ d²L / (d alpha_m d alpha_n),
+
+ for the standard Topeol profile
+
+ T(x) = (1 - x^alpha_m)^alpha_n.
+
+ The expression contains coupled logarithmic sensitivities in both
+ alpha_m and alpha_n arising from differentiation of the power-law
+ structure of the profile.
+
+ Warning
+ -------
+ The expression is undefined at x = 0 (log(x)) and may become singular
+ at x = 1 due to log(1 - x^alpha_m).
+ """
+
if Tstd is None:
Tstd = Topeol_std(x, alpha_m, alpha_n, beta_0)
res = 2 * beta_0 * Tstd - t
@@ -157,6 +451,49 @@ def d2Ldmdn(t, x, alpha_m, alpha_n, beta_0, Tstd=None):
def dLdn(t, x, alpha_m, alpha_n, beta_0, Tstd=None):
+ """
+ Evaluate the first derivative of the loss function with respect to
+ alpha_n.
+
+ Parameters
+ ----------
+ t : float or ndarray
+ Target values used in the loss function.
+ x : float or ndarray
+ Independent variable, typically a normalised coordinate in the interval [0, 1].
+ alpha_m : float
+ Inner exponent of the standard Topeol profile.
+ alpha_n : float
+ Outer exponent of the standard Topeol profile.
+ beta_0 : float
+ Amplitude parameter of the fitted profile.
+ Tstd : float or ndarray, optional
+ Precomputed value of Topeol_std(x, alpha_m, alpha_n, beta_0).
+ If not provided, it is evaluated internally.
+
+ Returns
+ -------
+ float or ndarray
+ First derivative of the loss function with respect to alpha_n.
+
+ Notes
+ -----
+ This function evaluates
+
+ dL / d alpha_n,
+
+ using the standard Topeol profile
+
+ T(x) = (1 - x^alpha_m)^alpha_n.
+
+ The dependence on alpha_n enters through a logarithmic term of the form
+ log(1 - x^alpha_m).
+
+ Warning
+ -------
+ The expression becomes singular when x = 1 since log(0) is undefined.
+ """
+
if Tstd is None:
Tstd = Topeol_std(x, alpha_m, alpha_n, beta_0)
res = t - beta_0 * Tstd
@@ -166,6 +503,50 @@ def dLdn(t, x, alpha_m, alpha_n, beta_0, Tstd=None):
def dLdm(t, x, alpha_m, alpha_n, beta_0, Tstd=None):
+ """
+ Evaluate the first derivative of the loss function with respect to
+ alpha_m.
+
+ Parameters
+ ----------
+ t : float or ndarray
+ Target values used in the loss function.
+ x : float or ndarray
+ Independent variable, typically a normalised coordinate in the interval [0, 1].
+ alpha_m : float
+ Inner exponent of the standard Topeol profile.
+ alpha_n : float
+ Outer exponent of the standard Topeol profile.
+ beta_0 : float
+ Amplitude parameter of the fitted profile.
+ Tstd : float or ndarray, optional
+ Precomputed value of Topeol_std(x, alpha_m, alpha_n, beta_0).
+ If not provided, it is evaluated internally.
+
+ Returns
+ -------
+ float or ndarray
+ First derivative of the loss function with respect to alpha_m.
+
+ Notes
+ -----
+ This function evaluates
+
+ dL / d alpha_m,
+
+ using the standard Topeol profile
+
+ T(x) = (1 - x^alpha_m)^alpha_n.
+
+ The dependence on alpha_m enters through logarithmic and power-law terms
+ involving log(x).
+
+ Warning
+ -------
+ The expression is undefined at x = 0 due to log(x), and may become
+ numerically unstable near x = 0 or x = 1.
+ """
+
if Tstd is None:
Tstd = Topeol_std(x, alpha_m, alpha_n, beta_0)
res = t - beta_0 * Tstd
@@ -175,6 +556,47 @@ def dLdm(t, x, alpha_m, alpha_n, beta_0, Tstd=None):
def dLdb(t, x, alpha_m, alpha_n, beta_0, Tstd=None):
+ """
+ Evaluate the first derivative of the loss function with respect to
+ beta_0.
+
+ Parameters
+ ----------
+ t : float or ndarray
+ Target values used in the loss function.
+ x : float or ndarray
+ Independent variable, typically a normalised coordinate in the interval [0, 1].
+ alpha_m : float
+ Inner exponent of the standard Topeol profile.
+ alpha_n : float
+ Outer exponent of the standard Topeol profile.
+ beta_0 : float
+ Amplitude parameter of the fitted profile.
+ Tstd : float or ndarray, optional
+ Precomputed value of Topeol_std(x, alpha_m, alpha_n, beta_0).
+ If not provided, it is evaluated internally.
+
+ Returns
+ -------
+ float or ndarray
+ First derivative of the loss function with respect to beta_0.
+
+ Notes
+ -----
+ This function evaluates
+
+ dL / d beta_0,
+
+ using the standard Topeol profile
+
+ T(x) = (1 - x^alpha_m)^alpha_n.
+
+ The derivative reflects a simple quadratic residual structure in beta_0.
+
+ Warning
+ -------
+ None.
+ """
if Tstd is None:
Tstd = Topeol_std(x, alpha_m, alpha_n, beta_0)
res = t - beta_0 * Tstd
@@ -183,6 +605,52 @@ def dLdb(t, x, alpha_m, alpha_n, beta_0, Tstd=None):
def dLdpars(tp, tf, x, alpha_m, alpha_n, beta_0, Tstd=None):
+ """
+ Compute the gradient of the total loss with respect to the model parameters.
+
+ The total loss is assumed to be the sum of two contributions:
+ one associated with ``tp`` and one associated with ``tf``, each depending
+ on the same Topeol profile but with different amplitude structure.
+
+ Parameters
+ ----------
+ tp : float or ndarray
+ First target dataset (e.g. pressure-like contribution).
+ tf : float or ndarray
+ Second target dataset (e.g. flux-function-like contribution).
+ x : float or ndarray
+ Independent variable, typically a normalised coordinate in the interval [0, 1].
+ alpha_m : float
+ Inner exponent of the standard Topeol profile.
+ alpha_n : float
+ Outer exponent of the standard Topeol profile.
+ beta_0 : float
+ Amplitude parameter of the fitted profile.
+ Tstd : float or ndarray, optional
+ Precomputed value of Topeol_std(x, alpha_m, alpha_n, beta_0).
+ If not provided, it is evaluated internally.
+
+ Returns
+ -------
+ ndarray, shape (3,)
+ Gradient of the total loss with respect to:
+ [alpha_m, alpha_n, beta_0].
+
+ Notes
+ -----
+ The gradient is constructed as the sum of two contributions:
+
+ 1. tp contribution:
+ - uses standard derivatives dLdm, dLdn, dLdb
+
+ 2. tf contribution:
+ - uses modified amplitude (1 - beta_0)
+ - flips sign of beta_0 derivative term
+
+ The final result is:
+ grad L = grad L(tp) + grad L(tf)
+ """
+
dLpdpars = np.zeros((len(tp), 3))
dLpdpars[:, 0] = dLdm(tp, x, alpha_m, alpha_n, beta_0, Tstd)
dLpdpars[:, 1] = dLdn(tp, x, alpha_m, alpha_n, beta_0, Tstd)
@@ -198,6 +666,53 @@ def dLdpars(tp, tf, x, alpha_m, alpha_n, beta_0, Tstd=None):
def d2Ldpars2(tp, tf, x, alpha_m, alpha_n, beta_0, Tstd=None):
+ """
+ Compute the Hessian matrix of the total loss with respect to the model parameters.
+
+ The total loss is assumed to be the sum of two contributions:
+ one associated with ``tp`` and one associated with ``tf``, each depending
+ on the same Topeol profile but with different amplitude structure.
+
+ Parameters
+ ----------
+ tp : float or ndarray
+ First target dataset (e.g. pressure-like contribution).
+ tf : float or ndarray
+ Second target dataset (e.g. flux-function-like contribution).
+ x : float or ndarray
+ Independent variable, typically a normalised coordinate in the interval [0, 1].
+ alpha_m : float
+ Inner exponent of the standard Topeol profile.
+ alpha_n : float
+ Outer exponent of the standard Topeol profile.
+ beta_0 : float
+ Amplitude parameter of the fitted profile.
+ Tstd : float or ndarray, optional
+ Precomputed value of Topeol_std(x, alpha_m, alpha_n, beta_0).
+ If not provided, it is evaluated internally.
+
+ Returns
+ -------
+ ndarray, shape (3, 3)
+ Hessian matrix of the total loss with respect to:
+ [alpha_m, alpha_n, beta_0].
+
+ Notes
+ -----
+ The Hessian is constructed as the sum of two contributions:
+
+ 1. tp contribution:
+ - uses second derivatives of the loss with respect to all parameter pairs
+ - evaluated using d2Ldm2, d2Ldn2, d2Ldb2, d2Ldmdn, d2Ldbdm, d2Ldbdn
+
+ 2. tf contribution:
+ - uses modified amplitude (1 - beta_0)
+ - includes sign flips in mixed derivatives involving beta_0 due to chain rule
+
+ The final Hessian is:
+ H = H(tp) + H(tf)
+ """
+
d2Lpdpars2 = np.zeros((3, 3))
d2Lpdpars2[0, 0] = np.sum(d2Ldm2(tp, x, alpha_m, alpha_n, beta_0, Tstd))
d2Lpdpars2[1, 1] = np.sum(d2Ldn2(tp, x, alpha_m, alpha_n, beta_0, Tstd))
@@ -229,6 +744,48 @@ def d2Ldpars2(tp, tf, x, alpha_m, alpha_n, beta_0, Tstd=None):
def Lpars(tp, tf, x, alpha_m, alpha_n, beta_0, Tstd=None):
+ """
+ Compute the total loss for the parameterised Topeol model.
+
+ The loss is defined as the sum of two squared-error contributions:
+ one for ``tp`` and one for ``tf``, using a shared profile.
+
+ Parameters
+ ----------
+ tp : float or ndarray
+ First target dataset (e.g. pressure-like contribution).
+ tf : float or ndarray
+ Second target dataset (e.g. flux-function-like contribution).
+ x : float or ndarray
+ Independent variable, typically a normalised coordinate in the interval [0, 1].
+ alpha_m : float
+ Inner exponent of the standard Topeol profile.
+ alpha_n : float
+ Outer exponent of the standard Topeol profile.
+ beta_0 : float
+ Amplitude parameter controlling the split between the two datasets.
+ Tstd : float or ndarray, optional
+ Precomputed value of Topeol_std(x, alpha_m, alpha_n, beta_0).
+ If not provided, it is evaluated internally.
+
+ Returns
+ -------
+ float
+ Scalar total loss value.
+
+ Notes
+ -----
+ The loss is
+
+ L = sum( (tp - beta_0 * T(x))^2 + (tf - (1 - beta_0) * T(x))^2 )
+
+ where
+
+ T(x) = (1 - x^alpha_m)^alpha_n.
+
+ The summation is performed over the input grid in ``x``.
+ """
+
if Tstd is None:
Tstd = Topeol_std(x, alpha_m, alpha_n, beta_0)
Lp = (tp - beta_0 * Tstd) ** 2
@@ -237,6 +794,54 @@ def Lpars(tp, tf, x, alpha_m, alpha_n, beta_0, Tstd=None):
def Topeol_opt_init(tp, tf):
+ """
+ Compute an initial guess for Topeol model scaling and amplitude split.
+
+ This routine provides a heuristic initialisation for optimisation by
+ normalising the input profiles and estimating a mixing parameter ``b0``
+ that balances contributions between ``tp`` and ``tf``.
+
+ Parameters
+ ----------
+ tp : ndarray
+ First input profile (e.g. pressure-like quantity).
+ tf : ndarray
+ Second input profile (e.g. flux-function-like quantity).
+
+ Returns
+ -------
+ tpn : ndarray
+ Normalised and rescaled version of ``tp`` using the estimated split.
+ tfn : ndarray
+ Normalised and rescaled version of ``tf`` using the estimated split.
+ b0 : float
+ Estimated initial amplitude split parameter in [0, 1].
+
+ Notes
+ -----
+ The initial estimate is computed as follows:
+
+ 1. A shared normalisation scale is defined using the larger of the
+ initial values of ``tp`` and ``tf``.
+ 2. A mask is applied where ``tf`` is positive to form a stable ratio.
+ 3. The mean ratio ``tp/tf`` over this region is used to estimate:
+
+ rr = mean(tp/tf)
+ b0 = rr / (1 + rr)
+
+ 4. The profiles are then rescaled consistently using ``b0`` and their
+ initial values.
+
+ Warning
+ -------
+ This heuristic assumes:
+ - ``tp[0]`` and ``tf[0]`` are non-zero
+ - There exists a region where ``tf > 0``
+ - Ratios are numerically stable in the masked region
+
+ Results may be unstable if these assumptions are violated.
+ """
+
tpn = tp / max(tp[0], tf[0])
tfn = tf / max(tp[0], tf[0])
@@ -250,6 +855,60 @@ def Topeol_opt_init(tp, tf):
def Topeol_opt_stepper(tp, tf, x, pars):
+ """
+ Perform one optimisation step for the Topeol parameter fit.
+
+ This routine computes a parameter update using either a Newton-like step
+ (based on the Hessian) or a fallback gradient-scaled step when the Hessian
+ is not sufficiently well-conditioned.
+
+ Parameters
+ ----------
+ tp : ndarray
+ First target dataset (e.g. pressure-like contribution).
+ tf : ndarray
+ Second target dataset (e.g. flux-function-like contribution).
+ x : ndarray
+ Independent variable, typically a normalised coordinate in the interval [0, 1].
+ pars : array-like, shape (3,)
+ Current parameter vector [alpha_m, alpha_n, beta_0].
+
+ Returns
+ -------
+ ndarray, shape (3,)
+ Updated parameter vector after one optimisation step.
+
+ Notes
+ -----
+ The update proceeds as follows:
+
+ 1. Compute the current profile:
+ T(x) = Topeol_std(x, alpha_m, alpha_n, beta_0)
+
+ 2. Compute gradient and Hessian:
+ grad = dLdpars(...)
+ H = d2Ldpars2(...)
+
+ 3. Check Hessian conditioning via eigenvalues:
+ - If sufficiently many positive eigenvalues are present,
+ solve Newton step:
+ dpars = -H^{-1} grad
+ - Otherwise, fall back to a scaled gradient step:
+ dpars = -L * grad / ||grad||
+
+ 4. Apply step limiting:
+ Prevent excessively large updates by clipping relative changes
+ in parameters.
+
+ Warning
+ -------
+ This is a heuristic optimizer:
+ - It is not guaranteed to converge globally.
+ - Hessian inversion may be unstable if ill-conditioned.
+ - Step limiting is ad hoc and may bias convergence.
+
+ """
+
alpha_m, alpha_n, beta_0 = pars
Tstd = Topeol_std(x, alpha_m, alpha_n, beta_0)
dLdp = dLdpars(tp, tf, x, alpha_m, alpha_n, beta_0, Tstd)
@@ -266,8 +925,57 @@ def Topeol_opt_stepper(tp, tf, x, pars):
return pars + dpars
-# the actual optimizer
def Topeol_opt(tp, tf, x, max_it, tol):
+ """
+ Optimise Topeol model parameters using an iterative stepper.
+
+ This function performs a fixed-point style optimisation of the parameters
+ [alpha_m, alpha_n, beta_0] using repeated calls to the Topeol_opt_stepper
+ until convergence or a maximum number of iterations is reached.
+
+ Parameters
+ ----------
+ tp : ndarray
+ First target dataset (e.g. pressure-like contribution).
+ tf : ndarray
+ Second target dataset (e.g. flux-function-like contribution).
+ x : ndarray
+ Independent variable, typically a normalised coordinate in the interval [0, 1].
+ max_it : int
+ Maximum number of optimisation iterations.
+ tol : float
+ Convergence tolerance. Iteration stops when all parameter updates
+ satisfy |Δpars| ≤ tol.
+
+ Returns
+ -------
+ ndarray, shape (3,)
+ Optimised parameters [alpha_m, alpha_n, beta_0].
+
+ Notes
+ -----
+ The optimisation procedure is:
+
+ 1. Initialise parameters using Topeol_opt_init:
+ tpn, tfn, b0 = Topeol_opt_init(tp, tf)
+
+ 2. Set initial guess:
+ pars = [2, 1, b0]
+
+ 3. Iterate:
+ pars_{k+1} = Topeol_opt_stepper(tpn, tfn, x, pars_k)
+
+ 4. Stop when:
+ max|pars_{k+1} - pars_k| ≤ tol
+ or when max_it is reached.
+
+ Warning
+ -------
+ - Convergence is not guaranteed.
+ - The algorithm depends strongly on the quality of the initial guess.
+ - If max_it is reached, the returned parameters may not be optimal.
+ """
+
tpn, tfn, b0 = Topeol_opt_init(tp, tf)
it = 0
pars = np.array([2, 1, b0])