Add calculation of force error with anisotropic pressure tensor

CHANGELOG.md

@@ -1,6 +1,9 @@
 Changelog
 =========
 
+- Adds ability to compute equilibria with anisotropic pressure. This includes a new
+profile, ``Equilibrium.anisotropy``, new compute quantity ``F_anisotropic``, and a new
+objective ``ForceBalanceAnisotropic``.
 - Refactors most of the optimizer subproblems to use JAX control flow, allowing them
 to run more efficiently on the GPU.
 - Adds ``'shear'`` as a compute quantity and ``Shear`` as an objective function.

desc/compute/_equil.py

@@ -14,7 +14,7 @@
 from desc.backend import jnp
 
 from .data_index import register_compute_fun
-from .utils import dot, surface_averages
+from .utils import cross, dot, surface_averages
 
 
 @register_compute_fun(
@@ -580,6 +580,30 @@ def _e_sup_helical_mag(params, transforms, profiles, data, **kwargs):
     return data
 
 
+@register_compute_fun(
+    name="F_anisotropic",
+    label="F_{anisotropic}",
+    units="N \\cdot m^{-3}",
+    units_long="Newtons / cubic meter",
+    description="Anisotropic force balance error",
+    dim=3,
+    params=[],
+    transforms={},
+    profiles=[],
+    coordinates="rtz",
+    data=["J", "B", "grad(beta_a)", "beta_a", "grad(|B|^2)", "grad(p)"],
+)
+def _F_anisotropic(params, transforms, profiles, data, **kwargs):
+    data["F_anisotropic"] = (
+        (1 - data["beta_a"]) * cross(data["J"], data["B"]).T
+        - dot(data["B"], data["grad(beta_a)"]) * data["B"].T / mu_0
+        - data["beta_a"] * data["grad(|B|^2)"].T / (2 * mu_0)
+        - data["grad(p)"].T
+    ).T
+
+    return data
+
+
 @register_compute_fun(
     name="W_B",
     label="W_B",

desc/compute/_profiles.py

@@ -427,6 +427,42 @@ def _p_r(params, transforms, profiles, data, **kwargs):
     return data
 
 
+@register_compute_fun(
+    name="p_t",
+    label="\\partial_{\\theta} p",
+    units="Pa",
+    units_long="Pascals",
+    description="Pressure, first poloidal derivative",
+    dim=1,
+    params=["p_l"],
+    transforms={},
+    profiles=["pressure"],
+    coordinates="rtz",
+    data=[],
+)
+def _p_t(params, transforms, profiles, data, **kwargs):
+    data["p_t"] = profiles["pressure"].compute(params["p_l"], dt=1)
+    return data
+
+
+@register_compute_fun(
+    name="p_z",
+    label="\\partial_{\\zeta} p",
+    units="Pa",
+    units_long="Pascals",
+    description="Pressure, first toroidal derivative",
+    dim=1,
+    params=["p_l"],
+    transforms={},
+    profiles=["pressure"],
+    coordinates="rtz",
+    data=[],
+)
+def _p_z(params, transforms, profiles, data, **kwargs):
+    data["p_z"] = profiles["pressure"].compute(params["p_l"], dz=1)
+    return data
+
+
 @register_compute_fun(
     name="grad(p)",
     label="\\nabla p",
@@ -438,10 +474,14 @@ def _p_r(params, transforms, profiles, data, **kwargs):
     transforms={},
     profiles=[],
     coordinates="rtz",
-    data=["p_r", "e^rho"],
+    data=["p_r", "p_t", "p_z", "e^rho", "e^theta", "e^zeta"],
 )
 def _gradp(params, transforms, profiles, data, **kwargs):
-    data["grad(p)"] = (data["p_r"] * data["e^rho"].T).T
+    data["grad(p)"] = (
+        data["p_r"] * data["e^rho"].T
+        + jnp.where(data["p_t"] == 0, 0, data["p_t"] * data["e^theta"].T)
+        + jnp.where(data["p_z"] == 0, 0, data["p_z"] * data["e^zeta"].T)
+    ).T
     return data
 
 
@@ -502,6 +542,112 @@ def _gradp_mag_vol(params, transforms, profiles, data, **kwargs):
     return data
 
 
+@register_compute_fun(
+    name="beta_a",
+    label="\\beta_a = \\mu_0 (p_{||} - p_{\\perp})/B^2",
+    units="~",
+    units_long="None",
+    description="Pressure anisotropy",
+    dim=1,
+    params=["a_lmn"],
+    transforms={},
+    profiles=["anisotropy"],
+    coordinates="rtz",
+    data=["0"],
+)
+def _beta_a(params, transforms, profiles, data, **kwargs):
+    if profiles["anisotropy"] is not None:
+        data["beta_a"] = profiles["anisotropy"].compute(params["a_lmn"], dr=0)
+    else:
+        data["beta_a"] = jnp.nan * data["0"]
+    return data
+
+
+@register_compute_fun(
+    name="beta_a_r",
+    label="\\partial_{\\rho} \\beta_a = \\mu_0 (p_{||} - p_{\\perp})/B^2",
+    units="~",
+    units_long="None",
+    description="Pressure anisotropy, first radial derivative",
+    dim=1,
+    params=["a_lmn"],
+    transforms={},
+    profiles=["anisotropy"],
+    coordinates="rtz",
+    data=["0"],
+)
+def _beta_a_r(params, transforms, profiles, data, **kwargs):
+    if profiles["anisotropy"] is not None:
+        data["beta_a_r"] = profiles["anisotropy"].compute(params["a_lmn"], dr=1)
+    else:
+        data["beta_a_r"] = jnp.nan * data["0"]
+    return data
+
+
+@register_compute_fun(
+    name="beta_a_t",
+    label="\\partial_{\\theta} \\beta_a = \\mu_0 (p_{||} - p_{\\perp})/B^2",
+    units="~",
+    units_long="None",
+    description="Pressure anisotropy, first poloidal derivative",
+    dim=1,
+    params=["a_lmn"],
+    transforms={},
+    profiles=["anisotropy"],
+    coordinates="rtz",
+    data=["0"],
+)
+def _beta_a_t(params, transforms, profiles, data, **kwargs):
+    if profiles["anisotropy"] is not None:
+        data["beta_a_t"] = profiles["anisotropy"].compute(params["a_lmn"], dt=1)
+    else:
+        data["beta_a_t"] = jnp.nan * data["0"]
+    return data
+
+
+@register_compute_fun(
+    name="beta_a_z",
+    label="\\partial_{\\zeta} \\beta_a = \\mu_0 (p_{||} - p_{\\perp})/B^2",
+    units="~",
+    units_long="None",
+    description="Pressure anisotropy, first toroidal derivative",
+    dim=1,
+    params=["a_lmn"],
+    transforms={},
+    profiles=["anisotropy"],
+    coordinates="rtz",
+    data=["0"],
+)
+def _beta_a_z(params, transforms, profiles, data, **kwargs):
+    if profiles["anisotropy"] is not None:
+        data["beta_a_z"] = profiles["anisotropy"].compute(params["a_lmn"], dz=1)
+    else:
+        data["beta_a_z"] = jnp.nan * data["0"]
+    return data
+
+
+@register_compute_fun(
+    name="grad(beta_a)",
+    label="\\nabla \\beta_a = \\nabla \\mu_0 (p_{||} - p_{\\perp})/B^2",
+    units="m^{-1}",
+    units_long="Inverse meters",
+    description="Pressure anisotropy gradient",
+    dim=3,
+    params=[],
+    transforms={},
+    profiles=[],
+    coordinates="rtz",
+    data=["beta_a_r", "beta_a_t", "beta_a_z", "e^rho", "e^theta", "e^zeta"],
+)
+def _gradbeta_a(params, transforms, profiles, data, **kwargs):
+    data["grad(beta_a)"] = (
+        data["beta_a_r"] * data["e^rho"].T
+        + data["beta_a_t"] * data["e^theta"].T
+        + data["beta_a_z"] * data["e^zeta"].T
+    ).T
+    return data
+
+
 @register_compute_fun(
     name="iota",
     label="\\iota",

desc/compute/utils.py

@@ -1295,12 +1295,12 @@ def body(i, mins):
     "ne_l",
     "Ti_l",
     "Zeff_l",
+    "a_lmn",
     "Ra_n",
     "Za_n",
     "Rb_lmn",
     "Zb_lmn",
 )
-
 # map from profile name to equilibrium parameter name
 profile_names = {
     "pressure": "p_l",
@@ -1310,4 +1310,5 @@ def body(i, mins):
     "electron_density": "ne_l",
     "ion_temperature": "Ti_l",
     "atomic_number": "Zeff_l",
+    "anisotropy": "a_lmn",
 }

desc/continuation.py

@@ -11,7 +11,7 @@
 from desc.objectives import get_equilibrium_objective, get_fixed_boundary_constraints
 from desc.optimize import Optimizer
 from desc.perturbations import get_deltas
-from desc.utils import Timer
+from desc.utils import Timer, errorif
 
 MIN_MRES_STEP = 1
 MIN_PRES_STEP = 0.1
@@ -92,7 +92,7 @@ def _solve_axisym(
 
         if ii > 0:
             eqi = eqfam[-1].copy()
-            # increase resolution of vacuum soln
+            # increase resolution of vacuum solution
             Mi = min(Mi + mres_step, M)
             Li = int(np.ceil(L / M) * Mi)
             L_gridi = np.ceil(L_grid / L * Li).astype(int)
@@ -516,11 +516,16 @@ def solve_continuation_automatic(  # noqa: C901
         final desired configuration,
 
     """
-    if eq.electron_temperature is not None:
-        raise NotImplementedError(
-            "Continuation method with kinetic profiles is not currently supported"
-        )
-
+    errorif(
+        eq.electron_temperature is not None,
+        NotImplementedError,
+        "Continuation method with kinetic profiles is not currently supported",
+    )
+    errorif(
+        eq.anisotropy is not None,
+        NotImplementedError,
+        "Continuation method with anisotropic pressure is not currently supported",
+    )
     timer = Timer()
     timer.start("Total time")
 
@@ -645,10 +650,16 @@ def solve_continuation(  # noqa: C901
         final desired configuration,
 
     """
-    if not all([eq.electron_temperature is None for eq in eqfam]):
-        raise NotImplementedError(
-            "Continuation method with kinetic profiles is not currently supported"
-        )
+    errorif(
+        not all([eq.electron_temperature is None for eq in eqfam]),
+        NotImplementedError,
+        "Continuation method with kinetic profiles is not currently supported",
+    )
+    errorif(
+        not all([eq.anisotropy is None for eq in eqfam]),
+        NotImplementedError,
+        "Continuation method with anisotropic pressure is not currently supported",
+    )
 
     timer = Timer()
     timer.start("Total time")

desc/equilibrium/equilibrium.py

@@ -97,6 +97,9 @@ class Equilibrium(IOAble):
     atomic_number : Profile or ndarray shape(k,2) (optional)
         Effective atomic number (Z_eff) profile or ndarray of mode numbers and spectral
         coefficients. Default is 1
+    anisotropy : Profile or ndarray
+        Anisotropic pressure profile or array of mode numbers and spectral coefficients.
+        Default is a PowerSeriesProfile with zero anisotropic pressure.
     surface: Surface or ndarray shape(k,5) (optional)
         Fixed boundary surface shape, as a Surface object or array of
         spectral mode numbers and coefficients of the form [l, m, n, R, Z].
@@ -135,6 +138,7 @@ class Equilibrium(IOAble):
         "_electron_density",
         "_ion_temperature",
         "_atomic_number",
+        "_anisotropy",
         "_spectral_indexing",
         "_bdry_mode",
         "_L_grid",
@@ -161,6 +165,7 @@ def __init__(
         electron_density=None,
         ion_temperature=None,
         atomic_number=None,
+        anisotropy=None,
         surface=None,
         axis=None,
         sym=None,
@@ -292,7 +297,7 @@ def __init__(
             "Cannot specify both iota and current profiles.",
         )
         errorif(
-            pressure is not None and use_kinetic,
+            ((pressure is not None) or (anisotropy is not None)) and use_kinetic,
             ValueError,
             "Cannot specify both pressure and kinetic profiles.",
         )
@@ -316,6 +321,7 @@ def __init__(
         self._ion_temperature = parse_profile(ion_temperature, "ion temperature")
         self._atomic_number = parse_profile(atomic_number, "atomic number")
         self._pressure = parse_profile(pressure, "pressure")
+        self._anisotropy = parse_profile(anisotropy, "anisotropy")
         self._iota = parse_profile(iota, "iota")
         self._current = parse_profile(current, "current")
 
@@ -328,6 +334,7 @@ def __init__(
             "electron_density",
             "ion_temperature",
             "atomic_number",
+            "anisotropy",
         ]:
             p = getattr(self, profile)
             if hasattr(p, "change_resolution"):
@@ -522,6 +529,7 @@ def change_resolution(
             "electron_density",
             "ion_temperature",
             "atomic_number",
+            "anisotropy",
         ]:
             p = getattr(self, profile)
             if hasattr(p, "change_resolution"):
@@ -1223,6 +1231,29 @@ def p_l(self, p_l):
         )
         self.pressure.params = p_l
 
+    @property
+    def anisotropy(self):
+        """Profile: Anisotropy profile."""
+        return self._anisotropy
+
+    @anisotropy.setter
+    def anisotropy(self, new):
+        self._anisotropy = parse_profile(new, "anisotropy")
+
+    @property
+    def a_lmn(self):
+        """ndarray: Coefficients of anisotropy profile."""
+        return np.empty(0) if self.anisotropy is None else self.anisotropy.params
+
+    @a_lmn.setter
+    def a_lmn(self, a_lmn):
+        errorif(
+            self.anisotropy is None,
+            ValueError,
+            "Attempt to set anisotropy on an equilibrium without anisotropy profile",
+        )
+        self.anisotropy.params = a_lmn
+
     @property
     def electron_temperature(self):
         """Profile: Electron temperature (eV) profile."""