diff --git a/CHANGELOG.md b/CHANGELOG.md index 139f8ba85d..5f1c02756b 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -3,6 +3,8 @@ Changelog New Features +- Adds ``SecondAdiabaticInvariantAlphaDerivative`` and ``SoftConnectivity`` objectives for omnigenity optimization targeting the second adiabatic invariant (bounce action) $J^*$ along Boozer magnetic field lines and enforcing clean single-well magnetic field structure. +- Introduces ``SplineZeta`` parameterization in ``desc.magnetic_fields`` for managing toroidal extrema curves $\zeta_{\rm min}(\alpha)$ and $\zeta_{\rm max}(\alpha)$ with optional stellarator symmetry. - Added warning for when ``deriv_mode="batched"`` is used in an ``ObjectiveFunction`` where one or more sub-objectives is using ``rev`` mode differentiation. Also adds more info about the derivative mode and Jacobian chunk sizes when building the objective with ``verbose>1``. Performance Improvements diff --git a/desc/compute/_omnigenity.py b/desc/compute/_omnigenity.py index 57f6ae85c5..2e09ca86f8 100644 --- a/desc/compute/_omnigenity.py +++ b/desc/compute/_omnigenity.py @@ -11,6 +11,8 @@ import functools +import jax +import numpy as np from interpax import interp1d from desc.backend import jnp, sign, vmap @@ -19,6 +21,9 @@ from ..utils import cross, dot, safediv from .data_index import register_compute_fun +SOFTPLUS_SHARPNESS = 100.0 +_trapz = getattr(jnp, "trapezoid", getattr(jnp, "trapz", None)) + @register_compute_fun( name="B_theta_mn", @@ -992,3 +997,371 @@ def _isodynamicity(params, transforms, profiles, data, **kwargs): dot(cross(data["b"], data["grad(|B|)"]), data["grad(psi)"]) / data["|B|^2"] ) return data + + +# --------------------------------------------------------------------------- +# Direct Second Adiabatic Invariant (J*) and Soft-Connectivity Kernels +# References: Chen et al., arXiv:2608.02418 (2026) +# --------------------------------------------------------------------------- + + +def _softplus_relu(x, beta=SOFTPLUS_SHARPNESS): + """Sharp, smooth approximation of ``max(x, 0)``.""" + return jnp.logaddexp(beta * x, 0.0) / beta + + +def _softplus_relu_sigmoid(x, beta=SOFTPLUS_SHARPNESS): + """Derivative of ``_softplus_relu``: ``sigmoid(beta * x)``.""" + return 1.0 / (1.0 + jnp.exp(-beta * x)) + + +def _reshape_surface_coefficients(grid, values): + """Reshape flattened per-surface Boozer coefficients.""" + return jnp.asarray(values).reshape((grid.num_rho, -1)) + + +def _boozer_B_star_from_t(B_min, B_max, t): + """Map normalized pitch samples ``t`` to surface-wise ``B_star`` values.""" + B_min = jnp.atleast_1d(jnp.asarray(B_min)) + B_max = jnp.atleast_1d(jnp.asarray(B_max)) + t = jnp.atleast_1d(jnp.asarray(t)) + return (1.0 - t[None, :]) * B_min[:, None] + t[None, :] * B_max[:, None] + + +def _smoothmax_logsumexp(x, axis, tau): + """Differentiable upper envelope approximating max(x, axis=axis).""" + tau = jnp.asarray(tau, dtype=x.dtype) + tau = jnp.maximum(tau, jnp.finfo(x.dtype).eps) + x_scaled = x / tau + x_max = jnp.max(x_scaled, axis=axis, keepdims=True) + lse = x_max + jnp.log(jnp.sum(jnp.exp(x_scaled - x_max), axis=axis, keepdims=True)) + return tau * lse + + +def _boozer_second_adiabatic_surface_alpha_deriv( + basis, + rho, + coeff_B, + iota, + alpha, + B_star, + zeta_min, + zeta_max, + nzeta, + softplus_sharpness, +): + """Analytical dJ/dalpha on a single flux surface via chain-rule integration. + + Computes the derivative integrand directly: + dJ/dalpha = integral [ df/dB * dB/dtheta_B ] dzeta + where f = sqrt(cutoff) / B and theta_B = alpha + iota*(zeta - zeta_min). + """ + zeta = jnp.linspace(zeta_min, zeta_max, nzeta) + theta = alpha[:, None] + iota * (zeta[None, :] - zeta_min) + rho2d = jnp.broadcast_to(rho, theta.shape) + zeta2d = jnp.broadcast_to(zeta[None, :], theta.shape) + nodes = jnp.stack((rho2d, theta, zeta2d), axis=-1).reshape((-1, 3)) + + mat = basis.evaluate(nodes) + mat_dt = basis.evaluate(nodes, derivatives=np.array([0, 1, 0])) + + B = (mat @ coeff_B).reshape((alpha.size, nzeta)) + dB_dt = (mat_dt @ coeff_B).reshape((alpha.size, nzeta)) + + B_star = jnp.atleast_1d(B_star) + arg = 1.0 - B[None] / B_star[:, None, None] + cutoff = _softplus_relu(arg, beta=softplus_sharpness) + sig = _softplus_relu_sigmoid(arg, beta=softplus_sharpness) + sqrt_c = jnp.sqrt(jnp.maximum(cutoff, 1e-30)) + + safe_sqrt_c = jnp.where(cutoff > 0, sqrt_c, 1e-30) + # df/dB for f = sqrt(cutoff) / B: + # d/dB[sqrt(c)/B] = (dc/dB) * (1/(2*B*sqrt(c))) - sqrt(c)/B**2, + # with dc/dB = sigmoid(beta*arg) * (-1/B_star). + df_dB = jnp.where( + cutoff > 0, + -sig / (2.0 * B_star[:, None, None] * safe_sqrt_c * B[None]) + - safe_sqrt_c / (B[None] ** 2), + 0.0, + ) + + integrand = df_dB * dB_dt[None] + return _trapz(integrand, zeta, axis=-1).transpose(1, 0) + + +def boozer_second_adiabatic_invariant_alpha_derivative_analytical( + basis, + rho, + iota, + coeff_B, + alpha, + B_star, + *, + nzeta=1000, + zeta_min=0.0, + zeta_max=None, + nfp=1, + softplus_sharpness=SOFTPLUS_SHARPNESS, +): + """Analytical dJ/dalpha via chain rule through the Boozer integral.""" + alpha = jnp.asarray(alpha) + B_star = jnp.asarray(B_star) + if B_star.ndim == 0: + B_star = B_star[None, None] + elif B_star.ndim == 1: + B_star = B_star[None, :] + B_star = jnp.broadcast_to(B_star, (rho.size, B_star.shape[-1])) + zeta_max = 2 * jnp.pi / nfp if zeta_max is None else zeta_max + + return vmap( + lambda r, it, cB, pi, a: _boozer_second_adiabatic_surface_alpha_deriv( + basis, + r, + cB, + it, + a, + pi, + zeta_min, + zeta_max, + nzeta, + softplus_sharpness, + ) + )(rho, iota, coeff_B, B_star, alpha) + + +def boozer_second_adiabatic_invariant_alpha_derivative_from_data( + grid, + basis, + data, + t, + *, + num_alpha=48, + nzeta=1000, + zeta_min=0.0, + zeta_max=None, + softplus_sharpness=SOFTPLUS_SHARPNESS, + soft_extrema_tau=0.1, +): + """Convenience wrapper for ``dJ/dalpha`` using per-surface Boozer data.""" + rho = jnp.asarray(grid.compress(grid.nodes[:, 0])) + iota = jnp.asarray(grid.compress(data["iota"])) + coeff_B = _reshape_surface_coefficients(grid, data["|B|_mn_B"]) + + nfp = grid.NFP + alpha0 = jnp.pi - iota * (jnp.pi / nfp) + alpha = jnp.linspace(-jnp.pi, 0.0, num_alpha, endpoint=False) + alpha0[:, None] + + A = basis.evaluate(grid.nodes[:: grid.num_rho]) + B_grid = coeff_B @ A.T + num_eval = B_grid.shape[1] + log_n = jnp.log(jnp.maximum(jnp.asarray(num_eval, dtype=B_grid.dtype), 2.0)) + B_range = jnp.max(B_grid, axis=1) - jnp.min(B_grid, axis=1) + B_range = jnp.maximum(B_range, 1e-30) + tau_eff = (soft_extrema_tau * B_range / log_n)[:, None] + B_max = _smoothmax_logsumexp(B_grid, axis=1, tau=tau_eff).squeeze(-1) + B_min = -_smoothmax_logsumexp(-B_grid, axis=1, tau=tau_eff).squeeze(-1) + B_star = _boozer_B_star_from_t(B_min, B_max, t)[0] + return boozer_second_adiabatic_invariant_alpha_derivative_analytical( + basis, + rho, + iota, + coeff_B, + alpha, + B_star, + nzeta=nzeta, + zeta_min=zeta_min, + zeta_max=zeta_max, + nfp=nfp, + softplus_sharpness=softplus_sharpness, + ) + + +def _boozer_soft_connectivity_surface( + basis, + rho, + coeff_B, + iota, + alpha, + nfp, + t, + reduced_alpha_knots, + zeta_min_knots, + zeta_max_knots, + sigmoid_sharpness, + spline_symmetry, +): + """Compute the structured soft-connectivity penalty on a single flux surface.""" + zeta_span = 2 * jnp.pi / nfp + alpha_next = alpha + iota * zeta_span + + reduced_alpha_knots = jnp.asarray(reduced_alpha_knots) + zeta_min_knots = jnp.clip(jnp.asarray(zeta_min_knots), 0.0, zeta_span) + + if spline_symmetry: + alpha_min_knots = jnp.concatenate( + [ + reduced_alpha_knots, + jnp.mod( + 2 * jnp.pi - reduced_alpha_knots - iota * zeta_span, + 2 * jnp.pi, + ), + ] + ) + zeta_min_knots = jnp.concatenate([zeta_min_knots, zeta_span - zeta_min_knots]) + else: + alpha_min_knots = reduced_alpha_knots + min_order = jnp.argsort(alpha_min_knots) + alpha_min_knots = alpha_min_knots[min_order] + zeta_min_knots = zeta_min_knots[min_order] + + zeta_min_alpha = interp1d( + alpha, alpha_min_knots, zeta_min_knots, method="cubic", period=2 * jnp.pi + ) + zeta_min_next = interp1d( + alpha_next, + alpha_min_knots, + zeta_min_knots, + method="cubic", + period=2 * jnp.pi, + ) + zeta_min_next = zeta_min_next + zeta_span + + if zeta_max_knots is not None: + zeta_max_knots = ( + jnp.mod(jnp.asarray(zeta_max_knots) + 0.5 * zeta_span, zeta_span) + - 0.5 * zeta_span + ) + if spline_symmetry: + alpha_max_knots = jnp.concatenate( + [ + reduced_alpha_knots, + jnp.mod(2 * jnp.pi - reduced_alpha_knots, 2 * jnp.pi), + ] + ) + zeta_max_knots = jnp.concatenate([zeta_max_knots, -zeta_max_knots]) + else: + alpha_max_knots = reduced_alpha_knots + max_order = jnp.argsort(alpha_max_knots) + alpha_max_knots = alpha_max_knots[max_order] + zeta_max_knots = zeta_max_knots[max_order] + + zeta_max_alpha = interp1d( + alpha, alpha_max_knots, zeta_max_knots, method="cubic", period=2 * jnp.pi + ) + zeta_max_next = interp1d( + alpha_next, + alpha_max_knots, + zeta_max_knots, + method="cubic", + period=2 * jnp.pi, + ) + zeta_max_next = zeta_max_next + zeta_span + else: + zeta_max_alpha = jnp.zeros(alpha.shape) + zeta_max_next = jnp.full(alpha.shape, zeta_span) + + zeta = (1.0 - t[None, :]) * zeta_max_alpha[:, None] + ( + t[None, :] * zeta_max_next[:, None] + ) + theta = alpha[:, None] + iota * zeta + rho2d = jnp.broadcast_to(rho, theta.shape) + nodes = jnp.stack((rho2d, theta, zeta), axis=-1).reshape((-1, 3)) + nt = t.size + + dB_dtheta = ( + basis.evaluate(nodes, derivatives=np.array([0, 1, 0])) @ coeff_B + ).reshape((alpha.size, nt)) + dB_dzeta = ( + basis.evaluate(nodes, derivatives=np.array([0, 0, 1])) @ coeff_B + ).reshape((alpha.size, nt)) + dB_dz_line = iota * dB_dtheta + dB_dzeta + + use_current_min = (zeta_min_alpha > zeta_max_alpha) & ( + zeta_min_alpha < zeta_max_next + ) + zeta_min_shifted = jnp.where(use_current_min, zeta_min_alpha, zeta_min_next) + delta = zeta - zeta_min_shifted[:, None] + sig = jax.nn.sigmoid(sigmoid_sharpness * delta) + penalty_left = _softplus_relu(dB_dz_line) + penalty_right = _softplus_relu(-dB_dz_line) + penalty = sig * penalty_right + (1.0 - sig) * penalty_left + return penalty + + +def boozer_soft_connectivity_penalty( + basis, + rho, + iota, + coeff_B, + alpha, + nfp, + t, + *, + reduced_alpha_knots, + zeta_min_knots, + zeta_max_knots=None, + sigmoid_sharpness=50.0, + spline_symmetry=True, +): + """Compute the soft-connectivity penalty over multiple flux surfaces.""" + nfp_arr = jnp.broadcast_to(jnp.asarray(nfp), rho.shape) + return vmap( + lambda r, it, cB, a, nf: _boozer_soft_connectivity_surface( + basis, + r, + cB, + it, + a, + nf, + t, + reduced_alpha_knots, + zeta_min_knots, + zeta_max_knots, + sigmoid_sharpness, + spline_symmetry, + ) + )(rho, iota, coeff_B, alpha, nfp_arr) + + +def boozer_soft_connectivity_penalty_from_data( + grid, + basis, + data, + t, + *, + reduced_alpha_knots, + zeta_min_knots, + zeta_max_knots=None, + num_alpha=50, + sigmoid_sharpness=50.0, + spline_symmetry=True, +): + """Convenience wrapper for the soft-connectivity penalty using per-surface Boozer data.""" + rho = jnp.asarray(grid.compress(grid.nodes[:, 0])) + iota = jnp.asarray(grid.compress(data["iota"])) + coeff_B = _reshape_surface_coefficients(grid, data["|B|_mn_B"]) + + nfp = grid.NFP + # alpha0 is the fixed point of the stellarator-symmetry mirror map + # alpha -> 2π − alpha − iota*span. With symmetric splines the penalty + # on the mirrored half is identical, so sampling one fundamental + # domain (length π) suffices. With symmetry=False the knots span the + # full [0, 2π) and every knot must be sampled: use the full period. + alpha0 = jnp.pi - iota * (jnp.pi / nfp) + alpha_span = jnp.pi if spline_symmetry else 2 * jnp.pi + alpha = jnp.linspace(-alpha_span, 0.0, num_alpha, endpoint=False) + alpha0[:, None] + + return boozer_soft_connectivity_penalty( + basis, + rho, + iota, + coeff_B, + alpha, + nfp, + t, + reduced_alpha_knots=reduced_alpha_knots, + zeta_min_knots=zeta_min_knots, + zeta_max_knots=zeta_max_knots, + sigmoid_sharpness=sigmoid_sharpness, + spline_symmetry=spline_symmetry, + ) diff --git a/desc/magnetic_fields/__init__.py b/desc/magnetic_fields/__init__.py index 49115a2a3d..d2775f0512 100644 --- a/desc/magnetic_fields/__init__.py +++ b/desc/magnetic_fields/__init__.py @@ -7,6 +7,7 @@ ScalarPotentialField, ScaledMagneticField, SplineMagneticField, + SplineZeta, SumMagneticField, ToroidalMagneticField, VectorPotentialField, diff --git a/desc/magnetic_fields/_core.py b/desc/magnetic_fields/_core.py index d09e67e6ad..6423933744 100644 --- a/desc/magnetic_fields/_core.py +++ b/desc/magnetic_fields/_core.py @@ -3259,3 +3259,203 @@ def helicity(self, helicity): and (int(helicity[1]) == helicity[1]) ) self._helicity = helicity + + +class SplineZeta(Optimizable, IOAble): + """Cubic spline control points defining zeta_min_Boozer(alpha) and zeta_max_Boozer(alpha). + + The optimizable parameters are the zeta_min and zeta_max knot values + that, together with fixed alpha_knots, define periodic cubic splines for + the Boozer-toroidal-angle positions of |B| minimum (low-field side) and + |B| maximum (high-field side) within one field period, as a function of + the field-line label alpha. + + With ``symmetry=True`` (default) stellarator symmetry is assumed: the + knots are the reduced free control points in (0, π) and the mirror knots + are expanded in the compute layer. With ``symmetry=False`` no + stellarator symmetry is assumed: the knots are the full set of free + control points spanning [0, 2π) and are used directly. + + Parameters + ---------- + n_control : int + Number of free spline control points. With ``symmetry=True`` they + lie in (0, π) and the full spline uses 2*n_control knots on + [0, 2π); with ``symmetry=False`` they span the full [0, 2π). + NFP : int + Number of field periods. + alpha_knots : ndarray, optional + Alpha positions of the free control points. Defaults to uniformly + spaced cell centers in (0, π) with ``symmetry=True`` (avoiding + duplicate mirrored knots at symmetry fixed points), or in [0, 2π) + with ``symmetry=False``. + zeta_min_knots : ndarray, optional + Zeta_min values at control points, shape (n_control,). + Defaults to π/NFP (center of one field period). + zeta_max_knots : ndarray, optional + Zeta_max values at control points, shape (n_control,). + Stored on the branch nearest the field-period boundary at ``ζ=0``, + i.e. wrapped to ``[-π/NFP, π/NFP)``. Defaults to 0.0. + symmetry : bool, optional + Whether to assume stellarator symmetry. Default True. With False, + the full-period knots are optimizable parameters and no mirror + expansion is applied in the soft-connectivity compute layer. + """ + + _io_attrs_ = [ + "_NFP", + "_n_control", + "_alpha_knots", + "_zeta_min_knots", + "_zeta_max_knots", + "_symmetry", + ] + _static_attrs = Optimizable._static_attrs + [ + "_NFP", + "_n_control", + "_alpha_knots", + "_symmetry", + ] + + def __init__( + self, + n_control=8, + NFP=1, + alpha_knots=None, + zeta_min_knots=None, + zeta_max_knots=None, + symmetry=True, + ): + self._NFP = int(NFP) + self._n_control = int(n_control) + self._symmetry = bool(symmetry) + if alpha_knots is None: + period = np.pi if self._symmetry else 2 * np.pi + alpha_knots = (np.arange(n_control) + 0.5) * period / n_control + self._alpha_knots = np.asarray(alpha_knots, dtype=float) + errorif( + self._alpha_knots.size != self._n_control, + ValueError, + "alpha_knots must have length n_control.", + ) + if zeta_min_knots is None: + zeta_min_knots = np.full(n_control, np.pi / NFP) + zeta_min_knots = np.asarray(zeta_min_knots, dtype=float) + errorif( + zeta_min_knots.size != self._n_control, + ValueError, + "zeta_min_knots must have length n_control.", + ) + self._zeta_min_knots = zeta_min_knots + if zeta_max_knots is None: + zeta_max_knots = np.zeros(n_control) + zeta_max_knots = np.asarray(zeta_max_knots, dtype=float) + errorif( + zeta_max_knots.size != self._n_control, + ValueError, + "zeta_max_knots must have length n_control.", + ) + self._zeta_max_knots = self._wrap_zeta_max_near_zero(zeta_max_knots) + + def _set_up(self): + """Validate serialized SplineZeta state after loading.""" + self._symmetry = bool(self._symmetry) + self._NFP = int(self._NFP) + self._n_control = int(self._n_control) + self._alpha_knots = np.asarray(self._alpha_knots, dtype=float) + errorif( + self._alpha_knots.size != self._n_control, + ValueError, + "Loaded SplineZeta alpha_knots must have length n_control.", + ) + self._zeta_min_knots = np.asarray(self._zeta_min_knots, dtype=float) + errorif( + self._zeta_min_knots.size != self._n_control, + ValueError, + "Loaded SplineZeta zeta_min_knots must have length n_control.", + ) + self._zeta_max_knots = self._wrap_zeta_max_near_zero(self._zeta_max_knots) + errorif( + self._zeta_max_knots.size != self._n_control, + ValueError, + "Loaded SplineZeta zeta_max_knots must have length n_control.", + ) + + def _wrap_zeta_max_near_zero(self, val): + """Wrap zeta_max values to the branch nearest the field-period boundary.""" + span = 2 * np.pi / self._NFP + return np.mod(np.asarray(val, dtype=float) + span / 2, span) - span / 2 + + def _wrap_zeta_max_near_zero_jax(self, val): + """JAX version of ``_wrap_zeta_max_near_zero`` for optimization params.""" + span = 2 * jnp.pi / self._NFP + return jnp.mod(jnp.asarray(val, dtype=float) + span / 2, span) - span / 2 + + def reduced_knots(self, params=None): + """Return the free knots used by the optimizer. + + With ``symmetry=True`` these are the reduced knots in (0, π) that the + compute layer mirrors; with ``symmetry=False`` they are the full + knots spanning [0, 2π), used directly. + """ + if params is None: + params = self.params_dict + return { + "reduced_alpha_knots": jnp.asarray(self._alpha_knots), + "zeta_min_knots": jnp.clip( + jnp.asarray(params["zeta_min_knots"]), 0.0, 2 * jnp.pi / self._NFP + ), + "zeta_max_knots": jnp.asarray(params["zeta_max_knots"]), + } + + @property + def symmetry(self): + """bool: Whether stellarator symmetry is assumed for the knot layout.""" + return self._symmetry + + @property + def NFP(self): + """int: Number of field periods.""" + return self._NFP + + @optimizable_parameter + @property + def zeta_min_knots(self): + """ndarray: zeta_min values at spline control points.""" + return self._zeta_min_knots + + @zeta_min_knots.setter + def zeta_min_knots(self, val): + val = np.asarray(val, dtype=float) + errorif( + val.size != self._n_control, + ValueError, + "zeta_min_knots must have length n_control.", + ) + self._zeta_min_knots = val + + @optimizable_parameter + @property + def zeta_max_knots(self): + """ndarray: zeta_max values at spline control points.""" + return self._zeta_max_knots + + @zeta_max_knots.setter + def zeta_max_knots(self, val): + val = np.asarray(val, dtype=float) + errorif( + val.size != self._n_control, + ValueError, + "zeta_max_knots must have length n_control.", + ) + self._zeta_max_knots = self._wrap_zeta_max_near_zero(val) + + @property + def alpha_knots_full(self): + """Full alpha knots spanning [0, 2π) for the periodic spline. + + With ``symmetry=False`` the stored knots already span [0, 2π). + """ + if not self._symmetry: + return self._alpha_knots + return np.concatenate([self._alpha_knots, self._alpha_knots + np.pi]) diff --git a/desc/objectives/__init__.py b/desc/objectives/__init__.py index 52ca059e98..8a861fb530 100644 --- a/desc/objectives/__init__.py +++ b/desc/objectives/__init__.py @@ -53,6 +53,8 @@ QuasisymmetryBoozer, QuasisymmetryTripleProduct, QuasisymmetryTwoTerm, + SecondAdiabaticInvariantAlphaDerivative, + SoftConnectivity, ) from ._power_balance import FusionPower, HeatingPowerISS04 from ._profiles import Pressure, RotationalTransform, Shear, ToroidalCurrent diff --git a/desc/objectives/_omnigenity.py b/desc/objectives/_omnigenity.py index 88f59ebab6..b13d7bf545 100644 --- a/desc/objectives/_omnigenity.py +++ b/desc/objectives/_omnigenity.py @@ -2,12 +2,19 @@ import warnings -from desc.backend import jnp +import numpy as np + +from desc.backend import jnp, vmap from desc.batching import vmap_chunked from desc.compute import get_profiles, get_transforms -from desc.compute._omnigenity import _omnigenity_mapping +from desc.compute._omnigenity import ( + SOFTPLUS_SHARPNESS, + _omnigenity_mapping, + boozer_second_adiabatic_invariant_alpha_derivative_from_data, + boozer_soft_connectivity_penalty_from_data, +) from desc.compute.utils import _compute as compute_fun -from desc.grid import LinearGrid +from desc.grid import LinearGrid, _Grid from desc.utils import Timer, errorif, warnif from desc.vmec_utils import ptolemy_linear_transform @@ -993,3 +1000,368 @@ def compute(self, params, constants=None): profiles=constants["profiles"], ) return data["isodynamicity"] + + +class SecondAdiabaticInvariantAlphaDerivative(_Objective): + r"""Boozer-space second adiabatic invariant derivative along alpha. + + Targets the omnigenity condition by directly penalizing the derivative + of the second adiabatic invariant (bounce action) :math:`J^*` with respect + to the field-line label :math:`\alpha`: + + .. math:: + + \frac{\partial J^*}{\partial \alpha} = 0, \quad + J^*(\psi, B_{\rm bounce}, \alpha) = \oint \sqrt{2m(E - \mu B)}\,d\ell + + A smooth `softplus` kernel is used to resolve the square-root endpoint + singularity at bounce points, enabling exact, stable reverse-mode + automatic differentiation. + + Parameters + ---------- + eq : Equilibrium + Equilibrium that will be optimized to satisfy the Objective. + grid : Grid, optional + Collocation grid used to compute Boozer coefficients and surface extrema. + Must be non-symmetric. Defaults to the boundary surface with + ``LinearGrid(rho=1.0, M=2*M_booz, N=2*N_booz, NFP=eq.NFP, sym=False)``. + M_booz : int, optional + Poloidal resolution of Boozer transformation. Default = 2 * eq.M. + N_booz : int, optional + Toroidal resolution of Boozer transformation. Default = 2 * eq.N. + num_alpha : int, optional + Number of alpha samples per surface. The actual alpha values are + computed dynamically from iota: + ``linspace(alpha0 - π, alpha0, N)`` where + ``alpha0 = π - iota * π/nfp``. Defaults to 48. + t : ndarray, optional + Normalized pitch samples used to define + ``B_star = (1-t) * B_min + t * B_max`` on each surface. + Defaults to 48 points in [0.05, 0.95]. + nzeta : int, optional + Number of Boozer toroidal samples used for the line integral. Default = 100. + zeta_min, zeta_max : float, optional + Boozer toroidal integration bounds. ``zeta_max`` defaults to one field period. + softplus_sharpness : float, optional + Sharpness of the smooth cutoff used in the second adiabatic invariant. + Default = 100.0. + soft_extrema_tau : float, optional + Temperature parameter for the soft Log-Sum-Exp extrema estimators. + Default = 0.1. + + References + ---------- + .. [1] Chen, H., Lu, Z., Xu, G., et al. (2026). "Direct Optimization + of Stellarator Omnigenity from the Second Adiabatic Invariant." + arXiv:2608.02418. + """ + + __doc__ = __doc__.rstrip() + collect_docs( + target_default="``target=0``.", bounds_default="``target=0``." + ) + + _coordinates = "r" + _units = "(1/T)" + _print_value_fmt = "dJ/dalpha: " + _static_attrs = _Objective._static_attrs + [ + "_data_keys", + "_grid", + "_nzeta", + "_num_alpha", + "_soft_extrema_tau", + "_softplus_sharpness", + "_t", + "_zeta_max", + "_zeta_min", + "M_booz", + "N_booz", + ] + + def __init__( + self, + eq, + target=None, + bounds=None, + weight=1, + normalize=True, + normalize_target=True, + loss_function=None, + deriv_mode="auto", + grid=None, + M_booz=None, + N_booz=None, + num_alpha=48, + t=None, + nzeta=100, + zeta_min=0.0, + zeta_max=None, + softplus_sharpness=SOFTPLUS_SHARPNESS, + soft_extrema_tau=0.1, + name="Second adiabatic invariant alpha derivative", + jac_chunk_size=None, + ): + if target is None and bounds is None: + target = 0 + self.M_booz = M_booz if M_booz is not None else 2 * eq.M + self.N_booz = N_booz if N_booz is not None else 2 * eq.N + self._grid = grid + self._num_alpha = num_alpha + self._t = np.linspace(0.05, 0.95, 48) if t is None else np.asarray(t) + self._nzeta = nzeta + self._zeta_min = zeta_min + self._zeta_max = zeta_max + self._softplus_sharpness = softplus_sharpness + self._soft_extrema_tau = soft_extrema_tau + self._jac_chunk_size = jac_chunk_size + + 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 and grid.""" + eq = self.things[0] + if self._grid is None: + grid = LinearGrid( + rho=np.array([1.0]), + M=2 * self.M_booz, + N=2 * self.N_booz, + NFP=eq.NFP, + sym=False, + ) + else: + grid = self._grid + + assert isinstance(grid, _Grid), f"Expected Grid, got {type(grid)}." + assert not grid.sym, "Grid must not be symmetric." + + self._data_keys = ["|B|_mn_B", "iota"] + timer = Timer() + if verbose > 0: + print("Precomputing transforms") + timer.start("Precomputing transforms") + + self._constants = { + "transforms": get_transforms( + self._data_keys, + obj=eq, + grid=grid, + M_booz=self.M_booz, + N_booz=self.N_booz, + ), + "profiles": get_profiles(self._data_keys, obj=eq, grid=grid), + "t": jnp.asarray(self._t), + } + timer.stop("Precomputing transforms") + if verbose > 1: + timer.disp("Precomputing transforms") + + if self._normalize: + scales = compute_scaling_factors(eq) + self._normalization = 1.0 / scales["B"] + + self._dim_f = grid.num_rho * self._num_alpha * self._t.size + super().build(use_jit=use_jit, verbose=verbose) + + def compute(self, params, constants=None): + """Compute ``dJ/dalpha`` on the configured rho/alpha/t grid.""" + constants = self._get_deprecated_constants(constants) + zeta_max = ( + 2 * np.pi / constants["transforms"]["grid"].NFP + if self._zeta_max is None + else self._zeta_max + ) + data = compute_fun( + "desc.equilibrium.equilibrium.Equilibrium", + self._data_keys, + params=params, + transforms=constants["transforms"], + profiles=constants["profiles"], + ) + dJ_dalpha = boozer_second_adiabatic_invariant_alpha_derivative_from_data( + constants["transforms"]["grid"], + constants["transforms"]["B"].basis, + data, + constants["t"], + num_alpha=self._num_alpha, + nzeta=self._nzeta, + zeta_min=self._zeta_min, + zeta_max=zeta_max, + softplus_sharpness=self._softplus_sharpness, + soft_extrema_tau=self._soft_extrema_tau, + ) + return jnp.ravel(dJ_dalpha) + + +class SoftConnectivity(_Objective): + r"""Soft-connectivity penalty for :math:`|B|` along Boozer field lines. + + Enforces clean, single-well magnetic field structure along field lines + without secondary local wells or trapped-branch bifurcations: + + .. math:: + + \mathcal{P}_{\rm conn} = \int \sigma(\Delta \zeta) \operatorname{softplus}(-\partial_\zeta |B|) + (1-\sigma(\Delta \zeta)) \operatorname{softplus}(\partial_\zeta |B|) + + Parameters + ---------- + eq : Equilibrium + Equilibrium that will be optimized to satisfy the Objective. + spline : SplineZeta + Spline parameterized control points defining the well extrema. + grid : Grid, optional + Collocation grid used to compute Boozer coefficients. + M_booz : int, optional + Poloidal resolution of Boozer transformation. Default = 2 * eq.M. + N_booz : int, optional + Toroidal resolution of Boozer transformation. Default = 2 * eq.N. + num_alpha : int, optional + Number of alpha samples per surface. Default = 48. + t : ndarray, optional + Normalized Boozer toroidal path samples in [0, 1]. Default = linspace(0, 1, 200). + sigmoid_sharpness : float, optional + Transition sharpness between left and right well sides. Default = 10.0. + + References + ---------- + .. [1] Chen, H., Lu, Z., Xu, G., et al. (2026). "Direct Optimization + of Stellarator Omnigenity from the Second Adiabatic Invariant." + arXiv:2608.02418. + """ + + __doc__ = __doc__.rstrip() + collect_docs( + target_default="``target=0``.", bounds_default="``target=0``." + ) + + _coordinates = "r" + _units = "~" + _print_value_fmt = "Soft connectivity penalty: " + _static_attrs = _Objective._static_attrs + [ + "_data_keys", + "_grid", + "_num_alpha", + "_sigmoid_sharpness", + "_t", + "M_booz", + "N_booz", + ] + + def __init__( + self, + eq, + spline, + target=None, + bounds=None, + weight=1, + normalize=False, + normalize_target=False, + loss_function=None, + deriv_mode="auto", + grid=None, + M_booz=None, + N_booz=None, + num_alpha=50, + t=None, + sigmoid_sharpness=50.0, + name="Soft connectivity penalty", + jac_chunk_size=None, + ): + if target is None and bounds is None: + target = 0 + self.M_booz = M_booz if M_booz is not None else 2 * eq.M + self.N_booz = N_booz if N_booz is not None else 2 * eq.N + self._grid = grid + self._num_alpha = num_alpha + self._t = np.linspace(0.0, 1.0, 200) if t is None else np.asarray(t) + self._sigmoid_sharpness = sigmoid_sharpness + self._jac_chunk_size = jac_chunk_size + + things = [eq, spline] + super().__init__( + things=things, + 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 and grid.""" + eq = self.things[0] + if self._grid is None: + grid = LinearGrid( + rho=np.array([1.0]), + M=2 * self.M_booz, + N=2 * self.N_booz, + NFP=eq.NFP, + sym=False, + ) + else: + grid = self._grid + + assert isinstance(grid, _Grid), f"Expected Grid, got {type(grid)}." + assert not grid.sym, "Grid must not be symmetric." + + self._data_keys = ["|B|_mn_B", "iota"] + timer = Timer() + if verbose > 0: + print("Precomputing transforms") + timer.start("Precomputing transforms") + + self._constants = { + "transforms": get_transforms( + self._data_keys, + obj=eq, + grid=grid, + M_booz=self.M_booz, + N_booz=self.N_booz, + ), + "profiles": get_profiles(self._data_keys, obj=eq, grid=grid), + "t": jnp.asarray(self._t), + } + timer.stop("Precomputing transforms") + if verbose > 1: + timer.disp("Precomputing transforms") + + self._dim_f = grid.num_rho * self._num_alpha * self._t.size + super().build(use_jit=use_jit, verbose=verbose) + + def compute(self, params_1, params_2=None, constants=None): + """Compute soft-connectivity penalty residuals.""" + constants = self._get_deprecated_constants(constants) + data = compute_fun( + "desc.equilibrium.equilibrium.Equilibrium", + self._data_keys, + params=params_1, + transforms=constants["transforms"], + profiles=constants["profiles"], + ) + spline = self.things[1] + spline_knots = spline.reduced_knots(params_2) + + penalty = boozer_soft_connectivity_penalty_from_data( + constants["transforms"]["grid"], + constants["transforms"]["B"].basis, + data, + constants["t"], + reduced_alpha_knots=spline_knots["reduced_alpha_knots"], + zeta_min_knots=spline_knots["zeta_min_knots"], + zeta_max_knots=spline_knots["zeta_max_knots"], + num_alpha=self._num_alpha, + sigmoid_sharpness=self._sigmoid_sharpness, + spline_symmetry=spline.symmetry, + ) + return jnp.ravel(penalty) diff --git a/tests/test_objective_funs.py b/tests/test_objective_funs.py index ab211c6599..086ecc77d0 100644 --- a/tests/test_objective_funs.py +++ b/tests/test_objective_funs.py @@ -87,6 +87,7 @@ QuasisymmetryTwoTerm, RotationalTransform, Shear, + SoftConnectivity, SurfaceCurrentRegularization, SurfaceQuadraticFlux, ToroidalCurrent, @@ -3327,6 +3328,7 @@ class TestComputeScalarResolution: PlasmaCoilSetMinDistance, PlasmaVesselDistance, QuadraticFlux, + SoftConnectivity, SurfaceQuadraticFlux, ToroidalFlux, SurfaceCurrentRegularization, @@ -3849,6 +3851,7 @@ class TestObjectiveNaNGrad: PlasmaCoilSetMinDistance, PlasmaVesselDistance, QuadraticFlux, + SoftConnectivity, SurfaceCurrentRegularization, SurfaceQuadraticFlux, ToroidalFlux, diff --git a/tests/test_second_adiabatic_invariant.py b/tests/test_second_adiabatic_invariant.py new file mode 100644 index 0000000000..0350332792 --- /dev/null +++ b/tests/test_second_adiabatic_invariant.py @@ -0,0 +1,372 @@ +"""Tests for SecondAdiabaticInvariantAlphaDerivative and SoftConnectivity objectives. + +References +---------- +.. [1] Chen, H., Lu, Z., Xu, G., et al. (2026). "Direct Optimization + of Stellarator Omnigenity from the Second Adiabatic Invariant." + arXiv:2608.02418. +""" + +from __future__ import annotations + +import numpy as np +import pytest + +from desc.backend import jnp +from desc.basis import DoubleFourierSeries +from desc.compute._omnigenity import ( + _smoothmax_logsumexp, + _softplus_relu, + _softplus_relu_sigmoid, + boozer_second_adiabatic_invariant_alpha_derivative_analytical, + boozer_soft_connectivity_penalty, +) +from desc.equilibrium import Equilibrium +from desc.examples import get +from desc.grid import LinearGrid +from desc.objectives import ( + ObjectiveFunction, + SecondAdiabaticInvariantAlphaDerivative, + SoftConnectivity, +) +from desc.optimize import Optimizer + + +class TestSecondAdiabaticInvariant: + """Test suite for direct second adiabatic invariant optimization.""" + + @pytest.mark.unit + def test_softplus_relu_and_sigmoid(self): + """Test mathematical properties of the softplus smoothing kernel.""" + x = np.linspace(-5.0, 5.0, 100) + beta = 50.0 + y = _softplus_relu(x, beta=beta) + dy = _softplus_relu_sigmoid(x, beta=beta) + + # For positive x, softplus(x) ≈ x + np.testing.assert_allclose(y[x > 0.5], x[x > 0.5], atol=1e-2) + # For negative x, softplus(x) ≈ 0 + np.testing.assert_allclose(y[x < -0.5], 0.0, atol=1e-2) + # Verify derivative bounds in [0, 1] + assert np.all(dy >= 0.0) and np.all(dy <= 1.0) + + @pytest.mark.unit + def test_smoothmax_logsumexp(self): + """Test Log-Sum-Exp differentiable extrema approximation.""" + data = np.array([[1.0, 3.0, 2.0], [5.0, 0.0, -1.0]]) + tau = 0.01 + smooth_max = _smoothmax_logsumexp(data, axis=1, tau=tau).squeeze() + exact_max = np.max(data, axis=1) + np.testing.assert_allclose(smooth_max, exact_max, atol=0.05) + + @pytest.mark.unit + def test_second_adiabatic_invariant_objective_build_and_compute(self): + """Test objective construction, compute, and residual evaluation.""" + eq = get("DSHAPE") + obj = SecondAdiabaticInvariantAlphaDerivative( + eq=eq, + num_alpha=16, + nzeta=32, + M_booz=4, + N_booz=0, + ) + obj.build() + residuals = obj.compute(eq.params_dict) + + assert np.all(np.isfinite(residuals)) + assert residuals.ndim == 1 + assert len(residuals) > 0 + + @pytest.mark.unit + def test_second_adiabatic_invariant_jacobian_and_grad(self): + """Test gradient and Jacobian derivative actions of SecondAdiabaticInvariant.""" + eq = get("DSHAPE") + obj_fun = ObjectiveFunction( + SecondAdiabaticInvariantAlphaDerivative( + eq=eq, + num_alpha=8, + nzeta=20, + M_booz=2, + N_booz=0, + ) + ) + obj_fun.build() + + x = obj_fun.x() + grad = obj_fun.grad(x) + assert np.all(np.isfinite(grad)) + assert not np.any(np.isnan(grad)) + assert grad.shape == (obj_fun.dim_x,) + + jac = obj_fun.jac_scaled(x) + assert np.all(np.isfinite(jac)) + assert not np.any(np.isnan(jac)) + assert jac.shape == (obj_fun.dim_f, obj_fun.dim_x) + + tangent = np.random.default_rng(42).normal(size=x.shape) + jvp_action = obj_fun.jvp_scaled(tangent, x) + assert np.all(np.isfinite(jvp_action)) + np.testing.assert_allclose(jvp_action, jac @ tangent, rtol=1e-8, atol=1e-8) + + cotangent = np.random.default_rng(43).normal(size=(obj_fun.dim_f,)) + vjp_action = obj_fun.vjp_scaled(cotangent, x) + assert np.all(np.isfinite(vjp_action)) + np.testing.assert_allclose(vjp_action, jac.T @ cotangent, rtol=1e-8, atol=1e-8) + + # Verify gradient-Jacobian consistency: grad = J^T * f + f = obj_fun.compute_scaled_error(x) + expected_grad = jac.T @ f + np.testing.assert_allclose(grad, expected_grad, rtol=1e-8, atol=1e-8) + + @pytest.mark.unit + def test_soft_connectivity_objective_build_and_compute(self): + """Test SoftConnectivity objective build and evaluation.""" + from desc.magnetic_fields import SplineZeta + + eq = get("DSHAPE") + spline = SplineZeta(n_control=4, NFP=eq.NFP) + obj = SoftConnectivity( + eq=eq, + spline=spline, + num_alpha=12, + M_booz=4, + N_booz=0, + ) + obj.build() + residuals = obj.compute(eq.params_dict, spline.params_dict) + + assert np.all(np.isfinite(residuals)) + assert residuals.ndim == 1 + assert len(residuals) > 0 + + @pytest.mark.unit + def test_spline_zeta_reduced_knots_and_symmetry(self): + """Test SplineZeta parameter extraction and symmetry behavior.""" + from desc.magnetic_fields import SplineZeta + + spline_sym = SplineZeta(n_control=4, NFP=2, symmetry=True) + assert spline_sym.symmetry is True + knots_sym = spline_sym.reduced_knots() + assert knots_sym["reduced_alpha_knots"].size == 4 + assert knots_sym["zeta_min_knots"].size == 4 + assert knots_sym["zeta_max_knots"].size == 4 + assert spline_sym.alpha_knots_full.size == 8 + + spline_asym = SplineZeta(n_control=4, NFP=2, symmetry=False) + assert spline_asym.symmetry is False + knots_asym = spline_asym.reduced_knots() + assert knots_asym["reduced_alpha_knots"].size == 4 + assert spline_asym.alpha_knots_full.size == 4 + + @pytest.mark.unit + def test_single_step_optimization(self): + """Test single-step equilibrium optimization using J* objective.""" + from desc.objectives import ForceBalance, get_fixed_boundary_constraints + + eq = get("DSHAPE") + objective = ObjectiveFunction( + SecondAdiabaticInvariantAlphaDerivative( + eq=eq, + num_alpha=8, + nzeta=20, + M_booz=2, + N_booz=0, + ) + ) + constraints = (ForceBalance(eq), *get_fixed_boundary_constraints(eq)) + optimizer = Optimizer("proximal-lsq-exact") + eq_opt, result = eq.optimize( + objective=objective, + constraints=constraints, + optimizer=optimizer, + maxiter=1, + verbose=0, + ) + assert result.success or result.nfev >= 1 + assert np.all(np.isfinite(eq_opt.R_lmn)) + assert np.all(np.isfinite(eq_opt.Z_lmn)) + + @pytest.mark.unit + def test_soft_connectivity_joint_objective_optimization(self): + """Test joint optimization with eq and SplineZeta.""" + from desc.magnetic_fields import SplineZeta + from desc.objectives import ( + FixParameters, + ForceBalance, + get_fixed_boundary_constraints, + ) + + eq = get("DSHAPE") + spline = SplineZeta(n_control=4, NFP=eq.NFP) + objective = ObjectiveFunction( + SoftConnectivity( + eq=eq, + spline=spline, + num_alpha=8, + M_booz=2, + N_booz=0, + ) + ) + constraints = ( + ForceBalance(eq), + *get_fixed_boundary_constraints(eq), + FixParameters(spline, {"zeta_max_knots": True}), + ) + optimizer = Optimizer("proximal-lsq-exact") + (eq_opt, spline_opt), result = optimizer.optimize( + things=(eq, spline), + objective=objective, + constraints=constraints, + maxiter=1, + verbose=0, + ) + assert result.success or result.nfev >= 1 + assert np.all(np.isfinite(eq_opt.R_lmn)) + assert np.all(np.isfinite(spline_opt.zeta_min_knots)) + + @pytest.mark.unit + def test_boozer_alpha_derivative_scalar_b_star(self): + """Test scalar B_star is promoted and matches the per-surface path.""" + basis = DoubleFourierSeries(M=2, N=0, NFP=1, sym=False) + rng = np.random.default_rng(0) + coeff_B = (1.0 + 0.1 * rng.normal(size=basis.num_modes))[None, :] + rho = np.array([1.0]) + iota = np.array([0.5]) + alpha = np.linspace(-np.pi, 0.0, 6, endpoint=False)[None, :] + + out_scalar = boozer_second_adiabatic_invariant_alpha_derivative_analytical( + basis, rho, iota, coeff_B, alpha, 1.0, nzeta=64, nfp=1 + ) + out_vector = boozer_second_adiabatic_invariant_alpha_derivative_analytical( + basis, rho, iota, coeff_B, alpha, np.array([1.0]), nzeta=64, nfp=1 + ) + + assert np.all(np.isfinite(out_scalar)) + np.testing.assert_allclose(out_scalar, out_vector, rtol=1e-10, atol=1e-12) + + @pytest.mark.unit + def test_soft_connectivity_penalty_asymmetric_spline(self): + """Test the compute-layer knot expansion with spline_symmetry=False.""" + basis = DoubleFourierSeries(M=2, N=0, NFP=1, sym=False) + rng = np.random.default_rng(1) + coeff_B = (1.0 + 0.1 * rng.normal(size=basis.num_modes))[None, :] + rho = np.array([1.0]) + iota = np.array([0.4]) + n_control = 4 + reduced_alpha_knots = (np.arange(n_control) + 0.5) * 2 * np.pi / n_control + zeta_min_knots = np.full(n_control, np.pi) + zeta_max_knots = np.zeros(n_control) + alpha = np.linspace(-2 * np.pi, 0.0, 12, endpoint=False)[None, :] + t = np.linspace(0.0, 1.0, 16) + + penalty = boozer_soft_connectivity_penalty( + basis, + rho, + iota, + coeff_B, + alpha, + 1, + t, + reduced_alpha_knots=reduced_alpha_knots, + zeta_min_knots=zeta_min_knots, + zeta_max_knots=zeta_max_knots, + spline_symmetry=False, + ) + + assert np.all(np.isfinite(penalty)) + assert penalty.shape == (rho.size, alpha.shape[1], t.size) + + @pytest.mark.unit + def test_soft_connectivity_penalty_no_zeta_max(self): + """Test the zeta_max_knots=None fallback of the connectivity penalty.""" + basis = DoubleFourierSeries(M=2, N=0, NFP=1, sym=False) + rng = np.random.default_rng(2) + coeff_B = (1.0 + 0.1 * rng.normal(size=basis.num_modes))[None, :] + rho = np.array([1.0]) + iota = np.array([0.4]) + n_control = 4 + reduced_alpha_knots = (np.arange(n_control) + 0.5) * np.pi / n_control + zeta_min_knots = np.full(n_control, np.pi) + alpha = np.linspace(-np.pi, 0.0, 12, endpoint=False)[None, :] + t = np.linspace(0.0, 1.0, 16) + + penalty = boozer_soft_connectivity_penalty( + basis, + rho, + iota, + coeff_B, + alpha, + 1, + t, + reduced_alpha_knots=reduced_alpha_knots, + zeta_min_knots=zeta_min_knots, + zeta_max_knots=None, + spline_symmetry=True, + ) + + assert np.all(np.isfinite(penalty)) + assert penalty.shape == (rho.size, alpha.shape[1], t.size) + + @pytest.mark.unit + def test_soft_connectivity_symmetry_false_objective(self): + """Test SoftConnectivity end-to-end with a non-symmetric SplineZeta.""" + from desc.magnetic_fields import SplineZeta + + eq = get("DSHAPE") + spline = SplineZeta(n_control=4, NFP=eq.NFP, symmetry=False) + obj = SoftConnectivity( + eq=eq, + spline=spline, + num_alpha=8, + t=np.linspace(0.0, 1.0, 16), + M_booz=2, + N_booz=0, + ) + obj.build() + residuals = obj.compute(eq.params_dict, spline.params_dict) + + assert np.all(np.isfinite(residuals)) + assert residuals.ndim == 1 + assert len(residuals) > 0 + + @pytest.mark.unit + def test_spline_zeta_save_load_roundtrip(self, tmpdir_factory): + """Test SplineZeta serialization roundtrip, which invokes _set_up.""" + from desc.io import load + from desc.magnetic_fields import SplineZeta + + spline = SplineZeta(n_control=4, NFP=2, symmetry=False) + tmpdir = tmpdir_factory.mktemp("test_spline_zeta_io") + spline.save(tmpdir.join("spline_zeta.h5")) + loaded = load(tmpdir.join("spline_zeta.h5")) + + assert loaded.symmetry == spline.symmetry + assert loaded.NFP == 2 + np.testing.assert_allclose(loaded.zeta_min_knots, spline.zeta_min_knots) + np.testing.assert_allclose(loaded.zeta_max_knots, spline.zeta_max_knots) + np.testing.assert_allclose(loaded._alpha_knots, spline._alpha_knots) + + @pytest.mark.unit + def test_spline_zeta_set_up_validation(self): + """Test _set_up raises ValueError on inconsistent loaded state.""" + from desc.magnetic_fields import SplineZeta + + for attr in ["_alpha_knots", "_zeta_min_knots", "_zeta_max_knots"]: + spline = SplineZeta(n_control=4, NFP=2) + setattr(spline, attr, np.zeros(5)) + with pytest.raises(ValueError): + spline._set_up() + + @pytest.mark.unit + def test_spline_zeta_wrap_zeta_max_jax(self): + """Test the JAX branch wrap matches the numpy implementation.""" + from desc.magnetic_fields import SplineZeta + + spline = SplineZeta(n_control=4, NFP=2) + vals = np.linspace(-3.0, 3.0, 25) + np.testing.assert_allclose( + np.asarray(spline._wrap_zeta_max_near_zero_jax(vals)), + spline._wrap_zeta_max_near_zero(vals), + rtol=1e-12, + atol=1e-12, + )