Merge branch 'master' into dp/remove-tor-flux-warn

.pre-commit-config.yaml

@@ -2,7 +2,7 @@ repos:
   - repo: https://github.com/psf/black
     rev: 24.3.0
     hooks:
-      - id: black
+      - id: black-jupyter
   - repo: https://github.com/pre-commit/pre-commit-hooks
     rev: v4.3.0
     hooks:

desc/compute/__init__.py

@@ -26,8 +26,6 @@
 # just need to import all the submodules here to register everything in the
 # data_index
 
-from desc.utils import flatten_list
-
 from . import (
     _basis_vectors,
     _bootstrap,

desc/compute/_bootstrap.py

@@ -13,8 +13,8 @@
 from scipy.special import roots_legendre
 
 from ..backend import fori_loop, jnp
+from ..integrals import surface_averages_map
 from .data_index import register_compute_fun
-from .utils import surface_averages_map
 
 
 @register_compute_fun(

desc/compute/_core.py

@@ -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",
@@ -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",
@@ -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",
@@ -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=[],
@@ -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",
@@ -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",
@@ -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",
@@ -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",
@@ -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",
@@ -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",