Reverse-Mode Differentiable Ideal-Ballooning Stability Solver

.gitignore

@@ -20,4 +20,6 @@ tests/benchmarks/.benchmarks/
 *.output
 *.pdf
 *.pyc
+*.swp
+*.swo
 .pymon

CHANGELOG.md

@@ -12,12 +12,14 @@ New Features
 - Changes ``ToroidalFlux`` objective to default using a 1D loop integral of the vector potential
 to compute the toroidal flux when possible, as opposed to a 2D surface integral of the magnetic field dotted with ``n_zeta``.
 - Allow specification of Nyquist spectrum maximum modenumbers when using ``VMECIO.save`` to save a DESC .h5 file as a VMEC-format wout file
+- Added and tested infinite-n ideal-ballooning stability solver implemented as a part of the BallooningStability Objective. DESC can use reverse-mode AD to now optimize equilibria against infinite-n ideal ballooning modes.
 - Add ``jac_chunk_size`` to ``ObjectiveFunction`` and ``_Objective`` to control the above chunk size for the ``fwd`` mode Jacobian calculation
   - if ``None``, the chunk size is equal to ``dim_x``, so no chunking is done
   - if an ``int``, this is the chunk size to be used.
   - if ``"auto"`` for the ``ObjectiveFunction``, will use a heuristic for the maximum ``jac_chunk_size`` needed to fit the jacobian calculation on the available device memory, according to the formula: ``max_jac_chunk_size = (desc_config.get("avail_mem") / estimated_memory_usage - 0.22)  / 0.85  * self.dim_x`` with ``estimated_memory_usage = 2.4e-7 * self.dim_f * self.dim_x + 1``
 - the ``ObjectiveFunction`` ``jac_chunk_size`` is used if ``deriv_mode="batched"``, and the ``_Objective`` ``jac_chunk_size`` will be used if ``deriv_mode="blocked"``
 
+
 Bug Fixes
 
 - Fixes bugs that occur when saving asymmetric equilibria as wout files

desc/compute/_core.py

@@ -1547,24 +1547,6 @@ def _alpha_t(params, transforms, profiles, data, **kwargs):
     return data
 
 
-@register_compute_fun(
-    name="alpha_tt",
-    label="\\partial_{\\theta \\theta} \\alpha",
-    units="~",
-    units_long="None",
-    description="Field line label, second derivative wrt poloidal coordinate",
-    dim=1,
-    params=[],
-    transforms={},
-    profiles=[],
-    coordinates="rtz",
-    data=["theta_PEST_tt", "phi_tt", "iota"],
-)
-def _alpha_tt(params, transforms, profiles, data, **kwargs):
-    data["alpha_tt"] = data["theta_PEST_tt"] - data["iota"] * data["phi_tt"]
-    return data
-
-
 @register_compute_fun(
     name="alpha_tz",
     label="\\partial_{\\theta \\zeta} \\alpha",
@@ -1601,12 +1583,30 @@ def _alpha_z(params, transforms, profiles, data, **kwargs):
     return data
 
 
+@register_compute_fun(
+    name="alpha_tt",
+    label="\\partial_{\\theta \\theta} \\alpha",
+    units="~",
+    units_long="None",
+    description="Field line label, second-order derivative wrt poloidal coordinate",
+    dim=1,
+    params=[],
+    transforms={},
+    profiles=[],
+    coordinates="rtz",
+    data=["phi_tt", "iota", "theta_PEST_tt"],
+)
+def _alpha_tt(params, transforms, profiles, data, **kwargs):
+    data["alpha_tt"] = data["theta_PEST_tt"] - data["iota"] * data["phi_tt"]
+    return data
+
+
 @register_compute_fun(
     name="alpha_zz",
     label="\\partial_{\\zeta \\zeta} \\alpha",
     units="~",
     units_long="None",
-    description="Field line label, second derivative wrt toroidal coordinate",
+    description="Field line label, second-order derivative wrt toroidal coordinate",
     dim=1,
     params=[],
     transforms={},
@@ -3105,8 +3105,8 @@ def _theta_PEST_r(params, transforms, profiles, data, **kwargs):
     label="\\partial_{\\rho \\theta} \\vartheta",
     units="rad",
     units_long="radians",
-    description="PEST straight field line poloidal angular coordinate, derivative wrt "
-    "radial and DESC poloidal coordinate",
+    description="PEST straight field line poloidal angular coordinate,"
+    "derivative wrt poloidal and radial coordinate",
     dim=1,
     params=[],
     transforms={},
@@ -3201,8 +3201,8 @@ def _theta_PEST_t(params, transforms, profiles, data, **kwargs):
     label="\\partial_{\\theta \\theta} \\vartheta",
     units="rad",
     units_long="radians",
-    description="PEST straight field line poloidal angular coordinate, second "
-    "derivative wrt poloidal coordinate",
+    description="PEST straight field line poloidal angular coordinate,"
+    "second derivative wrt poloidal coordinate",
     dim=1,
     params=[],
     transforms={},
@@ -3277,8 +3277,8 @@ def _theta_PEST_tzz(params, transforms, profiles, data, **kwargs):
     label="\\partial_{\\zeta} \\vartheta",
     units="rad",
     units_long="radians",
-    description="PEST straight field line poloidal angular coordinate, derivative wrt "
-    "toroidal coordinate",
+    description="PEST straight field line poloidal angular coordinate,"
+    " derivative wrt toroidal coordinate",
     dim=1,
     params=[],
     transforms={},

desc/compute/_metric.py

@@ -83,7 +83,7 @@ def _sqrtg_clebsch(params, transforms, profiles, data, **kwargs):
 
 @register_compute_fun(
     name="|e_theta x e_zeta|",
-    label="|\\mathbf{e}_{\\theta} \\times \\mathbf{e}_{\\zeta}|",
+    label="| \\mathbf{e}_{\\theta} \\times \\mathbf{e}_{\\zeta} |",
     units="m^{2}",
     units_long="square meters",
     description="2D Jacobian determinant for constant rho surface",
@@ -141,7 +141,7 @@ def _e_theta_x_e_zeta_r(params, transforms, profiles, data, **kwargs):
 
 @register_compute_fun(
     name="|e_theta x e_zeta|_rr",
-    label="\\partial_{\\rho\\rho} |\\mathbf{e}_{\\theta} \\times \\mathbf{e}_{\\zeta}|",
+    label="\\partial_{\\rho\\rho}|\\mathbf{e}_{\\theta}\\times\\mathbf{e}_{\\zeta}|",
     units="m^{2}",
     units_long="square meters",
     description="2D Jacobian determinant for constant rho surface"
@@ -180,7 +180,7 @@ def _e_theta_x_e_zeta_rr(params, transforms, profiles, data, **kwargs):
 
 @register_compute_fun(
     name="|e_theta x e_zeta|_z",
-    label="\\partial_{\\zeta}|e_{\\theta} \\times e_{\\zeta}|",
+    label="\\partial_{\\zeta}|\\mathbf{e}_{\\theta} \\times \\mathbf{e}_{\\zeta}|",
     units="m^{2}",
     units_long="square meters",
     description="2D Jacobian determinant for constant rho surface,"
@@ -1443,6 +1443,44 @@ def _g_sup_rr_r(params, transforms, profiles, data, **kwargs):
     return data
 
 
+@register_compute_fun(
+    name="g^rr_t",
+    label="\\partial_{\\theta} g^{\\rho \\rho}",
+    units="m^-2",
+    units_long="inverse square meters",
+    description="Radial/Radial element of contravariant metric tensor, "
+    + "first poloidal derivative",
+    dim=1,
+    params=[],
+    transforms={},
+    profiles=[],
+    coordinates="rtz",
+    data=["e^rho", "e^rho_t"],
+)
+def _g_sup_rr_t(params, transforms, profiles, data, **kwargs):
+    data["g^rr_t"] = 2 * dot(data["e^rho_t"], data["e^rho"])
+    return data
+
+
+@register_compute_fun(
+    name="g^rr_z",
+    label="\\partial_{\\zeta} g^{\\rho \\rho}",
+    units="m^-2",
+    units_long="inverse square meters",
+    description="Radial/Radial element of contravariant metric tensor, "
+    + "first toroidal derivative",
+    dim=1,
+    params=[],
+    transforms={},
+    profiles=[],
+    coordinates="rtz",
+    data=["e^rho", "e^rho_z"],
+)
+def _g_sup_rr_z(params, transforms, profiles, data, **kwargs):
+    data["g^rr_z"] = 2 * dot(data["e^rho_z"], data["e^rho"])
+    return data
+
+
 @register_compute_fun(
     name="g^rt_r",
     label="\\partial_{\\rho} g^{\\rho \\theta}",
@@ -1544,25 +1582,6 @@ def _g_sup_zz_r(params, transforms, profiles, data, **kwargs):
     return data
 
 
-@register_compute_fun(
-    name="g^rr_t",
-    label="\\partial_{\\theta} g^{\\rho \\rho}",
-    units="m^-2",
-    units_long="inverse square meters",
-    description="Radial/Radial element of contravariant metric tensor, "
-    + "first poloidal derivative",
-    dim=1,
-    params=[],
-    transforms={},
-    profiles=[],
-    coordinates="rtz",
-    data=["e^rho", "e^rho_t"],
-)
-def _g_sup_rr_t(params, transforms, profiles, data, **kwargs):
-    data["g^rr_t"] = 2 * dot(data["e^rho_t"], data["e^rho"])
-    return data
-
-
 @register_compute_fun(
     name="g^rt_t",
     label="\\partial_{\\theta} g^{\\rho \\theta}",
@@ -1664,25 +1683,6 @@ def _g_sup_zz_t(params, transforms, profiles, data, **kwargs):
     return data
 
 
-@register_compute_fun(
-    name="g^rr_z",
-    label="\\partial_{\\zeta} g^{\\rho \\rho}",
-    units="m^-2",
-    units_long="inverse square meters",
-    description="Radial/Radial element of contravariant metric tensor, "
-    + "first toroidal derivative",
-    dim=1,
-    params=[],
-    transforms={},
-    profiles=[],
-    coordinates="rtz",
-    data=["e^rho", "e^rho_z"],
-)
-def _g_sup_rr_z(params, transforms, profiles, data, **kwargs):
-    data["g^rr_z"] = 2 * dot(data["e^rho_z"], data["e^rho"])
-    return data
-
-
 @register_compute_fun(
     name="g^rt_z",
     label="\\partial_{\\zeta} g^{\\rho \\theta}",
@@ -1900,6 +1900,42 @@ def _gradzeta(params, transforms, profiles, data, **kwargs):
     return data
 
 
+@register_compute_fun(
+    name="g^aa",
+    label="g^{\\alpha \\alpha}",
+    units="m^{-2}",
+    units_long="inverse square meters",
+    description="Contravariant metric tensor grad alpha dot grad alpha",
+    dim=1,
+    params=[],
+    transforms={},
+    profiles=[],
+    coordinates="rtz",
+    data=["grad(alpha)"],
+)
+def _g_sup_aa(params, transforms, profiles, data, **kwargs):
+    data["g^aa"] = dot(data["grad(alpha)"], data["grad(alpha)"])
+    return data
+
+
+@register_compute_fun(
+    name="g^ra",
+    label="g^{\\rho \\alpha}",
+    units="m^{-2}",
+    units_long="inverse square meters",
+    description="Contravariant metric tensor grad rho dot grad alpha",
+    dim=1,
+    params=[],
+    transforms={},
+    profiles=[],
+    coordinates="rtz",
+    data=["grad(alpha)", "e^rho"],
+)
+def _g_sup_ra(params, transforms, profiles, data, **kwargs):
+    data["g^ra"] = dot(data["grad(alpha)"], data["e^rho"])
+    return data
+
+
 @register_compute_fun(
     name="gbdrift",
     # Exact definition of the magnetic drifts taken from
@@ -1967,10 +2003,13 @@ def _cvdrift(params, transforms, profiles, data, **kwargs):
     transforms={},
     profiles=[],
     coordinates="rtz",
-    data=["|B|^2", "b", "e^rho", "grad(|B|)"],
+    data=["rho", "|B|^2", "b", "e^rho", "grad(|B|)"],
 )
 def _cvdrift0(params, transforms, profiles, data, **kwargs):
     data["cvdrift0"] = (
-        1 / data["|B|^2"] * (dot(data["b"], cross(data["grad(|B|)"], data["e^rho"])))
+        2
+        * data["rho"]
+        / data["|B|^2"]
+        * (dot(data["b"], cross(data["grad(|B|)"], data["e^rho"])))
     )
     return data