Merge pull request #598 from PlasmaControl/basis_vectors

Generalize basis vectors for new toroidal angle. See #465
PlasmaControl · Aug 11, 2023 · 7e0dcf1 · 7e0dcf1
2 parents cc10a94 + 309c4ef
commit 7e0dcf1
Show file tree

Hide file tree

Showing 12 changed files with 3,690 additions and 542 deletions.
diff --git a/desc/compute/_basis_vectors.py b/desc/compute/_basis_vectors.py
diff --git a/desc/compute/_bootstrap.py b/desc/compute/_bootstrap.py
@@ -334,6 +334,7 @@ def j_dot_B_Redl(
         "Ti",
         "Ti_r",
         "Zeff",
+        "rho",
     ],
     helicity="helicity",
 )

diff --git a/desc/compute/_field.py b/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",

diff --git a/desc/compute/_geometry.py b/desc/compute/_geometry.py
@@ -257,7 +257,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]

diff --git a/desc/compute/_metric.py b/desc/compute/_metric.py
@@ -766,10 +766,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):
@@ -907,16 +907,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):
@@ -1005,15 +1002,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",
     ],
 )
@@ -1054,16 +1048,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):
@@ -1201,15 +1192,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",
     ],
 )
@@ -1299,15 +1287,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",
     ],
 )
@@ -1348,15 +1333,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",
     ],
 )
@@ -1495,15 +1477,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",
     ],
 )
@@ -1593,15 +1572,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",
     ],
 )

diff --git a/desc/compute/_profiles.py b/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",

diff --git a/docs/sort_compute_funs.py b/docs/sort_compute_funs.py
@@ -0,0 +1,39 @@
+"""Sort the compute functions and print to specified output.
+
+First command line argument is module name, e.g. "_basis_vectors".
+Second command line argument is output file to print to,
+e.g. "sorted_basis_vectors.txt".
+"""
+
+import inspect
+import re
+import sys
+
+import desc.compute
+
+preamble = ""
+
+# Gather all compute function source code and map 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]:
+        # everything in the module until the first compute function
+        # (import statements, etc.)
+        preamble = inspect.getsource(module).partition("@register_compute_fun")[0]
+        for function_name, fun in inspect.getmembers(module, inspect.isfunction):
+            source_code = inspect.getsource(fun)
+            matches = pattern.findall(source_code)
+            if matches:  # skip imported functions
+                # matches[0] is the thing this function says it computes, e.g. x
+                # while function_name is e.g. _x_ZernikeRZToroidalSection
+                if function_name in source_codes:
+                    raise ValueError("Can't sort when things have same name.")
+                source_codes[function_name] = source_code
+
+# Write functions sorted to file.
+with open(sys.argv[2], "w") as output_file:
+    output_file.write(preamble)
+    for function_name in sorted(source_codes):
+        output_file.write(source_codes[function_name])
+        output_file.write("\n\n")