Magnetic axis limits and new variables

desc/compute/_basis_vectors.py

@@ -322,57 +322,39 @@ def _e_sub_rho_rt(params, transforms, profiles, data, **kwargs):
 def _e_sub_rho_rrt(params, transforms, profiles, data, **kwargs):
     data["e_rho_rrt"] = jnp.array(
         [
-            -data["R_rt"] * data["omega_r"] ** 2
-            - data["R_t"] * data["omega_r"] * data["omega_rr"]
-            - 2 * data["R_r"] * data["omega_r"] * data["omega_rt"]
-            - 2
+            -3 * data["R_rt"] * data["omega_r"] ** 2
+            - 3 * data["R_t"] * data["omega_r"] * data["omega_rr"]
+            - 4 * data["R_r"] * data["omega_r"] * data["omega_rt"]
+            - 3
             * data["R"]
             * (
                 data["omega_rr"] * data["omega_rt"]
                 + data["omega_r"] * data["omega_rrt"]
             )
-            - data["omega_rt"]
-            * (2 * data["R_r"] * data["omega_r"] + data["R"] * data["omega_rr"])
+            - data["omega_rt"] * (2 * data["R_r"] * data["omega_r"])
             - data["omega_t"]
             * (
-                2 * data["R_rr"] * data["omega_r"]
+                3 * data["R_rr"] * data["omega_r"]
                 + 3 * data["R_r"] * data["omega_rr"]
                 + data["R"] * data["omega_rrr"]
             )
             + data["R_rrrt"]
-            - data["omega_r"]
-            * (
-                2 * data["omega_r"] * data["R_rt"]
-                + 2 * data["R_r"] * data["omega_rt"]
-                + data["omega_t"] * data["R_rr"]
-                + data["R_t"] * data["omega_rr"]
-                + data["R"]
-                * (-data["omega_t"] * data["omega_r"] ** 2 + data["omega_rrt"])
-            ),
-            2 * data["omega_rr"] * data["R_rt"]
-            + 2 * data["omega_r"] * data["R_rrt"]
-            + 2 * data["R_rr"] * data["omega_rt"]
+            + data["omega_r"] * data["R"] * data["omega_t"] * data["omega_r"] ** 2,
+            3 * data["omega_rr"] * data["R_rt"]
+            + 3 * data["omega_r"] * data["R_rrt"]
+            + 3 * data["R_rr"] * data["omega_rt"]
             + 2 * data["R_r"] * data["omega_rrt"]
-            + data["omega_rt"] * data["R_rr"]
             + data["omega_t"] * data["R_rrr"]
-            + data["R_rt"] * data["omega_rr"]
             + data["R_t"] * data["omega_rrr"]
             + data["R_r"]
-            * (-data["omega_t"] * data["omega_r"] ** 2 + data["omega_rrt"])
+            * (-3 * data["omega_t"] * data["omega_r"] ** 2 + data["omega_rrt"])
             + data["R"]
             * (
-                -data["omega_rt"] * data["omega_r"] ** 2
-                - data["omega_t"] * data["omega_r"] * data["omega_rr"]
+                -3 * data["omega_rt"] * data["omega_r"] ** 2
+                - 3 * data["omega_t"] * data["omega_r"] * data["omega_rr"]
                 + data["omega_rrrt"]
             )
-            + data["omega_r"]
-            * (
-                -data["R_t"] * data["omega_r"] ** 2
-                - 2 * data["R"] * data["omega_r"] * data["omega_rt"]
-                - data["omega_t"]
-                * (2 * data["R_r"] * data["omega_r"] + data["R"] * data["omega_rr"])
-                + data["R_rrt"]
-            ),
+            - data["R_t"] * data["omega_r"] ** 3,
             data["Z_rrrt"],
         ]
     ).T
@@ -418,8 +400,7 @@ def _e_sub_rho_rz(params, transforms, profiles, data, **kwargs):
             + data["R_rrz"],
             2 * data["omega_r"] * data["R_rz"]
             + 2 * data["R_r"] * data["omega_rz"]
-            + data["R_rr"]
-            + data["omega_z"] * data["R_rr"]
+            + (1 + data["omega_z"]) * data["R_rr"]
             + data["R_z"] * data["omega_rr"]
             - data["R"]
             * ((1 + data["omega_z"]) * data["omega_r"] ** 2 - data["omega_rrz"]),
@@ -469,7 +450,7 @@ def _e_sub_rho_rrz(params, transforms, profiles, data, **kwargs):
             * (1 + data["omega_z"])
             * (data["R_rr"] * data["omega_r"] + data["R_r"] * data["omega_rr"])
             - data["R_rz"] * data["omega_r"] ** 2
-            - data["R_z"] * 2 * data["omega_r"] * data["omega_rr"]
+            - 2 * data["R_z"] * data["omega_r"] * data["omega_rr"]
             - 2 * data["R_r"] * data["omega_r"] * data["omega_rz"]
             - 2
             * data["R"]
@@ -488,8 +469,7 @@ def _e_sub_rho_rrz(params, transforms, profiles, data, **kwargs):
             * (
                 2 * data["omega_r"] * data["R_rz"]
                 + 2 * data["R_r"] * data["omega_rz"]
-                + data["R_rr"]
-                + data["omega_z"] * data["R_rr"]
+                + (1 + data["omega_z"]) * data["R_rr"]
                 + data["R_z"] * data["omega_rr"]
                 - data["R"]
                 * ((1 + data["omega_z"]) * data["omega_r"] ** 2 - data["omega_rrz"])
@@ -498,17 +478,16 @@ def _e_sub_rho_rrz(params, transforms, profiles, data, **kwargs):
             + 2 * data["omega_r"] * data["R_rrz"]
             + 2 * data["R_rr"] * data["omega_rz"]
             + 2 * data["R_r"] * data["omega_rrz"]
-            + data["R_rrr"]
             + data["omega_rz"] * data["R_rr"]
-            + data["omega_z"] * data["R_rrr"]
+            + (1 + data["omega_z"]) * data["R_rrr"]
             + data["R_rz"] * data["omega_rr"]
             + data["R_z"] * data["omega_rrr"]
             - data["R_r"]
             * ((1 + data["omega_z"]) * data["omega_r"] ** 2 - data["omega_rrz"])
             - data["R"]
             * (
                 data["omega_rz"] * data["omega_r"] ** 2
-                + (1 + data["omega_z"]) * 2 * data["omega_r"] * data["omega_rr"]
+                + 2 * (1 + data["omega_z"]) * data["omega_r"] * data["omega_rr"]
                 - data["omega_rrrz"]
             )
             + data["omega_r"]

desc/compute/_equil.py

@@ -20,9 +20,18 @@
     profiles=[],
     coordinates="rtz",
     data=["sqrt(g)", "B_zeta_t", "B_theta_z"],
+    axis_limit_data=["sqrt(g)_r", "B_zeta_rt", "B_theta_rz"],
 )
 def _J_sup_rho(params, transforms, profiles, data, **kwargs):
     data["J^rho"] = (data["B_zeta_t"] - data["B_theta_z"]) / (mu_0 * data["sqrt(g)"])
+    if transforms["grid"].axis.size:
+        # Recall that B_zeta_t(r=0) = -psi_r sqrt(g)_t ||e_zeta|| / sqrt(g)^2,
+        # which is not equal to zero at the magnetic axis.
+        # Still, it can be shown J^rho is of the indeterminate form 0/0.
+        limit = (data["B_zeta_rt"] - data["B_theta_rz"]) / (mu_0 * data["sqrt(g)_r"])
+        data["J^rho"] = put(
+            data["J^rho"], transforms["grid"].axis, limit[transforms["grid"].axis]
+        )
     return data
 
 
@@ -76,7 +85,6 @@ def _J_sup_zeta(params, transforms, profiles, data, **kwargs):
     data=["J^rho", "J^theta", "J^zeta", "e_rho", "e_theta", "e_zeta"],
 )
 def _J(params, transforms, profiles, data, **kwargs):
-    # TODO: should be infinite because density at infinitesimal volume
     data["J"] = (
         data["J^rho"] * data["e_rho"].T
         + data["J^theta"] * data["e_theta"].T
@@ -424,6 +432,7 @@ def _Fmag_vol(params, transforms, profiles, data, **kwargs):
     data=["B^theta", "B^zeta", "e^theta", "e^zeta"],
 )
 def _e_helical(params, transforms, profiles, data, **kwargs):
+    # TODO: ask question
     data["e^helical"] = (
         data["B^zeta"] * data["e^theta"].T - data["B^theta"] * data["e^zeta"].T
     ).T

desc/compute/_field.py

@@ -12,6 +12,7 @@
     surface_integrals,
     surface_max,
     surface_min,
+    surface_integrals_map,
 )
 
 
@@ -2654,6 +2655,7 @@ def _B_rms(params, transforms, profiles, data, **kwargs):
     profiles=[],
     coordinates="r",
     data=["sqrt(g)", "|B|", "V_r(r)"],
+    axis_limit_data=["sqrt(g)_r", "V_rr(r)"],
 )
 def _B_fsa(params, transforms, profiles, data, **kwargs):
     data["<|B|>"] = surface_averages(
@@ -2662,6 +2664,16 @@ def _B_fsa(params, transforms, profiles, data, **kwargs):
         data["sqrt(g)"],
         denominator=data["V_r(r)"],
     )
+    if transforms["grid"].axis.size:
+        limit = surface_averages(
+            transforms["grid"],
+            data["|B|"],
+            data["sqrt(g)_r"],
+            denominator=data["V_rr(r)"],
+        )
+        data["<|B|>"] = put(
+            data["<|B|>"], transforms["grid"].axis, limit[transforms["grid"].axis]
+        )
     return data
 
 
@@ -2677,6 +2689,7 @@ def _B_fsa(params, transforms, profiles, data, **kwargs):
     profiles=[],
     coordinates="r",
     data=["sqrt(g)", "|B|^2", "V_r(r)"],
+    axis_limit_data=["sqrt(g)_r", "V_rr(r)"],
 )
 def _B2_fsa(params, transforms, profiles, data, **kwargs):
     data["<B^2>"] = surface_averages(
@@ -2685,6 +2698,16 @@ def _B2_fsa(params, transforms, profiles, data, **kwargs):
         data["sqrt(g)"],
         denominator=data["V_r(r)"],
     )
+    if transforms["grid"].axis.size:
+        limit = surface_averages(
+            transforms["grid"],
+            data["|B|^2"],
+            data["sqrt(g)_r"],
+            denominator=data["V_rr(r)"],
+        )
+        data["<B^2>"] = put(
+            data["<B^2>"], transforms["grid"].axis, limit[transforms["grid"].axis]
+        )
     return data
 
 
@@ -2700,6 +2723,7 @@ def _B2_fsa(params, transforms, profiles, data, **kwargs):
     profiles=[],
     coordinates="r",
     data=["sqrt(g)", "|B|", "V_r(r)"],
+    axis_limit_data=["sqrt(g)_r", "V_rr(r)"],
 )
 def _1_over_B_fsa(params, transforms, profiles, data, **kwargs):
     data["<1/|B|>"] = surface_averages(
@@ -2708,6 +2732,16 @@ def _1_over_B_fsa(params, transforms, profiles, data, **kwargs):
         data["sqrt(g)"],
         denominator=data["V_r(r)"],
     )
+    if transforms["grid"].axis.size:
+        limit = surface_averages(
+            transforms["grid"],
+            1 / data["|B|"],
+            data["sqrt(g)_r"],
+            denominator=data["V_rr(r)"],
+        )
+        data["<1/|B|>"] = put(
+            data["<1/|B|>"], transforms["grid"].axis, limit[transforms["grid"].axis]
+        )
     return data
 
 
@@ -2728,20 +2762,33 @@ def _1_over_B_fsa(params, transforms, profiles, data, **kwargs):
         "B",
         "B_r",
         "|B|^2",
-        "<B^2>",
         "V_r(r)",
         "V_rr(r)",
     ],
+    axis_limit_data=["sqrt(g)_rr", "V_rrr(r)"],
 )
 def _B2_fsa_r(params, transforms, profiles, data, **kwargs):
+    compute_surface_integrals = surface_integrals_map(transforms["grid"])
+    num = compute_surface_integrals(data["sqrt(g)"] * data["|B|^2"])
+    num_r = compute_surface_integrals(
+        data["sqrt(g)_r"] * data["|B|^2"]
+        + 2 * data["sqrt(g)"] * dot(data["B"], data["B_r"]),
+    )
     data["<B^2>_r"] = (
-        surface_integrals(
-            transforms["grid"],
-            data["sqrt(g)_r"] * data["|B|^2"]
-            + data["sqrt(g)"] * 2 * dot(data["B"], data["B_r"]),
+        num_r / data["V_r(r)"] - num * data["V_rr(r)"] / data["V_r(r)"] ** 2
+    )
+    if transforms["grid"].axis.size:
+        limit = (
+            compute_surface_integrals(
+                data["sqrt(g)_rr"] * data["|B|^2"]
+                + 4 * data["sqrt(g)_r"] * dot(data["B"], data["B_r"])
+            )
+            * data["V_rr(r)"]
+            - num_r * data["V_rrr(r)"]
+        ) / (2 * data["V_rr(r)"] ** 2)
+        data["<B^2>_r"] = put(
+            data["<B^2>_r"], transforms["grid"].axis, limit[transforms["grid"].axis]
         )
-        - data["V_rr(r)"] * data["<B^2>"]
-    ) / data["V_r(r)"]
     return data
 
 
@@ -3336,9 +3383,9 @@ def _kappa_g(params, transforms, profiles, data, **kwargs):
 )
 def _grad_B_vec(params, transforms, profiles, data, **kwargs):
     data["grad(B)"] = (
-        (data["B_r"][:, None, :] * data["e^rho"][:, :, None])
-        + (data["B_t"][:, None, :] * data["e^theta"][:, :, None])
-        + (data["B_z"][:, None, :] * data["e^zeta"][:, :, None])
+        (data["B_r"][:, jnp.newaxis, :] * data["e^rho"][:, :, jnp.newaxis])
+        + (data["B_t"][:, jnp.newaxis, :] * data["e^theta"][:, :, jnp.newaxis])
+        + (data["B_z"][:, jnp.newaxis, :] * data["e^zeta"][:, :, jnp.newaxis])
     )
     return data
 

desc/compute/_geometry.py

@@ -77,12 +77,29 @@ def _V_r_of_r(params, transforms, profiles, data, **kwargs):
     transforms={"grid": []},
     profiles=[],
     coordinates="r",
-    data=["sqrt(g)_r", "sqrt(g)"],
+    data=["sqrt(g)_r"],
 )
 def _V_rr_of_r(params, transforms, profiles, data, **kwargs):
-    data["V_rr(r)"] = surface_integrals(
-        transforms["grid"], data["sqrt(g)_r"] * jnp.sign(data["sqrt(g)"])
-    )
+    data["V_rr(r)"] = surface_integrals(transforms["grid"], data["sqrt(g)_r"])
+    return data
+
+
+@register_compute_fun(
+    name="V_rrr(r)",
+    label="\\partial_{\\rho\\rho\\rho} V(\\rho)",
+    units="m^{3}",
+    units_long="cubic meters",
+    description="Volume enclosed by flux surfaces, third derivative wrt radial "
+    + "coordinate",
+    dim=1,
+    params=[],
+    transforms={"grid": []},
+    profiles=[],
+    coordinates="r",
+    data=["sqrt(g)_rr"],
+)
+def _V_rrr_of_r(params, transforms, profiles, data, **kwargs):
+    data["V_rrr(r)"] = surface_integrals(transforms["grid"], data["sqrt(g)_rr"])
     return data
 
 

desc/compute/_qs.py

@@ -268,9 +268,8 @@ def _B_mn(params, transforms, profiles, data, **kwargs):
     norm = 2 ** (3 - jnp.sum((transforms["B"].basis.modes == 0), axis=1))
     data["|B|_mn"] = (
         norm  # 1 if m=n=0, 2 if m=0 or n=0, 4 if m!=0 and n!=0
-        * jnp.matmul(
-            transforms["B"].basis.evaluate(nodes).T, data["sqrt(g)_B"] * data["|B|"]
-        )
+        * transforms["B"].basis.evaluate(nodes).T
+        @ (data["sqrt(g)_B"] * data["|B|"])
         / transforms["B"].grid.num_nodes
     )
     return data

desc/compute/_stability.py

@@ -2,7 +2,7 @@
 
 from scipy.constants import mu_0
 
-from desc.backend import jnp
+from desc.backend import jnp, put
 
 from .data_index import register_compute_fun
 from .utils import dot, surface_integrals
@@ -25,6 +25,9 @@ def _D_shear(params, transforms, profiles, data, **kwargs):
     # Implements equations 4.16 in M. Landreman & R. Jorge (2020)
     # doi:10.1017/S002237782000121X.
     data["D_shear"] = (data["iota_r"] / (4 * jnp.pi * data["psi_r"])) ** 2
+    # TODO: limit at magnetic axis is
+    #  data["iota_rrr"] / (2 * (4 * jnp.pi * data["psi_rr"]) ** 2)
+    #  if iota_r = iota_rr = 0, and nan otherwise.
     return data
 
 
@@ -199,4 +202,7 @@ def _magnetic_well(params, transforms, profiles, data, **kwargs):
         * (2 * mu_0 * data["p_r"] + data["<B^2>_r"])
         / (data["V_r(r)"] * data["<B^2>"])
     )
+    if transforms["grid"].axis.size:
+        # coefficient of limit is V_r / V_rr = 0
+        data["magnetic well"] = put(data["magnetic well"], transforms["grid"].axis, 0)
     return data