Make changes in light of #1101 (comment)...

Some of the master compute data changes with this commit.
The changes which are not due to floating point differences or simply grid
resolution differences (recall the grid used in test_compute_everything has
L = 9, M=N=5 which is less than required for convergence to true integral
quantities on that equilbrium) are given below:

Not equal to tolerance rtol=1e-10, atol=1e-10
Parameterization: desc.equilibrium.equilibrium.Equilibrium. Name: A(z).
Mismatched elements: 660 / 660 (100%)
Max absolute difference: 0.6494644
Max relative difference: 0.81816076
 x: array([0.792591, 0.792591, 0.792591, 0.792591, 0.792591, 0.792591,
       0.792591, 0.792591, 0.792591, 0.792591, 0.792591, 0.792591,
       0.792591, 0.792591, 0.792591, 0.792591, 0.792591, 0.792591,...
 y: array([0.762978, 0.762978, 0.762978, 0.762978, 0.762978, 0.762978,
       0.762978, 0.762978, 0.762978, 0.762978, 0.762978, 0.762978,
       0.762978, 0.762978, 0.762978, 0.762978, 0.762978, 0.762978,...

Not equal to tolerance rtol=1e-10, atol=1e-10
Parameterization: desc.equilibrium.equilibrium.Equilibrium. Name: a_major/a_minor.
Mismatched elements: 660 / 660 (100%)
Max absolute difference: 11.33931768
Max relative difference: 4.99527722
 x: array([ 3.864152,  3.864152,  3.864152,  3.864152,  3.864152,  3.864152,
        3.864152,  3.864152,  3.864152,  3.864152,  3.864152,  3.864152,
        3.864152,  3.864152,  3.864152,  3.864152,  3.864152,  3.864152,...
 y: array([4.057173, 4.057173, 4.057173, 4.057173, 4.057173, 4.057173,
       4.057173, 4.057173, 4.057173, 4.057173, 4.057173, 4.057173,
       4.057173, 4.057173, 4.057173, 4.057173, 4.057173, 4.057173,...

NOTE that both quanties are incorrect, because in general we cannot compute
these quantities accuratly on grids that do not sample the full poloidal
domain.

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/_basis_vectors.py

@@ -1514,7 +1514,7 @@ def _e_sup_zeta_zz(params, transforms, profiles, data, **kwargs):
     parameterization=[
         "desc.equilibrium.equilibrium.Equilibrium",
         "desc.geometry.surface.FourierRZToroidalSurface",
-        "desc.geometry.surface.ZernikeRZToroidalSection",
+        "desc.geometry.core.Surface",
     ],
     basis="basis",
 )
@@ -3545,7 +3545,7 @@ def _n_zeta(params, transforms, profiles, data, **kwargs):
     parameterization=[
         "desc.equilibrium.equilibrium.Equilibrium",
         "desc.geometry.surface.FourierRZToroidalSurface",
-        "desc.geometry.surface.ZernikeRZToroidalSection",
+        "desc.geometry.core.Surface",
     ],
 )
 def _e_sub_theta_rp(params, transforms, profiles, data, **kwargs):

desc/compute/_geometry.py

@@ -9,10 +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
 
 
@@ -171,9 +172,8 @@ def _V_rrr_of_r(params, transforms, profiles, data, **kwargs):
         "desc.geometry.surface.FourierRZToroidalSurface",
     ],
     resolution_requirement="t",
-    # FIXME: Add source grid requirement once omega is nonzero.
 )
-def _A_of_z(params, transforms, profiles, data, **kwargs):
+def _A_of_z_FourierRZToroidalSurface(params, transforms, profiles, data, **kwargs):
     # 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
@@ -182,15 +182,14 @@ def _A_of_z(params, transforms, profiles, data, **kwargs):
     # 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]
     max_rho = jnp.max(data["rho"])
     data["A(z)"] = jnp.abs(
         line_integrals(
             transforms["grid"],
-            # FIXME: Integrand is not symmetric function, so this will fail
-            #   on symmetric grids.
             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.
@@ -218,17 +217,40 @@ def _A_of_z(params, transforms, profiles, data, **kwargs):
     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
 
 
@@ -424,6 +446,7 @@ def _R0_over_a(params, transforms, profiles, data, **kwargs):
     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``}

desc/objectives/_geometry.py

@@ -9,6 +9,7 @@
 from desc.grid import LinearGrid, QuadratureGrid
 from desc.utils import Timer, warnif
 
+from ..compute.geom_utils import errorif_sym
 from .normalization import compute_scaling_factors
 from .objective_funs import _Objective
 from .utils import softmin
@@ -118,6 +119,7 @@ def build(self, use_jit=True, verbose=1):
                     M=eq.M * 2,
                     N=eq.N * 2,
                     NFP=eq.NFP,
+                    sym=False,
                 )
         else:
             grid = self._grid
@@ -238,6 +240,8 @@ def __init__(
     ):
         if target is None and bounds is None:
             target = 1
+        if grid is not None:
+            errorif_sym(grid, name)
         self._grid = grid
         super().__init__(
             things=eq,
@@ -280,6 +284,7 @@ def build(self, use_jit=True, verbose=1):
                     M=eq.M * 2,
                     N=eq.N * 2,
                     NFP=eq.NFP,
+                    sym=False,
                 )
         else:
             grid = self._grid

docs/adding_compute_funs.rst

@@ -124,6 +124,46 @@ dependencies specified by the decorator. The function itself should do any calcu
 needed using these dependencies and then add the output to the ``data`` dictionary and
 return it. The key in the ``data`` dictionary should match the ``name`` of the quantity.
 
+Here is another example.
+::
+
+    @register_compute_fun(
+        name="perimeter(z)",
+        label="P(\\zeta)",
+        units="m",
+        units_long="meters",
+        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", "e_theta|r,p"],
+        parameterization=[
+            "desc.equilibrium.equilibrium.Equilibrium",
+            "desc.geometry.core.Surface",
+        ],
+        # Document that the grid should have resolution in theta for an accurate
+        # computation, as required for integration over theta.
+        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(
+                transforms["grid"],
+                safenorm(data["e_theta|r,p"], axis=-1),
+                line_label="theta",
+                fix_surface=("rho", max_rho),
+            )
+            # To approximate perimeter at ρ ~ 1, we scale by ρ⁻¹, assuming the integrand
+            # varies little from ρ = max_rho to ρ = 1.
+            / max_rho
+        )
+        return data
+
 Once a new quantity is added to the ``desc.compute`` module, there are two final steps involving the testing suite which must be checked.
 The first step is implementing the correct axis limit, or marking it as not finite or not implemented.
 We can check whether the axis limit currently evaluates as finite by computing the quantity on a grid with nodes at the axis.

tests/test_compute_funs.py

@@ -194,14 +194,15 @@ def test_elongation():
         modes_R=[[0, 0], [1, 0], [0, 1]],
         modes_Z=[[-1, 0], [0, -1]],
     )
-    grid = LinearGrid(rho=1, M=2 * surf3.M, N=surf3.N, NFP=surf3.NFP, sym=surf3.sym)
-    data1 = surf1.compute(["a_major/a_minor", "A"], grid=grid)
-    data2 = surf2.compute(["a_major/a_minor", "A"], grid=grid)
-    data3 = surf3.compute(["a_major/a_minor", "A"], grid=grid)
+    assert surf3.sym
+    grid = LinearGrid(rho=1, M=3 * surf3.M, N=surf3.N, NFP=surf3.NFP, sym=False)
+    data1 = surf1.compute(["a_major/a_minor"], grid=grid)
+    data2 = surf2.compute(["a_major/a_minor"], grid=grid)
+    data3 = surf3.compute(["a_major/a_minor"], grid=grid)
     # elongation approximation is less accurate as elongation increases
     np.testing.assert_allclose(1.0, data1["a_major/a_minor"])
-    np.testing.assert_allclose(2.0, data2["a_major/a_minor"], rtol=1e-3)
-    np.testing.assert_allclose(3.0, data3["a_major/a_minor"], rtol=1e-2)
+    np.testing.assert_allclose(2.0, data2["a_major/a_minor"], rtol=1e-4)
+    np.testing.assert_allclose(3.0, data3["a_major/a_minor"], rtol=1e-3)
 
 
 @pytest.mark.slow
@@ -1657,11 +1658,13 @@ def test_surface_equilibrium_geometry():
         for key in data:
             x = eq.compute(key)[key].max()  # max needed for elongation broadcasting
             y = eq.surface.compute(key)[key].max()
-            if key == "a_major/a_minor":
-                rtol, atol = 1e-2, 0  # need looser tol here bc of different grids
+            if key in ("A", "a", "R0", "R0/a", "a_major/a_minor"):
+                rtol, atol = 1e-3, 0  # need looser tol here bc of different grids
             else:
                 rtol, atol = 1e-8, 0
-            np.testing.assert_allclose(x, y, rtol=rtol, atol=atol, err_msg=name + key)
+            np.testing.assert_allclose(
+                x, y, rtol=rtol, atol=atol, err_msg=name + " " + key
+            )
         # compare at rho=1, where we expect the eq.compute and the
         # surface.compute to agree for these surface basis vectors
         grid = LinearGrid(rho=np.array(1.0), M=10, N=10, NFP=eq.NFP)