Merge branch 'master' into rc/toroidal_angle

.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/_equil.py

@@ -9,12 +9,14 @@
 expensive computations.
 """
 
+from interpax import interp1d
 from scipy.constants import mu_0
 
 from desc.backend import jnp
 
+from ..integrals import surface_averages
 from .data_index import register_compute_fun
-from .utils import cross, dot, safediv, safenorm, surface_averages
+from .utils import cross, dot, safediv, safenorm
 
 
 @register_compute_fun(
@@ -802,3 +804,42 @@ def _beta_volpol(params, transforms, profiles, data, **kwargs):
 def _beta_voltor(params, transforms, profiles, data, **kwargs):
     data["<beta_tor>_vol"] = jnp.abs(data["W_p"] / data["W_Btor"])
     return data
+
+
+@register_compute_fun(
+    name="P_ISS04",
+    label="P_{ISS04}",
+    units="W",
+    units_long="Watts",
+    description="Heating power required by the ISS04 energy confinement time scaling",
+    dim=0,
+    params=[],
+    transforms={"grid": []},
+    profiles=[],
+    coordinates="",
+    data=["a", "iota", "rho", "R0", "W_p", "<ne>_vol", "<|B|>_axis"],
+    method="str: Interpolation method. Default 'cubic'.",
+    H_ISS04="float: ISS04 confinement enhancement factor. Default 1.",
+)
+def _P_ISS04(params, transforms, profiles, data, **kwargs):
+    rho = transforms["grid"].compress(data["rho"], surface_label="rho")
+    iota = transforms["grid"].compress(data["iota"], surface_label="rho")
+    fx = {}
+    if "iota_r" in data:
+        fx["fx"] = transforms["grid"].compress(
+            data["iota_r"]
+        )  # noqa: unused dependency
+    iota_23 = interp1d(2 / 3, rho, iota, method=kwargs.get("method", "cubic"), **fx)
+    data["P_ISS04"] = 1e6 * (  # MW -> W
+        jnp.abs(data["W_p"] / 1e6)  # J -> MJ
+        / (
+            0.134
+            * data["a"] ** 2.28  # m
+            * data["R0"] ** 0.64  # m
+            * (data["<ne>_vol"] / 1e19) ** 0.54  # 1/m^3 -> 1e19/m^3
+            * data["<|B|>_axis"] ** 0.84  # T
+            * iota_23**0.41
+            * kwargs.get("H_ISS04", 1)
+        )
+    ) ** (1 / 0.39)
+    return data