From b17c627db56b96aa83c87f8a931343c25335d5cc Mon Sep 17 00:00:00 2001 From: Greta Hibbard Date: Mon, 2 Dec 2024 18:50:23 -0500 Subject: [PATCH] fixing #1288, adding opt test for effective ripple --- desc/objectives/_neoclassical.py | 36 +++++++--- tests/test_neoclassical.py | 115 +++++++++++++++++++++++++++++++ 2 files changed, 140 insertions(+), 11 deletions(-) diff --git a/desc/objectives/_neoclassical.py b/desc/objectives/_neoclassical.py index e609cc0d0a..7f8e8fb908 100644 --- a/desc/objectives/_neoclassical.py +++ b/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) diff --git a/tests/test_neoclassical.py b/tests/test_neoclassical.py index d517678cb9..af22a77a1b 100644 --- a/tests/test_neoclassical.py +++ b/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."""