Merge branch 'master' into dp/vmecio-asym

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

desc/compute/_profiles.py

@@ -1909,3 +1909,51 @@ def _shear(params, transforms, profiles, data, **kwargs):
         None,
     )
     return data
+
+
+@register_compute_fun(
+    name="<sigma*nu>",
+    label="\\langle\\sigma\\nu\\rangle",
+    units="m^3 \\cdot s^{-1}",
+    units_long="cubic meters / second",
+    description="Thermal reactivity from Bosch-Hale parameterization",
+    dim=1,
+    params=[],
+    transforms={"grid": []},
+    profiles=[],
+    coordinates="r",
+    data=["Ti"],
+    fuel="str: Fusion fuel, assuming a 50/50 mix. One of {'DT'}. Default is 'DT'.",
+)
+def _reactivity(params, transforms, profiles, data, **kwargs):
+    # Bosch and Hale. “Improved Formulas for Fusion Cross-Sections and Thermal
+    # Reactivities.” Nuclear Fusion 32 (April 1992): 611-631.
+    # https://doi.org/10.1088/0029-5515/32/4/I07.
+    coefficients = {
+        "DT": {
+            "B_G": 34.382,
+            "mc2": 1124656,
+            "C1": 1.17302e-9,
+            "C2": 1.51361e-2,
+            "C3": 7.51886e-2,
+            "C4": 4.60643e-3,
+            "C5": 1.35000e-2,
+            "C6": -1.06750e-4,
+            "C7": 1.36600e-5,
+        }
+    }
+    fuel = kwargs.get("fuel", "DT")
+    coeffs = coefficients.get(fuel)
+
+    T = data["Ti"] / 1e3  # keV
+    theta = T / (
+        1
+        - (T * (coeffs["C2"] + T * (coeffs["C4"] + T * coeffs["C6"])))
+        / (1 + T * (coeffs["C3"] + T * (coeffs["C5"] + T * coeffs["C7"])))
+    )
+    xi = (coeffs["B_G"] ** 2 / (4 * theta)) ** (1 / 3)
+    sigma_nu = (
+        coeffs["C1"] * theta * jnp.sqrt(xi / (coeffs["mc2"] * T**3)) * jnp.exp(-3 * xi)
+    )  # cm^3/s
+    data["<sigma*nu>"] = sigma_nu / 1e6  # m^3/s
+    return data

desc/objectives/__init__.py

@@ -11,7 +11,6 @@
     QuadraticFlux,
     ToroidalFlux,
 )
-from ._confinement import HeatingPowerISS04
 from ._equilibrium import (
     CurrentDensity,
     Energy,
@@ -39,6 +38,7 @@
     QuasisymmetryTripleProduct,
     QuasisymmetryTwoTerm,
 )
+from ._power_balance import FusionPower, HeatingPowerISS04
 from ._profiles import Pressure, RotationalTransform, Shear, ToroidalCurrent
 from ._stability import MagneticWell, MercierStability
 from .getters import (

desc/objectives/_confinement.py → desc/objectives/_power_balance.py

@@ -9,6 +9,188 @@
 from .objective_funs import _Objective
 
 
+class FusionPower(_Objective):
+    """Fusion power.
+
+    P = e E ∫ n^2 ⟨σν⟩ dV (W)
+
+    References
+    ----------
+    https://doi.org/10.1088/0029-5515/32/4/I07.
+    Improved Formulas for Fusion Cross-Sections and Thermal Reactivities.
+    H.-S. Bosch and G.M. Hale. Nucl. Fusion April 1992; 32 (4): 611-631.
+
+    Parameters
+    ----------
+    eq : Equilibrium
+        Equilibrium that will be optimized to satisfy the Objective.
+    target : {float, ndarray}, optional
+        Target value(s) of the objective. Only used if bounds is None.
+        Must be broadcastable to Objective.dim_f. Defaults to ``target=1e9``.
+    bounds : tuple of {float, ndarray}, optional
+        Lower and upper bounds on the objective. Overrides target.
+        Both bounds must be broadcastable to to Objective.dim_f.
+        Defaults to ``target=1e9``.
+    weight : {float, ndarray}, optional
+        Weighting to apply to the Objective, relative to other Objectives.
+        Must be broadcastable to to Objective.dim_f
+    normalize : bool, optional
+        Whether to compute the error in physical units or non-dimensionalize.
+    normalize_target : bool, optional
+        Whether target and bounds should be normalized before comparing to computed
+        values. If `normalize` is `True` and the target is in physical units,
+        this should also be set to True.
+    loss_function : {None, 'mean', 'min', 'max'}, optional
+        Loss function to apply to the objective values once computed. This loss function
+        is called on the raw compute value, before any shifting, scaling, or
+        normalization. Note: Has no effect for this objective.
+    deriv_mode : {"auto", "fwd", "rev"}
+        Specify how to compute jacobian matrix, either forward mode or reverse mode AD.
+        "auto" selects forward or reverse mode based on the size of the input and output
+        of the objective. Has no effect on self.grad or self.hess which always use
+        reverse mode and forward over reverse mode respectively.
+    fuel : str, optional
+        Fusion fuel, assuming a 50/50 mix. One of {'DT'}. Default = 'DT'.
+    grid : Grid, optional
+        Collocation grid used to compute the intermediate quantities.
+        Defaults to ``QuadratureGrid(eq.L_grid, eq.M_grid, eq.N_grid, eq.NFP)``.
+    name : str, optional
+        Name of the objective function.
+
+    """
+
+    _scalar = True
+    _units = "(W)"
+    _print_value_fmt = "Fusion power: {:10.3e} "
+
+    def __init__(
+        self,
+        eq,
+        target=None,
+        bounds=None,
+        weight=1,
+        normalize=True,
+        normalize_target=True,
+        loss_function=None,
+        deriv_mode="auto",
+        fuel="DT",
+        grid=None,
+        name="fusion power",
+    ):
+        errorif(
+            fuel not in ["DT"], ValueError, f"fuel must be one of ['DT'], got {fuel}."
+        )
+        if target is None and bounds is None:
+            target = 1e9
+        self._fuel = fuel
+        self._grid = grid
+        super().__init__(
+            things=eq,
+            target=target,
+            bounds=bounds,
+            weight=weight,
+            normalize=normalize,
+            normalize_target=normalize_target,
+            loss_function=loss_function,
+            deriv_mode=deriv_mode,
+            name=name,
+        )
+
+    def build(self, use_jit=True, verbose=1):
+        """Build constant arrays.
+
+        Parameters
+        ----------
+        use_jit : bool, optional
+            Whether to just-in-time compile the objective and derivatives.
+        verbose : int, optional
+            Level of output.
+
+        """
+        eq = self.things[0]
+        errorif(
+            eq.electron_density is None,
+            ValueError,
+            "Equilibrium must have an electron density profile.",
+        )
+        errorif(
+            eq.ion_temperature is None,
+            ValueError,
+            "Equilibrium must have an ion temperature profile.",
+        )
+        if self._grid is None:
+            self._grid = QuadratureGrid(
+                L=eq.L_grid,
+                M=eq.M_grid,
+                N=eq.N_grid,
+                NFP=eq.NFP,
+            )
+        self._dim_f = 1
+        self._data_keys = ["P_fusion"]
+
+        timer = Timer()
+        if verbose > 0:
+            print("Precomputing transforms")
+        timer.start("Precomputing transforms")
+
+        self._constants = {
+            "profiles": get_profiles(self._data_keys, obj=eq, grid=self._grid),
+            "transforms": get_transforms(self._data_keys, obj=eq, grid=self._grid),
+        }
+
+        timer.stop("Precomputing transforms")
+        if verbose > 1:
+            timer.disp("Precomputing transforms")
+
+        if self._normalize:
+            scales = compute_scaling_factors(eq)
+            self._normalization = scales["W_p"]
+
+        super().build(use_jit=use_jit, verbose=verbose)
+
+    def compute(self, params, constants=None):
+        """Compute fusion power.
+
+        Parameters
+        ----------
+        params : dict
+            Dictionary of equilibrium or surface degrees of freedom, eg
+            Equilibrium.params_dict
+        constants : dict
+            Dictionary of constant data, eg transforms, profiles etc. Defaults to
+            self.constants
+
+        Returns
+        -------
+        P : float
+            Fusion power (W).
+
+        """
+        if constants is None:
+            constants = self.constants
+        data = compute_fun(
+            "desc.equilibrium.equilibrium.Equilibrium",
+            self._data_keys,
+            params=params,
+            transforms=constants["transforms"],
+            profiles=constants["profiles"],
+            fuel=self.fuel,
+        )
+        return data["P_fusion"]
+
+    @property
+    def fuel(self):
+        """str: Fusion fuel, assuming a 50/50 mix. One of {'DT'}. Default = 'DT'."""
+        return self._fuel
+
+    @fuel.setter
+    def fuel(self, new):
+        errorif(
+            new not in ["DT"], ValueError, f"fuel must be one of ['DT'], got {new}."
+        )
+        self._fuel = new
+
+
 class HeatingPowerISS04(_Objective):
     """Heating power required by the ISS04 energy confinement time scaling.
 

tests/test_axis_limits.py

@@ -25,9 +25,9 @@
 # d²ψ/(dρ)² and 𝜕√𝑔/𝜕𝜌 are both finite nonzero at the magnetic axis.
 # Also, dⁿψ/(dρ)ⁿ for n > 3 is assumed zero everywhere.
 zero_limits = {"rho", "psi", "psi_r", "e_theta", "sqrt(g)", "B_t"}
-# "current Redl" and "P_ISS04" need special treatment because they are not defined for
-# all configurations (giving NaN values)
-not_continuous_limits = {"current Redl", "P_ISS04"}
+# These compute quantities require kinetic profiles, which are not defined for all
+# configurations (giving NaN values)
+not_continuous_limits = {"current Redl", "P_ISS04", "P_fusion", "<sigma*nu>"}
 not_finite_limits = {
     "D_Mercier",
     "D_geodesic",

tests/test_objective_funs.py

@@ -48,6 +48,7 @@
     Energy,
     ForceBalance,
     ForceBalanceAnisotropic,
+    FusionPower,
     GenericObjective,
     HeatingPowerISS04,
     Isodynamicity,
@@ -2118,6 +2119,7 @@ class TestComputeScalarResolution:
         CoilLength,
         CoilSetMinDistance,
         CoilTorsion,
+        FusionPower,
         GenericObjective,
         HeatingPowerISS04,
         Omnigenity,
@@ -2179,6 +2181,28 @@ def test_compute_scalar_resolution_bootstrap(self):
             f[i] = obj.compute_scalar(obj.x())
         np.testing.assert_allclose(f, f[-1], rtol=5e-2)
 
+    @pytest.mark.regression
+    def test_compute_scalar_resolution_fusion_power(self):
+        """FusionPower."""
+        eq = self.eq.copy()
+        eq.electron_density = PowerSeriesProfile([1e19, 0, -1e19])
+        eq.electron_temperature = PowerSeriesProfile([1e3, 0, -1e3])
+        eq.ion_temperature = PowerSeriesProfile([1e3, 0, -1e3])
+        eq.atomic_number = 1.0
+
+        f = np.zeros_like(self.res_array, dtype=float)
+        for i, res in enumerate(self.res_array):
+            grid = QuadratureGrid(
+                L=int(self.eq.L * res),
+                M=int(self.eq.M * res),
+                N=int(self.eq.N * res),
+                NFP=self.eq.NFP,
+            )
+            obj = ObjectiveFunction(FusionPower(eq=eq, grid=grid))
+            obj.build(verbose=0)
+            f[i] = obj.compute_scalar(obj.x())
+        np.testing.assert_allclose(f, f[-1], rtol=5e-2)
+
     @pytest.mark.regression
     def test_compute_scalar_resolution_heating_power(self):
         """HeatingPowerISS04."""
@@ -2470,6 +2494,7 @@ class TestObjectiveNaNGrad:
         CoilSetMinDistance,
         CoilTorsion,
         ForceBalanceAnisotropic,
+        FusionPower,
         HeatingPowerISS04,
         Omnigenity,
         PlasmaCoilSetMinDistance,
@@ -2519,6 +2544,22 @@ def test_objective_no_nangrad_bootstrap(self):
         g = obj.grad(obj.x(eq))
         assert not np.any(np.isnan(g)), "redl bootstrap"
 
+    @pytest.mark.unit
+    def test_objective_no_nangrad_fusion_power(self):
+        """FusionPower."""
+        eq = Equilibrium(
+            L=2,
+            M=2,
+            N=2,
+            electron_density=PowerSeriesProfile([1e19, 0, -1e19]),
+            electron_temperature=PowerSeriesProfile([1e3, 0, -1e3]),
+            current=PowerSeriesProfile([1, 0, -1]),
+        )
+        obj = ObjectiveFunction(FusionPower(eq))
+        obj.build()
+        g = obj.grad(obj.x(eq))
+        assert not np.any(np.isnan(g)), "fusion power"
+
     @pytest.mark.unit
     def test_objective_no_nangrad_heating_power(self):
         """HeatingPowerISS04."""