WIP: Consistent handling of COCOS conventions

hypnotoad/cases/circular.py

@@ -124,6 +124,8 @@ def __init__(self, settings=None, nonorthogonal_settings=None):
             # Core-only geometry: add connection so domain is periodic in y
             self.makeConnection("circular", 0, "circular", 0)
 
+        self.f_psi_sign = np.sign(self.psi_r(1.0) - self.psi_r(0.0))
+
     def makeRegion(self):
         """
         Create the EquilibriumRegion to build the grid around

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:
@@ -403,50 +410,83 @@ def __init__(
             ):
                 # if psi_outer is not beyond the last point of psi1D, no need to extend
                 # Exclude first point since duplicates last point in core
-                psiSOL = np.linspace(psi1D[-1], psi_outer, 50)[1:]
+                psi_lcfs = psi1D[-1]
+                psiSOL = np.linspace(psi_lcfs, psi_outer, 50)[1:]
                 psi1D = np.concatenate([psi1D, psiSOL])
 
                 # fpol constant in SOL
                 fpol1D = np.concatenate([fpol1D, np.full(psiSOL.shape, fpol1D[-1])])
 
-            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)])
+                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
+                    p_lcfs = pressure[-1]
+                    # p_SOL = p_lcfs * exp( (psi - psi_lcfs) * dpdpsi / p_lcfs)
+                    p_SOL = p_lcfs * np.exp((psiSOL - psi_lcfs) * dpdpsi / p_lcfs)
+                    pressure = np.concatenate([pressure, p_SOL])
+                    if pprime is not None:
+                        pprime = np.concatenate([pprime, dpdpsi * p_SOL / p_lcfs])
 
         self.magneticFunctionsFromGrid(
             R1D, Z1D, psi2D, self.user_options.psi_interpolation_method
         )
 
-        self.f_psi_sign = 1.0
-        if len(fpol1D) > 0:
-            # Spline for interpolation of f = R*Bt
+        # Note: radial label must be increasing
+        if psi1D[-1] < psi1D[0]:
+            self.f_psi_sign = -1.0
+        else:
+            self.f_psi_sign = 1.0
 
-            # Note: psi1D must be increasing
-            if psi1D[-1] < psi1D[0]:
-                self.f_psi_sign = -1.0
+        xcoord = self.f_psi_sign * psi1D
 
-            self.f_spl = interpolate.InterpolatedUnivariateSpline(
-                psi1D * self.f_psi_sign, fpol1D, ext=3
-            )
+        if len(fpol1D) > 0:
+
+            # Spline for interpolation of f = R*Bt
+            self.f_spl = interpolate.InterpolatedUnivariateSpline(xcoord, fpol1D, ext=3)
             # 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(
+                    xcoord, fprime, ext=1
+                )
+            else:
+                # fprime is derivative of fpol with respect to psi
+                # rather than increasing radial label xcoord
+                fpol_f_psi_sign_spl = interpolate.InterpolatedUnivariateSpline(
+                    xcoord, fpol1D * self.f_psi_sign, ext=3
+                )
+                self.fprime_spl = fpol_f_psi_sign_spl.derivative()
         else:
             self.f_spl = lambda psi: 0.0
             self.fprime_spl = lambda psi: 0.0
 
         # Optional pressure profile
         if pressure is not None:
             self.p_spl = interpolate.InterpolatedUnivariateSpline(
-                psi1D * self.f_psi_sign, pressure, ext=3
+                xcoord, pressure, ext=3
             )
+
+            if pprime is not None:
+                self.pprime_spl = interpolate.InterpolatedUnivariateSpline(
+                    xcoord, pprime, ext=1
+                )
+            else:
+                # pprime is derivative of pressure with respect to psi
+                # rather than increasing radial label xcoord
+
+                pressure_f_psi_sign_spl = interpolate.InterpolatedUnivariateSpline(
+                    xcoord, pressure * self.f_psi_sign, ext=3
+                )
+                self.pprime_spl = pressure_f_psi_sign_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")
@@ -816,7 +856,13 @@ def describeSingleNull(self):
             dpsidi_pf = (self.psi_sep[0] - self.psi_pf_upper) / self.user_options.nx_pf
 
         # Get the smallest absolute grid spacing for the separatrix
-        dpsidi_sep = min([dpsidi_sol, dpsidi_core, dpsidi_pf], key=abs)
+        # dpsidi_pf_minimum is introduced to avoid very narrow radial grid width
+        # around separatrix in core and SOL region and 0.3 is an emprical factor.
+
+        # dpsidi_sep = min([dpsidi_sol, dpsidi_core, dpsidi_pf], key=abs)
+        dpsidi_sep = min([dpsidi_sol, dpsidi_core], key=abs)  # set dx at separatrix
+        # dpsidi_pf_minimum = max([0.3*dpsidi_sep, 0.3*dpsidi_pf], key=abs)
+        dpsidi_pf_minimum = max([0.3 * dpsidi_sep, dpsidi_pf], key=abs)
 
         # decrease (assuming the factor is <1) the spacing around the separatrix by the
         # factor psi_spacing_separatrix_multiplier
@@ -859,6 +905,7 @@ def describeSingleNull(self):
                 "psi_start": self.psi_pf_lower,
                 "psi_end": psi_sep,
                 "grad_end": dpsidi_sep,
+                "grad_start": dpsidi_pf_minimum,
             }
 
             leg_regions = {
@@ -909,6 +956,7 @@ def describeSingleNull(self):
                 "psi_start": self.psi_pf_upper,
                 "psi_end": psi_sep,
                 "grad_end": dpsidi_sep,
+                "grad_start": dpsidi_pf_minimum,
             }
 
             leg_regions = {
@@ -1621,9 +1669,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
@@ -1687,6 +1740,13 @@ 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"""
+        if self.pprime_spl is None:
+            return None
+        return self.pprime_spl(psi * self.f_psi_sign)
+
     @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/cases/torpex.py

@@ -162,6 +162,9 @@ def __init__(self, equilibOptions, meshOptions):
             self.Rmax = float("inf")
             self.Zmin = -float("inf")
             self.Zmax = float("inf")
+
+            self.f_psi_sign = 1.0  # H.seto (dummy)
+
         elif "gfile" in equilibOptions:
             # load a g-file
             with open(equilibOptions["gfile"], "rt") as fh:
@@ -206,6 +209,9 @@ def __init__(self, equilibOptions, meshOptions):
                         "direction of plasma current should be clockwise to be "
                         "consistent with sign of grad(psi)"
                     )
+
+            self.f_psi_sign = numpy.sign(psi_bndry - psi_axis)  # H.seto (dummy)
+
             # index of a point close to the magnetic axis
             i_axis = numpy.searchsorted(R, R_axis)
             j_axis = numpy.searchsorted(Z, Z_axis)
@@ -294,6 +300,9 @@ def __init__(self, equilibOptions, meshOptions):
                     f"{eqfile['Z'][0, 0][ZindMid, RindMid]}"
                 )
             self.Bt_axis = Bt[ZindMid, RindMid]
+
+            self.f_psi_sign = numpy.sign(0.0 - psi[ZindMid, RindMid])  # H.seto (dummy)
+
         else:
             raise ValueError("Failed to initialise psi function from inputs")