diff --git a/desc/compute/_basis_vectors.py b/desc/compute/_basis_vectors.py index d260ddee31..44433252cb 100644 --- a/desc/compute/_basis_vectors.py +++ b/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 diff --git a/desc/compute/_geometry.py b/desc/compute/_geometry.py index d9d736aa95..36bb1ac32a 100644 --- a/desc/compute/_geometry.py +++ b/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,13 +417,14 @@ 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", @@ -415,15 +432,21 @@ def _R0_over_a(params, transforms, profiles, data, **kwargs): ) 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 diff --git a/tests/inputs/master_compute_data.pkl b/tests/inputs/master_compute_data.pkl index 836a76b534..b356ddfceb 100644 Binary files a/tests/inputs/master_compute_data.pkl and b/tests/inputs/master_compute_data.pkl differ diff --git a/tests/test_data_index.py b/tests/test_data_index.py index e3fd65a495..1eb3ea6639 100644 --- a/tests/test_data_index.py +++ b/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