From b9c904c608c5b05292203177baee0b7fa1d3396c Mon Sep 17 00:00:00 2001 From: Timothy Nunn Date: Tue, 13 Jun 2023 10:39:38 +0100 Subject: [PATCH] Extract jtor shape function into a class --- 01-freeboundary.py | 45 +++++++++----- docs/creating_equilibria.rst | 23 ++++--- freegs/jtor.py | 117 +++++++++++++++++++++++------------ freegs/test_jtor.py | 50 +++++++++++---- 4 files changed, 159 insertions(+), 76 deletions(-) diff --git a/01-freeboundary.py b/01-freeboundary.py index ffad543..54ab33a 100644 --- a/01-freeboundary.py +++ b/01-freeboundary.py @@ -8,20 +8,28 @@ tokamak = freegs.machine.TestTokamak() -eq = freegs.Equilibrium(tokamak=tokamak, - Rmin=0.1, Rmax=2.0, # Radial domain - Zmin=-1.0, Zmax=1.0, # Height range - nx=65, ny=65, # Number of grid points - boundary=freegs.boundary.freeBoundaryHagenow) # Boundary condition +eq = freegs.Equilibrium( + tokamak=tokamak, + Rmin=0.1, + Rmax=2.0, # Radial domain + Zmin=-1.0, + Zmax=1.0, # Height range + nx=65, + ny=65, # Number of grid points + boundary=freegs.boundary.freeBoundaryHagenow, +) # Boundary condition ######################################### # Plasma profiles -profiles = freegs.jtor.ConstrainPaxisIp(eq, - 1e3, # Plasma pressure on axis [Pascals] - 2e5, # Plasma current [Amps] - 2.0) # Vacuum f=R*Bt +profiles = freegs.jtor.ConstrainPaxisIpArbShape( + eq, + 1e3, # Plasma pressure on axis [Pascals] + 2e5, # Plasma current [Amps] + 2.0, # Vacuum f=R*Bt + shape_function=freegs.jtor.DoublePowerShapeFunction(alpha_m=1, alpha_n=2), +) ######################################### # Coil current constraints @@ -29,20 +37,24 @@ # Specify locations of the X-points # to use to constrain coil currents -xpoints = [(1.1, -0.6), # (R,Z) locations of X-points - (1.1, 0.8)] +xpoints = [ + (1.1, -0.6), # (R,Z) locations of X-points + (1.1, 0.8), +] -isoflux = [(1.1,-0.6, 1.1,0.6)] # (R1,Z1, R2,Z2) pair of locations +isoflux = [(1.1, -0.6, 1.1, 0.6)] # (R1,Z1, R2,Z2) pair of locations constrain = freegs.control.constrain(xpoints=xpoints, isoflux=isoflux) ######################################### # Nonlinear solve -freegs.solve(eq, # The equilibrium to adjust - profiles, # The toroidal current profile function - constrain, - show=True) # Constraint function to set coil currents +freegs.solve( + eq, # The equilibrium to adjust + profiles, # The toroidal current profile function + constrain, + show=True, +) # Constraint function to set coil currents # eq now contains the solution @@ -80,6 +92,7 @@ # Safety factor import matplotlib.pyplot as plt + plt.plot(*eq.q()) plt.xlabel(r"Normalised $\psi$") plt.ylabel("Safety factor") diff --git a/docs/creating_equilibria.rst b/docs/creating_equilibria.rst index 478bd1a..c77c6ff 100644 --- a/docs/creating_equilibria.rst +++ b/docs/creating_equilibria.rst @@ -399,13 +399,15 @@ Constrain pressure and current One of the most intuitive methods is to fix the shape of the plasma profiles, and adjust them to fix the pressure on the magnetic axis and total plasma current. -To do this, create a ``ConstrainPaxisIp`` profile object: +To do this, create a ``ConstrainBetapIpArbShape`` profile object: :: - profiles = freegs.jtor.ConstrainPaxisIp(1e4, # Pressure on axis [Pa] - 1e6, # Plasma current [Amps] - 1.0) # Vacuum f=R*Bt + profiles = freegs.jtor.ConstrainBetapIpArbShape(eq, + 1e4, # Pressure on axis [Pa] + 1e6, # Plasma current [Amps] + 1.0, # Vacuum f=R*Bt + shape_function=freegs.jtor.DoublePowerShapeFunction()) This sets the toroidal current to: @@ -415,7 +417,7 @@ This sets the toroidal current to: J_\phi = L \left[\beta_0 R + \left(1-\beta_0\right)/R\right] \left(1-\psi_n^{\alpha_m}\right)^{\alpha_n} where :math:`\psi_n` is the normalised poloidal flux, 0 on the magnetic axis and 1 on the plasma boundary/separatrix. -The constants which determine the profile shapes are :math:`\alpha_m = 1` and :math:`\alpha_n = 2`. These can be changed by specifying in the initialisation of ``ConstrainPaxisIp``. +The constants which determine the profile shapes are :math:`\alpha_m = 1` and :math:`\alpha_n = 2`. These can be changed by specifying in the initialisation of ``DoublePowerShapeFunction``. Any arbitrary shape function can be passed via the ``shape_function`` argument provided the object complies with the ``ProfileShapeFunction`` interface. The values of :math:`L` and :math:`\beta_0` are determined from the constraints: The pressure on axis is given by integrating the pressure gradient flux function @@ -448,12 +450,15 @@ This is the method used in `Y.M.Jeon 2015 `_, :: - profiles = freegs.jtor.ConstrainBetapIp(0.5, # Poloidal beta - 1e6, # Plasma current [Amps] - 1.0) # Vacuum f=R*Bt - + profiles = freegs.jtor.ConstrainBetapIpArbShape(eq, + 0.5, # Poloidal beta + 1e6, # Plasma current [Amps] + 1.0, # Vacuum f=R*Bt + shape_function=freegs.jtor.DoublePowerShapeFunction()) + By integrating over the plasma domain and combining the constraints on poloidal beta and plasma current, the values of :math:`L` and :math:`\beta_0` are found. + Feedback and shape control -------------------------- diff --git a/freegs/jtor.py b/freegs/jtor.py index 2e35bb8..cb74e04 100644 --- a/freegs/jtor.py +++ b/freegs/jtor.py @@ -18,6 +18,7 @@ You should have received a copy of the GNU Lesser General Public License along with FreeGS. If not, see . """ + from scipy.integrate import romb, quad # Romberg integration from . import critical from .gradshafranov import mu0 @@ -136,36 +137,57 @@ def fvac(self) -> float: pass -class ConstrainBetapIp(Profile): +class ProfileShapeFunction(abc.ABC): + """When instantiated, can be called with the normalised psi to + shape the profiles. """ - Constrain poloidal Beta and plasma current - This is the constraint used in - YoungMu Jeon arXiv:1503.03135 + @abc.abstractmethod + def __call__(self, psi_norm): + """ + psi_norm - normalised psi [0, 1] + """ + pass + + +class DoublePowerShapeFunction(ProfileShapeFunction): + def __init__(self, alpha_m=1.0, alpha_n=2.0) -> None: + # Check inputs + if alpha_m < 0: + raise ValueError("alpha_m must be positive") + if alpha_n < 0: + raise ValueError("alpha_n must be positive") + self.alpha_m = alpha_m + self.alpha_n = alpha_n + + def __call__(self, psi_norm): + return (1.0 - psi_norm**self.alpha_m) ** self.alpha_n + + +class ConstrainBetapIpArbShape(Profile): + """Constrain poloidal Beta and plasma current with an arbitrary + shaping function """ - def __init__(self, eq, betap, Ip, fvac, alpha_m=1.0, alpha_n=2.0, Raxis=1.0): + def __init__(self, eq, betap, Ip, fvac, shape_function=None, Raxis=1.0): """ betap - Poloidal beta Ip - Plasma current [Amps] fvac - Vacuum f = R*Bt - + shape_function - a callable which accepts a scalar normalised psi and array of alpha. + Defaults to `DoublePowerShapeFunction`. Raxis - R used in p' and ff' components """ - # Check inputs - if alpha_m < 0: - raise ValueError("alpha_m must be positive") - if alpha_n < 0: - raise ValueError("alpha_n must be positive") + if shape_function is None: + shape_function = DoublePowerShapeFunction() # Set parameters for later use self.betap = betap self.Ip = Ip self._fvac = fvac - self.alpha_m = alpha_m - self.alpha_n = alpha_n + self.shape_function = shape_function self.Raxis = Raxis self.eq = eq @@ -208,7 +230,7 @@ def Jtor(self, R, Z, psi, psi_bndry=None): psi_norm = (psi - psi_axis) / (psi_bndry - psi_axis) # Current profile shape - jtorshape = (1.0 - np.clip(psi_norm, 0.0, 1.0) ** self.alpha_m) ** self.alpha_n + jtorshape = self.shape_function(np.clip(psi_norm, 0.0, 1.0)) if mask is not None: # If there is a masking function (X-points, limiters) @@ -219,9 +241,7 @@ def Jtor(self, R, Z, psi, psi_bndry=None): # Need integral of jtorshape to calculate pressure # Note factor to convert from normalised psi integral def pshape(psinorm): - shapeintegral, _ = quad( - lambda x: (1.0 - x**self.alpha_m) ** self.alpha_n, psinorm, 1.0 - ) + shapeintegral, _ = quad(self.shape_function, psinorm, 1.0) shapeintegral *= psi_bndry - psi_axis return shapeintegral @@ -286,7 +306,7 @@ def pprime(self, pn): dp/dpsi as a function of normalised psi. 0 outside core Calculate pprimeshape inside the core only """ - shape = (1.0 - np.clip(pn, 0.0, 1.0) ** self.alpha_m) ** self.alpha_n + shape = self.shape_function(np.clip(pn, 0.0, 1.0)) return self.L * self.Beta0 / self.Raxis * shape def ffprime(self, pn): @@ -294,20 +314,37 @@ def ffprime(self, pn): f * df/dpsi as a function of normalised psi. 0 outside core. Calculate ffprimeshape inside the core only. """ - shape = (1.0 - np.clip(pn, 0.0, 1.0) ** self.alpha_m) ** self.alpha_n + shape = self.shape_function(np.clip(pn, 0.0, 1.0)) return mu0 * self.L * (1 - self.Beta0) * self.Raxis * shape def fvac(self): return self._fvac -class ConstrainPaxisIp(Profile): +class ConstrainBetapIp(ConstrainBetapIpArbShape): + """Constrain poloidal Beta and plasma current with the + double power shaping function. + + Provided for API backwards compatability. `ConstrainBetapIpArbShape` + provides more flexibility for profile shaping. + + This is the constraint used in + YoungMu Jeon arXiv:1503.03135 """ - Constrain pressure on axis and plasma current + def __init__(self, eq, betap, Ip, fvac, alpha_m=1.0, alpha_n=2.0, Raxis=1.0): + shape_function = DoublePowerShapeFunction(alpha_m, alpha_n) + super().__init__(eq, betap, Ip, fvac, shape_function, Raxis) + + +class ConstrainPaxisIpArbShape(Profile): """ + Constrain pressure on axis and plasma current with an + arbitrary shaping function. - def __init__(self, eq, paxis, Ip, fvac, alpha_m=1.0, alpha_n=2.0, Raxis=1.0): + """ + + def __init__(self, eq, paxis, Ip, fvac, shape_function=None, Raxis=1.0): """ paxis - Pressure at magnetic axis [Pa] Ip - Plasma current [Amps] @@ -316,18 +353,14 @@ def __init__(self, eq, paxis, Ip, fvac, alpha_m=1.0, alpha_n=2.0, Raxis=1.0): Raxis - R used in p' and ff' components """ - # Check inputs - if alpha_m < 0: - raise ValueError("alpha_m must be positive") - if alpha_n < 0: - raise ValueError("alpha_n must be positive") + if shape_function is None: + shape_function = DoublePowerShapeFunction() # Set parameters for later use self.paxis = paxis self.Ip = Ip self._fvac = fvac - self.alpha_m = alpha_m - self.alpha_n = alpha_n + self.shape_function = shape_function self.Raxis = Raxis self.eq = eq @@ -370,7 +403,7 @@ def Jtor(self, R, Z, psi, psi_bndry=None): psi_norm = (psi - psi_axis) / (psi_bndry - psi_axis) # Current profile shape - jtorshape = (1.0 - np.clip(psi_norm, 0.0, 1.0) ** self.alpha_m) ** self.alpha_n + jtorshape = self.shape_function(np.clip(psi_norm, 0.0, 1.0)) if mask is not None: # If there is a masking function (X-points, limiters) @@ -380,9 +413,7 @@ def Jtor(self, R, Z, psi, psi_bndry=None): # Need integral of jtorshape to calculate paxis # Note factor to convert from normalised psi integral - shapeintegral, _ = quad( - lambda x: (1.0 - x**self.alpha_m) ** self.alpha_n, 0.0, 1.0 - ) + shapeintegral, _ = quad(self.shape_function, 0.0, 1.0) shapeintegral *= psi_bndry - psi_axis # Pressure on axis is @@ -422,7 +453,7 @@ def pprime(self, pn): dp/dpsi as a function of normalised psi. 0 outside core Calculate pprimeshape inside the core only """ - shape = (1.0 - np.clip(pn, 0.0, 1.0) ** self.alpha_m) ** self.alpha_n + shape = self.shape_function(np.clip(pn, 0.0, 1.0)) return self.L * self.Beta0 / self.Raxis * shape def ffprime(self, pn): @@ -430,19 +461,29 @@ def ffprime(self, pn): f * df/dpsi as a function of normalised psi. 0 outside core. Calculate ffprimeshape inside the core only. """ - shape = (1.0 - np.clip(pn, 0.0, 1.0) ** self.alpha_m) ** self.alpha_n + shape = self.shape_function(np.clip(pn, 0.0, 1.0)) return mu0 * self.L * (1 - self.Beta0) * self.Raxis * shape def fvac(self): return self._fvac -class ProfilesPprimeFfprime(Profile): +class ConstrainPaxisIp(ConstrainPaxisIpArbShape): + """Constrain pressure on axis and plasma current + + Provided for API backwards compatability. `ConstrainPaxisIpArbShape` + provides more flexibility for profile shaping. """ - Specified profile functions p'(psi), ff'(psi) - Jtor = R*p' + ff'/(R*mu0) + def __init__(self, eq, paxis, Ip, fvac, alpha_m=1.0, alpha_n=2.0, Raxis=1.0): + shape_function = DoublePowerShapeFunction(alpha_m, alpha_n) + super().__init__(eq, paxis, Ip, fvac, shape_function, Raxis) + +class ProfilesPprimeFfprime(Profile): + """Specified profile functions p'(psi), ff'(psi) + + Jtor = R*p' + ff'/(R*mu0) """ def __init__(self, pprime_func, ffprime_func, fvac, p_func=None, f_func=None): diff --git a/freegs/test_jtor.py b/freegs/test_jtor.py index ecb15da..a2144a7 100644 --- a/freegs/test_jtor.py +++ b/freegs/test_jtor.py @@ -1,25 +1,49 @@ import numpy as np +import pytest from . import jtor from . import equilibrium -def test_psinorm_range(): +@pytest.mark.parametrize( + ("profile_class", "kwargs"), + ( + (jtor.ConstrainPaxisIp, {}), + (jtor.ConstrainBetapIp, {}), + (jtor.ConstrainPaxisIpArbShape, {}), + (jtor.ConstrainBetapIpArbShape, {}), + ( + jtor.ConstrainPaxisIpArbShape, + {"shape_function": jtor.DoublePowerShapeFunction()}, + ), + ( + jtor.ConstrainBetapIpArbShape, + {"shape_function": jtor.DoublePowerShapeFunction()}, + ), + ( + jtor.ConstrainPaxisIpArbShape, + {"shape_function": jtor.DoublePowerShapeFunction(1.0, 2.0)}, + ), + ( + jtor.ConstrainBetapIpArbShape, + {"shape_function": jtor.DoublePowerShapeFunction(1.0, 2.0)}, + ), + ), +) +def test_psinorm_range(profile_class, kwargs): """Test that the profiles produce finite values outside core""" eq = equilibrium.Equilibrium(Rmin=0.5, Rmax=1.5, Zmin=-1.0, Zmax=1.0, nx=33, ny=33) - for profiles in [ - jtor.ConstrainPaxisIp(eq, 1e3, 2e5, 2.0), - jtor.ConstrainBetapIp(eq, 1.0, 2e5, 2.0), - ]: - current_density = profiles.Jtor(eq.R, eq.Z, eq.psi()) - assert np.all(np.isfinite(current_density)) + profiles = profile_class(eq, 1e3, 2e5, 2.0, **kwargs) - assert profiles.pprime(1.0) == 0.0 - assert profiles.pprime(1.1) == 0.0 - assert np.isfinite(profiles.pprime(-0.32)) + current_density = profiles.Jtor(eq.R, eq.Z, eq.psi()) + assert np.all(np.isfinite(current_density)) - assert profiles.ffprime(1.0) == 0.0 - assert profiles.ffprime(1.1) == 0.0 - assert np.isfinite(profiles.ffprime(-0.32)) + assert profiles.pprime(1.0) == 0.0 + assert profiles.pprime(1.1) == 0.0 + assert np.isfinite(profiles.pprime(-0.32)) + + assert profiles.ffprime(1.0) == 0.0 + assert profiles.ffprime(1.1) == 0.0 + assert np.isfinite(profiles.ffprime(-0.32))