Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Fix math bug in surface computations #1175

Merged
merged 7 commits into from
Aug 14, 2024
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 0 additions & 1 deletion desc/compute/_basis_vectors.py
Original file line number Diff line number Diff line change
Expand Up @@ -1431,7 +1431,6 @@ def _e_sub_rho(params, transforms, profiles, data, **kwargs):
# At the magnetic axis, this function returns the multivalued map whose
# image is the set { 𝐞ᵨ | ρ=0 }.
data["e_rho"] = jnp.array([data["R_r"], data["R"] * data["omega_r"], data["Z_r"]]).T

return data


Expand Down
241 changes: 37 additions & 204 deletions desc/compute/_core.py
Original file line number Diff line number Diff line change
Expand Up @@ -22,12 +22,37 @@
"desc.geometry.core.Surface",
"desc.geometry.core.Curve",
],
aliases=["rho_t", "rho_z", "theta_r", "theta_z", "zeta_r", "zeta_t"],
)
def _0(params, transforms, profiles, data, **kwargs):
data["0"] = jnp.zeros(transforms["grid"].num_nodes)
return data


@register_compute_fun(
name="1",
label="1",
units="~",
units_long="None",
description="Ones",
dim=1,
params=[],
transforms={"grid": []},
profiles=[],
coordinates="rtz",
data=[],
parameterization=[
"desc.equilibrium.equilibrium.Equilibrium",
"desc.geometry.core.Surface",
"desc.geometry.core.Curve",
],
aliases=["rho_r", "theta_t", "zeta_z"],
)
def _1(params, transforms, profiles, data, **kwargs):
data["1"] = jnp.ones(transforms["grid"].num_nodes)
return data


@register_compute_fun(
name="R",
label="R",
Expand Down Expand Up @@ -2711,6 +2736,10 @@ def _phi(params, transforms, profiles, data, **kwargs):
profiles=[],
coordinates="rtz",
data=["omega_r"],
parameterization=[
"desc.equilibrium.equilibrium.Equilibrium",
"desc.geometry.core.Surface",
],
)
def _phi_r(params, transforms, profiles, data, **kwargs):
data["phi_r"] = data["omega_r"]
Expand Down Expand Up @@ -2785,6 +2814,10 @@ def _phi_rz(params, transforms, profiles, data, **kwargs):
profiles=[],
coordinates="rtz",
data=["omega_t"],
parameterization=[
"desc.equilibrium.equilibrium.Equilibrium",
"desc.geometry.core.Surface",
],
)
def _phi_t(params, transforms, profiles, data, **kwargs):
data["phi_t"] = data["omega_t"]
Expand Down Expand Up @@ -2841,6 +2874,10 @@ def _phi_tz(params, transforms, profiles, data, **kwargs):
profiles=[],
coordinates="rtz",
data=["omega_z"],
parameterization=[
"desc.equilibrium.equilibrium.Equilibrium",
"desc.geometry.core.Surface",
],
)
def _phi_z(params, transforms, profiles, data, **kwargs):
data["phi_z"] = 1 + data["omega_z"]
Expand Down Expand Up @@ -2890,75 +2927,6 @@ def _rho(params, transforms, profiles, data, **kwargs):
return data


@register_compute_fun(
name="rho_r",
label="\\partial_{\\rho} \\rho",
units="~",
units_long="None",
description="Radial coordinate, proportional to the square root "
+ "of the toroidal flux, derivative wrt radial coordinate",
dim=1,
params=[],
transforms={},
profiles=[],
coordinates="r",
data=["0"],
parameterization=[
"desc.equilibrium.equilibrium.Equilibrium",
"desc.geometry.core.Surface",
],
)
def _rho_r(params, transforms, profiles, data, **kwargs):
data["rho_r"] = jnp.ones_like(data["0"])
return data


@register_compute_fun(
name="rho_t",
label="\\partial_{\\theta} \\rho",
units="~",
units_long="None",
description="Radial coordinate, proportional to the square root "
"of the toroidal flux, derivative wrt poloidal coordinate",
dim=1,
params=[],
transforms={},
profiles=[],
coordinates="r",
data=["0"],
parameterization=[
"desc.equilibrium.equilibrium.Equilibrium",
"desc.geometry.core.Surface",
],
)
def _rho_t(params, transforms, profiles, data, **kwargs):
data["rho_t"] = data["0"]
return data


@register_compute_fun(
name="rho_z",
label="\\partial_{\\zeta} \\rho",
units="~",
units_long="None",
description="Radial coordinate, proportional to the square root "
"of the toroidal flux, derivative wrt toroidal coordinate",
dim=1,
params=[],
transforms={},
profiles=[],
coordinates="r",
data=["0"],
parameterization=[
"desc.equilibrium.equilibrium.Equilibrium",
"desc.geometry.core.Surface",
],
)
def _rho_z(params, transforms, profiles, data, **kwargs):
data["rho_z"] = data["0"]
return data


@register_compute_fun(
name="theta",
label="\\theta",
Expand Down Expand Up @@ -3113,75 +3081,6 @@ def _theta_PEST_zz(params, transforms, profiles, data, **kwargs):
return data


@register_compute_fun(
name="theta_r",
label="\\partial_{\\rho} \\theta",
units="rad",
units_long="radians",
description="Poloidal angular coordinate (geometric, not magnetic), "
"derivative wrt radial coordinate",
dim=1,
params=[],
transforms={},
profiles=[],
coordinates="t",
data=["0"],
parameterization=[
"desc.equilibrium.equilibrium.Equilibrium",
"desc.geometry.core.Surface",
],
)
def _theta_r(params, transforms, profiles, data, **kwargs):
data["theta_r"] = data["0"]
return data


@register_compute_fun(
name="theta_t",
label="\\partial_{\\theta} \\theta",
units="rad",
units_long="radians",
description="Poloidal angular coordinate (geometric, not magnetic), "
"derivative wrt poloidal coordinate",
dim=1,
params=[],
transforms={},
profiles=[],
coordinates="t",
data=["0"],
parameterization=[
"desc.equilibrium.equilibrium.Equilibrium",
"desc.geometry.core.Surface",
],
)
def _theta_t(params, transforms, profiles, data, **kwargs):
data["theta_t"] = jnp.ones_like(data["0"])
return data


@register_compute_fun(
name="theta_z",
label="\\partial_{\\zeta} \\theta",
units="rad",
units_long="radians",
description="Poloidal angular coordinate (geometric, not magnetic), "
"derivative wrt toroidal coordinate",
dim=1,
params=[],
transforms={},
profiles=[],
coordinates="t",
data=["0"],
parameterization=[
"desc.equilibrium.equilibrium.Equilibrium",
"desc.geometry.core.Surface",
],
)
def _theta_z(params, transforms, profiles, data, **kwargs):
data["theta_z"] = data["0"]
return data


@register_compute_fun(
name="zeta",
label="\\zeta",
Expand All @@ -3202,69 +3101,3 @@ def _theta_z(params, transforms, profiles, data, **kwargs):
def _zeta(params, transforms, profiles, data, **kwargs):
data["zeta"] = transforms["grid"].nodes[:, 2]
return data


@register_compute_fun(
name="zeta_r",
label="\\partial_{\\rho} \\zeta",
units="rad",
units_long="radians",
description="Toroidal angular coordinate derivative, wrt radial coordinate",
dim=1,
params=[],
transforms={},
profiles=[],
coordinates="z",
data=["0"],
parameterization=[
"desc.equilibrium.equilibrium.Equilibrium",
"desc.geometry.core.Surface",
],
)
def _zeta_r(params, transforms, profiles, data, **kwargs):
data["zeta_r"] = data["0"]
return data


@register_compute_fun(
name="zeta_t",
label="\\partial_{\\theta} \\zeta",
units="rad",
units_long="radians",
description="Toroidal angular coordinate, derivative wrt poloidal coordinate",
dim=1,
params=[],
transforms={},
profiles=[],
coordinates="z",
data=["0"],
parameterization=[
"desc.equilibrium.equilibrium.Equilibrium",
"desc.geometry.core.Surface",
],
)
def _zeta_t(params, transforms, profiles, data, **kwargs):
data["zeta_t"] = data["0"]
return data


@register_compute_fun(
name="zeta_z",
label="\\partial_{\\zeta} \\zeta",
units="rad",
units_long="radians",
description="Toroidal angular coordinate, derivative wrt toroidal coordinate",
dim=1,
params=[],
transforms={},
profiles=[],
coordinates="z",
data=["0"],
parameterization=[
"desc.equilibrium.equilibrium.Equilibrium",
"desc.geometry.core.Surface",
],
)
def _zeta_z(params, transforms, profiles, data, **kwargs):
data["zeta_z"] = jnp.ones_like(data["0"])
return data
2 changes: 1 addition & 1 deletion desc/compute/_field.py
Original file line number Diff line number Diff line change
Expand Up @@ -2935,7 +2935,7 @@ def _gradB2(params, transforms, profiles, data, **kwargs):

@register_compute_fun(
name="|grad(|B|^2)|/2mu0",
label="|\\nabla |B|^{2}/(2\\mu_0)|",
label="|\\nabla |B|^{2}|/(2\\mu_0)",
units="N \\cdot m^{-3}",
units_long="Newton / cubic meter",
description="Magnitude of magnetic pressure gradient",
Expand Down
4 changes: 2 additions & 2 deletions desc/compute/_profiles.py
Original file line number Diff line number Diff line change
Expand Up @@ -82,10 +82,10 @@ def _psi_r(params, transforms, profiles, data, **kwargs):
transforms={},
profiles=[],
coordinates="r",
data=["rho"],
data=["1"],
)
def _psi_rr(params, transforms, profiles, data, **kwargs):
data["psi_rr"] = params["Psi"] * jnp.ones_like(data["rho"]) / jnp.pi
data["psi_rr"] = data["1"] * params["Psi"] / jnp.pi
return data


Expand Down
60 changes: 0 additions & 60 deletions desc/compute/_surface.py
Original file line number Diff line number Diff line change
Expand Up @@ -118,66 +118,6 @@ def _phi_Surface(params, transforms, profiles, data, **kwargs):
return data


@register_compute_fun(
name="phi_r",
label="\\partial_{\\rho} \\phi",
units="rad",
units_long="radians",
description="Toroidal angle in lab frame, derivative wrt radial coordinate",
dim=1,
params=[],
transforms={},
profiles=[],
coordinates="rtz",
data=["e_rho"],
parameterization="desc.geometry.core.Surface",
)
def _phi_r_Surface(params, transforms, profiles, data, **kwargs):
coords_r = data["e_rho"]
data["phi_r"] = coords_r[:, 1]
return data


@register_compute_fun(
name="phi_t",
label="\\partial_{\\theta} \\phi",
units="rad",
units_long="radians",
description="Toroidal angle in lab frame, derivative wrt poloidal coordinate",
dim=1,
params=[],
transforms={},
profiles=[],
coordinates="rtz",
data=["e_theta"],
parameterization="desc.geometry.core.Surface",
)
def _phi_t_Surface(params, transforms, profiles, data, **kwargs):
coords_t = data["e_theta"]
data["phi_t"] = coords_t[:, 1]
return data


@register_compute_fun(
name="phi_z",
label="\\partial_{\\zeta} \\phi",
units="rad",
units_long="radians",
description="Toroidal angle in lab frame, derivative wrt toroidal coordinate",
dim=1,
params=[],
transforms={},
profiles=[],
coordinates="rtz",
data=["e_zeta"],
parameterization="desc.geometry.core.Surface",
)
def _phi_z_Surface(params, transforms, profiles, data, **kwargs):
coords_z = data["e_zeta"]
data["phi_z"] = coords_z[:, 1]
return data


@register_compute_fun(
name="Z",
label="Z",
Expand Down
2 changes: 2 additions & 0 deletions desc/compute/data_index.py
Original file line number Diff line number Diff line change
Expand Up @@ -132,6 +132,8 @@
source_grid_requirement = {}
if not isinstance(parameterization, (tuple, list)):
parameterization = [parameterization]
if not isinstance(aliases, (tuple, list)):
aliases = [aliases]

Check warning on line 136 in desc/compute/data_index.py

View check run for this annotation

Codecov / codecov/patch

desc/compute/data_index.py#L136

Added line #L136 was not covered by tests

deps = {
"params": params,
Expand Down
Binary file modified tests/inputs/master_compute_data_rpz.pkl
Binary file not shown.
Loading
Loading