Integrate on boundary to compute length scale quantities

desc/backend.py

@@ -371,6 +371,14 @@ def tangent_solve(g, y):
         )
         return x, (jnp.linalg.norm(res), niter)
 
+    def trapezoid(y, x=None, dx=1.0, axis=-1):
+        """Integrate along the given axis using the composite trapezoidal rule."""
+        if hasattr(jnp, "trapezoid"):
+            # https://github.com/google/jax/issues/20410
+            return jnp.trapezoid(y, x, dx, axis)
+        else:
+            return jax.scipy.integrate.trapezoid(y, x, dx, axis)
+
 
 # we can't really test the numpy backend stuff in automated testing, so we ignore it
 # for coverage purposes
@@ -644,6 +652,14 @@ def repeat(a, repeats, axis=None, total_repeat_length=None):
             out = out[:total_repeat_length]
         return out
 
+    def trapezoid(y, x=None, dx=1.0, axis=-1):
+        """Integrate along the given axis using the composite trapezoidal rule."""
+        if hasattr(np, "trapezoid"):
+            # https://github.com/numpy/numpy/issues/25586
+            return np.trapezoid(y, x, dx, axis)
+        else:
+            return np.trapz(y, x, dx, axis)
+
     def custom_jvp(fun, *args, **kwargs):
         """Dummy function for custom_jvp without JAX."""
         fun.defjvp = lambda *args, **kwargs: None

desc/compute/_geometry.py

@@ -9,9 +9,11 @@
 expensive computations.
 """
 
-from desc.backend import jnp
+from desc.backend import jnp, trapezoid
 
+from ..grid import QuadratureGrid
 from .data_index import register_compute_fun
+from .geom_utils import warnif_sym
 from .utils import cross, dot, line_integrals, safenorm, surface_integrals
 
 
@@ -26,11 +28,16 @@
     transforms={"grid": []},
     profiles=[],
     coordinates="",
-    data=["sqrt(g)"],
+    data=["sqrt(g)", "V(r)", "rho"],
     resolution_requirement="rtz",
 )
 def _V(params, transforms, profiles, data, **kwargs):
-    data["V"] = jnp.sum(data["sqrt(g)"] * transforms["grid"].weights)
+    if isinstance(transforms["grid"], QuadratureGrid):
+        data["V"] = jnp.sum(data["sqrt(g)"] * transforms["grid"].weights)
+    else:
+        # To approximate volume at ρ ~ 1, we scale by ρ⁻², assuming the integrand
+        # varies little from ρ = max_rho to ρ = 1 and a roughly circular cross-section.
+        data["V"] = jnp.max(data["V(r)"]) / jnp.max(data["rho"]) ** 2
     return data
 
 
@@ -39,27 +46,19 @@ def _V(params, transforms, profiles, data, **kwargs):
     label="V",
     units="m^{3}",
     units_long="cubic meters",
-    description="Volume",
+    description="Volume enclosed by surface, scaled by max(ρ)⁻²",
     dim=1,
     params=[],
     transforms={"grid": []},
     profiles=[],
     coordinates="",
-    data=["e_theta", "e_zeta", "x"],
+    data=["V(r)", "rho"],
     parameterization="desc.geometry.surface.FourierRZToroidalSurface",
-    resolution_requirement="tz",
 )
 def _V_FourierRZToroidalSurface(params, transforms, profiles, data, **kwargs):
-    # divergence theorem: integral(dV div [0, 0, Z]) = integral(dS dot [0, 0, Z])
-    data["V"] = jnp.max(  # take max in case there are multiple surfaces for some reason
-        jnp.abs(
-            surface_integrals(
-                transforms["grid"],
-                cross(data["e_theta"], data["e_zeta"])[:, 2] * data["x"][:, 2],
-                expand_out=False,
-            )
-        )
-    )
+    # To approximate volume at ρ ~ 1, we scale by ρ⁻², assuming the integrand
+    # varies little from ρ = max_rho to ρ = 1 and a roughly circular cross-section.
+    data["V"] = jnp.max(data["V(r)"]) / jnp.max(data["rho"]) ** 2
     return data
 
 
@@ -75,6 +74,10 @@ def _V_FourierRZToroidalSurface(params, transforms, profiles, data, **kwargs):
     profiles=[],
     coordinates="r",
     data=["e_theta", "e_zeta", "Z"],
+    parameterization=[
+        "desc.equilibrium.equilibrium.Equilibrium",
+        "desc.geometry.surface.FourierRZToroidalSurface",
+    ],
     resolution_requirement="tz",
 )
 def _V_of_r(params, transforms, profiles, data, **kwargs):
@@ -155,45 +158,20 @@ def _V_rrr_of_r(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=["|e_rho x e_theta|"],
+    data=["Z", "n_rho", "e_theta|r,p", "rho"],
     parameterization=[
         "desc.equilibrium.equilibrium.Equilibrium",
         "desc.geometry.surface.ZernikeRZToroidalSection",
+        "desc.geometry.surface.FourierRZToroidalSurface",
     ],
-    resolution_requirement="rt",
-)
-def _A_of_z(params, transforms, profiles, data, **kwargs):
-    data["A(z)"] = surface_integrals(
-        transforms["grid"],
-        data["|e_rho x e_theta|"],
-        surface_label="zeta",
-        expand_out=True,
-    )
-    return data
-
-
-@register_compute_fun(
-    name="A(z)",
-    label="A(\\zeta)",
-    units="m^{2}",
-    units_long="square meters",
-    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", "e_theta|r,p", "rho"],
-    parameterization=["desc.geometry.surface.FourierRZToroidalSurface"],
-    resolution_requirement="rt",  # just need max(rho) near 1
-    # FIXME: Add source grid requirement once omega is nonzero.
+    resolution_requirement="t",
 )
 def _A_of_z_FourierRZToroidalSurface(params, transforms, profiles, data, **kwargs):
     # Denote any vector v = [vᴿ, v^ϕ, vᶻ] with a tuple of its contravariant components.
@@ -204,6 +182,7 @@ def _A_of_z_FourierRZToroidalSurface(params, transforms, profiles, data, **kwarg
     # 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].
+    warnif_sym(transforms["grid"], "A(z)")
     n = data["n_rho"]
     n = n.at[:, 1].set(0)
     n = n / jnp.linalg.norm(n, axis=-1)[:, jnp.newaxis]
@@ -232,23 +211,46 @@ def _A_of_z_FourierRZToroidalSurface(params, transforms, profiles, data, **kwarg
     label="A",
     units="m^{2}",
     units_long="square meters",
-    description="Average cross-sectional area",
+    description="Average enclosed cross-sectional area, scaled by max(ρ)⁻²",
     dim=0,
     params=[],
     transforms={"grid": []},
     profiles=[],
     coordinates="",
-    data=["A(z)"],
+    data=["Z", "n_rho", "e_theta|r,p", "rho", "phi"],
     parameterization=[
-        "desc.equilibrium.equilibrium.Equilibrium",
         "desc.geometry.core.Surface",
+        "desc.equilibrium.equilibrium.Equilibrium",
     ],
-    resolution_requirement="z",
+    resolution_requirement="tz",
 )
 def _A(params, transforms, profiles, data, **kwargs):
-    data["A"] = jnp.mean(
-        transforms["grid"].compress(data["A(z)"], surface_label="zeta")
+    # 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)[:, jnp.newaxis]
+    max_rho = jnp.max(data["rho"])
+    A = jnp.abs(
+        line_integrals(
+            transforms["grid"],
+            data["Z"] * n[:, 2] * safenorm(data["e_theta|r,p"], axis=-1),
+            line_label="theta",
+            fix_surface=("rho", max_rho),
+            expand_out=False,
+        )
+        # 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
     )
+    phi = transforms["grid"].compress(data["phi"], "zeta")
+    data["A"] = jnp.squeeze(trapezoid(A, phi) / jnp.ptp(phi) if A.size > 1 else A)
     return data
 
 
@@ -275,25 +277,22 @@ def _A_of_r(params, transforms, profiles, data, **kwargs):
     label="S",
     units="m^{2}",
     units_long="square meters",
-    description="Surface area of outermost flux surface",
+    description="Surface area of outermost flux surface, scaled by max(ρ)⁻¹",
     dim=0,
     params=[],
     transforms={"grid": []},
     profiles=[],
     coordinates="",
-    data=["|e_theta x e_zeta|"],
+    data=["S(r)", "rho"],
     parameterization=[
         "desc.equilibrium.equilibrium.Equilibrium",
         "desc.geometry.surface.FourierRZToroidalSurface",
     ],
-    resolution_requirement="tz",
 )
 def _S(params, transforms, profiles, data, **kwargs):
-    data["S"] = jnp.max(
-        surface_integrals(
-            transforms["grid"], data["|e_theta x e_zeta|"], expand_out=False
-        )
-    )
+    # To approximate surface are at ρ ~ 1, we scale by ρ⁻¹, assuming the integrand
+    # varies little from ρ = max_rho to ρ = 1.
+    data["S"] = jnp.max(data["S(r)"]) / jnp.max(data["rho"])
     return data
 
 
@@ -309,6 +308,10 @@ def _S(params, transforms, profiles, data, **kwargs):
     profiles=[],
     coordinates="r",
     data=["|e_theta x e_zeta|"],
+    parameterization=[
+        "desc.equilibrium.equilibrium.Equilibrium",
+        "desc.geometry.surface.FourierRZToroidalSurface",
+    ],
     resolution_requirement="tz",
 )
 def _S_of_r(params, transforms, profiles, data, **kwargs):
@@ -440,9 +443,10 @@ def _R0_over_a(params, transforms, profiles, data, **kwargs):
         "desc.equilibrium.equilibrium.Equilibrium",
         "desc.geometry.core.Surface",
     ],
-    resolution_requirement="rt",  # just need max(rho) near 1
+    resolution_requirement="t",
 )
 def _perimeter_of_z(params, transforms, profiles, data, **kwargs):
+    warnif_sym(transforms["grid"], "perimeter(z)")
     max_rho = jnp.max(data["rho"])
     data["perimeter(z)"] = (
         line_integrals(

desc/compute/geom_utils.py

@@ -2,11 +2,40 @@
 
 import functools
 
+from termcolor import colored
+
 from desc.backend import jnp
 
+from ..utils import ResolutionWarning, errorif, warnif
 from .utils import safenorm, safenormalize
 
 
+def warnif_sym(grid, name):
+    """Warn if grid has truncated poloidal domain to [0, π] ⊂ [0, 2π)."""
+    warnif(
+        grid.sym,
+        ResolutionWarning,
+        msg=colored(
+            "This grid only samples the poloidal domain θ ∈ [0, π], "
+            f"but, in general, the full domain [0, 2π) is required to compute {name}.",
+            "yellow",
+        ),
+    )
+
+
+def errorif_sym(grid, name):
+    """Warn if grid has truncated poloidal domain to [0, π] ⊂ [0, 2π)."""
+    errorif(
+        grid.sym,
+        ResolutionWarning,
+        msg=colored(
+            "This grid only samples the poloidal domain θ ∈ [0, π], "
+            f"but, in general, the full domain [0, 2π) is required to compute {name}.",
+            "yellow",
+        ),
+    )
+
+
 def reflection_matrix(normal):
     """Matrix to reflect points across plane through origin with specified normal.
 

desc/compute/utils.py

@@ -903,7 +903,7 @@ def line_integrals(
         The coordinate curve to compute the integration over.
         To clarify, a theta (poloidal) curve is the intersection of a
         rho surface (flux surface) and zeta (toroidal) surface.
-    fix_surface : str, float
+    fix_surface : (str, float)
         A tuple of the form: label, value.
         ``fix_surface`` label should differ from ``line_label``.
         By default, ``fix_surface`` is chosen to be the flux surface at rho=1.

desc/grid.py

@@ -66,19 +66,19 @@ def _set_up(self):
             del self._unique_theta_idx
 
     def _enforce_symmetry(self):
-        """Enforce stellarator symmetry.
+        """Remove unnecessary nodes assuming poloidal symmetry.
 
-        1. Remove nodes with theta > pi.
-        2. Rescale theta spacing to preserve dtheta weight.
-            Need to rescale on each theta coordinate curve by a different factor.
-            dtheta should = 2π / number of nodes remaining on that theta curve.
-            Nodes on the symmetry line should not be rescaled.
+        1. Remove nodes with θ > π.
+        2. Rescale θ spacing to preserve dθ weight.
+           Need to rescale on each θ coordinate curve by a different factor.
+           dθ = 2π / number of nodes remaining on that θ curve.
+           Nodes on the symmetry line should not be rescaled.
 
         """
         if not self.sym:
             return
         # indices where poloidal coordinate is off the symmetry line of
-        # poloidal coord=0 or pi
+        # poloidal coord=0 or π
         off_sym_line_idx = self.nodes[:, 1] % np.pi != 0
         __, inverse, off_sym_line_per_rho_surf_count = np.unique(
             self.nodes[off_sym_line_idx, 0], return_inverse=True, return_counts=True
@@ -109,7 +109,7 @@ def _enforce_symmetry(self):
         # The first two assumptions let _per_poloidal_curve = _per_rho_surf.
         # The third assumption lets the scale factor be constant over a
         # particular theta curve, so that each node in the open interval
-        # (0, pi) has its spacing scaled up by the same factor.
+        # (0, π) has its spacing scaled up by the same factor.
         # Nodes at endpoints 0, π should not be scaled.
         scale = off_sym_line_per_rho_surf_count / (
             off_sym_line_per_rho_surf_count - to_delete_per_rho_surf_count
@@ -206,7 +206,13 @@ def NFP(self):
 
     @property
     def sym(self):
-        """bool: True for stellarator symmetry, False otherwise."""
+        """bool: True for poloidal symmetry, False otherwise.
+
+        If true, the poloidal domain of this grid is [0, π] ⊂ [0, 2π).
+        Note that this is distinct from stellarator symmetry.
+        Still, when stellarator symmetry exists, flux surface integrals and
+        volume integrals are invariant to this truncation.
+        """
         return self.__dict__.setdefault("_sym", False)
 
     @property
@@ -865,7 +871,11 @@ class LinearGrid(_Grid):
     NFP : int
         Number of field periods (Default = 1).
     sym : bool
-        True for stellarator symmetry, False otherwise (Default = False).
+        Whether to truncate the poloidal domain to [0, π] ⊂ [0, 2π).
+        Note that this is distinct from stellarator symmetry.
+        Still, when stellarator symmetry exists, flux surface integrals and
+        volume integrals are invariant to this truncation, so setting this flag
+        to true will reduce memory consumption. (Default = False).
     axis : bool
         True to include a point at rho=0 (default), False for rho[0] = rho[1]/2.
     endpoint : bool
@@ -1318,7 +1328,11 @@ class ConcentricGrid(_Grid):
     NFP : int
         number of field periods (Default = 1)
     sym : bool
-        True for stellarator symmetry, False otherwise (Default = False)
+        Whether to truncate the poloidal domain to [0, π] ⊂ [0, 2π).
+        Note that this is distinct from stellarator symmetry.
+        Still, when stellarator symmetry exists, flux surface integrals and
+        volume integrals are invariant to this truncation, so setting this flag
+        to true will reduce memory consumption. (Default = False).
     axis : bool
         True to include the magnetic axis, False otherwise (Default = False)
     node_pattern : {``'cheb1'``, ``'cheb2'``, ``'jacobi'``, ``linear``}