Skip to content

Commit

Permalink
Use max rho of surface to compute A(z) (#1086)
Browse files Browse the repository at this point in the history
  • Loading branch information
unalmis authored Jul 3, 2024
2 parents 4745460 + 5c3d5f0 commit fd2645a
Show file tree
Hide file tree
Showing 4 changed files with 74 additions and 15 deletions.
28 changes: 28 additions & 0 deletions desc/compute/_basis_vectors.py
Original file line number Diff line number Diff line change
Expand Up @@ -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",
)
Expand Down Expand Up @@ -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
49 changes: 36 additions & 13 deletions desc/compute/_geometry.py
Original file line number Diff line number Diff line change
Expand Up @@ -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(
Expand Down Expand Up @@ -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

Expand Down Expand Up @@ -401,29 +417,36 @@ 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",
],
)
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

Expand Down
Binary file modified tests/inputs/master_compute_data.pkl
Binary file not shown.
12 changes: 10 additions & 2 deletions tests/test_data_index.py
Original file line number Diff line number Diff line change
Expand Up @@ -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:
Expand Down Expand Up @@ -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
Expand All @@ -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"])
Expand All @@ -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

0 comments on commit fd2645a

Please sign in to comment.