Add compute implementation for surface geometry terms

desc/compute/_geometry.py

@@ -144,6 +144,65 @@ def _V_rrr_of_r(params, transforms, profiles, data, **kwargs):
     return data
 
 
+@register_compute_fun(
+    name="A(z)",
+    label="A(\\zeta)",
+    units="m^{2}",
+    units_long="square meters",
+    description="Cross-sectional area as function of zeta",
+    dim=1,
+    params=[],
+    transforms={"grid": []},
+    profiles=[],
+    coordinates="z",
+    data=["|e_rho x e_theta|"],
+    parameterization=[
+        "desc.equilibrium.equilibrium.Equilibrium",
+        "desc.geometry.surface.ZernikeRZToroidalSection",
+    ],
+)
+def _A_of_z(params, transforms, profiles, data, **kwargs):
+    data["A(z)"] = surface_integrals(
+        transforms["grid"],
+        jnp.abs(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="Cross-sectional area as function of zeta",
+    dim=1,
+    params=[],
+    transforms={"grid": []},
+    profiles=[],
+    coordinates="z",
+    data=["Z", "n_rho", "g_tt"],
+    parameterization=["desc.geometry.surface.FourierRZToroidalSurface"],
+)
+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
+    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"]),
+            line_label="theta",
+            fix_surface=("rho", 1.0),
+            expand_out=True,
+        )
+    )
+    return data
+
+
 @register_compute_fun(
     name="A",
     label="A",
@@ -155,20 +214,15 @@ def _V_rrr_of_r(params, transforms, profiles, data, **kwargs):
     transforms={"grid": []},
     profiles=[],
     coordinates="",
-    data=["|e_rho x e_theta|"],
+    data=["A(z)"],
     parameterization=[
         "desc.equilibrium.equilibrium.Equilibrium",
-        "desc.geometry.surface.ZernikeRZToroidalSection",
+        "desc.geometry.core.Surface",
     ],
 )
 def _A(params, transforms, profiles, data, **kwargs):
     data["A"] = jnp.mean(
-        surface_integrals(
-            transforms["grid"],
-            jnp.abs(data["|e_rho x e_theta|"]),
-            surface_label="zeta",
-            expand_out=False,
-        )
+        transforms["grid"].compress(data["A(z)"], surface_label="zeta")
     )
     return data
 
@@ -288,6 +342,10 @@ def _S_rr_of_r(params, transforms, profiles, data, **kwargs):
     profiles=[],
     coordinates="",
     data=["V", "A"],
+    parameterization=[
+        "desc.equilibrium.equilibrium.Equilibrium",
+        "desc.geometry.surface.FourierRZToroidalSurface",
+    ],
 )
 def _R0(params, transforms, profiles, data, **kwargs):
     data["R0"] = data["V"] / (2 * jnp.pi * data["A"])
@@ -306,6 +364,10 @@ def _R0(params, transforms, profiles, data, **kwargs):
     profiles=[],
     coordinates="",
     data=["A"],
+    parameterization=[
+        "desc.equilibrium.equilibrium.Equilibrium",
+        "desc.geometry.surface.FourierRZToroidalSurface",
+    ],
 )
 def _a(params, transforms, profiles, data, **kwargs):
     data["a"] = jnp.sqrt(data["A"] / jnp.pi)
@@ -324,44 +386,68 @@ def _a(params, transforms, profiles, data, **kwargs):
     profiles=[],
     coordinates="",
     data=["R0", "a"],
+    parameterization=[
+        "desc.equilibrium.equilibrium.Equilibrium",
+        "desc.geometry.surface.FourierRZToroidalSurface",
+    ],
 )
 def _R0_over_a(params, transforms, profiles, data, **kwargs):
     data["R0/a"] = data["R0"] / data["a"]
     return data
 
 
 @register_compute_fun(
-    name="a_major/a_minor",
-    label="a_{\\mathrm{major}} / a_{\\mathrm{minor}}",
-    units="~",
-    units_long="None",
-    description="Elongation at a toroidal cross-section",
+    name="perimeter(z)",
+    label="P(\\zeta)",
+    units="m",
+    units_long="meters",
+    description="Perimeter of cross section as function of zeta",
     dim=1,
     params=[],
     transforms={"grid": []},
     profiles=[],
     coordinates="z",
-    data=["rho", "sqrt(g)", "g_tt", "R"],
+    data=["rho", "g_tt"],
+    parameterization=[
+        "desc.equilibrium.equilibrium.Equilibrium",
+        "desc.geometry.core.Surface",
+    ],
 )
-def _a_major_over_a_minor(params, transforms, profiles, data, **kwargs):
+def _perimeter_of_z(params, transforms, profiles, data, **kwargs):
     max_rho = jnp.max(data["rho"])
-    P = (  # perimeter at rho=1
+    data["perimeter(z)"] = (  # perimeter at rho ~ 1
         line_integrals(
             transforms["grid"],
             jnp.sqrt(data["g_tt"]),
             line_label="theta",
             fix_surface=("rho", max_rho),
-            expand_out=False,
+            expand_out=True,
         )
-        / max_rho
-    )
-    # surface area
-    A = surface_integrals(
-        transforms["grid"],
-        jnp.abs(data["sqrt(g)"] / data["R"]),
-        surface_label="zeta",
-        expand_out=False,
+        / max_rho  # to account for quadrature grid not having nodes out to rho=1
     )
+    return data
+
+
+@register_compute_fun(
+    name="a_major/a_minor",
+    label="a_{\\mathrm{major}} / a_{\\mathrm{minor}}",
+    units="~",
+    units_long="None",
+    description="Elongation at a toroidal cross-section",
+    dim=1,
+    params=[],
+    transforms={"grid": []},
+    profiles=[],
+    coordinates="z",
+    data=["A(z)", "perimeter(z)"],
+    parameterization=[
+        "desc.equilibrium.equilibrium.Equilibrium",
+        "desc.geometry.core.Surface",
+    ],
+)
+def _a_major_over_a_minor(params, transforms, profiles, data, **kwargs):
+    A = transforms["grid"].compress(data["A(z)"], surface_label="zeta")
+    P = transforms["grid"].compress(data["perimeter(z)"], surface_label="zeta")
     # derived from Ramanujan approximation for the perimeter of an ellipse
     a = (  # semi-major radius
         jnp.sqrt(3)

desc/objectives/_geometry.py

@@ -16,15 +16,14 @@
 from .utils import softmin
 
 
-# TODO: add Surface parametrization to compute R0/a
-# so can use this objective with FourierRZToroidalSurface
 class AspectRatio(_Objective):
     """Aspect ratio = major radius / minor radius.
 
     Parameters
     ----------
-    eq : Equilibrium
-        Equilibrium that will be optimized to satisfy the Objective.
+    eq : Equilibrium or FourierRZToroidalSurface
+        Equilibrium or FourierRZToroidalSurface that
+        will be optimized to satisfy the Objective.
     target : {float, ndarray}, optional
         Target value(s) of the objective. Only used if bounds is None.
         Must be broadcastable to Objective.dim_f.
@@ -102,7 +101,23 @@
         """
         eq = self.things[0]
         if self._grid is None:
-            grid = QuadratureGrid(L=eq.L_grid, M=eq.M_grid, N=eq.N_grid, NFP=eq.NFP)
+            if hasattr(eq, "L_grid"):
+                grid = QuadratureGrid(
+                    L=eq.L_grid,
+                    M=eq.M_grid,
+                    N=eq.N_grid,
+                    NFP=eq.NFP,
+                )
+            else:
+                # if not an Equilibrium, is a Surface,
+                # has no radial resolution so just need
+                # the surface points
+                grid = LinearGrid(
+                    rho=1.0,
+                    M=eq.M * 2,
+                    N=eq.N * 2,
+                    NFP=eq.NFP,
+                )
         else:
             grid = self._grid
 
@@ -133,7 +148,8 @@
         Parameters
         ----------
         params : dict
-            Dictionary of equilibrium degrees of freedom, eg Equilibrium.params_dict
+            Dictionary of equilibrium or surface degrees of freedom, eg
+            Equilibrium.params_dict
         constants : dict
             Dictionary of constant data, eg transforms, profiles etc. Defaults to
             self.constants
@@ -147,7 +163,7 @@
         if constants is None:
             constants = self.constants
         data = compute_fun(
-            "desc.equilibrium.equilibrium.Equilibrium",
+            self.things[0],
             self._data_keys,
             params=params,
             transforms=constants["transforms"],
@@ -164,8 +180,9 @@
 
     Parameters
     ----------
-    eq : Equilibrium
-        Equilibrium that will be optimized to satisfy the Objective.
+    eq : Equilibrium or FourierRZToroidalSurface
+        Equilibrium or FourierRZToroidalSurface that
+        will be optimized to satisfy the Objective.
     target : {float, ndarray}, optional
         Target value(s) of the objective. Only used if bounds is None.
         Must be broadcastable to Objective.dim_f.
@@ -243,7 +260,23 @@
         """
         eq = self.things[0]
         if self._grid is None:
-            grid = QuadratureGrid(L=eq.L_grid, M=eq.M_grid, N=eq.N_grid, NFP=eq.NFP)
+            if hasattr(eq, "L_grid"):
+                grid = QuadratureGrid(
+                    L=eq.L_grid,
+                    M=eq.M_grid,
+                    N=eq.N_grid,
+                    NFP=eq.NFP,
+                )
+            else:
+                # if not an Equilibrium, is a Surface,
+                # has no radial resolution so just need
+                # the surface points
+                grid = LinearGrid(
+                    rho=1.0,
+                    M=eq.M * 2,
+                    N=eq.N * 2,
+                    NFP=eq.NFP,
+                )
         else:
             grid = self._grid
 
@@ -274,7 +307,8 @@
         Parameters
         ----------
         params : dict
-            Dictionary of equilibrium degrees of freedom, eg Equilibrium.params_dict
+            Dictionary of equilibrium or surface degrees of freedom,
+            eg Equilibrium.params_dict
         constants : dict
             Dictionary of constant data, eg transforms, profiles etc. Defaults to
             self.constants
@@ -288,7 +322,7 @@
         if constants is None:
             constants = self.constants
         data = compute_fun(
-            "desc.equilibrium.equilibrium.Equilibrium",
+            self.things[0],
             self._data_keys,
             params=params,
             transforms=constants["transforms"],

tests/test_compute_funs.py

@@ -1653,3 +1653,20 @@ def test_iota_components(HELIOTRON_vac):
     data = eq.compute(["iota", "iota current", "iota vacuum"], grid)
     np.testing.assert_allclose(data["iota"], data["iota vacuum"])
     np.testing.assert_allclose(data["iota current"], 0)
+
+
+@pytest.mark.unit
+def test_surface_equilibrium_geometry():
+    """Test that computing stuff from surface gives same result as equilibrium."""
+    names = ["DSHAPE", "HELIOTRON", "NCSX"]
+    data = ["A", "V", "a", "R0", "R0/a", "a_major/a_minor"]
+    for name in names:
+        eq = get(name)
+        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
+            else:
+                rtol, atol = 1e-8, 0
+            np.testing.assert_allclose(x, y, rtol=rtol, atol=atol, err_msg=name + key)

tests/test_objective_funs.py

@@ -152,6 +152,7 @@ def test(eq):
 
         test(Equilibrium(iota=PowerSeriesProfile(0)))
         test(Equilibrium(current=PowerSeriesProfile(0)))
+        test(Equilibrium(iota=PowerSeriesProfile(0)).surface)
 
     @pytest.mark.unit
     def test_elongation(self):
@@ -166,6 +167,7 @@ def test(eq):
             np.testing.assert_allclose(f_scaled, 2 * (1.3 / 0.7), rtol=5e-3)
 
         test(get("HELIOTRON"))
+        test(get("HELIOTRON").surface)
 
     @pytest.mark.unit
     def test_energy(self):