Save Jpar0 to output tokamak grids

hypnotoad/cases/tokamak.py

@@ -317,6 +317,8 @@ def __init__(
         psi1D,
         fpol1D,
         pressure=None,
+        fprime=None,
+        pprime=None,
         wall=None,
         psi_axis_gfile=None,
         psi_bdry_gfile=None,
@@ -340,6 +342,8 @@ def __init__(
         --------
 
         pressure[nf] = 1D array of pressure as a function of psi1D [Pa]
+        fprime[nf] = 1D array of df / dpsi
+        pprime[nf] = 1D array of dp / dpsi . If set then pressure must also be set.
 
         wall = [(R0,Z0), (R1, Z1), ...]
                A list of coordinate pairs, defining the vessel wall.
@@ -362,6 +366,9 @@ def __init__(
 
         """
 
+        if pprime is not None:
+            assert pressure is not None
+
         self.user_options = self.user_options_factory.create(settings)
 
         if self.user_options.reverse_current:
@@ -409,12 +416,22 @@ def __init__(
                 # fpol constant in SOL
                 fpol1D = np.concatenate([fpol1D, np.full(psiSOL.shape, fpol1D[-1])])
 
+                if fprime is not None:
+                    fprime = np.concatenate([fprime, np.full(psiSOL.shape, 0.0)])
+
             if pressure is not None:
                 # Use an exponential decay for the pressure, based on
                 # the value and gradient at the plasma edge
                 p0 = pressure[-1]
                 # p = p0 * exp( (psi - psi0) * dpdpsi / p0)
-                pressure = np.concatenate([pressure, p0 * np.exp(psiSOL * dpdpsi / p0)])
+                pressure = np.concatenate(
+                    [pressure, p0 * np.exp((psiSOL - psi1D[-1]) * dpdpsi / p0)]
+                )
+
+                if pprime is not None:
+                    pprime = np.concatenate(
+                        [pprime, dpdpsi * np.exp((psiSOL - psi1D[-1]) * dpdpsi / p0)]
+                    )
 
         self.magneticFunctionsFromGrid(
             R1D, Z1D, psi2D, self.user_options.psi_interpolation_method
@@ -434,7 +451,12 @@ def __init__(
             # ext=3 specifies that boundary values are used outside range
 
             # Spline representing the derivative of f
-            self.fprime_spl = self.f_spl.derivative()
+            if fprime is not None:
+                self.fprime_spl = interpolate.InterpolatedUnivariateSpline(
+                    psi1D * self.f_psi_sign, fprime, ext=3
+                )
+            else:
+                self.fprime_spl = self.f_spl.derivative()
         else:
             self.f_spl = lambda psi: 0.0
             self.fprime_spl = lambda psi: 0.0
@@ -444,9 +466,17 @@ def __init__(
             self.p_spl = interpolate.InterpolatedUnivariateSpline(
                 psi1D * self.f_psi_sign, pressure, ext=3
             )
+
+            if pprime is not None:
+                self.pprime_spl = interpolate.InterpolatedUnivariateSpline(
+                    psi1D * self.f_psi_sign, pprime, ext=3
+                )
+            else:
+                self.pprime_spl = self.p_spl.derivative()
         else:
             # If no pressure, then not output to grid file
             self.p_spl = None
+            self.pprime_spl = None
 
         # Find critical points (O- and X-points)
         R2D, Z2D = np.meshgrid(R1D, Z1D, indexing="ij")
@@ -476,6 +506,8 @@ def __init__(
 
         if len(xpoints) == 0:
             warnings.warn("No X-points found in TokamakEquilibrium input")
+            self.psi_bdry = psi_bdry_gfile
+            self.psi_bdry_gfile = psi_bdry_gfile
         else:
             self.psi_bdry = xpoints[0][2]  # Psi on primary X-point
             self.x_point = Point2D(xpoints[0][0], xpoints[0][1])
@@ -1621,9 +1653,14 @@ def createRegionObjects(self, all_regions, segments):
                     eqreg.pressure = lambda psi: self.pressure(
                         leg_psi + sign * abs(psi - leg_psi)
                     )
+                    # Set pprime and fprime to zero so Jpar0 = 0
+                    eqreg.pprime = lambda psi: 0.0
+                    eqreg.fprime = lambda psi: 0.0
                 else:
-                    # Core region, so use the core pressure
+                    # Core region, so use the core profiles
                     eqreg.pressure = self.pressure
+                    eqreg.pprime = self.pprime
+                    eqreg.fprime = self.fpolprime
 
             region_objects[name] = eqreg
         # The region objects need to be sorted, so that the
@@ -1677,8 +1714,16 @@ def fpol(self, psi):
 
     @Equilibrium.handleMultiLocationArray
     def fpolprime(self, psi):
-        """psi-derivative of fpol"""
-        return self.fprime_spl(psi * self.f_psi_sign)
+        """psi-derivative of fpol
+        Note: Zero outside core."""
+        fprime = self.fprime_spl(psi * self.f_psi_sign)
+        if self.psi_bdry is not None:
+            psinorm = (psi - self.psi_axis) / (self.psi_bdry - self.psi_axis)
+            if np.isscalar(psi) and psinorm > 1.0:
+                fprime = 0.0
+            else:
+                fprime[psinorm > 1.0] = 0.0
+        return fprime
 
     @Equilibrium.handleMultiLocationArray
     def pressure(self, psi):
@@ -1687,6 +1732,21 @@ def pressure(self, psi):
             return None
         return self.p_spl(psi * self.f_psi_sign)
 
+    @Equilibrium.handleMultiLocationArray
+    def pprime(self, psi):
+        """psi-derivative of plasma pressure
+        Note: Zero outside the core"""
+        if self.pprime_spl is None:
+            return None
+        pprime = self.pprime_spl(psi * self.f_psi_sign)
+        if self.psi_bdry is not None:
+            psinorm = (psi - self.psi_axis) / (self.psi_bdry - self.psi_axis)
+            if np.isscalar(psi) and psinorm > 1.0:
+                pprime = 0.0
+            else:
+                pprime[psinorm > 1.0] = 0.0
+        return pprime
+
     @property
     def Bt_axis(self):
         """Calculate toroidal field on axis"""
@@ -1751,6 +1811,8 @@ def read_geqdsk(
 
     pressure = data["pres"]
     fpol = data["fpol"]
+    fprime = data["ffprime"] / fpol
+    pprime = data["pprime"]
 
     result = TokamakEquilibrium(
         R1D,
@@ -1761,6 +1823,8 @@ def read_geqdsk(
         psi_bdry_gfile=psi_bdry_gfile,
         psi_axis_gfile=psi_axis_gfile,
         pressure=pressure,
+        fprime=fprime,
+        pprime=pprime,
         wall=wall,
         make_regions=make_regions,
         settings=settings,

hypnotoad/core/mesh.py

@@ -902,6 +902,27 @@ def geometry1(self):
 
         self.Bxy = numpy.sqrt(self.Bpxy**2 + self.Btxy**2)
 
+        if hasattr(
+            self.meshParent.equilibrium.regions[self.equilibriumRegion.name], "pprime"
+        ) and hasattr(
+            self.meshParent.equilibrium.regions[self.equilibriumRegion.name], "fprime"
+        ):
+            # Calculate parallel current density from p' and f'
+            # Note: COCOS +1 convention is assumed
+            # so mu0 * Jpar = -B f' - mu0 p' f / B
+
+            pprime = self.meshParent.equilibrium.regions[
+                self.equilibriumRegion.name
+            ].pprime(self.psixy)
+            fprime = self.meshParent.equilibrium.regions[
+                self.equilibriumRegion.name
+            ].fprime(self.psixy)
+
+            mu0 = 4e-7 * numpy.pi
+            self.Jpar0 = -(
+                self.Bxy * fprime / mu0 + self.Rxy * self.Btxy * pprime / self.Bxy
+            )
+
     def geometry2(self):
         """
         Continuation of geometry1(), but needs neighbours to have calculated Bp so called
@@ -3657,9 +3678,12 @@ def addFromRegionsXArray(name):
         addFromRegions("bxcvy")
         addFromRegions("bxcvz")
 
-        if hasattr(next(iter(self.equilibrium.regions.values())), "pressure"):
+        if hasattr(next(iter(self.regions.values())), "pressure"):
             addFromRegions("pressure")
 
+        if hasattr(next(iter(self.regions.values())), "Jpar0"):
+            addFromRegions("Jpar0")
+
         # Penalty mask
         self.penalty_mask = BoutArray(
             numpy.zeros((self.nx, self.ny)),

hypnotoad/geqdsk/_geqdsk.py

@@ -22,6 +22,7 @@
 
 from datetime import date
 from numpy import zeros, pi
+import numpy as np
 
 from ._fileutils import f2s, ChunkOutput, write_1d, write_2d, next_value
 
@@ -48,6 +49,11 @@ def write(data, fh, label=None, shot=None, time=None):
 
       psi           2D array (nx,ny) of poloidal flux
 
+      ffprime       1D array of f(psi) * f'(psi). If not present
+                    then this is calculated from fpol
+      pprime        1D array of p'(psi). If not present then
+                    this is calculated from pres
+
     fh - file handle
 
     label - Text label to put in the file
@@ -120,11 +126,7 @@ def write(data, fh, label=None, shot=None, time=None):
         f2s(data["zmagx"]) + f2s(0.0) + f2s(data["sibdry"]) + f2s(0.0) + f2s(0.0) + "\n"
     )
 
-    # SCENE has actual ff' and p' data so can use that
     # fill arrays
-    # Lukas Kripner (16/10/2018): uncommenting this, since you left there
-    # check for data existence bellow. This seems to as safer variant.
-    workk = zeros([nx])
 
     # Write arrays
     co = ChunkOutput(fh)
@@ -134,11 +136,29 @@ def write(data, fh, label=None, shot=None, time=None):
     if "ffprime" in data:
         write_1d(data["ffprime"], co)
     else:
-        write_1d(workk, co)
+        psi1D = np.linspace(data["simagx"], data["sibdry"], nx)
+        from scipy import interpolate
+
+        sign = -1.0 if psi1D[1] < psi1D[0] else 1.0
+        # spline fitting requires increasing X axis, so reverse if needed
+
+        fprime_spl = interpolate.InterpolatedUnivariateSpline(
+            sign * psi1D, sign * data["fpol"]
+        ).derivative()
+        ffprime = data["fpol"] * fprime_spl(sign * psi1D)
+        write_1d(ffprime, co)
+
     if "pprime" in data:
         write_1d(data["pprime"], co)
     else:
-        write_1d(workk, co)
+        psi1D = np.linspace(data["simagx"], data["sibdry"], nx)
+        from scipy import interpolate
+
+        sign = -1.0 if psi1D[1] < psi1D[0] else 1.0
+        pprime_spl = interpolate.InterpolatedUnivariateSpline(
+            sign * psi1D, sign * data["pres"]
+        ).derivative()
+        write_1d(pprime_spl(sign * psi1D), co)
 
     write_2d(data["psi"], co)
     write_1d(data["qpsi"], co)
@@ -195,6 +215,8 @@ def read(fh, cocos=1):
 
       fpol          1D array of f(psi)=R*Bt  [meter-Tesla]
       pres          1D array of p(psi) [Pascals]
+      ffprime       1D array of ff'(psi)
+      pprime        1D array of p'(psi)
       qpsi          1D array of q(psi)
 
       psi           2D array (nx,ny) of poloidal flux

integrated_tests/utils.py

@@ -22,6 +22,9 @@
     "penalty_mask",
     "closed_wall_R",
     "closed_wall_Z",
+    "Jpar0",
+    "Jpar0_xlow",
+    "Jpar0_ylow",
 ]