diff --git a/desc/backend.py b/desc/backend.py index 6e123f49e9..4fda28545f 100644 --- a/desc/backend.py +++ b/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 diff --git a/desc/compute/_basis_vectors.py b/desc/compute/_basis_vectors.py index d8f45a1280..44433252cb 100644 --- a/desc/compute/_basis_vectors.py +++ b/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): diff --git a/desc/compute/_geometry.py b/desc/compute/_geometry.py index 3fd612629a..1b7de579c0 100644 --- a/desc/compute/_geometry.py +++ b/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,6 +182,7 @@ 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] @@ -189,8 +190,6 @@ def _A_of_z(params, transforms, profiles, data, **kwargs): 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( diff --git a/desc/compute/geom_utils.py b/desc/compute/geom_utils.py index fc5e1dab83..8618a74410 100644 --- a/desc/compute/geom_utils.py +++ b/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. diff --git a/desc/compute/utils.py b/desc/compute/utils.py index d5176bc40d..c0d5cecba0 100644 --- a/desc/compute/utils.py +++ b/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. diff --git a/desc/grid.py b/desc/grid.py index aad2c86f33..225cb15c7c 100644 --- a/desc/grid.py +++ b/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``} diff --git a/desc/objectives/_geometry.py b/desc/objectives/_geometry.py index 664f80ff27..9b2fc1a3d7 100644 --- a/desc/objectives/_geometry.py +++ b/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 diff --git a/docs/adding_compute_funs.rst b/docs/adding_compute_funs.rst index 66e3c2c047..efe9f708d6 100644 --- a/docs/adding_compute_funs.rst +++ b/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. diff --git a/tests/inputs/master_compute_data.pkl b/tests/inputs/master_compute_data.pkl index b356ddfceb..3452586309 100644 Binary files a/tests/inputs/master_compute_data.pkl and b/tests/inputs/master_compute_data.pkl differ diff --git a/tests/test_compute_funs.py b/tests/test_compute_funs.py index 5f8d8e7cc7..27b757eb6a 100644 --- a/tests/test_compute_funs.py +++ b/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)