Fix bug in boozer coordinate computation for Omnigenous magnetic field (

#1166)

- [x] Move test compute everything to own file and exclude source grid
quantities. from commit 8a8d9de.
- [x] Remove xyz basis file in favor of correctness check to catch bug
#1088
- [x] Debug zeta_B computation. #1163 has been reproduced here (and on
master by Dario). This issue was resolved after fixing the computation
of zeta_B and theta_B for omnigenous fields.

desc/compute/_geometry.py

@@ -625,32 +625,14 @@ def _curvature_k1_rho(params, transforms, profiles, data, **kwargs):
     transforms={},
     profiles=[],
     coordinates="rtz",
-    data=["g_tt", "g_tz", "g_zz", "L_sff_rho", "M_sff_rho", "N_sff_rho"],
+    data=["curvature_k1_rho"],
     parameterization=[
         "desc.equilibrium.equilibrium.Equilibrium",
         "desc.geometry.surface.FourierRZToroidalSurface",
     ],
 )
 def _curvature_k2_rho(params, transforms, profiles, data, **kwargs):
-    # following notation from
-    # https://en.wikipedia.org/wiki/Parametric_surface
-    E = data["g_tt"]
-    F = data["g_tz"]
-    G = data["g_zz"]
-    L = data["L_sff_rho"]
-    M = data["M_sff_rho"]
-    N = data["N_sff_rho"]
-    a = E * G - F**2
-    b = 2 * F * M - L * G - E * N
-    c = L * N - M**2
-    r1 = (-b + jnp.sqrt(b**2 - 4 * a * c)) / (2 * a)
-    r2 = (-b - jnp.sqrt(b**2 - 4 * a * c)) / (2 * a)
-    # In the axis limit, the matrix of the first fundamental form is singular.
-    # The diagonal of the shape operator becomes unbounded,
-    # so the eigenvalues do not exist.
-    data["curvature_k1_rho"] = jnp.maximum(r1, r2)
-    data["curvature_k2_rho"] = jnp.minimum(r1, r2)
-    return data
+    return data  # noqa: unused dependency
 
 
 @register_compute_fun(
@@ -804,25 +786,10 @@ def _curvature_k1_theta(params, transforms, profiles, data, **kwargs):
     transforms={},
     profiles=[],
     coordinates="rtz",
-    data=["g_rr", "g_rz", "g_zz", "L_sff_theta", "M_sff_theta", "N_sff_theta"],
+    data=["curvature_k1_theta"],
 )
 def _curvature_k2_theta(params, transforms, profiles, data, **kwargs):
-    # following notation from
-    # https://en.wikipedia.org/wiki/Parametric_surface
-    E = data["g_zz"]
-    F = data["g_rz"]
-    G = data["g_rr"]
-    L = data["L_sff_theta"]
-    M = data["M_sff_theta"]
-    N = data["N_sff_theta"]
-    a = E * G - F**2
-    b = 2 * F * M - L * G - E * N
-    c = L * N - M**2
-    r1 = (-b + jnp.sqrt(b**2 - 4 * a * c)) / (2 * a)
-    r2 = (-b - jnp.sqrt(b**2 - 4 * a * c)) / (2 * a)
-    data["curvature_k1_theta"] = jnp.maximum(r1, r2)
-    data["curvature_k2_theta"] = jnp.minimum(r1, r2)
-    return data
+    return data  # noqa: unused dependency
 
 
 @register_compute_fun(
@@ -989,32 +956,14 @@ def _curvature_k1_zeta(params, transforms, profiles, data, **kwargs):
     transforms={},
     profiles=[],
     coordinates="rtz",
-    data=["g_rr", "g_rt", "g_tt", "L_sff_zeta", "M_sff_zeta", "N_sff_zeta"],
+    data=["curvature_k1_zeta"],
     parameterization=[
         "desc.equilibrium.equilibrium.Equilibrium",
         "desc.geometry.surface.ZernikeRZToroidalSection",
     ],
 )
 def _curvature_k2_zeta(params, transforms, profiles, data, **kwargs):
-    # following notation from
-    # https://en.wikipedia.org/wiki/Parametric_surface
-    E = data["g_rr"]
-    F = data["g_rt"]
-    G = data["g_tt"]
-    L = data["L_sff_zeta"]
-    M = data["M_sff_zeta"]
-    N = data["N_sff_zeta"]
-    a = E * G - F**2
-    b = 2 * F * M - L * G - E * N
-    c = L * N - M**2
-    r1 = (-b + jnp.sqrt(b**2 - 4 * a * c)) / (2 * a)
-    r2 = (-b - jnp.sqrt(b**2 - 4 * a * c)) / (2 * a)
-    # In the axis limit, the matrix of the first fundamental form is singular.
-    # The diagonal of the shape operator becomes unbounded,
-    # so the eigenvalues do not exist.
-    data["curvature_k1_zeta"] = jnp.maximum(r1, r2)
-    data["curvature_k2_zeta"] = jnp.minimum(r1, r2)
-    return data
+    return data  # noqa: unused dependency
 
 
 @register_compute_fun(

desc/compute/_omnigenity.py

@@ -429,22 +429,21 @@ def _omni_angle(params, transforms, profiles, data, **kwargs):
 
 @register_compute_fun(
     name="theta_B",
-    label="(\\theta_{B},\\zeta_{B})",
+    label="\\theta_{B}",
     units="rad",
     units_long="radians",
-    description="Boozer angular coordinates",
+    description="Boozer poloidal angle",
     dim=1,
     params=[],
     transforms={},
     profiles=[],
     coordinates="rtz",
     data=["alpha", "h"],
-    aliases=["zeta_B"],
     parameterization="desc.magnetic_fields._core.OmnigenousField",
     helicity="tuple: Type of quasisymmetry, (M,N). Default (1,0)",
     iota="float: Value of rotational transform on the Omnigenous surface. Default 1.0",
 )
-def _omni_map(params, transforms, profiles, data, **kwargs):
+def _omni_map_theta_B(params, transforms, profiles, data, **kwargs):
     M, N = kwargs.get("helicity", (1, 0))
     iota = kwargs.get("iota", 1)
 
@@ -476,6 +475,24 @@ def _omni_map(params, transforms, profiles, data, **kwargs):
     return data
 
 
+@register_compute_fun(
+    name="zeta_B",
+    label="\\zeta_{B}",
+    units="rad",
+    units_long="radians",
+    description="Boozer toroidal angle",
+    dim=1,
+    params=[],
+    transforms={},
+    profiles=[],
+    coordinates="rtz",
+    data=["theta_B"],
+    parameterization="desc.magnetic_fields._core.OmnigenousField",
+)
+def _omni_map_zeta_B(params, transforms, profiles, data, **kwargs):
+    return data  # noqa: unused dependency
+
+
 @register_compute_fun(
     name="|B|",
     label="|\\mathbf{B}|",

desc/compute/data_index.py

@@ -176,10 +176,11 @@ def _decorator(func):
                     if name in data_index[base_class]:
                         if p == data_index[base_class][name]["parameterization"]:
                             raise ValueError(
-                                f"Already registered function with parameterization {p} and name {name}."
+                                f"Already registered function with parameterization {p}"
+                                f" and name {name}."
                             )
-                        # if it was already registered from a parent class, we prefer
-                        # the child class.
+                        # if it was already registered from a parent class, we
+                        # prefer the child class.
                         inheritance_order = [base_class] + superclasses
                         if inheritance_order.index(p) > inheritance_order.index(
                             data_index[base_class][name]["parameterization"]

desc/plotting.py

@@ -2756,30 +2756,22 @@ def plot_boozer_surface(
 
     # default grids
     if grid_compute is None:
-        if eq_switch:
-            grid_kwargs = {
-                "rho": rho,
-                "M": 4 * thing.M,
-                "N": 4 * thing.N,
-                "NFP": thing.NFP,
-                "endpoint": False,
-            }
-        else:
-            grid_kwargs = {
-                "rho": rho,
-                "M": 50,
-                "N": 50,
-                "NFP": thing.NFP,
-                "endpoint": False,
-            }
+        # grid_compute only used for Equilibrium, not OmnigenousField
+        grid_kwargs = {
+            "rho": rho,
+            "M": 4 * getattr(thing, "M", 1),
+            "N": 4 * getattr(thing, "N", 1),
+            "NFP": thing.NFP,
+            "endpoint": False,
+        }
         grid_compute = _get_grid(**grid_kwargs)
     if grid_plot is None:
         grid_kwargs = {
             "rho": rho,
             "theta": 91,
             "zeta": 91,
             "NFP": thing.NFP,
-            "endpoint": True,
+            "endpoint": eq_switch,
         }
         grid_plot = _get_grid(**grid_kwargs)
 
@@ -2816,7 +2808,7 @@ def plot_boozer_surface(
             warnings.simplefilter("ignore")
             data = thing.compute(
                 ["theta_B", "zeta_B", "|B|"],
-                grid=grid_compute,
+                grid=grid_plot,
                 helicity=thing.helicity,
                 iota=iota,
             )

docs/adding_compute_funs.rst

@@ -155,8 +155,8 @@ if ``False`` is printed, then the limit of the quantity does not evaluate as fin
   The tests automatically detect this, so no further action is needed from developers in this case.
 
 
-The second step is to run the ``test_compute_everything`` test located in the ``tests/test_compute_funs.py`` file.
-This can be done with the command :console:`pytest -k test_compute_everything tests/test_compute_funs.py`.
+The second step is to run the ``test_compute_everything`` test located in the ``tests/test_compute_everything.py`` file.
+This can be done with the command :console:`pytest tests/test_compute_everything.py`.
 This test is a regression test to ensure that compute quantities in each new update of DESC do not differ significantly
 from previous versions of DESC.
 Since the new quantity did not exist in previous versions of DESC, one must run this test