Skip to content

Commit

Permalink
Add new surface quantities needed after merging #583...
Browse files Browse the repository at this point in the history 
- Switch test marked regression to unit for same reason as
  git commit 44f2877
- Remove unneded test for limits with fix iota
  • Loading branch information
@unalmis
unalmis committed Aug 9, 2023
1 parent b16400f commit 4605ad1
Show file tree
Hide file tree
Showing 7 changed files with 255 additions and 52 deletions.
4 changes: 4 additions & 0 deletions desc/compute/_basis_vectors.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -196,6 +196,10 @@ def _e_sup_theta(params, transforms, profiles, data, **kwargs):
profiles=[],
coordinates="rtz",
data=["e_rho", "e_zeta"],
parameterization=[
"desc.equilibrium.equilibrium.Equilibrium",
"desc.geometry.core.Surface",
],
)
def _e_sup_theta_times_sqrt_g(params, transforms, profiles, data, **kwargs):
data["e^theta*sqrt(g)"] = cross(data["e_zeta"], data["e_rho"])
Expand Down
16 changes: 16 additions & 0 deletions desc/compute/_metric.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -90,6 +90,10 @@ def _e_theta_x_e_zeta(params, transforms, profiles, data, **kwargs):
profiles=[],
coordinates="rtz",
data=["e_theta", "e_zeta", "e_theta_r", "e_zeta_r"],
parameterization=[
"desc.equilibrium.equilibrium.Equilibrium",
"desc.geometry.core.Surface",
],
)
def _e_theta_x_e_zeta_r(params, transforms, profiles, data, **kwargs):
a = cross(data["e_theta"], data["e_zeta"])
Expand Down Expand Up @@ -120,6 +124,10 @@ def _e_theta_x_e_zeta_r(params, transforms, profiles, data, **kwargs):
profiles=[],
coordinates="rtz",
data=["e_theta", "e_zeta", "e_theta_r", "e_zeta_r", "e_theta_rr", "e_zeta_rr"],
parameterization=[
"desc.equilibrium.equilibrium.Equilibrium",
"desc.geometry.core.Surface",
],
)
def _e_theta_x_e_zeta_rr(params, transforms, profiles, data, **kwargs):
a = cross(data["e_theta"], data["e_zeta"])
Expand Down Expand Up @@ -201,6 +209,10 @@ def _e_rho_x_e_theta(params, transforms, profiles, data, **kwargs):
profiles=[],
coordinates="rtz",
data=["e_rho", "e_theta", "e_rho_r", "e_theta_r"],
parameterization=[
"desc.equilibrium.equilibrium.Equilibrium",
"desc.geometry.core.Surface",
],
)
def _e_rho_x_e_theta_r(params, transforms, profiles, data, **kwargs):
a = cross(data["e_rho"], data["e_theta"])
Expand Down Expand Up @@ -231,6 +243,10 @@ def _e_rho_x_e_theta_r(params, transforms, profiles, data, **kwargs):
profiles=[],
coordinates="rtz",
data=["e_rho", "e_theta", "e_rho_r", "e_theta_r", "e_rho_rr", "e_theta_rr"],
parameterization=[
"desc.equilibrium.equilibrium.Equilibrium",
"desc.geometry.core.Surface",
],
)
def _e_rho_x_e_theta_rr(params, transforms, profiles, data, **kwargs):
a = cross(data["e_rho"], data["e_theta"])
Expand Down
158 changes: 158 additions & 0 deletions desc/compute/_surface.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -141,6 +141,30 @@ def _e_rho_r_FourierRZToroidalSurface(params, transforms, profiles, data, **kwar
return data


@register_compute_fun(
name="e_rho_rr",
label="\\partial_{\\rho \\rho} \\mathbf{e}_{\\rho}",
units="m",
units_long="meters",
description="Covariant radial basis vector,"
" second derivative wrt radial coordinate",
dim=3,
params=[],
transforms={
"grid": [],
},
profiles=[],
coordinates="tz",
data=[],
parameterization="desc.geometry.surface.FourierRZToroidalSurface",
basis="basis",
)
def _e_rho_rr_FourierRZToroidalSurface(params, transforms, profiles, data, **kwargs):
coords = jnp.zeros((transforms["grid"].num_nodes, 3))
data["e_rho_rr"] = coords
return data


@register_compute_fun(
name="e_rho_t",
label="\\partial_{\\theta} \\mathbf{e}_{\\rho}",
Expand Down Expand Up @@ -210,6 +234,30 @@ def _e_theta_r_FourierRZToroidalSurface(params, transforms, profiles, data, **kw
return data


@register_compute_fun(
name="e_theta_rr",
label="\\partial_{\\rho \\rho} \\mathbf{e}_{\\theta}",
units="m",
units_long="meters",
description="Covariant poloidal basis vector,"
" second derivative wrt radial coordinate",
dim=3,
params=[],
transforms={
"grid": [],
},
profiles=[],
coordinates="tz",
data=[],
parameterization="desc.geometry.surface.FourierRZToroidalSurface",
basis="basis",
)
def _e_theta_rr_FourierRZToroidalSurface(params, transforms, profiles, data, **kwargs):
coords = jnp.zeros((transforms["grid"].num_nodes, 3))
data["e_theta_rr"] = coords
return data


@register_compute_fun(
name="e_theta_t",
label="\\partial_{\\theta} \\mathbf{e}_{\\theta}",
Expand Down Expand Up @@ -293,6 +341,30 @@ def _e_zeta_r_FourierRZToroidalSurface(params, transforms, profiles, data, **kwa
return data


@register_compute_fun(
name="e_zeta_rr",
label="\\partial_{\\rho \\rho} \\mathbf{e}_{\\zeta}",
units="m",
units_long="meters",
description="Covariant toroidal basis vector,"
" second derivative wrt radial coordinate",
dim=3,
params=[],
transforms={
"grid": [],
},
profiles=[],
coordinates="tz",
data=[],
parameterization="desc.geometry.surface.FourierRZToroidalSurface",
basis="basis",
)
def _e_zeta_rr_FourierRZToroidalSurface(params, transforms, profiles, data, **kwargs):
coords = jnp.zeros((transforms["grid"].num_nodes, 3))
data["e_zeta_rr"] = coords
return data


@register_compute_fun(
name="e_zeta_t",
label="\\partial_{\\theta} \\mathbf{e}_{\\zeta}",
Expand Down Expand Up @@ -498,6 +570,37 @@ def _e_rho_r_ZernikeRZToroidalSection(params, transforms, profiles, data, **kwar
return data


@register_compute_fun(
name="e_rho_rr",
label="\\partial_{\\rho \\rho} \\mathbf{e}_{\\rho}",
units="m",
units_long="meters",
description="Covariant radial basis vector,"
" second derivative wrt radial coordinate",
dim=3,
params=["R_lmn", "Z_lmn"],
transforms={
"R": [[3, 0, 0]],
"Z": [[3, 0, 0]],
"grid": [],
},
profiles=[],
coordinates="rt",
data=[],
parameterization="desc.geometry.surface.ZernikeRZToroidalSection",
basis="basis",
)
def _e_rho_rr_ZernikeRZToroidalSection(params, transforms, profiles, data, **kwargs):
R = transforms["R"].transform(params["R_lmn"], dr=3)
Z = transforms["Z"].transform(params["Z_lmn"], dr=3)
phi = jnp.zeros(transforms["grid"].num_nodes)
coords = jnp.stack([R, phi, Z], axis=1)
if kwargs.get("basis", "rpz").lower() == "xyz":
coords = rpz2xyz(coords)
data["e_rho_rr"] = coords
return data


@register_compute_fun(
name="e_rho_t",
label="\\partial_{\\theta} \\mathbf{e}_{\\rho}",
Expand Down Expand Up @@ -581,6 +684,37 @@ def _e_theta_r_ZernikeRZToroidalSection(params, transforms, profiles, data, **kw
return data


@register_compute_fun(
name="e_theta_rr",
label="\\partial_{\\rho \\rho} \\mathbf{e}_{\\theta}",
units="m",
units_long="meters",
description="Covariant poloidal basis vector,"
" second derivative wrt radial coordinate",
dim=3,
params=["R_lmn", "Z_lmn"],
transforms={
"R": [[2, 1, 0]],
"Z": [[2, 1, 0]],
"grid": [],
},
profiles=[],
coordinates="rt",
data=[],
parameterization="desc.geometry.surface.ZernikeRZToroidalSection",
basis="basis",
)
def _e_theta_rr_ZernikeRZToroidalSection(params, transforms, profiles, data, **kwargs):
R = transforms["R"].transform(params["R_lmn"], dr=2, dt=1)
Z = transforms["Z"].transform(params["Z_lmn"], dr=2, dt=1)
phi = jnp.zeros(transforms["grid"].num_nodes)
coords = jnp.stack([R, phi, Z], axis=1)
if kwargs.get("basis", "rpz").lower() == "xyz":
coords = rpz2xyz(coords)
data["e_theta_rr"] = coords
return data


@register_compute_fun(
name="e_theta_t",
label="\\partial_{\\theta} \\mathbf{e}_{\\theta}",
Expand Down Expand Up @@ -657,6 +791,30 @@ def _e_zeta_r_ZernikeRZToroidalSection(params, transforms, profiles, data, **kwa
return data


@register_compute_fun(
name="e_zeta_rr",
label="\\partial_{\\rho \\rho} \\mathbf{e}_{\\zeta}",
units="m",
units_long="meters",
description="Covariant toroidal basis vector,"
" second derivative wrt radial coordinate",
dim=3,
params=[],
transforms={
"grid": [],
},
profiles=[],
coordinates="rt",
data=[],
parameterization="desc.geometry.surface.ZernikeRZToroidalSection",
basis="basis",
)
def _e_zeta_rr_ZernikeRZToroidalSection(params, transforms, profiles, data, **kwargs):
coords = jnp.zeros((transforms["grid"].num_nodes, 3))
data["e_zeta_rr"] = coords
return data


@register_compute_fun(
name="e_zeta_t",
label="\\partial_{\\theta} \\mathbf{e}_{\\zeta}",
Expand Down
2 changes: 1 addition & 1 deletion desc/compute/data_index.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -46,7 +46,7 @@ def register_compute_fun(
a flux function, etc.
data : list of str
Names of other items in the data index needed to compute qty.
parameterization: str
parameterization: str or list of str
Name of desc types the method is valid for. eg 'desc.geometry.FourierXYZCurve'
or `desc.equilibrium.Equilibrium`.
axis_limit_data : list of str
Expand Down
Loading

0 comments on commit 4605ad1

Please sign in to comment.