Skip to content

Commit

Permalink
Merge branch 'master' into rg/parallel-testing
Browse files Browse the repository at this point in the history
  • Loading branch information
YigitElma authored Aug 28, 2024
2 parents e7e040b + 9f5cb55 commit 9a7a222
Show file tree
Hide file tree
Showing 7 changed files with 454 additions and 170 deletions.
250 changes: 249 additions & 1 deletion desc/compute/_core.py
Original file line number Diff line number Diff line change
Expand Up @@ -2764,6 +2764,25 @@ def _phi_rr(params, transforms, profiles, data, **kwargs):
return data


@register_compute_fun(
name="phi_rrz",
label="\\partial_{\\rho \\rho \\zeta} \\phi",
units="rad",
units_long="radians",
description="Toroidal angle in lab frame, second derivative wrt radial coordinate "
"and first wrt DESC toroidal coordinate",
dim=1,
params=[],
transforms={},
profiles=[],
coordinates="rtz",
data=["omega_rrz"],
)
def _phi_rrz(params, transforms, profiles, data, **kwargs):
data["phi_rrz"] = data["omega_rrz"]
return data


@register_compute_fun(
name="phi_rt",
label="\\partial_{\\rho \\theta} \\phi",
Expand All @@ -2783,6 +2802,25 @@ def _phi_rt(params, transforms, profiles, data, **kwargs):
return data


@register_compute_fun(
name="phi_rtz",
label="\\partial_{\\rho \\theta \\zeta} \\phi",
units="rad",
units_long="radians",
description="Toroidal angle in lab frame, third derivative wrt radial, "
"poloidal, and toroidal coordinates",
dim=1,
params=[],
transforms={},
profiles=[],
coordinates="rtz",
data=["omega_rtz"],
)
def _phi_rtz(params, transforms, profiles, data, **kwargs):
data["phi_rtz"] = data["omega_rtz"]
return data


@register_compute_fun(
name="phi_rz",
label="\\partial_{\\rho \\zeta} \\phi",
Expand All @@ -2802,6 +2840,25 @@ def _phi_rz(params, transforms, profiles, data, **kwargs):
return data


@register_compute_fun(
name="phi_rzz",
label="\\partial_{\\rho \\zeta \\zeta} \\phi",
units="rad",
units_long="radians",
description="Toroidal angle in lab frame, first derivative wrt radial and "
"second derivative wrt DESC toroidal coordinate",
dim=1,
params=[],
transforms={},
profiles=[],
coordinates="rtz",
data=["omega_rzz"],
)
def _phi_rzz(params, transforms, profiles, data, **kwargs):
data["phi_rzz"] = data["omega_rzz"]
return data


@register_compute_fun(
name="phi_t",
label="\\partial_{\\theta} \\phi",
Expand Down Expand Up @@ -2843,12 +2900,31 @@ def _phi_tt(params, transforms, profiles, data, **kwargs):
return data


@register_compute_fun(
name="phi_ttz",
label="\\partial_{\\theta \\theta \\zeta} \\phi",
units="rad",
units_long="radians",
description="Toroidal angle in lab frame, second derivative wrt poloidal "
"coordinate and first derivative wrt toroidal coordinate",
dim=1,
params=[],
transforms={},
profiles=[],
coordinates="rtz",
data=["omega_ttz"],
)
def _phi_ttz(params, transforms, profiles, data, **kwargs):
data["phi_ttz"] = data["omega_ttz"]
return data


@register_compute_fun(
name="phi_tz",
label="\\partial_{\\theta \\zeta} \\phi",
units="rad",
units_long="radians",
description="Toroidal angle in lab frame, second derivative wrt poloidal and "
description="Toroidal angle in lab frame, derivative wrt poloidal and "
"toroidal coordinate",
dim=1,
params=[],
Expand All @@ -2862,6 +2938,25 @@ def _phi_tz(params, transforms, profiles, data, **kwargs):
return data


@register_compute_fun(
name="phi_tzz",
label="\\partial_{\\theta \\zeta \\zeta} \\phi",
units="rad",
units_long="radians",
description="Toroidal angle in lab frame, derivative wrt poloidal coordinate and "
"second derivative wrt toroidal coordinate",
dim=1,
params=[],
transforms={},
profiles=[],
coordinates="rtz",
data=["omega_tzz"],
)
def _phi_tzz(params, transforms, profiles, data, **kwargs):
data["phi_tzz"] = data["omega_tzz"]
return data


@register_compute_fun(
name="phi_z",
label="\\partial_{\\zeta} \\phi",
Expand Down Expand Up @@ -2903,6 +2998,25 @@ def _phi_zz(params, transforms, profiles, data, **kwargs):
return data


@register_compute_fun(
name="phi_zzz",
label="\\partial_{\\zeta \\zeta \\zeta} \\phi",
units="rad",
units_long="radians",
description="Toroidal angle in lab frame, third derivative wrt toroidal "
"coordinate",
dim=1,
params=[],
transforms={},
profiles=[],
coordinates="rtz",
data=["omega_zzz"],
)
def _phi_zzz(params, transforms, profiles, data, **kwargs):
data["phi_zzz"] = data["omega_zzz"]
return data


@register_compute_fun(
name="rho",
label="\\rho",
Expand Down Expand Up @@ -2986,6 +3100,83 @@ def _theta_PEST_r(params, transforms, profiles, data, **kwargs):
return data


@register_compute_fun(
name="theta_PEST_rt",
label="\\partial_{\\rho \\theta} \\vartheta",
units="rad",
units_long="radians",
description="PEST straight field line poloidal angular coordinate, derivative wrt "
"radial and DESC poloidal coordinate",
dim=1,
params=[],
transforms={},
profiles=[],
coordinates="rtz",
data=["lambda_rt"],
)
def _theta_PEST_rt(params, transforms, profiles, data, **kwargs):
data["theta_PEST_rt"] = data["lambda_rt"]
return data


@register_compute_fun(
name="theta_PEST_rrt",
label="\\partial_{\\rho \\rho \\theta} \\vartheta",
units="rad",
units_long="radians",
description="PEST straight field line poloidal angular coordinate, second "
"derivative wrt radial coordinate and first derivative wrt DESC poloidal "
"coordinate",
dim=1,
params=[],
transforms={},
profiles=[],
coordinates="rtz",
data=["lambda_rrt"],
)
def _theta_PEST_rrt(params, transforms, profiles, data, **kwargs):
data["theta_PEST_rrt"] = data["lambda_rrt"]
return data


@register_compute_fun(
name="theta_PEST_rtz",
label="\\partial_{\\rho \\theta \\zeta} \\vartheta",
units="rad",
units_long="radians",
description="PEST straight field line poloidal angular coordinate, derivative wrt "
"radial and DESC poloidal and toroidal coordinates",
dim=1,
params=[],
transforms={},
profiles=[],
coordinates="rtz",
data=["lambda_rtz"],
)
def _theta_PEST_rtz(params, transforms, profiles, data, **kwargs):
data["theta_PEST_rtz"] = data["lambda_rtz"]
return data


@register_compute_fun(
name="theta_PEST_rtt",
label="\\partial_{\\rho \\theta \\theta} \\vartheta",
units="rad",
units_long="radians",
description="PEST straight field line poloidal angular coordinate, derivative wrt "
"radial coordinate once and DESC poloidal coordinate twice",
dim=1,
params=[],
transforms={},
profiles=[],
coordinates="rtz",
data=["lambda_rtt"],
)
def _theta_PEST_rtt(params, transforms, profiles, data, **kwargs):
data["theta_PEST_rtt"] = data["lambda_rtt"]
return data


@register_compute_fun(
name="theta_PEST_t",
label="\\partial_{\\theta} \\vartheta",
Expand Down Expand Up @@ -3024,6 +3215,25 @@ def _theta_PEST_tt(params, transforms, profiles, data, **kwargs):
return data


@register_compute_fun(
name="theta_PEST_ttt",
label="\\partial_{\\theta \\theta \\theta} \\vartheta",
units="rad",
units_long="radians",
description="PEST straight field line poloidal angular coordinate, third "
"derivative wrt poloidal coordinate",
dim=1,
params=[],
transforms={},
profiles=[],
coordinates="rtz",
data=["lambda_ttt"],
)
def _theta_PEST_ttt(params, transforms, profiles, data, **kwargs):
data["theta_PEST_ttt"] = data["lambda_ttt"]
return data


@register_compute_fun(
name="theta_PEST_tz",
label="\\partial_{\\theta \\zeta} \\vartheta",
Expand All @@ -3043,6 +3253,25 @@ def _theta_PEST_tz(params, transforms, profiles, data, **kwargs):
return data


@register_compute_fun(
name="theta_PEST_tzz",
label="\\partial_{\\theta \\zeta \\zeta} \\vartheta",
units="rad",
units_long="radians",
description="PEST straight field line poloidal angular coordinate, derivative wrt "
"poloidal coordinate once and toroidal coordinate twice",
dim=1,
params=[],
transforms={},
profiles=[],
coordinates="rtz",
data=["lambda_tzz"],
)
def _theta_PEST_tzz(params, transforms, profiles, data, **kwargs):
data["theta_PEST_tzz"] = data["lambda_tzz"]
return data


@register_compute_fun(
name="theta_PEST_z",
label="\\partial_{\\zeta} \\vartheta",
Expand Down Expand Up @@ -3081,6 +3310,25 @@ def _theta_PEST_zz(params, transforms, profiles, data, **kwargs):
return data


@register_compute_fun(
name="theta_PEST_ttz",
label="\\partial_{\\theta \\theta \\zeta} \\vartheta",
units="rad",
units_long="radians",
description="PEST straight field line poloidal angular coordinate, second "
"derivative wrt poloidal coordinate and derivative wrt toroidal coordinate",
dim=1,
params=[],
transforms={},
profiles=[],
coordinates="rtz",
data=["lambda_ttz"],
)
def _theta_PEST_ttz(params, transforms, profiles, data, **kwargs):
data["theta_PEST_ttz"] = data["lambda_ttz"]
return data


@register_compute_fun(
name="zeta",
label="\\zeta",
Expand Down
Loading

0 comments on commit 9a7a222

Please sign in to comment.