From 60653d8d5304d36fab4b0cad4731dd194867a091 Mon Sep 17 00:00:00 2001 From: Timothy Nunn Date: Mon, 19 Jan 2026 13:15:44 +0000 Subject: [PATCH 1/7] Add option for different initial psi guess using parabolic Jtor assumption --- freegs4e/equilibrium.py | 296 +++++++++++++++++++++++++++++----------- 1 file changed, 213 insertions(+), 83 deletions(-) diff --git a/freegs4e/equilibrium.py b/freegs4e/equilibrium.py index 3a91ec2..175a8be 100644 --- a/freegs4e/equilibrium.py +++ b/freegs4e/equilibrium.py @@ -81,8 +81,9 @@ def __init__( boundary : callable The boundary condition function. Can be either `freeBoundary` or `fixedBoundary`. - psi : np.array - Initial guess for plasma flux [Webers/2pi]. If `None`, default initial guess used. + psi : np.array | None | callable + Initial guess for plasma flux [Webers/2pi]. If `None`, default initial Gaussian-shaped + guess used. If a callable is provided it is called with the equilibrium to get the initial guess. current : float Plasma current [A]. order : int @@ -121,16 +122,6 @@ def __init__( # assign boundary function self._applyBoundary = boundary - # assign initial guess for plasma flux (if None) - if psi is None: - # Starting guess for psi - psi = self.create_psi_plasma_default() - self.gpars = np.array([0.5, 0.5, 0, 2]) - - if psi.shape != (nx, ny): - raise ValueError("Shape of psi must match grid size (nx, ny).") - self.plasma_psi = psi - # generate Greens function mappings (used # in self.psi() to speed up calculations) self._pgreen = tokamak.createPsiGreens(self.R, self.Z) @@ -157,78 +148,18 @@ def __init__( nx, ny, generator, nlevels=1, ncycle=1, niter=2, direct=True ) - def create_psi_plasma_default( - self, adaptive_centre=False, gpars=(0.5, 0.5, 0, 2) - ): - """ - Generate a Gaussian-shaped initial guess for the plasma flux. - - The flux function ψ(R, Z) is given by: - - ψ(R, Z) = exp(C - (|x - xc|^p + |y - yc|^p)) - - where (xc, yc) is the center, C is a shift, and p is a power term. - The function ensures the flux is zero on the boundaries. - - Parameters - ---------- - adaptive_centre : bool, optional - If True, the plasma core position (xc, yc) is adjusted dynamically. - If False (default), it remains fixed. - gpars : tuple (xc, yc, C, p) - Parameters defining the flux shape: - - xc, yc : float - Coordinates of the flux center. - - C : float - Shift parameter in the exponent. - - p : float - Power of the distance term. - - Returns - ------- - np.array - The computed flux values on the given (R, Z) grid. - """ - - # define unit meshgrid - xx, yy = meshgrid( - linspace(0, 1, self.nx), linspace(0, 1, self.ny), indexing="ij" - ) - - # finds approximate (weighted) centre of the core plasma - if adaptive_centre == True: - ntot = np.sum(self.mask_inside_limiter) - xc = ( - np.sum( - self.mask_inside_limiter - * linspace(0, 1, self.nx)[:, np.newaxis] - ) - / ntot - ) - yc = ( - np.sum( - self.mask_inside_limiter - * linspace(0, 1, self.ny)[np.newaxis, :] - ) - / ntot - ) - # else sets centre according to gpars - else: - xc, yc = gpars[:2] - - # generate the plasma flux - psi = exp( - gpars[2] - - ((np.abs(xx - xc)) ** gpars[3] + (np.abs(yy - yc)) ** gpars[3]) - ) - - # set zero flux on boundary - psi[0, :] = 0.0 - psi[:, 0] = 0.0 - psi[-1, :] = 0.0 - psi[:, -1] = 0.0 + # assign initial guess for plasma flux (if None) + if psi is None: + # Starting guess for psi using Gaussian-shape + psi_guess_callable = PsiGuessGaussian() + psi = psi_guess_callable(self) + self.gpars = np.array(psi_guess_callable.gpars) + elif callable(psi): + psi = psi(self) - return psi + if psi.shape != (nx, ny): + raise ValueError("Shape of psi must match grid size (nx, ny).") + self.plasma_psi = psi def setSolverVcycle(self, nlevels=1, ncycle=1, niter=1, direct=True): """ @@ -2840,6 +2771,205 @@ def ellipse_points(R0, Z0, A, B, N=360): return np.column_stack((R, Z)) +class PsiGuessGaussian: + def __init__( + self, + adaptive_centre: bool = False, + gpars: tuple[float, ...] = (0.5, 0.5, 0, 2), + ): + """Class to generate a Gaussian-shaped initial guess for the plasma flux. + + The flux function ψ(R, Z) is given by: + + ψ(R, Z) = exp(C - (|x - xc|^p + |y - yc|^p)) + + where (xc, yc) is the center, C is a shift, and p is a power term. + The function ensures the flux is zero on the boundaries. + + Parameters + ---------- + adaptive_centre : bool, optional + If True, the plasma core position (xc, yc) is adjusted dynamically. + If False (default), it remains fixed. + gpars : tuple (xc, yc, C, p) + Parameters defining the flux shape: + - xc, yc : float + Coordinates of the flux center. + - C : float + Shift parameter in the exponent. + - p : float + Power of the distance term. + """ + self._adaptive_centre = adaptive_centre + self.gpars = gpars + + def __call__(self, eq): + """Calculate the initial Gaussian-shaped psi guess. + + Returns + ------- + np.array + The computed flux values on the given (R, Z) grid. + """ + + # define unit meshgrid + xx, yy = meshgrid( + linspace(0, 1, eq.nx), linspace(0, 1, eq.ny), indexing="ij" + ) + + # finds approximate (weighted) centre of the core plasma + if self._adaptive_centre: + ntot = np.sum(eq.mask_inside_limiter) + xc = ( + np.sum( + eq.mask_inside_limiter + * linspace(0, 1, eq.nx)[:, np.newaxis] + ) + / ntot + ) + yc = ( + np.sum( + eq.mask_inside_limiter + * linspace(0, 1, eq.ny)[np.newaxis, :] + ) + / ntot + ) + # else sets centre according to gpars + else: + xc, yc = self.gpars[:2] + + # generate the plasma flux + psi = exp( + self.gpars[2] + - ( + (np.abs(xx - xc)) ** self.gpars[3] + + (np.abs(yy - yc)) ** self.gpars[3] + ) + ) + + # set zero flux on boundary + psi[0, :] = 0.0 + psi[:, 0] = 0.0 + psi[-1, :] = 0.0 + psi[:, -1] = 0.0 + + return psi + + +class PsiGuessParabolicJtor: + def __init__(self, h, k, a, b, Ip): + """A class to generate the initial psi guess from an elliptic description + of the plasma boundary and the plasma current. + + We assume a plasma core with a separatrix that can be described by an ellipse. From this, + we create an elliptic distribution function E(R, Z) such that E(h, k) = 0 where (h, k) is the + centre of the ellipse and E(R, Z) = 1 for all points on the separatrix. + + We assume a parabolic distribution of current and calculate Jtor using equation A.7 from [1]: + Jtor = Jhat(1-E(R,Z)^i)^j + Jhat is a constant found to match the calculated plasma current to the desired plasma current (Ip). + + Finally, we apply a linear Grad-Shafranov step to find our estimate of the initial plasma psi. + + Parameters + ---------- + h : float + The radial centre of the ellipse [m]. + k : float + The vertical centre of the ellipse [m]. + a : float + Half the length of the radial axis of the ellipse (normally the semi-major axis). + b : float + Half the length of the vertical axis of the ellipse (normally the semi-minor axis). + Ip : float + The desired plasma current [A]. + + References + ---------- + [1] Wai, J. T., et al. "Feedforward equilibrium trajectory optimization with GSPulse." arXiv preprint arXiv:2506.21760 (2025). + """ + self._h = h + self._k = k + self._a = a + self._b = b + self._Ip = Ip + + @classmethod + def from_xpoint_and_midplane(cls, Rx, Zx, Ri, Ro, Zm, Ip): + h = (Rx + Ri + Ro) / 3 + k = Zm + a = (Ro - Ri) / 2 + b = Zx - Zm + + return cls(h, k, a, b, Ip) + + def elliptic_distribution_function(self, R, Z): + """The elliptic distribution function (E(R, Z)) for the ellipse represented by this class. + + Parameters + ---------- + R : float | ndarray + The radial position(s) to evaluate the distribution function at. + Z : float | ndarray + The vertical position(s) to evaluate the distribution function at. + """ + return (((R - self._h) ** 2) / self._a**2) + ( + ((Z - self._k) ** 2) / self._b**2 + ) + + def parabolic_jtor(self, eq, Ip, *, a=2, b=2): + """Calculate the Jtor assuming a parabolic current distribution for the given equilibrium. + + Solves the equation: + Jtor = Jhat*(1-E(R,Z)^a)^b + where Jhat is calculated according to the desired Ip. + + Parameters + ---------- + eq : Equilibrium + The equilibrium object to calculate the initial psi for + Ip : float + The desired plasma current. + """ + x = self.elliptic_distribution_function(eq.R, eq.Z) + + # x > 1 indicates this grid point is outside of the core + # because E(R, Z) = 1 at the separatrix + outside_core_mask = x > 1 + # Set x = 1 outside the core to ensure the plasma current + # is only distributed within the core + # Jtor outside plasma = (1 - 1**a)**b = 0 + x[outside_core_mask] = 1.0 + + dR = eq.R[1, 0] - eq.R[0, 0] + dZ = eq.Z[0, 1] - eq.Z[0, 0] + + Jtor_normalised = (1 - x**a) ** b + Jtor_integrated = np.sum(Jtor_normalised[~outside_core_mask]) * dR * dZ + + Jhat = Ip / Jtor_integrated + + return Jhat * Jtor_normalised * ~outside_core_mask + + def __call__(self, eq: Equilibrium): + """Calculates the initial psi guess by approximating the Jtor and + applying a linear Grad-Shafranov step. + """ + # ensure equilibrium has attribute plasma_psi which will be overwritten later + eq.plasma_psi = np.zeros((eq.nx, eq.ny)) + + jtor_guess = self.parabolic_jtor(eq, self._Ip) + psi = eq.callSolver(psi=eq.plasma_psi, rhs=-mu0 * eq.R * jtor_guess) + + # set zero flux on boundary + psi[0, :] = 0.0 + psi[:, 0] = 0.0 + psi[-1, :] = 0.0 + psi[:, -1] = 0.0 + + return psi + + if __name__ == "__main__": # Test the different spline interpolation routines From 5ef733aad50b82c72d58fe668439b2d6df95c270 Mon Sep 17 00:00:00 2001 From: Timothy Nunn Date: Tue, 20 Jan 2026 10:52:14 +0000 Subject: [PATCH 2/7] Add option to calculate elliptic distribution function from separatrix constraints --- freegs4e/equilibrium.py | 66 +++++++++++++++++++++++++++++++++++++++-- 1 file changed, 64 insertions(+), 2 deletions(-) diff --git a/freegs4e/equilibrium.py b/freegs4e/equilibrium.py index 175a8be..ce9657b 100644 --- a/freegs4e/equilibrium.py +++ b/freegs4e/equilibrium.py @@ -2856,6 +2856,9 @@ def __call__(self, eq): return psi +from scipy.optimize import least_squares + + class PsiGuessParabolicJtor: def __init__(self, h, k, a, b, Ip): """A class to generate the initial psi guess from an elliptic description @@ -2896,6 +2899,13 @@ def __init__(self, h, k, a, b, Ip): @classmethod def from_xpoint_and_midplane(cls, Rx, Zx, Ri, Ro, Zm, Ip): + """Finds the elliptic distribution function that produces a boundary + according to the desired X-point location (Rx, Zx), inner midplane (Ri, Zm), + and outer midplane (Ro, Zm). + + This assumes that the centroid of the ellipse (where the distribution function + is 0) occurs radially at the average of Rx, Ri, Ro. + """ h = (Rx + Ri + Ro) / 3 k = Zm a = (Ro - Ri) / 2 @@ -2903,6 +2913,58 @@ def from_xpoint_and_midplane(cls, Rx, Zx, Ri, Ro, Zm, Ip): return cls(h, k, a, b, Ip) + @classmethod + def from_separatrix_points(cls, Ip: float, *points): + """Finds the elliptic distribution function that best fits the + provided points that describe the desired core's separatrix. + + Parameters + ---------- + Ip : float + The desired plasma current [A]. + points : tuple[tuple[float, float], ...] + A list of (R, Z) locations on the separatrix. + + Notes + ----- + This method requires at least 4 isoflux points + """ + if len(points) < 4: + raise ValueError( + "At least 4 isoflux points required to approximate the ellipse shape parameters" + ) + + points_np = np.array(points) + Rpoints = points_np[:, 0] + Zpoints = points_np[:, 1] + + def _func(args): + h, k, a, b = args.squeeze().tolist() + + return cls._elliptic_function(Rpoints, Zpoints, h, k, a, b) - 1.0 + + result = least_squares(_func, [1.0] * 4, bounds=(0.0, np.inf)) + + return cls(*result.x.tolist(), Ip) + + @staticmethod + def _elliptic_function(R, Z, h, k, a, b): + """Evaluates the elliptic function: + E(R, Z) = (R-h)^2/a^2 + (Z-k)^2/b^2 + + Parameters + ---------- + h : float + The radial centre of the ellipse. + k : float + The vertical centre of the ellipse. + a : float + Half the length of the radial axis of the ellipse (normally the semi-major axis). + b : float + Half the length of the vertical axis of the ellipse (normally the semi-minor axis). + """ + return (((R - h) ** 2) / a**2) + (((Z - k) ** 2) / b**2) + def elliptic_distribution_function(self, R, Z): """The elliptic distribution function (E(R, Z)) for the ellipse represented by this class. @@ -2913,8 +2975,8 @@ def elliptic_distribution_function(self, R, Z): Z : float | ndarray The vertical position(s) to evaluate the distribution function at. """ - return (((R - self._h) ** 2) / self._a**2) + ( - ((Z - self._k) ** 2) / self._b**2 + return self._elliptic_function( + R, Z, self._h, self._k, self._a, self._b ) def parabolic_jtor(self, eq, Ip, *, a=2, b=2): From 9f1fd67fcaf5ab9cfc2c842cd66cbfa32054f767 Mon Sep 17 00:00:00 2001 From: Timothy Nunn Date: Tue, 20 Jan 2026 16:29:25 +0000 Subject: [PATCH 3/7] Add option to scale psi guess from PsiGuessParabolicJtor --- freegs4e/equilibrium.py | 25 ++++++++++++++----------- 1 file changed, 14 insertions(+), 11 deletions(-) diff --git a/freegs4e/equilibrium.py b/freegs4e/equilibrium.py index ce9657b..3003f9f 100644 --- a/freegs4e/equilibrium.py +++ b/freegs4e/equilibrium.py @@ -25,9 +25,10 @@ import matplotlib.pyplot as plt import numpy as np import shapely as sh -from numpy import array, exp, linspace, meshgrid, pi +from numpy import exp, linspace, meshgrid, pi from scipy import interpolate from scipy.integrate import cumulative_trapezoid +from scipy.optimize import least_squares from scipy.spatial.distance import pdist, squareform from . import critical, machine, multigrid, polygons # multigrid solver @@ -2856,11 +2857,8 @@ def __call__(self, eq): return psi -from scipy.optimize import least_squares - - class PsiGuessParabolicJtor: - def __init__(self, h, k, a, b, Ip): + def __init__(self, h, k, a, b, Ip, *, psi_scale=1.0): """A class to generate the initial psi guess from an elliptic description of the plasma boundary and the plasma current. @@ -2886,6 +2884,8 @@ def __init__(self, h, k, a, b, Ip): Half the length of the vertical axis of the ellipse (normally the semi-minor axis). Ip : float The desired plasma current [A]. + psi_scale : float + Multiplies the psi estimate when calling this class before returning. References ---------- @@ -2896,9 +2896,12 @@ def __init__(self, h, k, a, b, Ip): self._a = a self._b = b self._Ip = Ip + self._psi_scale = psi_scale @classmethod - def from_xpoint_and_midplane(cls, Rx, Zx, Ri, Ro, Zm, Ip): + def from_xpoint_and_midplane( + cls, Rx, Zx, Ri, Ro, Zm, Ip, *, psi_scale=1.0 + ): """Finds the elliptic distribution function that produces a boundary according to the desired X-point location (Rx, Zx), inner midplane (Ri, Zm), and outer midplane (Ro, Zm). @@ -2911,10 +2914,10 @@ def from_xpoint_and_midplane(cls, Rx, Zx, Ri, Ro, Zm, Ip): a = (Ro - Ri) / 2 b = Zx - Zm - return cls(h, k, a, b, Ip) + return cls(h, k, a, b, Ip, psi_scale=psi_scale) @classmethod - def from_separatrix_points(cls, Ip: float, *points): + def from_separatrix_points(cls, Ip: float, *points, psi_scale=1.0): """Finds the elliptic distribution function that best fits the provided points that describe the desired core's separatrix. @@ -2945,7 +2948,7 @@ def _func(args): result = least_squares(_func, [1.0] * 4, bounds=(0.0, np.inf)) - return cls(*result.x.tolist(), Ip) + return cls(*result.x.tolist(), Ip, psi_scale=psi_scale) @staticmethod def _elliptic_function(R, Z, h, k, a, b): @@ -2979,7 +2982,7 @@ def elliptic_distribution_function(self, R, Z): R, Z, self._h, self._k, self._a, self._b ) - def parabolic_jtor(self, eq, Ip, *, a=2, b=2): + def parabolic_jtor(self, eq, Ip, *, a=2.0, b=2.0): """Calculate the Jtor assuming a parabolic current distribution for the given equilibrium. Solves the equation: @@ -3029,7 +3032,7 @@ def __call__(self, eq: Equilibrium): psi[-1, :] = 0.0 psi[:, -1] = 0.0 - return psi + return self._psi_scale * psi if __name__ == "__main__": From 4fa2402c9a311227c2c1253c107f5fdd533c383a Mon Sep 17 00:00:00 2001 From: Timothy Nunn Date: Fri, 6 Feb 2026 15:43:07 +0000 Subject: [PATCH 4/7] Add create_psi_plasma_default back to Equilibrium for backwards compatability --- freegs4e/equilibrium.py | 8 ++++++++ 1 file changed, 8 insertions(+) diff --git a/freegs4e/equilibrium.py b/freegs4e/equilibrium.py index 3003f9f..723e33d 100644 --- a/freegs4e/equilibrium.py +++ b/freegs4e/equilibrium.py @@ -162,6 +162,14 @@ def __init__( raise ValueError("Shape of psi must match grid size (nx, ny).") self.plasma_psi = psi + def create_psi_plasma_default( + self, adaptive_centre=False, gpars=(0.5, 0.5, 0, 2) + ): + """This method is kept for backwards compatibility and will generate a Gaussian-shaped + initial guess for the plasma flux. It is NOT called by this class anymore. + """ + return PsiGuessGaussian(adaptive_centre, gpars) + def setSolverVcycle(self, nlevels=1, ncycle=1, niter=1, direct=True): """ Sets a new linear solver based on the multigrid scheme. From 48a0c2bdb2fe131eca490b2245808869aeaa650c Mon Sep 17 00:00:00 2001 From: Timothy Nunn Date: Mon, 9 Feb 2026 12:12:45 +0000 Subject: [PATCH 5/7] Update Jtor powers to equal Wai --- freegs4e/equilibrium.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/freegs4e/equilibrium.py b/freegs4e/equilibrium.py index 723e33d..5270a89 100644 --- a/freegs4e/equilibrium.py +++ b/freegs4e/equilibrium.py @@ -168,7 +168,7 @@ def create_psi_plasma_default( """This method is kept for backwards compatibility and will generate a Gaussian-shaped initial guess for the plasma flux. It is NOT called by this class anymore. """ - return PsiGuessGaussian(adaptive_centre, gpars) + return PsiGuessGaussian(adaptive_centre, gpars)(self) def setSolverVcycle(self, nlevels=1, ncycle=1, niter=1, direct=True): """ @@ -2990,7 +2990,7 @@ def elliptic_distribution_function(self, R, Z): R, Z, self._h, self._k, self._a, self._b ) - def parabolic_jtor(self, eq, Ip, *, a=2.0, b=2.0): + def parabolic_jtor(self, eq, Ip, *, a=1.87, b=1.5): """Calculate the Jtor assuming a parabolic current distribution for the given equilibrium. Solves the equation: From a07329e1d83df1fb7b7e9a35b65b5530ceaf9600 Mon Sep 17 00:00:00 2001 From: kpentland Date: Tue, 3 Mar 2026 13:49:59 +0000 Subject: [PATCH 6/7] updating RaiseValue error if psi_plasma grids do not match --- freegs4e/equilibrium.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/freegs4e/equilibrium.py b/freegs4e/equilibrium.py index 5270a89..93a534b 100644 --- a/freegs4e/equilibrium.py +++ b/freegs4e/equilibrium.py @@ -159,7 +159,9 @@ def __init__( psi = psi(self) if psi.shape != (nx, ny): - raise ValueError("Shape of psi must match grid size (nx, ny).") + raise ValueError( + f"Shape of psi grid {psi.shape} must match grid size ({nx}, {ny})." + ) self.plasma_psi = psi def create_psi_plasma_default( From 247945c40a646ce4cf0c3fee92c318e240316964 Mon Sep 17 00:00:00 2001 From: Timothy Nunn Date: Fri, 6 Mar 2026 13:58:49 +0000 Subject: [PATCH 7/7] Update docstrings to improve clarity of ParabolicJtor initial psi --- freegs4e/equilibrium.py | 18 +++++++++--------- 1 file changed, 9 insertions(+), 9 deletions(-) diff --git a/freegs4e/equilibrium.py b/freegs4e/equilibrium.py index 93a534b..c436a31 100644 --- a/freegs4e/equilibrium.py +++ b/freegs4e/equilibrium.py @@ -84,7 +84,8 @@ def __init__( or `fixedBoundary`. psi : np.array | None | callable Initial guess for plasma flux [Webers/2pi]. If `None`, default initial Gaussian-shaped - guess used. If a callable is provided it is called with the equilibrium to get the initial guess. + guess used. If a callable is provided, it is called with the equilibrium to get the initial guess. + The array (returned by the callable) must have the shape (nx, ny). current : float Plasma current [A]. order : int @@ -2992,11 +2993,11 @@ def elliptic_distribution_function(self, R, Z): R, Z, self._h, self._k, self._a, self._b ) - def parabolic_jtor(self, eq, Ip, *, a=1.87, b=1.5): + def parabolic_jtor(self, eq, Ip, *, i=1.87, j=1.5): """Calculate the Jtor assuming a parabolic current distribution for the given equilibrium. Solves the equation: - Jtor = Jhat*(1-E(R,Z)^a)^b + Jtor = Jhat*(1-E(R,Z)^i)^j where Jhat is calculated according to the desired Ip. Parameters @@ -3013,14 +3014,13 @@ def parabolic_jtor(self, eq, Ip, *, a=1.87, b=1.5): outside_core_mask = x > 1 # Set x = 1 outside the core to ensure the plasma current # is only distributed within the core - # Jtor outside plasma = (1 - 1**a)**b = 0 + # Jtor outside plasma = (1 - 1**i)**j = 0 x[outside_core_mask] = 1.0 - dR = eq.R[1, 0] - eq.R[0, 0] - dZ = eq.Z[0, 1] - eq.Z[0, 0] - - Jtor_normalised = (1 - x**a) ** b - Jtor_integrated = np.sum(Jtor_normalised[~outside_core_mask]) * dR * dZ + Jtor_normalised = (1 - x**i) ** j + Jtor_integrated = ( + np.sum(Jtor_normalised[~outside_core_mask]) * eq.dR * eq.dZ + ) Jhat = Ip / Jtor_integrated