Merge remote-tracking branch 'origin' into dp/poincare

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",

desc/compute/_equil.py

@@ -10,7 +10,7 @@
 """
 
 from interpax import interp1d
-from scipy.constants import mu_0
+from scipy.constants import elementary_charge, mu_0
 
 from desc.backend import jnp
 
@@ -843,3 +843,33 @@ def _P_ISS04(params, transforms, profiles, data, **kwargs):
         )
     ) ** (1 / 0.39)
     return data
+
+
+@register_compute_fun(
+    name="P_fusion",
+    label="P_{fusion}",
+    units="W",
+    units_long="Watts",
+    description="Fusion power",
+    dim=0,
+    params=[],
+    transforms={"grid": []},
+    profiles=[],
+    coordinates="",
+    data=["ni", "<sigma*nu>", "sqrt(g)"],
+    resolution_requirement="rtz",
+    fuel="str: Fusion fuel, assuming a 50/50 mix. One of {'DT'}. Default is 'DT'.",
+)
+def _P_fusion(params, transforms, profiles, data, **kwargs):
+    energies = {"DT": 3.52e6 + 14.06e6}  # eV
+    fuel = kwargs.get("fuel", "DT")
+    energy = energies.get(fuel)
+
+    reaction_rate = jnp.sum(
+        data["ni"] ** 2
+        * data["<sigma*nu>"]
+        * data["sqrt(g)"]
+        * transforms["grid"].weights
+    )  # reactions/s
+    data["P_fusion"] = reaction_rate * energy * elementary_charge  # J/s
+    return data