Fix Redl consistency with spline profiles (#1651)

- Adds error for incorrect grid with Redl, allow Redl to use
SplineProfile,
- Update Redl tutorial to include a spline profile optimization

Resolves #1650

---------

Co-authored-by: Yigit Gunsur Elmacioglu <[email protected]>

Author: dpanici

CHANGELOG.md

@@ -2,9 +2,13 @@ Changelog
 =========
 
 New Features
+
 - Adds a new utility function ``desc.compat.contract_equilibrium`` which takes in an ``Equilibrium`` object and an argument ``inner_rho``, and returns a new ``Equilibrium`` with original ``Equilibrium``'s ``inner_rho`` flux surface as its boundary.
 Optionally can also contract the profiles of the original ``Equilibrium`` so that the new ``Equilibrium``'s profiles match the original's in real space.
 - Adds second-order NAE constraints, accessible by passing ``order=2`` to ``desc.objectives.get_NAE_constraints``.
+- Adds error for incorrect grids in ``desc.objectives.BootstrapRedlConsistency`` and when computing ``current Redl`` and ``<J*B> Redl`` compute quantities
+- Allows Redl compute quantities to use SplineProfile
+- Updates Redl bootstrap current consistency tutorial to include a ``SplineProfile`` optimization
 - Adds automatically generated header file showing date the input file was created with `desc.vmec.VMECIO.write_vmec_input`
 
 

desc/compute/_bootstrap.py

@@ -14,6 +14,7 @@
 
 from ..backend import fori_loop, jnp
 from ..integrals.surface_integral import surface_averages_map
+from ..profiles import PowerSeriesProfile
 from .data_index import register_compute_fun
 
 
@@ -421,20 +422,21 @@ def _current_Redl(params, transforms, profiles, data, **kwargs):
         * transforms["grid"].compress(data["<J*B> Redl"])
         / transforms["grid"].compress(data["<|B|^2>"])
     )
-    degree = kwargs.get(
-        "degree",
-        min(
-            (
-                profiles["current"].basis.L
-                if profiles["current"] is not None
-                else transforms["grid"].num_rho - 1
+    if isinstance(profiles["current"], PowerSeriesProfile):
+        degree = kwargs.get(
+            "degree",
+            min(
+                profiles["current"].basis.L,
+                transforms["grid"].num_rho - 1,
             ),
-            transforms["grid"].num_rho - 1,
-        ),
-    )
+        )
+    else:
+        degree = kwargs.get("degree", transforms["grid"].num_rho - 1)
+
     XX = jnp.vander(rho, degree + 1)[:, :-1]  # remove constant term
     c_l_r = jnp.pad(jnp.linalg.lstsq(XX, current_r)[0], (0, 1))  # manual polyfit
     c_l = jnp.polyint(c_l_r)
     current = jnp.polyval(c_l, rho)
+
     data["current Redl"] = transforms["grid"].expand(current)
     return data

desc/equilibrium/equilibrium.py

@@ -858,6 +858,46 @@ def compute(  # noqa: C901
             msg="must pass in a Grid object for argument grid!"
             f" instead got type {type(grid)}",
         )
+        # a check for Redl to prevent computing on-axis or
+        # at a rho=1.0 point where profiles vanish
+        if (
+            any(["Redl" in name for name in names])
+            and self.electron_density is not None
+        ):
+            warnif(
+                grid.axis.size,
+                UserWarning,
+                "Redl formula is undefined at rho=0, "
+                "but grid has grid points at rho=0, note that on-axis"
+                "current will be NaN.",
+            )
+            rho = grid.compress(grid.nodes[:, 0], "rho")
+
+            # check if profiles may go to zero
+            # if they are exactly zero this would cause NaNs since the profiles
+            # vanish.
+            warnif(
+                np.any(np.isclose(self.electron_density(rho), 0.0, atol=1e-8)),
+                UserWarning,
+                "Redl formula is undefined where kinetic profiles vanish, "
+                "but given electron density vanishes at at least one provided"
+                "rho grid point.",
+            )
+            warnif(
+                np.any(np.isclose(self.electron_temperature(rho), 0.0, atol=1e-8)),
+                UserWarning,
+                "Redl formula is undefined where kinetic profiles vanish, "
+                "but given electron temperature vanishes at at least one provided"
+                "rho grid point.",
+            )
+            warnif(
+                np.any(np.isclose(self.ion_temperature(rho), 0.0, atol=1e-8)),
+                UserWarning,
+                "Redl formula is undefined where kinetic profiles vanish, "
+                "but given ion temperature vanishes at at least one provided"
+                "rho grid point.",
+            )
+
         if grid.coordinates != "rtz":
             inbasis = {
                 "r": "rho",

desc/objectives/_bootstrap.py

@@ -32,13 +32,20 @@ class BootstrapRedlConsistency(_Objective):
     eq : Equilibrium
         Equilibrium that will be optimized to satisfy the Objective.
     grid : Grid, optional
-        Collocation grid containing the nodes to evaluate at.
-        Defaults to ``LinearGrid(eq.L_grid, eq.M_grid, eq.N_grid)``
+        Collocation grid containing the nodes to evaluate at. Requires poloidal and
+        toroidal resolution as the objective must compute flux surface averages.
+        Defaults to
+        ``grid = LinearGrid(M=eq.M_grid,N=eq.N_grid,NFP=eq.NFP,``
+        ``sym=eq.sym,rho=np.linspace(1 / eq.L, 1, eq.L) - 1 / (2 * eq.L),)``
     helicity : tuple, optional
         Type of quasi-symmetry (M, N). Default = quasi-axisymmetry (1, 0).
         First entry must be M=1. Second entry is the toroidal mode number N,
-        used for evaluating the Redl bootstrap current formula. Set to 0 for axisymmetry
-        or quasi-axisymmetry; set to +/-NFP for quasi-helical symmetry.
+        used for evaluating the Redl bootstrap current formula. Set to 0 for
+        axisymmetry or quasi-axisymmetry; set to +/-NFP for quasi-helical symmetry.
+    degree : int, optional
+        The `degree` kwarg to pass to the `<J*B>_Redl` compute call, which
+        controls the degree of polynomial fit to the Redl current derivative
+        before it is integrated.
 
     """
 
@@ -62,6 +69,7 @@ def __init__(
         deriv_mode="auto",
         grid=None,
         helicity=(1, 0),
+        degree=None,
         name="Bootstrap current self-consistency (Redl)",
         jac_chunk_size=None,
     ):
@@ -71,6 +79,7 @@ def __init__(
         assert helicity[0] == 1, "Redl bootstrap current model assumes helicity[0] == 1"
         self._grid = grid
         self.helicity = helicity
+        self._degree = degree
         super().__init__(
             things=eq,
             target=target,
@@ -125,6 +134,12 @@ def build(self, use_jit=True, verbose=1):
             ValueError,
             "Helicity toroidal mode number should be 0 (QA) or +/- NFP (QH)",
         )
+        rho = grid.nodes[grid.unique_rho_idx, 0]
+        errorif(
+            grid.axis.size,
+            ValueError,
+            "Redl formula is undefined at rho=0, but grid has grid points at rho=0",
+        )
 
         self._dim_f = grid.num_rho
         self._data_keys = ["<J*B>", "<J*B> Redl"]
@@ -146,10 +161,34 @@ def build(self, use_jit=True, verbose=1):
             "Bootstrap current calculation requires an ion temperature profile.",
         )
 
+        rho = grid.compress(grid.nodes[:, 0], "rho")
+
+        # check if profiles may go to zero
+        # if they are exactly zero this would cause NaNs since the profiles
+        # vanish.
+        errorif(
+            np.any(np.isclose(eq.electron_density(rho), 0.0, atol=1e-8)),
+            ValueError,
+            "Redl formula is undefined where kinetic profiles vanish, "
+            "but given electron density vanishes at at least one provided"
+            "rho grid point.",
+        )
+        errorif(
+            np.any(np.isclose(eq.electron_temperature(rho), 0.0, atol=1e-8)),
+            ValueError,
+            "Redl formula is undefined where kinetic profiles vanish, "
+            "but given electron temperature vanishes at at least one provided"
+            "rho grid point.",
+        )
+        errorif(
+            np.any(np.isclose(eq.ion_temperature(rho), 0.0, atol=1e-8)),
+            ValueError,
+            "Redl formula is undefined where kinetic profiles vanish, "
+            "but given ion temperature vanishes at at least one provided"
+            "rho grid point.",
+        )
         # Try to catch cases in which density or temperatures are specified in the
         # wrong units. Densities should be ~ 10^20, temperatures are ~ 10^3.
-        rho = eq.compute("rho", grid=grid)["rho"]
-
         warnif(
             jnp.any(eq.electron_density(rho) > 1e22),
             UserWarning,
@@ -166,7 +205,9 @@ def build(self, use_jit=True, verbose=1):
             UserWarning,
             "Ion temperature is surprisingly high. It should have units of eV",
         )
-        # Profiles may go to 0 at rho=1, so exclude the last few grid points from lower
+
+        # Profiles may go to 0 at rho=1 (and we've already checked if our
+        # grid has points there), so exclude the last few grid points from lower
         # bounds:
         rho = rho[rho < 0.85]
         warnif(
@@ -197,6 +238,7 @@ def build(self, use_jit=True, verbose=1):
             "transforms": transforms,
             "profiles": profiles,
             "helicity": self._helicity,
+            "degree": self._degree,
         }
 
         timer.stop("Precomputing transforms")
@@ -235,6 +277,7 @@ def compute(self, params, constants=None):
             transforms=constants["transforms"],
             profiles=constants["profiles"],
             helicity=constants["helicity"],
+            degree=constants["degree"],
         )
         return constants["transforms"]["grid"].compress(
             data["<J*B>"] - data["<J*B> Redl"]