fixing #1288, adding opt test for effective ripple

desc/objectives/_neoclassical.py

@@ -132,6 +132,11 @@ def __init__(
 
         rho, alpha = np.atleast_1d(rho, alpha)
         self._dim_f = rho.size
+
+        zeta = np.linspace(0, 2 * np.pi * num_transit, knots_per_transit * num_transit)
+
+        grid = LinearGrid(rho=rho, theta=alpha, zeta=zeta)
+
         self._keys_1dr = [
             "iota",
             "iota_r",
@@ -141,12 +146,8 @@ def __init__(
             "R0",  # TODO: GitHub PR #1094
         ]
         self._constants = {
+            "grid": grid,
             "quad_weights": 1,
-            "rho": rho,
-            "alpha": alpha,
-            "zeta": np.linspace(
-                0, 2 * np.pi * num_transit, knots_per_transit * num_transit
-            ),
             "quad": chebgauss2(num_quad),
         }
         self._hyperparameters = {
@@ -168,7 +169,7 @@ def __init__(
             jac_chunk_size=jac_chunk_size,
         )
 
-    def build(self, use_jit=True, verbose=1):
+    def build(self, use_jit=True, verbose=1, constants=None):
         """Build constant arrays.
 
         Parameters
@@ -179,12 +180,20 @@ def build(self, use_jit=True, verbose=1):
             Level of output.
 
         """
+        if constants is None:
+            constants = self.constants
+        rtzgrid = constants["grid"]
+
         eq = self.things[0]
         self._grid_1dr = LinearGrid(
-            rho=self._constants["rho"], M=eq.M_grid, N=eq.N_grid, NFP=eq.NFP, sym=eq.sym
+            rho=rtzgrid.nodes[rtzgrid.unique_rho_idx, 0],
+            M=eq.M_grid,
+            N=eq.N_grid,
+            NFP=eq.NFP,
+            sym=eq.sym,
         )
         self._target, self._bounds = _parse_callable_target_bounds(
-            self._target, self._bounds, self._constants["rho"]
+            self._target, self._bounds, rtzgrid.nodes[rtzgrid.unique_rho_idx, 0]
         )
 
         timer = Timer()
@@ -226,6 +235,9 @@ def compute(self, params, constants=None):
         if constants is None:
             constants = self.constants
         eq = self.things[0]
+
+        rtzgrid = constants["grid"]
+
         # TODO: compute all deps of effective ripple here
         data = compute_fun(
             eq,
@@ -235,14 +247,16 @@ def compute(self, params, constants=None):
             constants["profiles"],
         )
         # TODO: interpolate all deps to this grid with fft utilities from fourier bounce
+
         grid = eq._get_rtz_grid(
-            constants["rho"],
-            constants["alpha"],
-            constants["zeta"],
+            rtzgrid.nodes[rtzgrid.unique_rho_idx, 0],
+            rtzgrid.nodes[rtzgrid.unique_theta_idx, 1],
+            rtzgrid.nodes[rtzgrid.unique_zeta_idx, 2],
             coordinates="raz",
             iota=self._grid_1dr.compress(data["iota"]),
             params=params,
         )
+
         data = {
             key: (
                 grid.copy_data_from_other(data[key], self._grid_1dr)

tests/test_neoclassical.py

@@ -7,9 +7,23 @@
 import pytest
 from tests.test_plotting import tol_1d
 
+import desc
 from desc.equilibrium.coords import get_rtz_grid
 from desc.examples import get
 from desc.grid import LinearGrid
+from desc.objectives import (
+    AspectRatio,
+    EffectiveRipple,
+    FixBoundaryR,
+    FixBoundaryZ,
+    FixIota,
+    FixPressure,
+    FixPsi,
+    ForceBalance,
+    GenericObjective,
+    ObjectiveFunction,
+)
+from desc.optimize import Optimizer
 from desc.utils import setdefault
 from desc.vmec import VMECIO
 
@@ -77,6 +91,107 @@ def test_effective_ripple():
     return fig
 
 
+def test_effective_ripple_opt():
+    """Test that an optimizatin with EffectiveRipple works without failing."""
+    eq = desc.examples.get("HELIOTRON")
+
+    # Flux surfaces on which to evaluate ballooning stability
+    surfaces = [1.0]
+
+    # Number of toroidal transits of the field line
+    num_transit = 3
+
+    # Number of point along a field line per transit
+    knots_per_transit = 64
+
+    # Determine which modes to unfix
+    k = 2
+
+    print("\n---------------------------------------")
+    print(f"Optimizing boundary modes M, N <= {k}")
+    print("---------------------------------------")
+
+    modes_R = np.vstack(
+        (
+            [0, 0, 0],
+            eq.surface.R_basis.modes[
+                np.max(np.abs(eq.surface.R_basis.modes), 1) > k, :
+            ],
+        )
+    )
+    modes_Z = eq.surface.Z_basis.modes[
+        np.max(np.abs(eq.surface.Z_basis.modes), 1) > k, :
+    ]
+    constraints = (
+        ForceBalance(eq=eq),
+        FixBoundaryR(eq=eq, modes=modes_R),
+        FixBoundaryZ(eq=eq, modes=modes_Z),
+        FixPressure(eq=eq),
+        FixIota(eq=eq),
+        FixPsi(eq=eq),
+    )
+
+    Curvature_grid = LinearGrid(
+        M=2 * int(eq.M),
+        N=2 * int(eq.N),
+        rho=np.array([1.0]),
+        NFP=eq.NFP,
+        sym=True,
+        axis=False,
+    )
+
+    objective = ObjectiveFunction(
+        (
+            EffectiveRipple(
+                eq=eq,
+                rho=np.array(surfaces),
+                alpha=0,
+                num_transit=num_transit,
+                knots_per_transit=knots_per_transit,
+                num_quad=17,
+                weight=1e6,
+                deriv_mode="rev",
+            ),
+            AspectRatio(
+                eq=eq,
+                bounds=(8, 11),
+                weight=1e3,
+            ),
+            GenericObjective(
+                f="curvature_k2_rho",
+                thing=eq,
+                grid=Curvature_grid,
+                bounds=(-128, 10),
+                weight=2e3,
+            ),
+        )
+    )
+
+    optimizer = Optimizer("proximal-lsq-exact")
+    (eq,), _ = optimizer.optimize(
+        eq,
+        objective,
+        constraints,
+        ftol=1e-4,
+        xtol=1e-6,
+        gtol=1e-6,
+        maxiter=2,  # increase maxiter to 50 for a better result
+        verbose=3,
+        options={"initial_trust_ratio": 2e-3},
+    )
+    (eq,), _ = optimizer.optimize(
+        eq,
+        objective,
+        constraints,
+        ftol=1e-4,
+        xtol=1e-6,
+        gtol=1e-6,
+        maxiter=2,  # increase maxiter to 50 for a better result
+        verbose=3,
+        options={"initial_trust_ratio": 2e-3},
+    )
+
+
 class NeoIO:
     """Class to interface with NEO."""