Fix bugs in compute quantities

desc/compute/_basis_vectors.py

@@ -94,7 +94,7 @@ def _e_sup_rho_r(params, transforms, profiles, data, **kwargs):
     name="e^rho_rr",
     label="\\partial_{\\rho\\rho} \\mathbf{e}^{\\rho}",
     units="m^{-1}",
-    units_long="inverse square meters",
+    units_long="inverse meters",
     description="Contravariant Radial basis vector, 2nd derivative"
     " wrt radial coordinate",
     dim=3,
@@ -144,7 +144,7 @@ def _e_sup_rho_rr(params, transforms, profiles, data, **kwargs):
     name="e^rho_rt",
     label="\\partial_{\\rho\\theta} \\mathbf{e}^{\\rho}",
     units="m^{-1}",
-    units_long="inverse square meters",
+    units_long="inverse meters",
     description="Contravariant Radial basis vector, derivative"
     " wrt radial and poloidal coordinate",
     dim=3,
@@ -199,7 +199,7 @@ def _e_sup_rho_rt(params, transforms, profiles, data, **kwargs):
     name="e^rho_rz",
     label="\\partial_{\\rho\\zeta} \\mathbf{e}^{\\rho}",
     units="m^{-1}",
-    units_long="inverse square meters",
+    units_long="inverse meters",
     description="Contravariant Radial basis vector, derivative"
     " wrt radial and toroidal coordinate",
     dim=3,
@@ -337,7 +337,7 @@ def _e_sup_rho_p_PEST(params, transforms, profiles, data, **kwargs):
     name="e^rho_tt",
     label="\\partial_{\\theta\\theta} \\mathbf{e}^{\\rho}",
     units="m^{-1}",
-    units_long="inverse square meters",
+    units_long="inverse meters",
     description="Contravariant Radial basis vector, 2nd derivative"
     " wrt poloidal coordinate",
     dim=3,
@@ -387,7 +387,7 @@ def _e_sup_rho_tt(params, transforms, profiles, data, **kwargs):
     name="e^rho_tz",
     label="\\partial_{\\theta\\zeta} \\mathbf{e}^{\\rho}",
     units="m^{-1}",
-    units_long="inverse square meters",
+    units_long="inverse meters",
     description="Contravariant Radial basis vector, derivative"
     " wrt poloidal and toroidal coordinate",
     dim=3,
@@ -484,7 +484,7 @@ def _e_sup_rho_z(params, transforms, profiles, data, **kwargs):
     name="e^rho_zz",
     label="\\partial_{\\zeta\\zeta} \\mathbf{e}^{\\rho}",
     units="m^{-1}",
-    units_long="inverse square meters",
+    units_long="inverse meters",
     description="Contravariant Radial basis vector, 2nd derivative"
     " wrt toroidal coordinate",
     dim=3,
@@ -731,7 +731,7 @@ def _e_sup_theta_r(params, transforms, profiles, data, **kwargs):
     name="e^theta_rr",
     label="\\partial_{\\rho\\rho} \\mathbf{e}^{\\theta}",
     units="m^{-1}",
-    units_long="inverse square meters",
+    units_long="inverse meters",
     description="Contravariant Poloidal basis vector, 2nd derivative"
     " wrt radial coordinate",
     dim=3,
@@ -781,7 +781,7 @@ def _e_sup_theta_rr(params, transforms, profiles, data, **kwargs):
     name="e^theta_rt",
     label="\\partial_{\\rho\\theta} \\mathbf{e}^{\\theta}",
     units="m^{-1}",
-    units_long="inverse square meters",
+    units_long="inverse meters",
     description="Contravariant Poloidal basis vector, derivative"
     " wrt radial and poloidal coordinate",
     dim=3,
@@ -836,7 +836,7 @@ def _e_sup_theta_rt(params, transforms, profiles, data, **kwargs):
     name="e^theta_rz",
     label="\\partial_{\\rho\\zeta} \\mathbf{e}^{\\theta}",
     units="m^{-1}",
-    units_long="inverse square meters",
+    units_long="inverse meters",
     description="Contravariant Poloidal basis vector, derivative"
     " wrt radial and toroidal coordinate",
     dim=3,
@@ -919,7 +919,7 @@ def _e_sup_theta_t(params, transforms, profiles, data, **kwargs):
     name="e^theta_tt",
     label="\\partial_{\\theta\\theta} \\mathbf{e}^{\\theta}",
     units="m^{-1}",
-    units_long="inverse square meters",
+    units_long="inverse meters",
     description="Contravariant Poloidal basis vector, 2nd derivative"
     " wrt poloidal coordinate",
     dim=3,
@@ -969,7 +969,7 @@ def _e_sup_theta_tt(params, transforms, profiles, data, **kwargs):
     name="e^theta_tz",
     label="\\partial_{\\theta\\zeta} \\mathbf{e}^{\\theta}",
     units="m^{-1}",
-    units_long="inverse square meters",
+    units_long="inverse meters",
     description="Contravariant Poloidal basis vector, derivative"
     " wrt poloidal and toroidal coordinate",
     dim=3,
@@ -1052,7 +1052,7 @@ def _e_sup_theta_z(params, transforms, profiles, data, **kwargs):
     name="e^theta_zz",
     label="\\partial_{\\zeta\\zeta} \\mathbf{e}^{\\theta}",
     units="m^{-1}",
-    units_long="inverse square meters",
+    units_long="inverse meters",
     description="Contravariant Poloidal basis vector, 2nd derivative"
     " wrt toroidal coordinate",
     dim=3,
@@ -1158,7 +1158,7 @@ def _e_sup_zeta_r(params, transforms, profiles, data, **kwargs):
     name="e^zeta_rr",
     label="\\partial_{\\rho\\rho} \\mathbf{e}^{\\zeta}",
     units="m^{-1}",
-    units_long="inverse square meters",
+    units_long="inverse meters",
     description="Contravariant Toroidal basis vector, 2nd derivative"
     " wrt radial coordinate",
     dim=3,
@@ -1208,7 +1208,7 @@ def _e_sup_zeta_rr(params, transforms, profiles, data, **kwargs):
     name="e^zeta_rt",
     label="\\partial_{\\rho\\theta} \\mathbf{e}^{\\zeta}",
     units="m^{-1}",
-    units_long="inverse square meters",
+    units_long="inverse meters",
     description="Contravariant Toroidal basis vector, derivative"
     " wrt radial and poloidal coordinate",
     dim=3,
@@ -1263,7 +1263,7 @@ def _e_sup_zeta_rt(params, transforms, profiles, data, **kwargs):
     name="e^zeta_rz",
     label="\\partial_{\\rho\\zeta} \\mathbf{e}^{\\zeta}",
     units="m^{-1}",
-    units_long="inverse square meters",
+    units_long="inverse meters",
     description="Contravariant Toroidal basis vector, derivative"
     " wrt radial and toroidal coordinate",
     dim=3,
@@ -1404,7 +1404,7 @@ def _e_sup_zeta_p_PEST(params, transforms, profiles, data, **kwargs):
     name="e^zeta_tt",
     label="\\partial_{\\theta\\theta} \\mathbf{e}^{\\zeta}",
     units="m^{-1}",
-    units_long="inverse square meters",
+    units_long="inverse meters",
     description="Contravariant Toroidal basis vector, 2nd derivative"
     " wrt poloidal coordinate",
     dim=3,
@@ -1454,7 +1454,7 @@ def _e_sup_zeta_tt(params, transforms, profiles, data, **kwargs):
     name="e^zeta_tz",
     label="\\partial_{\\theta\\zeta} \\mathbf{e}^{\\zeta}",
     units="m^{-1}",
-    units_long="inverse square meters",
+    units_long="inverse meters",
     description="Contravariant Toroidal basis vector, derivative"
     " wrt poloidal and toroidal coordinate",
     dim=3,
@@ -1552,7 +1552,7 @@ def _e_sup_zeta_z(params, transforms, profiles, data, **kwargs):
     name="e^zeta_zz",
     label="\\partial_{\\zeta\\zeta} \\mathbf{e}^{\\zeta}",
     units="m^{-1}",
-    units_long="inverse square meters",
+    units_long="inverse meters",
     description="Contravariant Toroidal basis vector, 2nd derivative"
     " wrt toroidal coordinate",
     dim=3,
@@ -1769,7 +1769,7 @@ def _e_sub_rho_rrr(params, transforms, profiles, data, **kwargs):
             4 * data["R_rrr"] * data["omega_r"]
             + 6 * data["R_rr"] * data["omega_rr"]
             - 4 * data["R_r"] * (data["omega_r"] ** 3 - data["omega_rrr"])
-            + data["R_r"]
+            + data["R"]
             * (data["omega_rrrr"] - 6 * data["omega_r"] ** 2 * data["omega_rr"]),
             data["Z_rrrr"],
         ]
@@ -2624,8 +2624,8 @@ def _e_sub_theta(params, transforms, profiles, data, **kwargs):
 @register_compute_fun(
     name="e_theta/sqrt(g)",
     label="\\mathbf{e}_{\\theta} / \\sqrt{g}",
-    units="m",
-    units_long="meters",
+    units="m^{-2}",
+    units_long="inverse square meters",
     description="Covariant Poloidal basis vector divided by 3-D volume Jacobian",
     dim=3,
     params=[],
@@ -2778,7 +2778,7 @@ def _e_sub_theta_rtt(params, transforms, profiles, data, **kwargs):
                 data["R_r"] * data["omega_tt"]
                 + data["R_tt"] * data["omega_r"]
                 + 2 * data["R_t"] * data["omega_rt"]
-                + data["R"] * data["omega_rt"]
+                + data["R"] * data["omega_rtt"]
             )
             - data["omega_r"]
             * (3 * data["R_t"] * data["omega_tt"] + data["R"] * data["omega_ttt"])
@@ -3336,10 +3336,7 @@ def _e_sub_zeta_rzz(params, transforms, profiles, data, **kwargs):
                 - data["omega_zzz"]
             )
             - data["R"]
-            * (
-                3 * data["omega_rz"] * (1 + data["omega_z"] * (1 + data["omega_z"]))
-                - data["omega_rzzz"]
-            )
+            * (3 * data["omega_rz"] * (1 + data["omega_z"]) ** 2 - data["omega_rzzz"])
             + data["omega_r"]
             * (
                 -3 * data["R_z"] * (1 + data["omega_z"]) ** 2

desc/compute/_bootstrap.py

@@ -118,17 +118,19 @@ def compute_J_dot_B_Redl(geom_data, profile_data, helicity_N=None):
         Dictionary containing the data described above.
     profile_data : dict
         Dictionary containing the data described above.
-    helicity_N : int
+    helicity_N : int, optional
         Set to 0 for quasi-axisymmetry, or +/- NFP for quasi-helical symmetry.
         This quantity is used to apply the quasisymmetry isomorphism to map the
         collisionality and bootstrap current from the tokamak expressions to
-        quasi-helical symmetry.
+        quasi-helical symmetry. Defaults to 0.
 
     Returns
     -------
     J_dot_B_data : dict
         Dictionary containing the computed data listed above.
     """
+    helicity_N = 0 if helicity_N is None else helicity_N
+
     G = geom_data["G"]
     R = geom_data["R"]
     iota = geom_data["iota"]
@@ -147,6 +149,7 @@ def compute_J_dot_B_Redl(geom_data, profile_data, helicity_N=None):
     # problem is avoided by adding a tiny number here:
     Zeff = jnp.maximum(1 + 1.0e-14, profile_data["Zeff"])
     d_ne_d_s = profile_data["ne_r"] / (2 * rho)
+    d_ni_d_s = profile_data["ni_r"] / (2 * rho)
     d_Te_d_s = profile_data["Te_r"] / (2 * rho)
     d_Ti_d_s = profile_data["Ti_r"] / (2 * rho)
 
@@ -254,9 +257,8 @@ def compute_J_dot_B_Redl(geom_data, profile_data, helicity_N=None):
     dnds_term = (
         -G
         * elementary_charge
-        * (ne * Te + ni * Ti)
         * L31
-        * (d_ne_d_s / ne)
+        * (Te * d_ne_d_s + Ti * d_ni_d_s)
         / (psi_edge * (iota - helicity_N))
     )
     dTeds_term = (
@@ -288,6 +290,7 @@ def compute_J_dot_B_Redl(geom_data, profile_data, helicity_N=None):
     J_dot_B_data["Te"] = Te
     J_dot_B_data["Ti"] = Ti
     J_dot_B_data["d_ne_d_s"] = d_ne_d_s
+    J_dot_B_data["d_ni_d_s"] = d_ni_d_s
     J_dot_B_data["d_Te_d_s"] = d_Te_d_s
     J_dot_B_data["d_Ti_d_s"] = d_Ti_d_s
     J_dot_B_data["ln_Lambda_e"] = ln_Lambda_e

desc/compute/_equil.py

@@ -205,7 +205,7 @@ def _J_sup_zeta_v_PEST(params, transforms, profiles, data, **kwargs):
 # TODO: Generalize for a general zeta before #568
 @register_compute_fun(
     name="(J^zeta_p)|PEST",
-    label="\\partial_{\\phi}|_{\\rho, \\phi} J^{\\zeta}",
+    label="\\partial_{\\phi}|_{\\rho, \\vartheta} J^{\\zeta}",
     units="A \\cdot m^{-3}",
     units_long="Amperes / cubic meter",
     description="Contravariant PEST toroidal component of plasma current density"
@@ -253,7 +253,7 @@ def _J_sup_theta_t(params, transforms, profiles, data, **kwargs):
 
 @register_compute_fun(
     name="J^theta_z",
-    label="\\partial_{\\theta} J^{\\theta}",
+    label="\\partial_{\\zeta} J^{\\theta}",
     units="A \\cdot m^{-3}",
     units_long="Amperes / cubic meter",
     description="Derivative of Contravariant poloidal component of plasma current"
@@ -825,7 +825,7 @@ def _Fmag_vol(params, transforms, profiles, data, **kwargs):
 
 @register_compute_fun(
     name="e^helical",
-    label="B^{\\theta} \\nabla \\zeta - B^{\\zeta} \\nabla \\theta",
+    label="B^{\\zeta} \\nabla \\theta - B^{\\theta} \\nabla \\zeta",
     units="T \\cdot m^{-2}",
     units_long="Tesla / square meter",
     description="Helical basis vector",
@@ -845,9 +845,9 @@ def _e_sup_helical(params, transforms, profiles, data, **kwargs):
 
 @register_compute_fun(
     name="e^helical*sqrt(g)",
-    label="\\sqrt{g}(B^{\\theta} \\nabla \\zeta - B^{\\zeta} \\nabla \\theta)",
-    units="T \\cdot m^{2}",
-    units_long="Tesla * square meter",
+    label="\\sqrt{g}(B^{\\zeta} \\nabla \\theta - B^{\\theta} \\nabla \\zeta)",
+    units="T \\cdot m",
+    units_long="Tesla * meter",
     description="Helical basis vector weighted by 3-D volume Jacobian",
     dim=3,
     params=[],
@@ -866,7 +866,7 @@ def _e_sup_helical_times_sqrt_g(params, transforms, profiles, data, **kwargs):
 
 @register_compute_fun(
     name="|e^helical|",
-    label="|B^{\\theta} \\nabla \\zeta - B^{\\zeta} \\nabla \\theta|",
+    label="|B^{\\zeta} \\nabla \\theta - B^{\\theta} \\nabla \\zeta|",
     units="T \\cdot m^{-2}",
     units_long="Tesla / square meter",
     description="Magnitude of helical basis vector",
@@ -884,9 +884,9 @@ def _e_sup_helical_mag(params, transforms, profiles, data, **kwargs):
 
 @register_compute_fun(
     name="|e^helical*sqrt(g)|",
-    label="|\\sqrt{g}(B^{\\theta} \\nabla \\zeta - B^{\\zeta} \\nabla \\theta)|",
-    units="T \\cdot m^{2}",
-    units_long="Tesla * square meter",
+    label="|\\sqrt{g}(B^{\\zeta} \\nabla \\theta - B^{\\theta} \\nabla \\zeta)|",
+    units="T \\cdot m",
+    units_long="Tesla * meter",
     description="Magnitude of helical basis vector weighted by 3-D volume Jacobian",
     dim=1,
     params=[],

desc/compute/_field.py

@@ -3485,7 +3485,7 @@ def _B_dot_grad_B_mag_vol(params, transforms, profiles, data, **kwargs):
 
 @register_compute_fun(
     name="B*grad(|B|)",
-    label="\\mathbf{B} \\cdot \\nabla B",
+    label="\\mathbf{B} \\cdot \\nabla |\\mathbf{B}|",
     units="T^2 \\cdot m^{-1}",
     units_long="Tesla squared / meters",
     description="",
@@ -3505,7 +3505,7 @@ def _B_dot_gradB(params, transforms, profiles, data, **kwargs):
 
 @register_compute_fun(
     name="(B*grad(|B|))_r",
-    label="\\partial_{\\theta} (\\mathbf{B} \\cdot \\nabla B)",
+    label="\\partial_{\\rho} (\\mathbf{B} \\cdot \\nabla |\\mathbf{B}|)",
     units="T^2 \\cdot m^{-1}",
     units_long="Tesla squared / meters",
     description="",
@@ -3537,7 +3537,7 @@ def _B_dot_gradB_r(params, transforms, profiles, data, **kwargs):
 
 @register_compute_fun(
     name="(B*grad(|B|))_t",
-    label="\\partial_{\\theta} (\\mathbf{B} \\cdot \\nabla B)",
+    label="\\partial_{\\theta} (\\mathbf{B} \\cdot \\nabla |\\mathbf{B}|)",
     units="T^2 \\cdot m^{-1}",
     units_long="Tesla squared / meters",
     description="",
@@ -3569,7 +3569,7 @@ def _B_dot_gradB_t(params, transforms, profiles, data, **kwargs):
 
 @register_compute_fun(
     name="(B*grad(|B|))_z",
-    label="\\partial_{\\zeta} (\\mathbf{B} \\cdot \\nabla B)",
+    label="\\partial_{\\zeta} (\\mathbf{B} \\cdot \\nabla |\\mathbf{B}|)",
     units="T^2 \\cdot m^{-1}",
     units_long="Tesla squared / meters",
     description="",