Basis vectors

desc/compute/_bootstrap.py

@@ -189,17 +189,15 @@
 
     # Redl eq (14):
     X32e = f_t / (
-        (
-            1
-            + 0.23 * (1 - 0.96 * f_t) * jnp.sqrt(nu_e) / jnp.sqrt(Zeff)
-            + 0.13
-            * (1 - 0.38 * f_t)
-            * nu_e
-            / (Zeff * Zeff)
-            * (
-                jnp.sqrt(1 + 2 * jnp.sqrt(Zeff - 1))
-                + f_t * f_t * jnp.sqrt((0.075 + 0.25 * (Zeff - 1) ** 2) * nu_e)
-            )
+        1
+        + 0.23 * (1 - 0.96 * f_t) * jnp.sqrt(nu_e) / jnp.sqrt(Zeff)
+        + 0.13
+        * (1 - 0.38 * f_t)
+        * nu_e
+        / (Zeff * Zeff)
+        * (
+            jnp.sqrt(1 + 2 * jnp.sqrt(Zeff - 1))
+            + f_t * f_t * jnp.sqrt((0.075 + 0.25 * (Zeff - 1) ** 2) * nu_e)
         )
     )
 
@@ -280,8 +278,8 @@
     J_dot_B = dnds_term + dTeds_term + dTids_term
 
     # Store all results in the J_dot_B_data dictionary:
-    nu_e_star = nu_e
-    nu_i_star = nu_i
+    nu_e_star = nu_e  # noqa: F841
+    nu_i_star = nu_i  # noqa: F841
     variables = [
         "rho",
         "ne",
@@ -342,6 +340,7 @@
         "Ti",
         "Ti_r",
         "Zeff",
+        "rho",
     ],
     helicity="helicity",
 )

desc/compute/_field.py

@@ -203,9 +203,9 @@ def _B0_r(params, transforms, profiles, data, **kwargs):
     data=["B0", "B0_r", "iota", "iota_r", "lambda_z", "lambda_rz"],
 )
 def _B_sup_theta_r(params, transforms, profiles, data, **kwargs):
-    data["B^theta_r"] = data["B0_r"] * (data["iota"] - data["lambda_z"]) + data[
-        "B0"
-    ] * (data["iota_r"] - data["lambda_rz"])
+    data["B^theta_r"] = data["B0_r"] * (data["iota"] - data["lambda_z"]) + (
+        data["B0"] * (data["iota_r"] - data["lambda_rz"])
+    )
     return data
 
 
@@ -2223,7 +2223,7 @@ def _B_mag_tz(params, transforms, profiles, data, **kwargs):
         "B^zeta_r",
         "B_theta_r",
         "B_zeta_r",
-        "B^theta_t",
+        "B^theta_z",
         "B^zeta_z",
         "B_theta_z",
         "B_zeta_z",

desc/compute/_geometry.py

@@ -194,7 +194,7 @@ def _R0_over_a(params, transforms, profiles, data, **kwargs):
     transforms={"grid": []},
     profiles=[],
     coordinates="",
-    data=["sqrt(g)", "g_tt"],
+    data=["sqrt(g)", "g_tt", "R"],
 )
 def _a_major_over_a_minor(params, transforms, profiles, data, **kwargs):
     max_rho = transforms["grid"].nodes[transforms["grid"].unique_rho_idx[-1], 0]

desc/compute/_metric.py

@@ -722,10 +722,10 @@ def _g_sup_tz(params, transforms, profiles, data, **kwargs):
         "e^rho",
         "sqrt(g)_r",
         "sqrt(g)",
-        "e_rho",
         "e_theta",
-        "e_rho_r",
         "e_theta_r",
+        "e_zeta_r",
+        "e_zeta",
     ],
 )
 def _g_sup_rr_r(params, transforms, profiles, data, **kwargs):
@@ -863,16 +863,13 @@ def _g_sup_rz_r(params, transforms, profiles, data, **kwargs):
     profiles=[],
     coordinates="rtz",
     data=[
-        "e^theta",
         "e^theta",
         "sqrt(g)_r",
         "sqrt(g)",
-        "e_theta",
-        "e_theta",
         "e_rho",
-        "e_theta_r",
-        "e_theta_r",
         "e_rho_r",
+        "e_zeta_r",
+        "e_zeta",
     ],
 )
 def _g_sup_tt_r(params, transforms, profiles, data, **kwargs):
@@ -961,15 +958,12 @@ def _g_sup_tz_r(params, transforms, profiles, data, **kwargs):
     profiles=[],
     coordinates="rtz",
     data=[
-        "e^zeta",
         "e^zeta",
         "sqrt(g)_r",
         "sqrt(g)",
-        "e_zeta",
-        "e_zeta",
+        "e_theta",
         "e_rho",
-        "e_zeta_r",
-        "e_zeta_r",
+        "e_theta_r",
         "e_rho_r",
     ],
 )
@@ -1010,16 +1004,13 @@ def _g_sup_zz_r(params, transforms, profiles, data, **kwargs):
     profiles=[],
     coordinates="rtz",
     data=[
-        "e^rho",
         "e^rho",
         "sqrt(g)_t",
         "sqrt(g)",
-        "e_rho",
-        "e_rho",
         "e_theta",
-        "e_rho_t",
-        "e_rho_t",
         "e_theta_t",
+        "e_zeta_t",
+        "e_zeta",
     ],
 )
 def _g_sup_rr_t(params, transforms, profiles, data, **kwargs):
@@ -1157,15 +1148,12 @@ def _g_sup_rz_t(params, transforms, profiles, data, **kwargs):
     profiles=[],
     coordinates="rtz",
     data=[
-        "e^theta",
         "e^theta",
         "sqrt(g)_t",
         "sqrt(g)",
-        "e_theta",
-        "e_theta",
+        "e_zeta",
         "e_rho",
-        "e_theta_t",
-        "e_theta_t",
+        "e_zeta_t",
         "e_rho_t",
     ],
 )
@@ -1255,15 +1243,12 @@ def _g_sup_tz_t(params, transforms, profiles, data, **kwargs):
     profiles=[],
     coordinates="rtz",
     data=[
-        "e^zeta",
         "e^zeta",
         "sqrt(g)_t",
         "sqrt(g)",
-        "e_zeta",
-        "e_zeta",
+        "e_theta",
         "e_rho",
-        "e_zeta_t",
-        "e_zeta_t",
+        "e_theta_t",
         "e_rho_t",
     ],
 )
@@ -1304,15 +1289,12 @@ def _g_sup_zz_t(params, transforms, profiles, data, **kwargs):
     profiles=[],
     coordinates="rtz",
     data=[
-        "e^rho",
         "e^rho",
         "sqrt(g)_z",
         "sqrt(g)",
-        "e_rho",
-        "e_rho",
+        "e_zeta",
         "e_theta",
-        "e_rho_z",
-        "e_rho_z",
+        "e_zeta_z",
         "e_theta_z",
     ],
 )
@@ -1451,15 +1433,12 @@ def _g_sup_rz_z(params, transforms, profiles, data, **kwargs):
     profiles=[],
     coordinates="rtz",
     data=[
-        "e^theta",
         "e^theta",
         "sqrt(g)_z",
         "sqrt(g)",
-        "e_theta",
-        "e_theta",
+        "e_zeta",
         "e_rho",
-        "e_theta_z",
-        "e_theta_z",
+        "e_zeta_z",
         "e_rho_z",
     ],
 )
@@ -1549,15 +1528,12 @@ def _g_sup_tz_z(params, transforms, profiles, data, **kwargs):
     profiles=[],
     coordinates="rtz",
     data=[
-        "e^zeta",
         "e^zeta",
         "sqrt(g)_z",
         "sqrt(g)",
-        "e_zeta",
-        "e_zeta",
+        "e_theta",
         "e_rho",
-        "e_zeta_z",
-        "e_zeta_z",
+        "e_theta_z",
         "e_rho_z",
     ],
 )

desc/compute/_profiles.py

@@ -68,7 +68,7 @@ def _psi_rr(params, transforms, profiles, data, **kwargs):
     units_long="Webers",
     description="Toroidal flux (normalized by 2pi), third radial derivative",
     dim=1,
-    params=["Psi"],
+    params=[],
     transforms={},
     profiles=[],
     coordinates="r",
@@ -509,9 +509,9 @@ def _iota(params, transforms, profiles, data, **kwargs):
             current_term = put(
                 current_term, transforms["grid"].axis, limit[transforms["grid"].axis]
             )
-        data["iota"] = (current_term + data["iota_zero_current_num"]) / data[
-            "iota_zero_current_den"
-        ]
+        data["iota"] = (current_term + data["iota_zero_current_num"]) / (
+            data["iota_zero_current_den"]
+        )
     return data
 
 
@@ -530,7 +530,6 @@ def _iota(params, transforms, profiles, data, **kwargs):
         "iota",
         "psi_r",
         "psi_rr",
-        "iota_zero_current_num",
         "iota_zero_current_num_r",
         "iota_zero_current_den",
         "iota_zero_current_den_r",
@@ -582,8 +581,6 @@ def _iota_r(params, transforms, profiles, data, **kwargs):
         "psi_r",
         "psi_rr",
         "psi_rrr",
-        "iota_zero_current_num",
-        "iota_zero_current_num_r",
         "iota_zero_current_num_rr",
         "iota_zero_current_den",
         "iota_zero_current_den_r",

docs/sort_compute_funs.py

@@ -0,0 +1,31 @@
+"""Sort the compute functions and print to specified output.
+
+First command line argument is module name, e.g. "_basis_vectors".
+Second command line output is file to print to,
+e.g. "sorted_basis_vectors.txt".
+"""
+
+import inspect
+import re
+import sys
+
+import desc.compute
+
+# Gather all compute function source code and map quantity name to source code.
+source_codes = {}
+pattern = re.compile(r"(?<=name=)[^,]+")
+for module_name, module in inspect.getmembers(desc.compute, inspect.ismodule):
+    if module_name == sys.argv[1]:
+        for _, fun in inspect.getmembers(module, inspect.isfunction):
+            source_code = inspect.getsource(fun)
+            # quantities that this function computes
+            matches = pattern.findall(source_code)
+            if matches:  # skip imported functions
+                source_codes[matches[0]] = source_code
+
+# Write compute functions sorted by name to file.
+with open(sys.argv[2], "w") as output_file:
+    for name in sorted(source_codes):
+        output_file.write("\n")
+        output_file.write(source_codes[name])
+        output_file.write("\n")