Merge branch 'master' into better_warning

desc/compute/_basis_vectors.py

@@ -1514,6 +1514,7 @@ def _e_sup_zeta_zz(params, transforms, profiles, data, **kwargs):
     parameterization=[
         "desc.equilibrium.equilibrium.Equilibrium",
         "desc.geometry.surface.FourierRZToroidalSurface",
+        "desc.geometry.core.Surface",
     ],
     basis="basis",
 )
@@ -3523,3 +3524,30 @@ 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=(
+        "Covariant poloidal basis vector in (ρ,θ,ϕ) coordinates. "
+        "ϕ increases counterclockwise when viewed from above "
+        "(cylindrical R,ϕ plane with Z out of page)."
+    ),
+    dim=3,
+    params=[],
+    transforms={},
+    profiles=[],
+    coordinates="rtz",
+    data=["e_theta", "e_phi", "phi_t"],
+    parameterization=[
+        "desc.equilibrium.equilibrium.Equilibrium",
+        "desc.geometry.surface.FourierRZToroidalSurface",
+        "desc.geometry.core.Surface",
+    ],
+)
+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,29 +176,45 @@ 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="Area of enclosed cross-section (enclosed constant phi surface), "
+    "scaled by max(ρ)⁻², 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", "rho"],
     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).
+    # In this geometry, the divergence operator on a polar basis vector is
+    # div = ([∂_R, ∂_ϕ, ∂_Z] ⊗ [1, 0, 1]) dot .
+    # ∫ dA div v = ∫ dℓ n dot v
+    # where n is the unit normal such that n dot e_θ|ρ,ϕ = 0 and n dot e_ϕ|R,Z = 0,
+    # and 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]
+    n = n / jnp.linalg.norm(n, axis=-1)[:, jnp.newaxis]
+    max_rho = jnp.max(data["rho"])
     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 phi.
+            #  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", max_rho),
             expand_out=True,
         )
+        # To approximate area at ρ ~ 1, we scale by ρ⁻², assuming the integrand
+        # varies little from ρ = max_rho to ρ = 1 and a roughly circular cross-section.
+        / max_rho**2
     )
     return data
 
@@ -401,29 +417,36 @@ def _R0_over_a(params, transforms, profiles, data, **kwargs):
     label="P(\\zeta)",
     units="m",
     units_long="meters",
-    description="Perimeter of cross section as function of zeta",
+    description="Perimeter of enclosed cross-section (enclosed constant phi surface), "
+    "scaled by max(ρ)⁻¹, as function of zeta",
     dim=1,
     params=[],
     transforms={"grid": []},
     profiles=[],
     coordinates="z",
-    data=["rho", "g_tt"],
+    data=["rho", "e_theta|r,p"],
     parameterization=[
         "desc.equilibrium.equilibrium.Equilibrium",
         "desc.geometry.core.Surface",
     ],
 )
 def _perimeter_of_z(params, transforms, profiles, data, **kwargs):
     max_rho = jnp.max(data["rho"])
-    data["perimeter(z)"] = (  # perimeter at rho ~ 1
+    data["perimeter(z)"] = (
         line_integrals(
             transforms["grid"],
-            jnp.sqrt(data["g_tt"]),
+            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 phi.
+            #  Should be simple once we have coordinate mapping and source grid
+            #  logic from GitHub pull request #1024.
             line_label="theta",
             fix_surface=("rho", max_rho),
             expand_out=True,
         )
-        / max_rho  # to account for quadrature grid not having nodes out to rho=1
+        # To approximate perimeter at ρ ~ 1, we scale by ρ⁻¹, assuming the integrand
+        # varies little from ρ = max_rho to ρ = 1.
+        / max_rho
     )
     return data
 

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:
@@ -74,7 +75,9 @@ 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):
+                # JITed functions are not functions according to inspect,
+                # so just check if callable.
+                for _, fun in inspect.getmembers(module, callable):
                     # quantities that this function computes
                     names = self.get_matches(fun, pattern_names)
                     # dependencies queried in source code of this function
@@ -97,7 +100,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 +113,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