diff --git a/17-spline-profiles.py b/17-spline-profiles.py new file mode 100644 index 0000000..1d15664 --- /dev/null +++ b/17-spline-profiles.py @@ -0,0 +1,141 @@ +#!/usr/bin/env python + +import freegs +import numpy as np +import matplotlib.pyplot as plt + +######################################### +# Create the machine, which specifies coil locations +# and equilibrium, specifying the domain to solve over + +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 + + +######################################### +# Plasma profiles + +profiles = freegs.jtor.ConstrainBetapIp(eq, + 3.214806e-02, # Poloidal beta + 2e5, # Plasma current [Amps] + 2.0) # Vacuum f=R*Bt + +######################################### +# Coil current constraints +# +# 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)] + +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 + +# eq now contains the solution + +print("Done!") + +print("Plasma current: %e Amps" % (eq.plasmaCurrent())) +print("Plasma pressure on axis: %e Pascals" % (eq.pressure(0.0))) +print("Poloidal beta: %e" % (eq.poloidalBeta())) + +# Now, extract the pprime and ffprime profiles from the solution +psi_n_data = np.linspace(0.0,1.0,101,endpoint=True) +pprime_data = eq.pprime(psi_n_data) +ffprime_data = eq.ffprime(psi_n_data) + +# Next, solve again, this time using the above data as a placeholder for bespoke profile data. +# In user-specific implementations this data may be obtained elsewhere, e.g. from MDSplus or +# from the output of another code. + +print('Using splined profiles for pprime and ffprime') + +######################################### +# Create the machine, which specifies coil locations +# and equilibrium, specifying the domain to solve over + +tokamak2 = freegs.machine.TestTokamak() + +eq2 = freegs.Equilibrium(tokamak=tokamak2, + 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 + +profiles2 = freegs.jtor.BetapIpConstrainedSplineProfiles(eq2, + 3.214806e-02, # Poloidal beta + 2e5, # Plasma current [Amps] + 1.0, # Raxis [m], + psi_n_data, # Spline data for normalised psi + pprime_data, # Spline data for pprime + ffprime_data, # Spline data for ffprime + 2.0) # Vacuum f=R*Bt + +######################################### +# Coil current constraints +# +# 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)] + +isoflux = [(1.1,-0.6, 1.1,0.6)] # (R1,Z1, R2,Z2) pair of locations + +constrain2 = freegs.control.constrain(xpoints=xpoints, isoflux=isoflux) + +######################################### +# Nonlinear solve + +freegs.solve(eq2, # The equilibrium to adjust + profiles2, # The toroidal current profile function + constrain2, + show=True) # Constraint function to set coil currents + +# eq2 now contains the solution + +print("Done!") + +print("Plasma current: %e Amps" % (eq2.plasmaCurrent())) +print("Plasma pressure on axis: %e Pascals" % (eq2.pressure(0.0))) +print("Poloidal beta: %e" % (eq2.poloidalBeta())) + +# Compare results - parameterised vs splined profiles +fig, ax = plt.subplots(1,3) + +ax[0].contour(eq.R,eq.Z,eq.psi(),levels=[eq.psi_bndry],colors='r') +ax[0].contour(eq2.R,eq2.Z,eq2.psi(),levels=[eq2.psi_bndry],colors='b') +ax[0].set_aspect('equal') +ax[0].set_xlabel('R (m)') +ax[0].set_ylabel('Z (m)') + +ax[1].plot(psi_n_data,eq.pprime(psi_n_data),color='r') +ax[1].plot(psi_n_data,eq2.pprime(psi_n_data),color='b') +ax[1].set_xlabel(r'$\psi_{N}$') +ax[1].set_ylabel(r'pprime($\psi_{N}$)') + +ax[2].plot(psi_n_data,eq.ffprime(psi_n_data),color='r') +ax[2].plot(psi_n_data,eq2.ffprime(psi_n_data),color='b') +ax[2].set_xlabel(r'$\psi_{N}$') +ax[2].set_ylabel(r'ffprime($\psi_{N}$)') + +plt.show() \ No newline at end of file diff --git a/freegs/critical.py b/freegs/critical.py index f79248c..11127f8 100644 --- a/freegs/critical.py +++ b/freegs/critical.py @@ -405,7 +405,7 @@ def find_separatrix( opoint, xpoint = find_critical(eq.R, eq.Z, psi) psinorm = (psi - opoint[0][2]) / (eq.psi_bndry - opoint[0][2]) - + psifunc = interpolate.RectBivariateSpline(eq.R[:, 0], eq.Z[0, :], psinorm) r0, z0 = opoint[0][0:2] diff --git a/freegs/equilibrium.py b/freegs/equilibrium.py index 4bf5eeb..756e7bd 100644 --- a/freegs/equilibrium.py +++ b/freegs/equilibrium.py @@ -10,7 +10,6 @@ from .boundary import fixedBoundary, freeBoundary from . import critical - from . import polygons # Operators which define the G-S equation @@ -18,12 +17,9 @@ # Multigrid solver from . import multigrid - from . import machine - import matplotlib.pyplot as plt - class Equilibrium: """ Represents the equilibrium state, including @@ -126,7 +122,6 @@ def __init__( self._current = current # Plasma current self.Jtor = None - self._updatePlasmaPsi(psi) # Needs to be after _pgreen # Create the solver @@ -430,6 +425,21 @@ def separatrix(self, npoints=360): :, 0:2 ] + def psi_surfRZ(self, psiN=0.995, npoints=360): + """ + Returns the R,Z of a flux surface specified by a value of psiN. This flux surface is closed on itself. + """ + + surf = critical.find_separatrix(self, opoint=None, xpoint=None, ntheta=npoints, psi=None, axis=None, psival=psiN) + + Rsurf = [point[0] for point in surf] + Zsurf = [point[1] for point in surf] + + Rsurf.append(Rsurf[0]) + Zsurf.append(Zsurf[0]) + + return np.array(Rsurf), np.array(Zsurf) + def solve(self, profiles, Jtor=None, psi=None, psi_bndry=None): """ Calculate the plasma equilibrium given new profiles @@ -626,7 +636,7 @@ def _updatePlasmaPsi(self, plasma_psi): self.psi_func = interpolate.RectBivariateSpline( self.R[:, 0], self.Z[0, :], plasma_psi ) - + # Update the plasma axis and boundary flux as well as mask self._updateBoundaryPsi() @@ -1056,19 +1066,35 @@ def internalInductance(self, npoints=360): integral = romb(romb(B_polvals_2 * dV)) return 2 * integral / (mu0 * mu0 * R_geo * Ip * Ip) - def poloidalBeta(self): - """Calculate plasma poloidal beta by integrating the thermal pressure - and poloidal magnetic field pressure over the plasma volume.""" + def flux_surface_averaged_Bpol2(self, psiN=0.995, npoints=360): + """ + Calculates the flux surface averaged value of the square of the poloidal field. + """ - R = self.R - Z = self.Z + # Get R, Z points of the flux surface + Rsurf, Zsurf = self.psi_surfRZ(psiN=psiN,npoints=npoints) - # Produce array of Bpol in (R,Z) - B_polvals_2 = self.Br(R, Z) ** 2 + self.Bz(R, Z) ** 2 + # Get the poloidal field + Bpol_surf = self.Bpol(Rsurf,Zsurf) - dR = R[1, 0] - R[0, 0] - dZ = Z[0, 1] - Z[0, 0] - dV = 2.0 * np.pi * R * dR * dZ + # Get the square of the poloidal field + Bpol_surf2 = Bpol_surf**2.0 + + # Get dl along the surface + dl = np.sqrt(np.diff(Rsurf)**2.0 + np.diff(Zsurf)**2.0) + dl = np.insert(dl,0,0.0) + + # Get l along the surface + l = np.cumsum(dl) + + # Calculate the flux surface averaged quantity + return np.trapz(x=l, y=Bpol_surf2 * Bpol_surf) / np.trapz(x=l, y=np.ones(np.size(l)) * Bpol_surf) + + def poloidalBeta(self): + """Return the poloidal beta. + + betaP = 2 * mu0 *

/ <> + """ # Normalised psi psi_norm = self.psiN() @@ -1076,12 +1102,10 @@ def poloidalBeta(self): # Plasma pressure pressure = self.pressure(psi_norm) - if self.mask is not None: # Only include points in the core - dV *= self.mask + volume_averaged_pressure = self.calc_volume_averaged(pressure) + line_averaged_Bpol2_lcfs = self.flux_surface_averaged_Bpol2(psiN=1.0) - pressure_integral = romb(romb(pressure * dV)) - field_integral_pol = romb(romb(B_polvals_2 * dV)) - return 2 * mu0 * pressure_integral / field_integral_pol + return (2.0 * mu0 * volume_averaged_pressure) / line_averaged_Bpol2_lcfs def poloidalBeta2(self): """Return the poloidal beta @@ -1250,6 +1274,26 @@ def qcyl(self): return val + def calc_volume_integrated(self,field): + """ + Calculates the volume integral of the input field. + """ + + dV = 2.0 * np.pi * self.R * self.dR *self.dZ + + if self.mask is not None: # Only include points in the core + dV *= self.mask + + return romb(romb(field * dV)) + + def calc_volume_averaged(self,field): + """ + Calculates the volume average of the input field. + """ + + volume_integrated_field = self.calc_volume_integrated(field) + + return volume_integrated_field / self.plasmaVolume() def refine(eq, nx=None, ny=None): """ @@ -1291,7 +1335,6 @@ def refine(eq, nx=None, ny=None): return result - def coarsen(eq): """ Reduce grid resolution, returning a new equilibrium @@ -1319,7 +1362,6 @@ def coarsen(eq): return result - def newDomain(eq, Rmin=None, Rmax=None, Zmin=None, Zmax=None, nx=None, ny=None): """Creates a new Equilibrium, solving in a different domain. The domain size (Rmin, Rmax, Zmin, Zmax) and resolution (nx,ny) @@ -1372,7 +1414,6 @@ def newDomain(eq, Rmin=None, Rmax=None, Zmin=None, Zmax=None, nx=None, ny=None): return result - if __name__ == "__main__": # Test the different spline interpolation routines diff --git a/freegs/jtor.py b/freegs/jtor.py index 2e35bb8..2f5a9e0 100644 --- a/freegs/jtor.py +++ b/freegs/jtor.py @@ -19,6 +19,7 @@ along with FreeGS. If not, see . """ from scipy.integrate import romb, quad # Romberg integration +from scipy.interpolate import interp1d from . import critical from .gradshafranov import mu0 @@ -26,7 +27,6 @@ import numpy as np import abc - class Profile(abc.ABC): """ Base class from which profiles classes can inherit @@ -46,10 +46,22 @@ def pressure(self, psinorm, out=None): """ if not hasattr(psinorm, "shape"): + # Assume a single value - val, _ = quad(self.pprime, psinorm, 1.0) + if hasattr(self,'psi_n_points'): + + # Splined profiles - use prepared spline of pprime integral + val = self.L * (self.Beta0 / self.Raxis) * self.pprime_int_spline(psinorm) + + else: + + # Parabolic profiles - integrate pprime + val, _ = quad(self.pprime, psinorm, 1.0) + # Convert from integral in normalised psi to integral in psi - return val * (self.psi_axis - self.psi_bndry) + val *= -(self.psi_bndry - self.psi_axis) + + return val # Assume a NumPy array @@ -63,9 +75,19 @@ def pressure(self, psinorm, out=None): raise ValueError("Input and output arrays of different lengths") for i in range(len(pvals)): - val, _ = quad(self.pprime, pvals[i], 1.0) + + if hasattr(self,'psi_n_points'): + + # Splined profiles - use prepared spline of pprime integral + val = self.L * self.Beta0 / self.Raxis *self.pprime_int_spline(pvals[i]) + + else: + + # Parabolic profiles - integrate pprime + val, _ = quad(self.pprime, pvals[i], 1.0) + # Convert from integral in normalised psi to integral in psi - val *= self.psi_axis - self.psi_bndry + val *= -(self.psi_bndry - self.psi_axis) ovals[i] = val return reshape(ovals, psinorm.shape) @@ -76,12 +98,21 @@ def fpol(self, psinorm, out=None): """ - if not hasattr(psinorm, "__len__"): - # Assume a single value + if not hasattr(psinorm, "shape"): + + # Assume a single value + if hasattr(self,'psi_n_points'): + + # Splined profiles - use prepared spline of ffprime integral + val = self.L * (1-self.Beta0) * self.Raxis *self.ffprime_int_spline(psinorm) + + else: + + # Parabolic profiles - integrate ffprime + val, _ = quad(self.ffprime, psinorm, 1.0) - val, _ = quad(self.ffprime, psinorm, 1.0) # Convert from integral in normalised psi to integral in psi - val *= self.psi_axis - self.psi_bndry + val *= -(self.psi_bndry - self.psi_axis) # ffprime = 0.5*d/dpsi(f^2) # Apply boundary condition at psinorm=1 val = fvac**2 @@ -99,9 +130,19 @@ def fpol(self, psinorm, out=None): if len(pvals) != len(ovals): raise ValueError("Input and output arrays of different lengths") for i in range(len(pvals)): - val, _ = quad(self.ffprime, pvals[i], 1.0) + + if hasattr(self,'psi_n_points'): + + # Splined profiles - use prepared spline of ffprime integral + val = self.L * (1-self.Beta0) * self.Raxis *self.ffprime_int_spline(pvals[i]) + + else: + + # Parabolic profiles - integrate ffprime + val, _ = quad(self.ffprime, pvals[i], 1.0) + # Convert from integral in normalised psi to integral in psi - val *= self.psi_axis - self.psi_bndry + val *= -(self.psi_bndry - self.psi_axis) # ffprime = 0.5*d/dpsi(f^2) # Apply boundary condition at psinorm=1 val = fvac**2 @@ -135,7 +176,6 @@ def fvac(self) -> float: """Return f = R*Bt in vacuum""" pass - class ConstrainBetapIp(Profile): """ Constrain poloidal Beta and plasma current @@ -178,11 +218,6 @@ def Jtor(self, R, Z, psi, psi_bndry=None): L and Beta0 are parameters which are set by constraints """ - # Intermediary update of the plasma - # boundary and axis flux - self.eq._updateBoundaryPsi(psi) - psi_bndry = self.eq.psi_bndry - # Analyse the equilibrium, finding O- and X-points opt, xpt = critical.find_critical(R, Z, psi) if not opt: @@ -220,7 +255,7 @@ def Jtor(self, R, Z, psi, psi_bndry=None): # 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 + lambda x: (1.0 - x ** self.alpha_m) ** self.alpha_n, psinorm, 1.0 ) shapeintegral *= psi_bndry - psi_axis return shapeintegral @@ -240,18 +275,13 @@ def pshape(psinorm): pfunc *= mask # Integrate over plasma - # betap = (2mu0) * (int(p)RdRdZ)/(int(B_poloidal**2)RdRdZ) - # = - (2L*Beta0*mu0/Raxis) * (pfunc*RdRdZ)/((int(B_poloidal**2)RdRdZ)) - - # Produce array of Bpol in (R,Z) for core plasma - B_polvals_2 = self.eq.Br(R, Z) ** 2 + self.eq.Bz(R, Z) ** 2 - if mask is not None: - B_polvals_2 *= mask + # betap = (2 mu0) * volume_av(p) / (flux_surf_av(B_poloidal**2)) + # = - (2 mu0 * L * Beta0 / Raxis) * volume_av(pfunc) / (flux_surf_av(B_poloidal**2)) - p_int = romb(romb(pfunc * R)) * dR * dZ - b_int = romb(romb(B_polvals_2 * R)) * dR * dZ + p_int = self.eq.calc_volume_averaged(pfunc) + b_int = self.eq.flux_surface_averaged_Bpol2(psiN=1.0) - # self.betap = - (2*LBeta0*mu0/ self.Raxis) * (p_int/b_int) + # self.betap = - (2 mu0 * L * Beta0 / Raxis) * (p_int/b_int) LBeta0 = (b_int / p_int) * (-self.betap * self.Raxis) / (2 * mu0) # Integrate current components @@ -268,13 +298,14 @@ def pshape(psinorm): L = self.Ip / I_R - LBeta0 * (IR / I_R - 1) Beta0 = LBeta0 / L - # print("Constraints: L = %e, Beta0 = %e" % (L, Beta0)) + print("Constraints: L = %e, Beta0 = %e" % (L, Beta0)) # Toroidal current Jtor = L * (Beta0 * R / self.Raxis + (1 - Beta0) * self.Raxis / R) * jtorshape self.L = L self.Beta0 = Beta0 + self.psi_bndry = psi_bndry self.psi_axis = psi_axis @@ -340,11 +371,6 @@ def Jtor(self, R, Z, psi, psi_bndry=None): L and Beta0 are parameters which are set by constraints """ - # Intermediary update of the plasma - # boundary and axis flux - self.eq._updateBoundaryPsi(psi) - psi_bndry = self.eq.psi_bndry - # Analyse the equilibrium, finding O- and X-points opt, xpt = critical.find_critical(R, Z, psi) if not opt: @@ -381,7 +407,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 + lambda x: (1.0 - x ** self.alpha_m) ** self.alpha_n, 0.0, 1.0 ) shapeintegral *= psi_bndry - psi_axis @@ -390,6 +416,8 @@ def Jtor(self, R, Z, psi, psi_bndry=None): # paxis = - (L*Beta0/Raxis) * shapeintegral # + LBeta0 = -self.paxis * self.Raxis / shapeintegral + # Integrate current components IR = romb(romb(jtorshape * R / self.Raxis)) * dR * dZ I_R = romb(romb(jtorshape * self.Raxis / R)) * dR * dZ @@ -400,12 +428,10 @@ def Jtor(self, R, Z, psi, psi_bndry=None): # = L*Beta0*(IR - I_R) + L*I_R # - LBeta0 = -self.paxis * self.Raxis / shapeintegral - L = self.Ip / I_R - LBeta0 * (IR / I_R - 1) Beta0 = LBeta0 / L - # print("Constraints: L = %e, Beta0 = %e" % (L, Beta0)) + print("Constraints: L = %e, Beta0 = %e" % (L, Beta0)) # Toroidal current Jtor = L * (Beta0 * R / self.Raxis + (1 - Beta0) * self.Raxis / R) * jtorshape @@ -424,7 +450,7 @@ def pprime(self, pn): """ shape = (1.0 - np.clip(pn, 0.0, 1.0) ** self.alpha_m) ** self.alpha_n return self.L * self.Beta0 / self.Raxis * shape - + def ffprime(self, pn): """ f * df/dpsi as a function of normalised psi. 0 outside core. @@ -436,8 +462,1032 @@ def ffprime(self, pn): def fvac(self): return self._fvac +class BetapIpConstrainedSplineProfiles(Profile): + """ + BetaP and Ip-constrained custom (splined) internal plasma profiles. + + """ + + def __init__(self, eq=None, betap=None, Ip=None, Raxis=None, psi_n=None, pprime=None, ffprime=None, fvac=None): + """ + eq - Equilibrium object + betap - Poloidal beta + Ip - Plasma current [Amps] + Raxis - R used in p' and ff' components + psi_n - Normalised (0,1) poloidal flux used to defined the profiles + pprime - Pressure gradient - dp/dpsi + ffprime - f*dfpol/dpsi + fvac - Vacuum f = R*Bt + + """ + + # Check inputs + if eq is None: + raise ValueError("No equilibrium object provided") + if betap is None: + raise ValueError("No betap value provided") + if Ip is None: + raise ValueError("No plasma current value provided") + if Raxis is None: + raise ValueError("No Raxis value provided") + if psi_n is None: + raise ValueError("No psi_n data provided") + if pprime is None: + raise ValueError("No pprime data provided") + if ffprime is None: + raise ValueError("No ffprime data provided") + if fvac is None: + raise ValueError("No fvac data provided") + + # Set values for later use + self.eq = eq + self.betap = betap + self.Ip = Ip + self.Raxis = Raxis + self.psi_n_points = psi_n + self.pprime_points = pprime + self.ffprime_points = ffprime + self._fvac = fvac + + # Create 1D splines for the internal profiles - these will be like jtorshape + self.pprime_spline = interp1d(self.psi_n_points,self.pprime_points,kind='linear',fill_value='extrapolate',bounds_error=False) + self.ffprime_spline = interp1d(self.psi_n_points,self.ffprime_points,kind='linear',fill_value='extrapolate',bounds_error=False) + + # Create 1D splines for the integral of pprime, ffprime + pn_points = np.linspace(0.0,1.0,100,endpoint=True) + + def pprime_int_func(pn): + val, _ = quad(self.pprime_spline,pn,1.0) + return val + + pprime_int_vals = [] + for pn in pn_points: + pprime_int_vals.append(pprime_int_func(pn)) + + pprime_int_vals = np.asarray(pprime_int_vals) + + self.pprime_int_spline = interp1d(pn_points,pprime_int_vals,kind='linear',fill_value='extrapolate',bounds_error=False) + + def ffprime_int_func(pn): + val, _ = quad(self.ffprime_spline,pn,1.0) + return val + + ffprime_int_vals = [] + for pn in pn_points: + ffprime_int_vals.append(ffprime_int_func(pn)) + + ffprime_int_vals = np.asarray(ffprime_int_vals) + + self.ffprime_int_spline = interp1d(pn_points,ffprime_int_vals,kind='linear',fill_value='extrapolate',bounds_error=False) + + def Jtor(self, R, Z, psi, psi_bndry=None): + """Calculate toroidal plasma current + + Jtor = R*pprime + ffprime/(R * mu0) + """ + + # Analyse the equilibrium, finding O- and X-points + opt, xpt = critical.find_critical(R, Z, psi) + if not opt: + raise ValueError("No O-points found!") + psi_axis = opt[0][2] + + if psi_bndry is not None: + mask = critical.core_mask(R, Z, psi, opt, xpt, psi_bndry) + elif xpt: + psi_bndry = xpt[0][2] + mask = critical.core_mask(R, Z, psi, opt, xpt) + else: + # No X-points + psi_bndry = psi[0, 0] + mask = None + + dR = R[1, 0] - R[0, 0] + dZ = Z[0, 1] - Z[0, 0] + + # Calculate normalised psi. + # 0 = magnetic axis + # 1 = plasma boundary + psi_norm = (psi - psi_axis) / (psi_bndry - psi_axis) + + pprime_shape = self.pprime_spline(psi_norm) + ffprime_shape = self.ffprime_spline(psi_norm) + + if mask is not None: + pprime_shape *= mask + ffprime_shape *= mask + + # Now apply constraints to define constants + + # Need integral of pprime_shape to calculate pressure + # as p(psinorm) = - (L*Beta0/Raxis) * pshape(psinorm) + + def pshape(psinorm): + shapeintegral = self.pprime_int_spline(psinorm) + shapeintegral *= psi_bndry - psi_axis + return shapeintegral + + nx, ny = psi_norm.shape + pfunc = zeros((nx, ny)) + for i in range(1, nx - 1): + for j in range(1, ny - 1): + if (psi_norm[i, j] >= 0.0) and (psi_norm[i, j] < 1.0): + pfunc[i, j] = pshape(psi_norm[i, j]) + if mask is not None: + pfunc *= mask + + # Integrate over plasma + # betap = (2 mu0) * volume_av(p) / (flux_surf_av(B_poloidal**2)) + # = - (2 mu0 * L * Beta0 / Raxis) * volume_av(pfunc) / (flux_surf_av(B_poloidal**2)) + p_int = self.eq.calc_volume_averaged(pfunc) + b_int = self.eq.flux_surface_averaged_Bpol2(psiN=1.0) + + # self.betap = - (2 mu0 * L * Beta0 / Raxis) * (p_int/b_int) + LBeta0 = (b_int / p_int) * (-self.betap * self.Raxis) / (2 * mu0) + + # Integrate current components + IR = romb(romb(pprime_shape * R/self.Raxis)) * dR*dZ # pprime component + I_R = romb(romb(ffprime_shape * self.Raxis/(R*mu0))) * dR*dZ # ffprime component + + # Toroidal plasma current Ip is + # + # Ip = L * (Beta0 * IR + (1-Beta0)*I_R) + # = L*Beta0*(IR - I_R) + L*I_R + # + # L = self.Ip / ( (Beta0*IR) + ((1.0-Beta0)*(I_R)) ) + + L = self.Ip/I_R - LBeta0*(IR/I_R - 1) + Beta0 = LBeta0 / L + + print("Constraints: L = %e, Beta0 = %e" % (L, Beta0)) + + # Toroidal current + Jtor = L *( (pprime_shape * Beta0 * R / self.Raxis) + ((1 - Beta0) * self.Raxis * ffprime_shape/ (R * mu0)) ) + + self.L = L + self.Beta0 = Beta0 + self.psi_bndry = psi_bndry + self.psi_axis = psi_axis + + return Jtor + + # Profile functions + def pprime(self, pn): + """ + dp/dpsi as a function of normalised psi. 0 outside core. + Calculate pprimeshape inside the core only + """ + shape = self.pprime_spline(pn) + return self.L * self.Beta0 / self.Raxis * shape + + def ffprime(self, pn): + """ + f * df/dpsi as a function of normalised psi. 0 outside core. + Calculate ffprimeshape inside the core only. + """ + shape = self.ffprime_spline(pn) + return self.L * (1 - self.Beta0) * self.Raxis * shape + + def fvac(self): + return self._fvac + +class PaxisIpConstrainedSplineProfiles(Profile): + """ + Paxis and Ip-constrained custom (splined) internal plasma profiles. + + """ + + def __init__(self, eq=None, paxis=None, Ip=None, Raxis=None, psi_n=None, pprime=None, ffprime=None, fvac=None): + """ + eq - Equilibrium object + paxis - Pressure at magnetic axis [Pa] + Ip - Plasma current [Amps] + Raxis - R used in p' and ff' components + psi_n - Normalised (0,1) poloidal flux used to defined the profiles + pprime - Pressure gradient - dp/dpsi + ffprime - f*dfpol/dpsi + fvac - Vacuum f = R*Bt + + """ + + # Check inputs + if eq is None: + raise ValueError("No equilibrium object provided") + if paxis is None: + raise ValueError("No paxis value provided") + if Ip is None: + raise ValueError("No plasma current value provided") + if Raxis is None: + raise ValueError("No Raxis value provided") + if psi_n is None: + raise ValueError("No psi_n data provided") + if pprime is None: + raise ValueError("No pprime data provided") + if ffprime is None: + raise ValueError("No ffprime data provided") + if fvac is None: + raise ValueError("No fvac data provided") + + # Set values for later use + self.eq = eq + self.paxis = paxis + self.Ip = Ip + self.Raxis = Raxis + self.psi_n_points = psi_n + self.pprime_points = pprime + self.ffprime_points = ffprime + self._fvac = fvac + + # Create 1D splines for the internal profiles - these will be like jtorshape + self.pprime_spline = interp1d(self.psi_n_points,self.pprime_points,kind='linear',fill_value='extrapolate',bounds_error=False) + self.ffprime_spline = interp1d(self.psi_n_points,self.ffprime_points,kind='linear',fill_value='extrapolate',bounds_error=False) + + # Create 1D splines for the integral of pprime, ffprime + pn_points = np.linspace(0.0,1.0,100,endpoint=True) + + def pprime_int_func(pn): + val, _ = quad(self.pprime_spline,pn,1.0) + return val + + pprime_int_vals = [] + for pn in pn_points: + pprime_int_vals.append(pprime_int_func(pn)) + + pprime_int_vals = np.asarray(pprime_int_vals) + + self.pprime_int_spline = interp1d(pn_points,pprime_int_vals,kind='linear',fill_value='extrapolate',bounds_error=False) + + def ffprime_int_func(pn): + val, _ = quad(self.ffprime_spline,pn,1.0) + return val + + ffprime_int_vals = [] + for pn in pn_points: + ffprime_int_vals.append(ffprime_int_func(pn)) + + ffprime_int_vals = np.asarray(ffprime_int_vals) + + self.ffprime_int_spline = interp1d(pn_points,ffprime_int_vals,kind='linear',fill_value='extrapolate',bounds_error=False) + + def Jtor(self, R, Z, psi, psi_bndry=None): + """Calculate toroidal plasma current + + Jtor = R*pprime + ffprime/(R * mu0) + """ + + # Analyse the equilibrium, finding O- and X-points + opt, xpt = critical.find_critical(R, Z, psi) + if not opt: + raise ValueError("No O-points found!") + psi_axis = opt[0][2] + + if psi_bndry is not None: + mask = critical.core_mask(R, Z, psi, opt, xpt, psi_bndry) + elif xpt: + psi_bndry = xpt[0][2] + mask = critical.core_mask(R, Z, psi, opt, xpt) + else: + # No X-points + psi_bndry = psi[0, 0] + mask = None + + dR = R[1, 0] - R[0, 0] + dZ = Z[0, 1] - Z[0, 0] + + # Calculate normalised psi. + # 0 = magnetic axis + # 1 = plasma boundary + psi_norm = (psi - psi_axis) / (psi_bndry - psi_axis) + + pprime_shape = self.pprime_spline(psi_norm) + ffprime_shape = self.ffprime_spline(psi_norm) + + if mask is not None: + pprime_shape *= mask + ffprime_shape *= mask + + # Now apply constraints to define constants + + # Need integral of pprime_shape to calculate pressure + # as p(psinorm) = - (L*Beta0/Raxis) * pshape(psinorm) + + shapeintegral = self.pprime_int_spline(0.0) + shapeintegral *= psi_bndry - psi_axis + + # Pressure on axis is + # + # paxis = - (L*Beta0/Raxis) * shapeintegral + # + + LBeta0 = -self.paxis * self.Raxis / shapeintegral + + # Integrate current components + IR = romb(romb(pprime_shape * R/self.Raxis)) * dR*dZ # pprime component + I_R = romb(romb(ffprime_shape * self.Raxis/(R*mu0))) * dR*dZ # ffprime component + + # Toroidal plasma current Ip is + # + # Ip = L * (Beta0 * IR + (1-Beta0)*I_R) + # = L*Beta0*(IR - I_R) + L*I_R + # + # L = self.Ip / ( (Beta0*IR) + ((1.0-Beta0)*(I_R)) ) + + L = self.Ip/I_R - LBeta0*(IR/I_R - 1) + Beta0 = LBeta0 / L + + print("Constraints: L = %e, Beta0 = %e" % (L, Beta0)) + + # Toroidal current + Jtor = L *( (pprime_shape * Beta0 * R / self.Raxis) + ((1 - Beta0) * self.Raxis * ffprime_shape/ (R * mu0)) ) + + self.L = L + self.Beta0 = Beta0 + + self.psi_bndry = psi_bndry + self.psi_axis = psi_axis + + return Jtor + + # Profile functions + def pprime(self, pn): + """ + dp/dpsi as a function of normalised psi. 0 outside core. + Calculate pprimeshape inside the core only + """ + shape = self.pprime_spline(pn) + return self.L * self.Beta0 / self.Raxis * shape + + def ffprime(self, pn): + """ + f * df/dpsi as a function of normalised psi. 0 outside core. + Calculate ffprimeshape inside the core only. + """ + shape = self.ffprime_spline(pn) + return self.L * (1 - self.Beta0) * self.Raxis * shape + + def fvac(self): + return self._fvac + +class PprimeIpConstrainedSplineProfiles(Profile): + """ + Pprime and Ip-constrained custom (splined) internal plasma profiles. + + """ -class ProfilesPprimeFfprime(Profile): + def __init__(self, eq=None, Ip=None, Raxis=None, psi_n=None, pprime=None, ffprime=None, fvac=None): + """ + eq - Equilibrium object + Ip - Plasma current [Amps] + Raxis - R used in p' and ff' components + psi_n - Normalised (0,1) poloidal flux used to defined the profiles + pprime - Pressure gradient - dp/dpsi + ffprime - f*dfpol/dpsi + fvac - Vacuum f = R*Bt + + """ + + # Check inputs + if eq is None: + raise ValueError("No equilibrium object provided") + if Ip is None: + raise ValueError("No plasma current value provided") + if Raxis is None: + raise ValueError("No Raxis value provided") + if psi_n is None: + raise ValueError("No psi_n data provided") + if pprime is None: + raise ValueError("No pprime data provided") + if ffprime is None: + raise ValueError("No ffprime data provided") + if fvac is None: + raise ValueError("No fvac data provided") + + # Set values for later use + self.eq = eq + self.Ip = Ip + self.Raxis = Raxis + self.psi_n_points = psi_n + self.pprime_points = pprime + self.ffprime_points = ffprime + self._fvac = fvac + + # Create 1D splines for the internal profiles - these will be like jtorshape + self.pprime_spline = interp1d(self.psi_n_points,self.pprime_points,kind='linear',fill_value='extrapolate',bounds_error=False) + self.ffprime_spline = interp1d(self.psi_n_points,self.ffprime_points,kind='linear',fill_value='extrapolate',bounds_error=False) + + # Create 1D splines for the integral of pprime, ffprime + pn_points = np.linspace(0.0,1.0,100,endpoint=True) + + def pprime_int_func(pn): + val, _ = quad(self.pprime_spline,pn,1.0) + return val + + pprime_int_vals = [] + for pn in pn_points: + pprime_int_vals.append(pprime_int_func(pn)) + + pprime_int_vals = np.asarray(pprime_int_vals) + + self.pprime_int_spline = interp1d(pn_points,pprime_int_vals,kind='linear',fill_value='extrapolate',bounds_error=False) + + def ffprime_int_func(pn): + val, _ = quad(self.ffprime_spline,pn,1.0) + return val + + ffprime_int_vals = [] + for pn in pn_points: + ffprime_int_vals.append(ffprime_int_func(pn)) + + ffprime_int_vals = np.asarray(ffprime_int_vals) + + self.ffprime_int_spline = interp1d(pn_points,ffprime_int_vals,kind='linear',fill_value='extrapolate',bounds_error=False) + + def Jtor(self, R, Z, psi, psi_bndry=None): + """Calculate toroidal plasma current + + Jtor = R*pprime + ffprime/(R * mu0) + """ + + # Analyse the equilibrium, finding O- and X-points + opt, xpt = critical.find_critical(R, Z, psi) + if not opt: + raise ValueError("No O-points found!") + psi_axis = opt[0][2] + + if psi_bndry is not None: + mask = critical.core_mask(R, Z, psi, opt, xpt, psi_bndry) + elif xpt: + psi_bndry = xpt[0][2] + mask = critical.core_mask(R, Z, psi, opt, xpt) + else: + # No X-points + psi_bndry = psi[0, 0] + mask = None + + dR = R[1, 0] - R[0, 0] + dZ = Z[0, 1] - Z[0, 0] + + # Calculate normalised psi. + # 0 = magnetic axis + # 1 = plasma boundary + psi_norm = (psi - psi_axis) / (psi_bndry - psi_axis) + + pprime_shape = self.pprime_spline(psi_norm) + ffprime_shape = self.ffprime_spline(psi_norm) + + if mask is not None: + pprime_shape *= mask + ffprime_shape *= mask + + # Now apply constraints to define constants + + # Constraining pprime = pprime_spline + # pprime(psinorm) = (LBeta0 / Raxis) * pprime_spline(psinorm) + # Hence LBeta0 = Raxis + LBeta0 = self.Raxis + + # Integrate current components + IR = romb(romb(pprime_shape * R/self.Raxis)) * dR*dZ # pprime component + I_R = romb(romb(ffprime_shape * self.Raxis/(R*mu0))) * dR*dZ # ffprime component + + # Toroidal plasma current Ip is + # + # Ip = L * (Beta0 * IR + (1-Beta0)*I_R) + # = L*Beta0*(IR - I_R) + L*I_R + # + # L = self.Ip / ( (Beta0*IR) + ((1.0-Beta0)*(I_R)) ) + + L = self.Ip/I_R - LBeta0*(IR/I_R - 1) + Beta0 = LBeta0 / L + + print("Constraints: L = %e, Beta0 = %e" % (L, Beta0)) + + # Toroidal current + Jtor = L *( (pprime_shape * Beta0 * R / self.Raxis) + ((1 - Beta0) * self.Raxis * ffprime_shape/ (R * mu0)) ) + + self.L = L + self.Beta0 = Beta0 + + self.psi_bndry = psi_bndry + self.psi_axis = psi_axis + + return Jtor + + # Profile functions + def pprime(self, pn): + """ + dp/dpsi as a function of normalised psi. 0 outside core. + Calculate pprimeshape inside the core only + """ + shape = self.pprime_spline(pn) + return self.L * self.Beta0 / self.Raxis * shape + + def ffprime(self, pn): + """ + f * df/dpsi as a function of normalised psi. 0 outside core. + Calculate ffprimeshape inside the core only. + """ + shape = self.ffprime_spline(pn) + return self.L * (1 - self.Beta0) * self.Raxis * shape + + def fvac(self): + return self._fvac + +class BetapFfprimeConstrainedSplineProfiles(Profile): + """ + BetaP and Ffprime-constrained custom (splined) internal plasma profiles. + + """ + + def __init__(self, eq=None, betap=None, Raxis=None, psi_n=None, pprime=None, ffprime=None, fvac=None): + """ + eq - Equilibrium object + betap - Poloidal beta + Raxis - R used in p' and ff' components + psi_n - Normalised (0,1) poloidal flux used to defined the profiles + pprime - Pressure gradient - dp/dpsi + ffprime - f*dfpol/dpsi + fvac - Vacuum f = R*Bt + + """ + + # Check inputs + if eq is None: + raise ValueError("No equilibrium object provided") + if betap is None: + raise ValueError("No betap value provided") + if Raxis is None: + raise ValueError("No Raxis value provided") + if psi_n is None: + raise ValueError("No psi_n data provided") + if pprime is None: + raise ValueError("No pprime data provided") + if ffprime is None: + raise ValueError("No ffprime data provided") + if fvac is None: + raise ValueError("No fvac data provided") + + # Set values for later use + self.eq = eq + self.betap = betap + self.Raxis = Raxis + self.psi_n_points = psi_n + self.pprime_points = pprime + self.ffprime_points = ffprime + self._fvac = fvac + + # Create 1D splines for the internal profiles - these will be like jtorshape + self.pprime_spline = interp1d(self.psi_n_points,self.pprime_points,kind='linear',fill_value='extrapolate',bounds_error=False) + self.ffprime_spline = interp1d(self.psi_n_points,self.ffprime_points,kind='linear',fill_value='extrapolate',bounds_error=False) + + # Create 1D splines for the integral of pprime, ffprime + pn_points = np.linspace(0.0,1.0,100,endpoint=True) + + def pprime_int_func(pn): + val, _ = quad(self.pprime_spline,pn,1.0) + return val + + pprime_int_vals = [] + for pn in pn_points: + pprime_int_vals.append(pprime_int_func(pn)) + + pprime_int_vals = np.asarray(pprime_int_vals) + + self.pprime_int_spline = interp1d(pn_points,pprime_int_vals,kind='linear',fill_value='extrapolate',bounds_error=False) + + def ffprime_int_func(pn): + val, _ = quad(self.ffprime_spline,pn,1.0) + return val + + ffprime_int_vals = [] + for pn in pn_points: + ffprime_int_vals.append(ffprime_int_func(pn)) + + ffprime_int_vals = np.asarray(ffprime_int_vals) + + self.ffprime_int_spline = interp1d(pn_points,ffprime_int_vals,kind='linear',fill_value='extrapolate',bounds_error=False) + + def Jtor(self, R, Z, psi, psi_bndry=None): + """Calculate toroidal plasma current + + Jtor = R*pprime + ffprime/(R * mu0) + """ + + # Analyse the equilibrium, finding O- and X-points + opt, xpt = critical.find_critical(R, Z, psi) + if not opt: + raise ValueError("No O-points found!") + psi_axis = opt[0][2] + + if psi_bndry is not None: + mask = critical.core_mask(R, Z, psi, opt, xpt, psi_bndry) + elif xpt: + psi_bndry = xpt[0][2] + mask = critical.core_mask(R, Z, psi, opt, xpt) + else: + # No X-points + psi_bndry = psi[0, 0] + mask = None + + dR = R[1, 0] - R[0, 0] + dZ = Z[0, 1] - Z[0, 0] + + # Calculate normalised psi. + # 0 = magnetic axis + # 1 = plasma boundary + psi_norm = (psi - psi_axis) / (psi_bndry - psi_axis) + + pprime_shape = self.pprime_spline(psi_norm) + ffprime_shape = self.ffprime_spline(psi_norm) + + if mask is not None: + pprime_shape *= mask + ffprime_shape *= mask + + # Now apply constraints to define constants + + # Need integral of pprime_shape to calculate pressure + # as p(psinorm) = - (L*Beta0/Raxis) * pshape(psinorm) + + def pshape(psinorm): + shapeintegral = self.pprime_int_spline(psinorm) + shapeintegral *= psi_bndry - psi_axis + return shapeintegral + + nx, ny = psi_norm.shape + pfunc = zeros((nx, ny)) + for i in range(1, nx - 1): + for j in range(1, ny - 1): + if (psi_norm[i, j] >= 0.0) and (psi_norm[i, j] < 1.0): + pfunc[i, j] = pshape(psi_norm[i, j]) + if mask is not None: + pfunc *= mask + + # Integrate over plasma + # betap = (2 mu0) * volume_av(p) / (flux_surf_av(B_poloidal**2)) + # = - (2 mu0 * L * Beta0 / Raxis) * volume_av(pfunc) / (flux_surf_av(B_poloidal**2)) + + p_int = self.eq.calc_volume_averaged(pfunc) + b_int = self.eq.flux_surface_averaged_Bpol2(psiN=1.0) + + # self.betap = - (2 mu0 * L * Beta0 / Raxis) * (p_int/b_int) + LBeta0 = (b_int / p_int) * (-self.betap * self.Raxis) / (2 * mu0) + + # Constrain Ffprime = ffprime_spline + # Ffprime = L*(1-Beta0)*Raxis*ffprime_spline + # L*(1-Beta0)*Raxis = 1 + # L = LBeta0 + (1.0/Raxis) + + L = LBeta0 + (1.0/self.Raxis) + Beta0 = LBeta0 / L + + print("Constraints: L = %e, Beta0 = %e" % (L, Beta0)) + + # Toroidal current + Jtor = L *( (pprime_shape * Beta0 * R / self.Raxis) + ((1 - Beta0) * self.Raxis * ffprime_shape/ (R * mu0)) ) + + self.L = L + self.Beta0 = Beta0 + + self.psi_bndry = psi_bndry + self.psi_axis = psi_axis + + return Jtor + + # Profile functions + def pprime(self, pn): + """ + dp/dpsi as a function of normalised psi. 0 outside core. + Calculate pprimeshape inside the core only + """ + shape = self.pprime_spline(pn) + return self.L * self.Beta0 / self.Raxis * shape + + def ffprime(self, pn): + """ + f * df/dpsi as a function of normalised psi. 0 outside core. + Calculate ffprimeshape inside the core only. + """ + shape = self.ffprime_spline(pn) + return self.L * (1 - self.Beta0) * self.Raxis * shape + + def fvac(self): + return self._fvac + +class PaxisFfprimeConstrainedSplineProfiles(Profile): + """ + Paxis and Ffprime-constrained custom (splined) internal plasma profiles. + + """ + + def __init__(self, eq=None, paxis=None, Raxis=None, psi_n=None, pprime=None, ffprime=None, fvac=None): + """ + eq - Equilibrium object + betap - Poloidal beta + Raxis - R used in p' and ff' components + psi_n - Normalised (0,1) poloidal flux used to defined the profiles + pprime - Pressure gradient - dp/dpsi + ffprime - f*dfpol/dpsi + fvac - Vacuum f = R*Bt + + """ + + # Check inputs + if eq is None: + raise ValueError("No equilibrium object provided") + if paxis is None: + raise ValueError("No paxis value provided") + if Raxis is None: + raise ValueError("No Raxis value provided") + if psi_n is None: + raise ValueError("No psi_n data provided") + if pprime is None: + raise ValueError("No pprime data provided") + if ffprime is None: + raise ValueError("No ffprime data provided") + if fvac is None: + raise ValueError("No fvac data provided") + + # Set values for later use + self.eq = eq + self.paxis = paxis + self.Raxis = Raxis + self.psi_n_points = psi_n + self.pprime_points = pprime + self.ffprime_points = ffprime + self._fvac = fvac + + # Create 1D splines for the internal profiles - these will be like jtorshape + self.pprime_spline = interp1d(self.psi_n_points,self.pprime_points,kind='linear',fill_value='extrapolate',bounds_error=False) + self.ffprime_spline = interp1d(self.psi_n_points,self.ffprime_points,kind='linear',fill_value='extrapolate',bounds_error=False) + + # Create 1D splines for the integral of pprime, ffprime + pn_points = np.linspace(0.0,1.0,100,endpoint=True) + + def pprime_int_func(pn): + val, _ = quad(self.pprime_spline,pn,1.0) + return val + + pprime_int_vals = [] + for pn in pn_points: + pprime_int_vals.append(pprime_int_func(pn)) + + pprime_int_vals = np.asarray(pprime_int_vals) + + self.pprime_int_spline = interp1d(pn_points,pprime_int_vals,kind='linear',fill_value='extrapolate',bounds_error=False) + + def ffprime_int_func(pn): + val, _ = quad(self.ffprime_spline,pn,1.0) + return val + + ffprime_int_vals = [] + for pn in pn_points: + ffprime_int_vals.append(ffprime_int_func(pn)) + + ffprime_int_vals = np.asarray(ffprime_int_vals) + + self.ffprime_int_spline = interp1d(pn_points,ffprime_int_vals,kind='linear',fill_value='extrapolate',bounds_error=False) + + def Jtor(self, R, Z, psi, psi_bndry=None): + """Calculate toroidal plasma current + + Jtor = R*pprime + ffprime/(R * mu0) + """ + + # Analyse the equilibrium, finding O- and X-points + opt, xpt = critical.find_critical(R, Z, psi) + if not opt: + raise ValueError("No O-points found!") + psi_axis = opt[0][2] + + if psi_bndry is not None: + mask = critical.core_mask(R, Z, psi, opt, xpt, psi_bndry) + elif xpt: + psi_bndry = xpt[0][2] + mask = critical.core_mask(R, Z, psi, opt, xpt) + else: + # No X-points + psi_bndry = psi[0, 0] + mask = None + + dR = R[1, 0] - R[0, 0] + dZ = Z[0, 1] - Z[0, 0] + + # Calculate normalised psi. + # 0 = magnetic axis + # 1 = plasma boundary + psi_norm = (psi - psi_axis) / (psi_bndry - psi_axis) + + pprime_shape = self.pprime_spline(psi_norm) + ffprime_shape = self.ffprime_spline(psi_norm) + + if mask is not None: + pprime_shape *= mask + ffprime_shape *= mask + + # Now apply constraints to define constants + + # Need integral of pprime_shape to calculate pressure + # as p(psinorm) = - (L*Beta0/Raxis) * pshape(psinorm) + + shapeintegral = self.pprime_int_spline(0.0) + shapeintegral *= psi_bndry - psi_axis + + # Pressure on axis is + # + # paxis = - (L*Beta0/Raxis) * shapeintegral + # + + LBeta0 = -self.paxis * self.Raxis / shapeintegral + + # Constrain Ffprime = ffprime_spline + # Ffprime = L*(1-Beta0)*Raxis*ffprime_spline + # L*(1-Beta0)*Raxis = 1 + # L = LBeta0 + (1.0/Raxis) + + L = LBeta0 + (1.0/self.Raxis) + Beta0 = LBeta0 / L + + print("Constraints: L = %e, Beta0 = %e" % (L, Beta0)) + + # Toroidal current + Jtor = L *( (pprime_shape * Beta0 * R / self.Raxis) + ((1 - Beta0) * self.Raxis * ffprime_shape/ (R * mu0)) ) + + self.L = L + self.Beta0 = Beta0 + + self.psi_bndry = psi_bndry + self.psi_axis = psi_axis + + return Jtor + + # Profile functions + def pprime(self, pn): + """ + dp/dpsi as a function of normalised psi. 0 outside core. + Calculate pprimeshape inside the core only + """ + shape = self.pprime_spline(pn) + return self.L * self.Beta0 / self.Raxis * shape + + def ffprime(self, pn): + """ + f * df/dpsi as a function of normalised psi. 0 outside core. + Calculate ffprimeshape inside the core only. + """ + shape = self.ffprime_spline(pn) + return self.L * (1 - self.Beta0) * self.Raxis * shape + + def fvac(self): + return self._fvac + +class PprimeFfprimeConstrainedSplineProfiles(Profile): + """ + Pprime and Ffprime-constrained custom (splined) internal plasma profiles. + + """ + + def __init__(self, eq=None, Raxis=None, psi_n=None, pprime=None, ffprime=None, fvac=None): + """ + eq - Equilibrium object + Raxis - R used in p' and ff' components + psi_n - Normalised (0,1) poloidal flux used to defined the profiles + pprime - Pressure gradient - dp/dpsi + ffprime - f*dfpol/dpsi + fvac - Vacuum f = R*Bt + + """ + + # Check inputs + if eq is None: + raise ValueError("No equilibrium object provided") + if Raxis is None: + raise ValueError("No Raxis value provided") + if psi_n is None: + raise ValueError("No psi_n data provided") + if pprime is None: + raise ValueError("No pprime data provided") + if ffprime is None: + raise ValueError("No ffprime data provided") + if fvac is None: + raise ValueError("No fvac data provided") + + # Set values for later use + self.eq = eq + self.Raxis = Raxis + self.psi_n_points = psi_n + self.pprime_points = pprime + self.ffprime_points = ffprime + self._fvac = fvac + + # Create 1D splines for the internal profiles - these will be like jtorshape + self.pprime_spline = interp1d(self.psi_n_points,self.pprime_points,kind='linear',fill_value='extrapolate',bounds_error=False) + self.ffprime_spline = interp1d(self.psi_n_points,self.ffprime_points,kind='linear',fill_value='extrapolate',bounds_error=False) + + # Create 1D splines for the integral of pprime, ffprime + pn_points = np.linspace(0.0,1.0,100,endpoint=True) + + def pprime_int_func(pn): + val, _ = quad(self.pprime_spline,pn,1.0) + return val + + pprime_int_vals = [] + for pn in pn_points: + pprime_int_vals.append(pprime_int_func(pn)) + + pprime_int_vals = np.asarray(pprime_int_vals) + + self.pprime_int_spline = interp1d(pn_points,pprime_int_vals,kind='linear',fill_value='extrapolate',bounds_error=False) + + def ffprime_int_func(pn): + val, _ = quad(self.ffprime_spline,pn,1.0) + return val + + ffprime_int_vals = [] + for pn in pn_points: + ffprime_int_vals.append(ffprime_int_func(pn)) + + ffprime_int_vals = np.asarray(ffprime_int_vals) + + self.ffprime_int_spline = interp1d(pn_points,ffprime_int_vals,kind='linear',fill_value='extrapolate',bounds_error=False) + + def Jtor(self, R, Z, psi, psi_bndry=None): + """Calculate toroidal plasma current + + Jtor = R*pprime + ffprime/(R * mu0) + """ + + # Analyse the equilibrium, finding O- and X-points + opt, xpt = critical.find_critical(R, Z, psi) + if not opt: + raise ValueError("No O-points found!") + psi_axis = opt[0][2] + + if psi_bndry is not None: + mask = critical.core_mask(R, Z, psi, opt, xpt, psi_bndry) + elif xpt: + psi_bndry = xpt[0][2] + mask = critical.core_mask(R, Z, psi, opt, xpt) + else: + # No X-points + psi_bndry = psi[0, 0] + mask = None + + dR = R[1, 0] - R[0, 0] + dZ = Z[0, 1] - Z[0, 0] + + # Calculate normalised psi. + # 0 = magnetic axis + # 1 = plasma boundary + psi_norm = (psi - psi_axis) / (psi_bndry - psi_axis) + + pprime_shape = self.pprime_spline(psi_norm) + ffprime_shape = self.ffprime_spline(psi_norm) + + if mask is not None: + pprime_shape *= mask + ffprime_shape *= mask + + # Now apply constraints to define constants + + LBeta0 = self.Raxis + L = self.Raxis + (1.0/self.Raxis) + Beta0 = LBeta0/L + + print("Constraints: L = %e, Beta0 = %e" % (L, Beta0)) + + # Toroidal current + Jtor = L *( (pprime_shape * Beta0 * R / self.Raxis) + ((1 - Beta0) * self.Raxis * ffprime_shape/ (R * mu0)) ) + + self.L = L + self.Beta0 = Beta0 + + self.psi_bndry = psi_bndry + self.psi_axis = psi_axis + + return Jtor + + # Profile functions + def pprime(self, pn): + """ + dp/dpsi as a function of normalised psi. 0 outside core. + Calculate pprimeshape inside the core only + """ + shape = self.pprime_spline(pn) + return self.L * self.Beta0 / self.Raxis * shape + + def ffprime(self, pn): + """ + f * df/dpsi as a function of normalised psi. 0 outside core. + Calculate ffprimeshape inside the core only. + """ + shape = self.ffprime_spline(pn) + return self.L * (1 - self.Beta0) * self.Raxis * shape + + def fvac(self): + return self._fvac + +class ProfilesPprimeFfprime: """ Specified profile functions p'(psi), ff'(psi) @@ -529,4 +1579,4 @@ def pprime(self, psinorm): return self._pprime(psinorm) def ffprime(self, psinorm): - return self._ffprime(psinorm) + return self._ffprime(psinorm) \ No newline at end of file