Use max rho of surface to compute A(z)

desc/compute/_basis_vectors.py

@@ -3523,3 +3523,25 @@ def _n_zeta(params, transforms, profiles, data, **kwargs):
         ).T,
     )
     return data
+
+
+@register_compute_fun(
+    name="e_theta|r,p",
+    label="\\mathbf{e}_{\\theta} |_{\\rho, \\phi}",
+    units="m",
+    units_long="meters",
+    description="Tangent vector along boundary of constant phi toroidal surface",
+    dim=3,
+    params=[],
+    transforms={},
+    profiles=[],
+    coordinates="rtz",
+    data=["e_theta", "e_phi", "phi_t"],
+    parameterization=[
+        "desc.equilibrium.equilibrium.Equilibrium",
+        "desc.geometry.surface.FourierRZToroidalSurface",
+    ],
+)
+def _e_sub_theta_rp(params, transforms, profiles, data, **kwargs):
+    data["e_theta|r,p"] = data["e_theta"] - (data["e_phi"].T * data["phi_t"]).T
+    return data

desc/compute/_geometry.py

@@ -12,7 +12,7 @@
 from desc.backend import jnp
 
 from .data_index import register_compute_fun
-from .utils import cross, dot, line_integrals, surface_integrals
+from .utils import cross, dot, line_integrals, safenorm, surface_integrals
 
 
 @register_compute_fun(
@@ -176,27 +176,38 @@ def _A_of_z(params, transforms, profiles, data, **kwargs):
     label="A(\\zeta)",
     units="m^{2}",
     units_long="square meters",
-    description="Cross-sectional area as function of zeta",
+    description="Enclosed cross-sectional (constant phi surface) area, "
+    "as function of zeta",
     dim=1,
     params=[],
     transforms={"grid": []},
     profiles=[],
     coordinates="z",
-    data=["Z", "n_rho", "g_tt"],
+    data=["Z", "n_rho", "e_theta|r,p"],
     parameterization=["desc.geometry.surface.FourierRZToroidalSurface"],
+    # FIXME: Add source grid requirement once omega is nonzero.
 )
 def _A_of_z_FourierRZToroidalSurface(params, transforms, profiles, data, **kwargs):
-    # divergence theorem: integral(dA div [0, 0, Z]) = integral(ds n dot [0, 0, Z])
-    # but we need only the part of n in the R,Z plane
+    # Denote any vector v = [vᴿ, v^ϕ, vᶻ] with a tuple of its contravariant components.
+    # We use a 2D divergence theorem over constant ϕ toroidal surface (i.e. R, Z plane).
+    # Denote the divergence operator div by [∂_R, ∂_ϕ, ∂_Z] ⊗ [1, 0, 1].
+    # ∫ dS div (v ⊗ [1, 0, 1]) = ∫ dℓ n dot (v ⊗ [1, 0, 1])
+    # where n is the unit normal such that n dot e_θ|ρ,ϕ = 0 and n dot e_ϕ|R,Z = 0,
+    # where the labels following | denote those coordinates are fixed.
+    # Now choose v = [0, 0, Z] and n in the direction (e_θ|ρ,ζ × e_ζ|ρ,θ) ⊗ [1, 0, 1].
     n = data["n_rho"]
     n = n.at[:, 1].set(0)
     n = n / jnp.linalg.norm(n, axis=-1)[:, None]
     data["A(z)"] = jnp.abs(
         line_integrals(
             transforms["grid"],
-            data["Z"] * n[:, 2] * jnp.sqrt(data["g_tt"]),
+            data["Z"] * n[:, 2] * safenorm(data["e_theta|r,p"], axis=-1),
+            # FIXME: Works currently for omega = zero, but for nonzero omega
+            #  we need to integrate over theta at constant zeta.
+            #  Should be simple once we have coordinate mapping and source grid
+            #  logic from GitHub pull request #1024.
             line_label="theta",
-            fix_surface=("rho", 1.0),
+            fix_surface=("rho", jnp.max(transforms["grid"].nodes[:, 0])),
             expand_out=True,
         )
     )

tests/test_data_index.py

@@ -8,6 +8,7 @@
 import desc.compute
 from desc.compute import data_index
 from desc.compute.data_index import _class_inheritance
+from desc.utils import errorif
 
 
 class TestDataIndex:
@@ -39,6 +40,11 @@ def get_parameterization(fun, default="desc.equilibrium.equilibrium.Equilibrium"
         matches.discard("")
         return matches if matches else {default}
 
+    @staticmethod
+    def _is_function(func):
+        # JITed functions are not functions according to inspect.
+        return inspect.isfunction(func) or callable(func)
+
     @pytest.mark.unit
     def test_data_index_deps(self):
         """Ensure developers do not add extra (or forget needed) dependencies.
@@ -74,7 +80,7 @@ def test_data_index_deps(self):
         pattern_params = re.compile(r"params\[(.*?)]")
         for module_name, module in inspect.getmembers(desc.compute, inspect.ismodule):
             if module_name[0] == "_":
-                for _, fun in inspect.getmembers(module, inspect.isfunction):
+                for _, fun in inspect.getmembers(module, self._is_function):
                     # quantities that this function computes
                     names = self.get_matches(fun, pattern_names)
                     # dependencies queried in source code of this function
@@ -97,7 +103,6 @@ def test_data_index_deps(self):
 
         for p in data_index:
             for name, val in data_index[p].items():
-                print(name)
                 err_msg = f"Parameterization: {p}. Name: {name}."
                 deps = val["dependencies"]
                 data = set(deps["data"])
@@ -111,6 +116,12 @@ def test_data_index_deps(self):
                 assert len(profiles) == len(deps["profiles"]), err_msg
                 assert len(params) == len(deps["params"]), err_msg
                 # assert correct dependencies are queried
+                errorif(
+                    name not in queried_deps[p],
+                    AssertionError,
+                    "Did you reuse the function name (i.e. def_...) for"
+                    f" '{name}' for some other quantity?",
+                )
                 assert queried_deps[p][name]["data"] == data | axis_limit_data, err_msg
                 assert queried_deps[p][name]["profiles"] == profiles, err_msg
                 assert queried_deps[p][name]["params"] == params, err_msg