diff --git a/desc/compute/_fast_ion.py b/desc/compute/_fast_ion.py index fc41df694d..ba0cf49fe4 100644 --- a/desc/compute/_fast_ion.py +++ b/desc/compute/_fast_ion.py @@ -119,9 +119,9 @@ def _Gamma_c(params, transforms, profiles, data, **kwargs): num_pitch, pitch_batch_size, surf_batch_size, - quad, nufft_eps, spline, + quad, vander, ) = Bounce2D._defaults(-2, grid, **kwargs) @@ -270,9 +270,9 @@ def _little_gamma_c_Nemov(params, transforms, profiles, data, **kwargs): num_pitch, _, _, - quad, nufft_eps, spline, + quad, vander, ) = Bounce2D._defaults(-2, grid, **kwargs) @@ -382,9 +382,9 @@ def _Gamma_c_Velasco(params, transforms, profiles, data, **kwargs): num_pitch, pitch_batch_size, surf_batch_size, - quad, nufft_eps, spline, + quad, vander, ) = Bounce2D._defaults(-1, grid, **kwargs) diff --git a/desc/compute/_neoclassical.py b/desc/compute/_neoclassical.py index 52769035bf..bd7747a1b0 100644 --- a/desc/compute/_neoclassical.py +++ b/desc/compute/_neoclassical.py @@ -20,7 +20,10 @@ If the option ``spline`` is ``True``, the bounce points are found with 8th order accuracy in this parameter. If the option ``spline`` is ``False``, then the bounce points are found with spectral accuracy in this parameter. - A reference value for the ``spline`` option is 100. + A reference value for the ``spline=True`` option is + ``grid.NFP*(grid.num_theta+grid.num_zeta)//2``. + A reference value for the ``spline=False`` option is + ``(grid.num_theta+grid.num_zeta)//2``. An error of ε in a bounce point manifests 𝒪(ε¹ᐧ⁵) error in bounce integrals with (v_∥)¹ and @@ -52,6 +55,8 @@ A tighter upper bound than ``num_well=(Aι+C)*num_transit`` is preferable. The ``check_points`` or ``plot`` methods in ``desc.integrals.Bounce2D`` are useful to select a reasonable value. + + This is the most important parameter to specify for performance. """, "num_quad": """int : Resolution for quadrature of bounce integrals. @@ -70,18 +75,18 @@ If given ``None``, then ``surf_batch_size`` is ``grid.num_rho``. Default is ``1``. Only consider increasing if ``pitch_batch_size`` is ``None``. """, - "quad": """tuple[jnp.ndarray] : - Used to compute bounce integrals. - Quadrature points xₖ and weights wₖ for the - approximate evaluation of the integral ∫₋₁¹ f(x) dx ≈ ∑ₖ wₖ f(xₖ). - """, "nufft_eps": """float : Precision requested for interpolation with non-uniform fast Fourier transform (NUFFT). If less than ``1e-14`` then NUFFT will not be used. """, "spline": """bool : - Whether to use cubic splines to compute bounce points instead of - Chebyshev series. Default is ``True``. + Whether to use cubic splines to compute initial guess for bounce points + instead of Chebyshev series. Default is ``True``. + """, + "quad": """tuple[jnp.ndarray] : + Used to compute bounce integrals. + Quadrature points xₖ and weights wₖ for the + approximate evaluation of the integral ∫₋₁¹ f(x) dx ≈ ∑ₖ wₖ f(xₖ). """, "_vander": """dict[str,jnp.ndarray] : Precomputed transform matrix "dct spline". @@ -199,9 +204,9 @@ def _epsilon_32(params, transforms, profiles, data, **kwargs): num_pitch, pitch_batch_size, surf_batch_size, - quad, nufft_eps, spline, + quad, vander, ) = Bounce2D._defaults(1, grid, **kwargs) diff --git a/desc/equilibrium/equilibrium.py b/desc/equilibrium/equilibrium.py index 004dd275fa..1b139b182d 100644 --- a/desc/equilibrium/equilibrium.py +++ b/desc/equilibrium/equilibrium.py @@ -946,6 +946,8 @@ def compute( # noqa: C901 inbasis=[inbasis[char] for char in grid.coordinates], outbasis=("rho", "theta", "zeta"), period=grid.period, + tol=kwargs.pop("tol", 1e-6), + maxiter=kwargs.pop("maxiter", 30), ) grid = Grid( nodes=rtz_nodes, diff --git a/desc/external/neo.py b/desc/external/neo.py index ba12ecd6fa..08595666b2 100644 --- a/desc/external/neo.py +++ b/desc/external/neo.py @@ -26,8 +26,8 @@ def __init__(self, name, eq, ns=256, M_booz=None, N_booz=None): self.eq = eq self.ns = ns # number of surfaces - self.M_booz = setdefault(M_booz, 5 * eq.M + 1) - self.N_booz = setdefault(N_booz, 5 * eq.N) + self.M_booz = setdefault(M_booz, 6 * eq.M + 1) + self.N_booz = setdefault(N_booz, 6 * eq.N) @staticmethod def read(name): @@ -39,12 +39,36 @@ def read(name): neo_eps[~good] = np.interp(neo_rho[~good], neo_rho[good], neo_eps[good]) return neo_rho, neo_eps - def write(self, **kwargs): - """Write neo input file.""" + def write(self, *, theta_n=300, phi_n=300, num_pitch=150, nstep_per=500, **kwargs): + """Write neo input file. + + Parameters + ---------- + theta_n, phi_n : int + Spline resolution. + num_pitch : int + Resolution for quadrature over velocity coordinate. + Default is 150. + nstep_per : int + Runge-Kutta steps to perform within each field period. Default is 500. + + NEO employs a constant step size (2π/NFP/``nstep_per``), + explicit Runge-Kutta scheme which has algebraic convergence of order 3/2 + for the bounce integrals of interest. The accuracy of the bounce points is + also only found with first order accuracy in (2π/NFP/``nstep_per``) so the + algebraic convergence is capped at 3/2 independent of the integrator. + + """ print(f"Writing VMEC wout to {self.vmec_file}") VMECIO.save(self.eq, self.vmec_file, surfs=self.ns, verbose=0) self._write_booz() - self._write_neo(**kwargs) + self._write_neo( + theta_n=theta_n, + phi_n=phi_n, + num_pitch=num_pitch, + nstep_per=nstep_per, + **kwargs, + ) def _write_booz(self): print(f"Writing boozer output file to {self.booz_file}") @@ -59,17 +83,18 @@ def _write_booz(self): def _write_neo( self, - theta_n=200, - phi_n=200, - num_pitch=150, + theta_n, + phi_n, + num_pitch, + nstep_per, multra=1, acc_req=0.01, nbins=100, - nstep_per=75, nstep_min=500, nstep_max=2000, verbose=2, ): + """Write NEO file.""" print(f"Writing NEO input file to {self.neo_in_file}") f = open(self.neo_in_file, "w") diff --git a/desc/integrals/__init__.py b/desc/integrals/__init__.py index fd793806ef..6bb01b77f7 100644 --- a/desc/integrals/__init__.py +++ b/desc/integrals/__init__.py @@ -1,6 +1,5 @@ """Classes for function integration.""" -from ._bounce_utils import fast_chebyshev, fast_cubic_spline from ._interp_utils import nufft1d2r, nufft2d2r from .bounce_integral import Bounce1D, Bounce2D from .singularities import ( diff --git a/desc/integrals/_bounce_utils.py b/desc/integrals/_bounce_utils.py index 24071dbe0f..fee1609022 100644 --- a/desc/integrals/_bounce_utils.py +++ b/desc/integrals/_bounce_utils.py @@ -26,6 +26,7 @@ from desc.backend import dct, ifft, jax, jnp from desc.integrals._interp_utils import ( + _JF_BUG, _eps, chebder, nufft1d2r, @@ -39,7 +40,7 @@ _sentinel = -1e5 -def bounce_points(pitch_inv, knots, B, num_well=-1): +def bounce_points(pitch_inv, knots, B, num_well=-1, return_mask=False): """Compute the bounce points given 1D spline of B and pitch λ. Parameters @@ -68,10 +69,12 @@ def bounce_points(pitch_inv, knots, B, num_well=-1): If there were fewer wells detected along a field line than the size of the last axis of the returned arrays, then that axis is padded with zero. + return_mask : bool + Whether to return the mask ``z1= 0) z1 = (dB_dz <= 0) & mask z2 = (dB_dz >= 0) & epigraph_and(mask, dB_dz) + del dB_dz # Transform out of local power basis expansion. intersect = flatten_mat(intersect + knots[:-1, None]) z1 = take_mask(intersect, z1, size=num_well, fill_value=_sentinel) z2 = take_mask(intersect, z2, size=num_well, fill_value=_sentinel) + del intersect mask = (z1 > _sentinel) & (z2 > _sentinel) # Set to zero so integration is over set of measure zero # and basis functions are faster to evaluate in downstream routines. z1 = jnp.where(mask, z1, 0.0) z2 = jnp.where(mask, z2, 0.0) - return z1, z2, mask + return (z1, z2, mask) if return_mask else (z1, z2) -def _newton(o, pitch_inv, z1, z2, mask, nufft_eps=1e-10): +def _newton(o, pitch_inv, z1, z2, mask, nufft_eps): """Newton step using maps used in the quadrature. An error of ε in a bounce point manifests @@ -132,7 +137,8 @@ def _newton(o, pitch_inv, z1, z2, mask, nufft_eps=1e-10): Shape (num ρ, num α, num pitch, num well). Subset of points to refine. nufft_eps : float - Desired error ε of the bounce points. + Desired error ε of the returned bounce points. + Should satisfy ε < εᵢₙ² where εᵢₙ is the error of the input points. Returns ------- @@ -170,7 +176,11 @@ def _newton(o, pitch_inv, z1, z2, mask, nufft_eps=1e-10): (0, 2 * jnp.pi / o._NFP), vec=True, eps=nufft_eps, - mask=flatten_mat(jnp.broadcast_to(mask[..., None, :, :], shape), 4), + mask=( + None + if _JF_BUG + else flatten_mat(jnp.broadcast_to(mask[..., None, :, :], shape), 4) + ), ) B, dB_dz, dB_dt = ( B.reshape(3, *shape) @@ -189,19 +199,30 @@ def _newton(o, pitch_inv, z1, z2, mask, nufft_eps=1e-10): return z[..., 0, :, :], z[..., 1, :, :] -@partial(jax.custom_jvp, nondiff_argnames=("num_well", "nufft_eps")) -def regular_points(o, pitch_inv, num_well, nufft_eps): - """Bounce points then newton, with regularized jvp.""" +@partial(jax.custom_jvp, nondiff_argnums=(2,)) +def regular_points(o, pitch_inv, num_well): + """Bounce points then Newton, with regularized jvp. + + Parameters + ---------- + o : Bounce2D + Object instance. + pitch_inv : jnp.ndarray + Shape broadcasts with (num ρ, num α, num pitch). + num_well : int + The usual suspect. + + """ return _newton( o, pitch_inv, - *bounce_points(pitch_inv, o._c["knots"], o._c["B(z)"], num_well), - nufft_eps, + *bounce_points(pitch_inv, o._c["knots"], o._c["B(z)"], num_well, True), + min(o._nufft_eps, 1e-10), ) @regular_points.defjvp -def regular_points_jvp(num_well, nufft_eps, primals, tangents): +def regular_points_jvp(num_well, primals, tangents): """Implicit function theorem with regularization. Regularization used to smooth the discretized system so that it recognizes @@ -210,8 +231,7 @@ def regular_points_jvp(num_well, nufft_eps, primals, tangents): References ---------- - Spectrally accurate, reverse-mode differentiable bounce-averaging algorithm - and its applications. Kaya Unalmis et al. Journal of Plasma Physics. + See supplementary information in DESC/publications/unalmis2025. """ # Cannot mix primals and tangents; see https://github.com/jax-ml/jax/issues/36319. @@ -219,7 +239,7 @@ def regular_points_jvp(num_well, nufft_eps, primals, tangents): o, p = primals do, dp = tangents - z1, z2 = regular_points(o, p, num_well, nufft_eps) + z1, z2 = regular_points(o, p, num_well) shape = (*z1.shape[:-2], 2, *z1.shape[-2:]) @@ -240,6 +260,7 @@ def regular_points_jvp(num_well, nufft_eps, primals, tangents): mask = (z1 < z2)[..., None, :, :] + nufft_eps = min(o._nufft_eps, 1e-10) dB_dz = nufft2d2r( z, t, @@ -250,7 +271,7 @@ def regular_points_jvp(num_well, nufft_eps, primals, tangents): (0, 2 * jnp.pi / o._NFP), vec=True, eps=nufft_eps, - mask=flatten_mat(jnp.broadcast_to(mask, shape), 4), + mask=None if _JF_BUG else flatten_mat(jnp.broadcast_to(mask, shape), 4), ) dB_do = nufft2d2r( z, @@ -258,7 +279,7 @@ def regular_points_jvp(num_well, nufft_eps, primals, tangents): do._c["|B|"].squeeze(-3), (0, 2 * jnp.pi / o._NFP), eps=nufft_eps, - mask=flatten_mat(jnp.broadcast_to(mask, shape), 4), + mask=None if _JF_BUG else flatten_mat(jnp.broadcast_to(mask, shape), 4), ).reshape(shape) dB_dz, dB_dt = ( @@ -688,7 +709,7 @@ def get_alphas(alpha, iota, num_transit, NFP): return alpha + iota * (2 * jnp.pi / NFP) * jnp.arange(num_transit * NFP) -def theta_on_fieldlines(angle, iota, alpha, num_transit, NFP): +def theta_on_fieldlines(angle, iota, alpha, num_transit, NFP, *, X_min=24): """Parameterize θ on field lines α. Parameters @@ -706,6 +727,15 @@ def theta_on_fieldlines(angle, iota, alpha, num_transit, NFP): Number of toroidal transits to follow field line. NFP : int Number of field periods. + X_min : int + See notes section. This parameter should never be changed. + It is included in the function signature for code optics only. + It is the number below which we short-circuit convergence to enforce + continuity by removing a discontinuity which is near machine precision + due to exponential convergence. This is not a hack; it has rigorous + mathematical justification regardless of the size of the removed + discontinuity, and does not bias the output beyond that of more + floating point operations in finite-precision. Returns ------- @@ -744,15 +774,12 @@ def theta_on_fieldlines(angle, iota, alpha, num_transit, NFP): (ϑ, NFP ζ) coordinates, then f(ϑ(α=α₀, ζ), ζ) will sample the approximation to F(α=α₀ ± ε, ζ) with ε → 0 as f converges to F. - This property was mentioned because parameterizing the stream map in (α, ζ) enables - partial summation. However, the small discontinuity due to discretization error - between branch cuts is undesirable as it can give significant error to the singular - integrals whose integration boundary is near a branch cut. If we were using splines - instead of pseudo-spectral methods to interpolate then we would have to account - for this. - """ + X = angle.shape[-2] + Y = truncate_rule(angle.shape[-1]) num_alpha = alpha.size + domain = (0, 2 * jnp.pi / NFP) + # peeling off field lines alpha = get_alphas(alpha, iota, num_transit, NFP) if angle.ndim == 2: @@ -762,16 +789,20 @@ def theta_on_fieldlines(angle, iota, alpha, num_transit, NFP): # (since this avoids modding on more points later and keeps θ bounded). alpha %= 2 * jnp.pi - domain = (0, 2 * jnp.pi / NFP) - Y = truncate_rule(angle.shape[-1]) delta = ( FourierChebyshevSeries(angle, domain, truncate=Y) .compute_cheb(alpha) .swapaxes(0, -3) ) alpha = alpha.swapaxes(0, -2) - delta = delta.at[..., 0].add(alpha) + delta = delta.at[..., 0].add(alpha) # This is now θ = α + δ. assert delta.shape == (*angle.shape[:-2], num_alpha, num_transit * NFP, Y) + + if X < X_min: + # This is needed as our algorithm assumes continuity of |B| along field + # lines when gathering bounce points. This is always true physically. + delta = PiecewiseChebyshevSeries.stitch(delta) + return PiecewiseChebyshevSeries(delta, domain) @@ -1019,7 +1050,19 @@ def broadcast_for_bounce(pitch_inv): def truncate_rule(Y): - """Truncation of Chebyshev series to reduce spectral aliasing.""" + """Truncation of Chebyshev series to reduce spectral aliasing. + + The truncation will remove aliasing error at the shortest wavelengths + where the signal to noise ratio is lowest. + We truncate with a 7/8 rule in the toroidal direction so that the Lebesgue + constant is ~ (4/π²) log(Y) when the data is corrupted by ≤ 10⁻⁸ noise from + the inexact Newton solve. The Lebesgue constant discussed here is the one in + L Mason, J.C. & Handscomb, David C. 2002 Chebyshev Polynomials, chapter 5. + This is useful since we evaluate the series on a much denser grid than the + Newton solve grid; and therefore, prefer all discretization error is from + the error of the projection rather than the interpolation. + + """ return max(1, 7 * Y // 8) @@ -1049,9 +1092,11 @@ def Y_B_rule(grid, spline): Parameters ---------- grid : Grid - Grid. + Tensor-product grid in (ρ, θ, ζ) with uniformly spaced nodes + (θ, ζ) ∈ [0, 2π) × [0, 2π/NFP). spline : bool - Whether |B| will be approximated with cubic spline or Chebyshev. + Whether to use cubic splines to compute initial guess for bounce points + instead of Chebyshev series. Returns ------- @@ -1066,17 +1111,21 @@ def Y_B_rule(grid, spline): return (Y_B * grid.NFP) if spline else Y_B -def num_well_rule(num_transit, NFP, Y_B=None): +def num_well_rule(num_transit, NFP, mins_per_transit=None): """Guess upper bound for number of wells based on spectrum. This should be loose enough that it is equivalent to ``num_well=None``, but more performant. """ num_well = num_transit * (20 + NFP) - return num_well if Y_B is None else min(num_well, num_transit * Y_B) + return ( + num_well + if mins_per_transit is None + else min(num_well, num_transit * mins_per_transit) + ) -def get_vander(grid, Y, Y_B, NFP): +def get_vander_spline(grid, Y, Y_B, NFP): """Builds Vandermonde matrices for objectives.""" Y_trunc = truncate_rule(Y) Y_B, num_z = round_up_rule(Y_B, NFP, grid.num_zeta == 1) diff --git a/desc/integrals/_interp_utils.py b/desc/integrals/_interp_utils.py index 1b47682ec1..b77a32e7ea 100644 --- a/desc/integrals/_interp_utils.py +++ b/desc/integrals/_interp_utils.py @@ -8,13 +8,11 @@ try: from jax_finufft import nufft2, options - can_use_nufft = True except (ImportError, ModuleNotFoundError): warnings.warn( "jax_finufft is not installed. NUFFT functions will not be available.", UserWarning, ) - can_use_nufft = False except Exception as e: error_str = str(e) # This error will probably happen pretty often, we skip it to prevent breaking @@ -35,10 +33,17 @@ f"will not be available: {e}", UserWarning, ) - can_use_nufft = False from desc.backend import jax, jnp +_JF_BUG = True +"""https://github.com/flatironinstitute/jax-finufft/issues/158. + + Wait for jax-finufft to merge + https://github.com/flatironinstitute/jax-finufft/pull/216 + then bump min version and set this to False. +""" + def nufft1d2r(x, f, domain=(0, 2 * jnp.pi), vec=False, eps=1e-6): """Non-uniform 1D real fast Fourier transform of second type. @@ -98,6 +103,7 @@ def nufft2d2r( vec=False, eps=1e-6, mask=None, + fill_value=None, ): """Non-uniform 2D real fast Fourier transform of second type. @@ -137,6 +143,13 @@ def nufft2d2r( of ``(x),(x),(b,f0,f1)->(b,x)``. eps : float Precision requested. Default is ``1e-6``. + mask : jnp.ndarray, optional + Boolean mask of points to interpolate to. Should have same shape as ``x0`` + and ``x1``. This does nothing until the merge of + https://github.com/flatironinstitute/jax-finufft/pull/216. + fill_value : float + Value to pad array where the mask is false. + Default is 0.0. Returns ------- @@ -160,19 +173,16 @@ def nufft2d2r( s = s[..., jnp.newaxis, :] if vec else s f = jnp.fft.ifftshift(f, rfft_axis) - opts = options.Opts(modeord=1) - - # TODO: Delete this if block after - # https://github.com/flatironinstitute/jax-finufft/pull/216 is merged - # and then bump min version requirement. - JF_BUG = True - if JF_BUG: - # https://github.com/flatironinstitute/jax-finufft/pull/159 + if _JF_BUG: opts = options.Opts(modeord=0) f = jnp.fft.fftshift(f, (-2, -1)) return (nufft2(f, x0, x1, iflag=1, eps=eps, opts=opts) * s).real - return (nufft2(f, x0, x1, points_mask=mask, iflag=1, eps=eps, opts=opts) * s).real + opts = options.Opts(modeord=1) + f = (nufft2(f, x0, x1, points_mask=mask, iflag=1, eps=eps, opts=opts) * s).real + if mask is not None and fill_value is not None: + f = jnp.where(mask[..., jnp.newaxis, :] if vec else mask, f, fill_value) + return f # Warning: method must be specified as keyword argument. @@ -261,7 +271,7 @@ def _filter_distinct(r, sentinel, eps): _eps = max(jnp.finfo(jnp.array(1.0).dtype).eps, 2.5e-12) -@partial(jax.custom_jvp, nondiff_argnames=("sort", "sentinel", "eps", "distinct")) +@partial(jax.custom_jvp, nondiff_argnums=(4, 5, 6, 7)) def polyroot_vec( c, k=0.0, @@ -360,8 +370,7 @@ def _polyroot_vec_jvp(sort, sentinel, eps, distinct, primals, tangents): References ---------- - Spectrally accurate, reverse-mode differentiable bounce-averaging algorithm - and its applications. Kaya Unalmis et al. Journal of Plasma Physics. + See supplementary information in DESC/publications/unalmis2025. """ c, k, a_min, a_max = primals diff --git a/desc/integrals/bounce_integral.py b/desc/integrals/bounce_integral.py index d14f209bc9..db0a491c13 100644 --- a/desc/integrals/bounce_integral.py +++ b/desc/integrals/bounce_integral.py @@ -3,7 +3,7 @@ import warnings from abc import ABC, abstractmethod -from equinox import Module +import equinox as eqx from interpax import CubicHermiteSpline, PPoly from interpax_fft import ( FourierChebyshevSeries, @@ -36,7 +36,7 @@ fast_chebyshev, fast_cubic_spline, get_mins, - get_vander, + get_vander_spline, mmt_for_bounce, move, num_well_rule, @@ -46,8 +46,8 @@ theta_on_fieldlines, ) from desc.integrals._interp_utils import ( + _JF_BUG, _eps, - can_use_nufft, interp1d_Hermite_vec, interp1d_vec, nufft2d2r, @@ -70,10 +70,11 @@ flatten_mat, parse_argname_change, setdefault, + warnif, ) -class Bounce(Module, ABC): +class Bounce(eqx.Module, ABC): """Abstract class for bounce integrals.""" @staticmethod @@ -92,9 +93,9 @@ def get_pitch_inv_quad(min_B, max_B, num_pitch, simp=True): If ``True``, then the pitch angles are chosen so that the quadrature over the velocity coordinate of 1/λ is done with Simpson’s 1/3 in the interior completed by an open midpoint scheme near the boundary such - that an accuracy of fourth order is preserved. If the integrand is not - sufficiently smooth then quadrature accuracy reduces to third order. - If ``False``, then an open midpoint scheme is returned. + that an accuracy of fourth order is preserved. + If ``False``, then an open midpoint scheme is returned, which + is only recommended for plotting purposes. Returns ------- @@ -207,7 +208,10 @@ class Bounce2D(Bounce): If the option ``spline`` is ``True``, the bounce points are found with 8th order accuracy in this parameter. If the option ``spline`` is ``False``, then the bounce points are found with spectral accuracy in this parameter. - A reference value for the ``spline`` option is 100. + A reference value for the ``spline=True`` option is + ``grid.NFP*(grid.num_theta+grid.num_zeta)//2``. + A reference value for the ``spline=False`` option is + ``(grid.num_theta+grid.num_zeta)//2``. An error of ε in a bounce point manifests 𝒪(ε¹ᐧ⁵) error in bounce integrals with (v_∥)¹ and @@ -252,14 +256,15 @@ class Bounce2D(Bounce): Lref : float Optional. Reference length scale for normalization. spline : bool - Whether to use cubic splines to compute bounce points instead of - Chebyshev series. Default is ``True``. + Whether to use cubic splines to compute initial guess for bounce points + instead of Chebyshev series. Default is ``True``. check : bool Flag for debugging. Must be false for JAX transformations. """ - required_names = ["B^zeta", "|B|", "iota"] + required_names = ["B^zeta", "|B|", "iota"] # TODO(#2152) + """Required keys in the ``data`` dictionary given to the ``__init__`` method.""" _quad: tuple[jax.Array] _NFP: int @@ -268,6 +273,7 @@ class Bounce2D(Bounce): _modes_t: jax.Array _c: dict[str, jax.Array] _theta: PiecewiseChebyshevSeries + _nufft_eps: float = eqx.field(static=True) def __init__( self, @@ -322,6 +328,8 @@ def __init__( iota, alpha = jnp.atleast_1d(iota, alpha) self._theta = theta_on_fieldlines(angle, iota, alpha, num_transit, grid.NFP) + self._nufft_eps = float(nufft_eps) + if Y_B is None: Y_B = Y_B_rule(grid, spline) if spline: @@ -333,11 +341,15 @@ def __init__( self._modes_t, self._modes_z, self._NFP, - nufft_eps, + self._nufft_eps, vander_t=vander.get("dct spline", None), check=check, ) else: + warnif( + Y_B > grid.num_theta + grid.num_zeta, + msg="Unnecessarily high resolution for Y_B with spline=False.", + ) self._c["B(z)"] = fast_chebyshev( self._theta, self._c["|B|"], @@ -348,8 +360,8 @@ def __init__( ) @staticmethod - def _objective_build(obj, names, eta): - """Default build for bounce integrals objectives. + def _build(obj, names, eta): + """Builds the objective, selecting default values if they were not specified. Examples -------- @@ -384,14 +396,22 @@ def _objective_build(obj, names, eta): Y_B = obj._hyperparam["Y_B"] if Y_B is None: - Y_B = Y_B_rule(obj._grid, spline=True) + Y_B = Y_B_rule(obj._grid, obj._hyperparam["spline"]) obj._hyperparam["Y_B"] = Y_B if obj._hyperparam["num_well"] is None: obj._hyperparam["num_well"] = num_well_rule( - obj._hyperparam["num_transit"], eq.NFP, Y_B + obj._hyperparam["num_transit"], + eq.NFP, + # Due to legacy reasons Y_B is resolution over full transit + # if spline is true. + Y_B if obj._hyperparam["spline"] else (Y_B * eq.NFP), ) - obj._constants["_vander"] = get_vander(obj._grid, Y, Y_B, eq.NFP) + obj._constants["_vander"] = ( + get_vander_spline(obj._grid, Y, Y_B, eq.NFP) + if obj._hyperparam["spline"] + else {} + ) num_quad = obj._hyperparam.pop("num_quad") if eta == 1: @@ -484,7 +504,12 @@ def _defaults(eta, grid, **kwargs): num_transit = kwargs.get("num_transit", 20) Y_B = kwargs.get("Y_B", Y_B_rule(grid, spline)) - num_well = kwargs.get("num_well", num_well_rule(num_transit, grid.NFP, Y_B)) + num_well = kwargs.get( + "num_well", + # Due to legacy reasons Y_B is resolution over full transit + # if spline is true. + num_well_rule(num_transit, grid.NFP, Y_B if spline else (Y_B * grid.NFP)), + ) return ( angle, @@ -495,9 +520,9 @@ def _defaults(eta, grid, **kwargs): num_pitch, pitch_batch_size, surf_batch_size, - quad, nufft_eps, spline, + quad, vander, ) @@ -630,7 +655,7 @@ def fourier(f): @staticmethod def compute_theta( eq, - X=16, + X=32, Y=32, rho=jnp.array([1.0]), iota=None, @@ -649,7 +674,7 @@ def compute_theta( @staticmethod def angle( eq, - X=16, + X=32, Y=32, rho=jnp.array([1.0]), iota=None, @@ -668,10 +693,12 @@ def angle( X : int Poloidal Fourier grid resolution to interpolate the angle. Preferably rounded down to power of 2. + Default is 32. Y : int Toroidal Chebyshev grid resolution over a single field period to interpolate the angle. Preferably rounded down to power of 2. + Default is 32. rho : float or jnp.ndarray Shape (num ρ, ). Flux surfaces labels in [0, 1] on which to compute. @@ -795,22 +822,20 @@ def points(self, pitch_inv, num_well=None): num_well = num_well_rule(self._theta.X // self._NFP, self._NFP) if isinstance(self._c["B(z)"], PiecewiseChebyshevSeries): + # Skip Newton update since these points are exponentially accurate. z1, z2 = self._c["B(z)"].intersect1d( - self._swap_pitch(pitch_inv), _eps, num_well + self._swap_pitch(pitch_inv), num_intersect=num_well, eps=_eps ) z1 = move(z1) z2 = move(z2) return z1, z2 pitch_inv = broadcast_for_bounce(pitch_inv) - if can_use_nufft: - return regular_points(self, pitch_inv, num_well, 1e-10) - - # Newton update has only been implemented for nuffts; contributions welcome. - z1, z2, _ = bounce_points( - pitch_inv, self._c["knots"], self._c["B(z)"], num_well - ) - return z1, z2 + if self._nufft_eps < 1e-14: + # FIXME: Newton update has only been implemented for nuffts; contributions + # welcome : copy logic in desc/equilibrium/coords._map_poloidal_coordinates. + return bounce_points(pitch_inv, self._c["knots"], self._c["B(z)"], num_well) + return regular_points(self, pitch_inv, num_well) def check_points(self, points, pitch_inv, *, plot=True, **kwargs): """Check that bounce points are computed correctly. @@ -832,7 +857,6 @@ def check_points(self, points, pitch_inv, *, plot=True, **kwargs): Whether to plot the field lines and bounce points of the given pitch angles. kwargs : dict Keyword arguments into - ``desc/integrals/_bounce_utils.py::PiecewiseChebyshevSeries.plot1d`` or ``desc/integrals/_bounce_utils.py::plot_ppoly``. Returns @@ -954,7 +978,6 @@ def integrate( if points is None: points = self.points(pitch_inv, num_well) - z1, z2 = points pitch = 1 / pitch_inv # to broadcast with (..., num pitch, num well, num quad) @@ -963,27 +986,17 @@ def integrate( elif jnp.ndim(pitch) > 1: pitch = pitch[:, None, :, None, None] - z = bijection_from_disc(x, z1[..., None], z2[..., None]) - if nufft_eps < 1e-14: - data = self._nummt(z, data, low_ram) + data = self._nummt(x, *points, data, low_ram) else: - data = self._nufft( - z, - data, - nufft_eps, - low_ram, - mask=flatten_mat(jnp.broadcast_to((z1 < z2)[..., None], z.shape), 4), - sentinel=0.5 * jnp.min(pitch_inv), - ) + data = self._nufft(x, *points, data, low_ram, nufft_eps, pitch_inv) data["|e_zeta|r,a|"] = data["|B|"] / jnp.abs(data["B^zeta"]) - data["zeta"] = z # Strictly increasing ζ knots enforces dζ > 0. # To retain dℓ = |B|/(B⋅∇ζ) dζ > 0 after fixing dζ > 0, we require # B⋅∇ζ > 0. This is equivalent to changing the sign of ∇ζ # or (∂ℓ/∂ζ)|ρ,a. Recall dζ = ∇ζ⋅dR ⇔ 1 = ∇ζ⋅(e_ζ|ρ,a). - cov = grad_bijection_from_disc(z1, z2) + cov = grad_bijection_from_disc(*points) result = [ (f(data, data["|B|"], pitch) * data["|e_zeta|r,a|"]).dot(w) * cov for f in integrand @@ -991,7 +1004,7 @@ def integrate( if check: check_interp( - z, + data["zeta"], jnp.reciprocal(data["|e_zeta|r,a|"]), data["|B|"], [data[k] for k in data if k not in ("zeta", "|e_zeta|r,a|", "|B|")], @@ -1001,25 +1014,30 @@ def integrate( return result[0] if len(result) == 1 else result - def _nufft(self, z, data, eps, low_ram, mask, sentinel): - shape = z.shape - z = flatten_mat(z, 3) + def _nufft(self, x, z1, z2, data, low_ram, eps, pitch_inv): + shape = (*z1.shape, x.size) + + z = flatten_mat(bijection_from_disc(x, z1[..., None], z2[..., None]), 3) t = flatten_mat(self._theta.eval1d(z, loop=low_ram)) z = flatten_mat(z) - c = jnp.where( - mask[:, None] if mask.ndim == 2 else mask, - nufft2d2r( - z, - t, - jnp.concatenate( - [*data.values(), self._c["B^zeta"], self._c["|B|"]], -3 - ), - (0, 2 * jnp.pi / self._NFP), - vec=True, - eps=eps, - mask=mask, - ), - sentinel, # replace junk zeros to avoid nans + # t and z have shape (num rho surfaces, num points on each surface) + # or just ( num points on each surface). + + if _JF_BUG: + mask = fill_value = None + else: + mask = flatten_mat(jnp.broadcast_to((z1 < z2)[..., None], shape), 4) + fill_value = 0.5 * jnp.min(pitch_inv) + + c = nufft2d2r( + z, + t, + jnp.concatenate([*data.values(), self._c["B^zeta"], self._c["|B|"]], -3), + (0, 2 * jnp.pi / self._NFP), + vec=True, + eps=eps, + mask=mask, + fill_value=fill_value, ) c = ( c.reshape(len(data) + 2, *shape) @@ -1027,18 +1045,45 @@ def _nufft(self, z, data, eps, low_ram, mask, sentinel): # reshape before swap to avoid memory copy else c.reshape(shape[0], len(data) + 2, *shape[1:]).swapaxes(0, 1) ) - return dict(zip([*data.keys(), "B^zeta", "|B|"], c)) - def _nummt(self, z, data, low_ram): - t = self._theta.eval1d(flatten_mat(z, 3), loop=low_ram).reshape(z.shape) - t = jnp.exp(1j * self._modes_t * t[..., None]) - z = jnp.exp(1j * self._modes_z * z[..., None]) + data = dict(zip([*data.keys(), "B^zeta", "|B|"], c)) + data["zeta"] = z.reshape(shape) + return data + + def _nummt(self, x, z1, z2, data, low_ram): + shape = (*z1.shape, x.size) + + zeta = bijection_from_disc(x, z1[..., None], z2[..., None]) + z = flatten_mat(zeta, 3) + t = self._theta.eval1d(z, loop=low_ram).reshape(*shape, 1) + + if isinstance(self._c["B(z)"], PiecewiseChebyshevSeries): + # Using the same |B| that gave bounce points increases correlation + # in discretization error in an open neighboorhood around the + # current point in the optimization landscape, and hence removes + # noise from optimization derivatives. For example, the difference + # between the auto derivative and a 4 point finite difference + # stencil at the optimal step size is reduced from 10³ to 10⁻⁶ for + # Γ_c (which has the singular weight 1/v_∥). + # Also uses less memory due to the dimension reduction. + B = self._c["B(z)"].eval1d(z, loop=low_ram).reshape(shape) + + z = zeta[..., None] + z = jnp.exp(1j * self._modes_z * z) + t = jnp.exp(1j * self._modes_t * t) data = {name: mmt_for_bounce(z, t, c) for name, c in data.items()} data["B^zeta"] = mmt_for_bounce(z, t, self._c["B^zeta"]) - data["|B|"] = mmt_for_bounce(z, t, self._c["|B|"]) + if isinstance(self._c["B(z)"], PiecewiseChebyshevSeries): + data["|B|"] = B + else: + data["|B|"] = mmt_for_bounce(z, t, self._c["|B|"]) + data["zeta"] = zeta + return data - def interp_to_argmin(self, f, points, *, nufft_eps=1e-6, is_fourier=False): + def interp_to_argmin( + self, f, points, *, nufft_eps=1e-6, is_fourier=False, **kwargs + ): """Interpolate ``f`` to the deepest point pⱼ in magnetic well j. Parameters @@ -1069,24 +1114,19 @@ def interp_to_argmin(self, f, points, *, nufft_eps=1e-6, is_fourier=False): ``f`` interpolated to the deepest point between ``points``. """ - errorif( - isinstance(self._c["B(z)"], PiecewiseChebyshevSeries), - NotImplementedError, - "Must choose Bounce2D(spline=True) for this feature.", - ) if not is_fourier: f = Bounce2D.fourier(f) num_transit = self._theta.X // self._NFP - K_z = self._num_z // 2 * self._NFP - mins, B_mins = get_mins( - self._c["knots"], - self._c["B(z)"], - num_mins=num_transit * max(K_z, 5), # liberal heuristic - # We set fill value to 0 since we chose our coordinates - # such that all bounce points are at ζ >= 0; and therefore, - # junk values in B_mins cannot be selected in argmin. - fill_value=0.0, + K_z = max(self._num_z // 2 * self._NFP, self._num_t // 2, 5) # liberal + num_mins = kwargs.get("num_mins", num_transit * K_z) + # We set fill value to 0 since we chose our coordinates + # such that all bounce points are at ζ >= 0; and therefore, + # junk values in B_mins cannot be selected in argmin. + mins, B_mins = ( + self._c["B(z)"].extrema1d(1, num_mins, fill_value=0.0, eps=_eps) + if isinstance(self._c["B(z)"], PiecewiseChebyshevSeries) + else get_mins(self._c["knots"], self._c["B(z)"], num_mins, fill_value=0.0) ) t = self._theta.eval1d(mins) @@ -1200,9 +1240,6 @@ def plot(self, l, m, pitch_inv=None, **kwargs): Shape (num pitch, ). Optional, 1/λ values whose corresponding bounce points on the field line specified by Clebsch coordinate ρ(l), α(m) will be plotted. - kwargs - Keyword arguments into - ``desc/integrals/_bounce_utils.py::PiecewiseChebyshevSeries.plot1d``. Returns ------- @@ -1226,7 +1263,7 @@ def plot(self, l, m, pitch_inv=None, **kwargs): B = B[m] B = PiecewiseChebyshevSeries(B, domain) if pitch_inv is not None: - kwargs["z1"], kwargs["z2"] = B.intersect1d(pitch_inv, _eps) + kwargs["z1"], kwargs["z2"] = B.intersect1d(pitch_inv, eps=_eps) kwargs["k"] = pitch_inv return B.plot1d(B.cheb, **kwargs) @@ -1235,9 +1272,7 @@ def plot(self, l, m, pitch_inv=None, **kwargs): if B.ndim == 3: B = B[m] if pitch_inv is not None: - kwargs["z1"], kwargs["z2"], _ = bounce_points( - pitch_inv, self._c["knots"], B - ) + kwargs["z1"], kwargs["z2"] = bounce_points(pitch_inv, self._c["knots"], B) kwargs["k"] = pitch_inv return plot_ppoly(PPoly(B.T, self._c["knots"]), **kwargs) @@ -1253,9 +1288,6 @@ def plot_theta(self, l, m, **kwargs): m : int Index into the ``alpha`` array supplied to make this object. The alpha value corresponds to ``alpha[m]``. - kwargs - Keyword arguments into - ``desc/integrals/_bounce_utils.py::PiecewiseChebyshevSeries.plot1d``. Returns ------- @@ -1449,6 +1481,7 @@ class Bounce1D(Bounce): """ required_names = ["B^zeta", "B^zeta_z|r,a", "|B|", "|B|_z|r,a"] + """Required keys in the ``data`` dictionary given to the ``__init__`` method.""" _quad: tuple[jax.Array] _data: dict[str, jax.Array] @@ -1633,10 +1666,9 @@ def points(self, pitch_inv, num_well=None): line and pitch, is padded with zero. """ - z1, z2, _ = bounce_points( + return bounce_points( broadcast_for_bounce(pitch_inv), self._zeta, self._B, num_well ) - return z1, z2 def check_points(self, points, pitch_inv, *, plot=True, **kwargs): """Check that bounce points are computed correctly. @@ -1860,7 +1892,7 @@ def plot(self, l, m, pitch_inv=None, **kwargs): jnp.ndim(pitch_inv) > 1, msg=f"Got pitch_inv.ndim={jnp.ndim(pitch_inv)}, but expected 1.", ) - kwargs["z1"], kwargs["z2"], _ = bounce_points(pitch_inv, self._zeta, B) + kwargs["z1"], kwargs["z2"] = bounce_points(pitch_inv, self._zeta, B) kwargs["k"] = pitch_inv fig, ax = plot_ppoly( PPoly(B.T, self._zeta), **set_default_plot_kwargs(kwargs, l, m) diff --git a/desc/objectives/_fast_ion.py b/desc/objectives/_fast_ion.py index afc92f3585..94983756b8 100644 --- a/desc/objectives/_fast_ion.py +++ b/desc/objectives/_fast_ion.py @@ -8,15 +8,14 @@ from desc.grid import LinearGrid from desc.integrals._bounce_utils import Y_B_rule, num_well_rule from desc.integrals.bounce_integral import Bounce2D -from desc.utils import parse_argname_change, setdefault, warnif +from desc.utils import setdefault, warnif from ..integrals.quad_utils import ( automorphism_sin, get_quadrature, grad_automorphism_sin, ) -from ._neoclassical import _bounce_overwrite -from .objective_funs import _Objective, collect_docs +from .objective_funs import _Objective, collect_docs, doc_bounce from .utils import _parse_callable_target_bounds @@ -50,87 +49,15 @@ class GammaC(_Objective): [3] Spectrally accurate, reverse-mode differentiable bounce-averaging algorithm and its applications. Kaya Unalmis et al. Journal of Plasma Physics. - Notes - ----- - Performance will improve significantly by resolving these GitHub issues. - * ``1206`` Upsample data above midplane to full grid assuming stellarator symmetry - * ``1034`` Optimizers/objectives with auxiliary output + """ - Parameters - ---------- - eq : Equilibrium - ``Equilibrium`` to be optimized. - grid : Grid - Tensor-product grid in (ρ, θ, ζ) with uniformly spaced nodes - (θ, ζ) ∈ [0, 2π) × [0, 2π/NFP). - Number of poloidal and toroidal nodes preferably rounded down to powers of two. - Determines the flux surfaces to compute on and resolution of FFTs. - Default grid samples the boundary surface at ρ=1. - X : int - Poloidal Fourier grid resolution to interpolate the angle. - Preferably rounded down to power of 2. - Y : int - Toroidal Chebyshev grid resolution over a single field period - to interpolate the angle. - Preferably rounded down to power of 2. - Y_B : int - Desired resolution for algorithm to compute bounce points. - The bounce points are found with 8th order accuracy in this parameter. - A reference value is 100. - - An error of ε in a bounce point manifests - 𝒪(ε¹ᐧ⁵) error in bounce integrals with (v_∥)¹ and - 𝒪(ε⁰ᐧ⁵) error in bounce integrals with (v_∥)⁻¹. - alpha : jnp.ndarray - Shape (num alpha, ). - Starting field line poloidal labels. - Default is single field line. To compute a surface average - on a rational surface, it is necessary to average over multiple - field lines until the surface is covered sufficiently. - num_transit : int - Number of toroidal transits to follow field line. - In an axisymmetric device, field line integration over a single poloidal - transit is sufficient to capture a surface average. For a 3D - configuration, more transits will approximate surface averages on an - irrational magnetic surface better, with diminishing returns. - num_well : int - Maximum number of wells to detect for each pitch and field line. - Giving ``-1`` will detect all wells but due to current limitations in - JAX this will have worse performance. - Specifying a number that tightly upper bounds the number of wells will - increase performance. In general, an upper bound on the number of wells - per toroidal transit is ``Aι+C`` where ``A``, ``C`` are the poloidal and - toroidal Fourier resolution of B, respectively, in straight-field line - PEST coordinates, and ι is the rotational transform normalized by 2π. - A tighter upper bound than ``num_well=(Aι+C)*num_transit`` is preferable. - The ``check_points`` or ``plot`` methods in ``desc.integrals.Bounce2D`` - are useful to select a reasonable value. - - This is the most important parameter to specify for performance. - num_quad : int - Resolution for quadrature of bounce integrals. Default is 32. - num_pitch : int - Resolution for quadrature over velocity coordinate. Default is 65. - pitch_batch_size : int - Number of pitch values with which to compute simultaneously. - If given ``None``, then ``pitch_batch_size`` is ``num_pitch``. - Default is ``num_pitch``. - surf_batch_size : int - Number of flux surfaces with which to compute simultaneously. - If given ``None``, then ``surf_batch_size`` is ``grid.num_rho``. - Default is ``1``. Only consider increasing if ``pitch_batch_size`` is ``None``. - nufft_eps : float - Precision requested for interpolation with non-uniform fast Fourier - transform (NUFFT). If less than ``1e-14`` then NUFFT will not be used. - use_bounce1d : bool - Set to ``True`` to use ``Bounce1D`` instead of ``Bounce2D``, - basically replacing some pseudo-spectral methods with local splines. - This can be efficient if ``num_transit`` and ``alpha.size`` are small, - depending on hardware and hardware features used by the JIT compiler. - If ``True``, then parameters ``X``, ``Y``, ``nufft_eps`` are ignored. + __doc__ = ( + __doc__.rstrip() + + doc_bounce + + """ Nemov : bool Whether to use the Γ_c as defined by Nemov et al. or Velasco et al. - Default is Nemov. Set to ``False`` to use Velascos's. + Default is Nemov. Set to ``False`` to use Velasco's. Nemov's Γ_c converges to a finite nonzero value in the infinity limit of the number of toroidal transits. Velasco's expression has a secular @@ -139,15 +66,13 @@ class GammaC(_Objective): integrals. At finite resolution, an optimization using Velasco's metric may need to be evaluated by measuring decrease in Γ_c at a fixed number of toroidal transits. - - """ - - __doc__ = __doc__.rstrip() + collect_docs( - target_default="``target=0``.", - bounds_default="``target=0``.", - normalize_detail=" Note: Has no effect for this objective.", - normalize_target_detail=" Note: Has no effect for this objective.", - overwrite=_bounce_overwrite, + """.rstrip() + + collect_docs( + target_default="``target=0``.", + bounds_default="``target=0``.", + normalize_detail=" Note: Has no effect for this objective.", + normalize_target_detail=" Note: Has no effect for this objective.", + ) ) _static_attrs = _Objective._static_attrs + [ @@ -175,7 +100,7 @@ def __init__( jac_chunk_size=None, name="Gamma_c", grid=None, - X=16, + X=32, Y=32, Y_B=None, alpha=jnp.array([0.0]), @@ -186,6 +111,7 @@ def __init__( pitch_batch_size=None, surf_batch_size=1, nufft_eps=1e-7, + spline=True, use_bounce1d=False, Nemov=True, **kwargs, @@ -204,9 +130,7 @@ def __init__( if target is None and bounds is None: target = 0.0 - self._use_bounce1d = parse_argname_change( - use_bounce1d, kwargs, "spline", "use_bounce1d" - ) + self._use_bounce1d = use_bounce1d self._grid = grid self._constants = {"quad_weights": 1.0, "alpha": alpha} self._hyperparam = { @@ -220,12 +144,14 @@ def __init__( "pitch_batch_size": pitch_batch_size, "surf_batch_size": surf_batch_size, "nufft_eps": nufft_eps, + "spline": spline, } if use_bounce1d: self._hyperparam.pop("X") self._hyperparam.pop("Y") self._hyperparam.pop("pitch_batch_size") self._hyperparam.pop("nufft_eps") + self._hyperparam.pop("spline") self._key = "Gamma_c" if Nemov else "Gamma_c Velasco" @@ -256,7 +182,7 @@ def build(self, use_jit=True, verbose=1): if self._use_bounce1d: return self._build_bounce1d(use_jit, verbose) - Bounce2D._objective_build( + Bounce2D._build( self, names=self._key, eta={"Gamma_c": -2, "Gamma_c Velasco": -1}[self._key] ) super().build(use_jit=use_jit, verbose=verbose) diff --git a/desc/objectives/_neoclassical.py b/desc/objectives/_neoclassical.py index 3f45d92ecb..f4c3cf9c62 100644 --- a/desc/objectives/_neoclassical.py +++ b/desc/objectives/_neoclassical.py @@ -6,25 +6,12 @@ from desc.grid import LinearGrid from desc.integrals._bounce_utils import Y_B_rule, num_well_rule from desc.integrals.bounce_integral import Bounce2D -from desc.utils import parse_argname_change, setdefault, warnif +from desc.utils import setdefault, warnif from ..integrals.quad_utils import chebgauss2 -from .objective_funs import _Objective, collect_docs +from .objective_funs import _Objective, collect_docs, doc_bounce from .utils import _parse_callable_target_bounds -_bounce_overwrite = {"deriv_mode": """ - 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. - - The default mode of ``auto`` will likely choose ``rev`` for this objective. - In ``rev`` mode, reducing the value of the parameter ``pitch_batch_size`` does - not reduce memory consumption, so it is recommended to retain the default unless - you have explicitly requested to use ``fwd`` mode. - """} - class EffectiveRipple(_Objective): """Proxy for neoclassical transport in the banana regime. @@ -46,94 +33,17 @@ class EffectiveRipple(_Objective): [2] Spectrally accurate, reverse-mode differentiable bounce-averaging algorithm and its applications. Kaya Unalmis et al. Journal of Plasma Physics. - Notes - ----- - Performance will improve significantly by resolving these GitHub issues. - * ``1206`` Upsample data above midplane to full grid assuming stellarator symmetry - * ``1034`` Optimizers/objectives with auxiliary output - - - Parameters - ---------- - eq : Equilibrium - ``Equilibrium`` to be optimized. - grid : Grid - Tensor-product grid in (ρ, θ, ζ) with uniformly spaced nodes - (θ, ζ) ∈ [0, 2π) × [0, 2π/NFP). - Number of poloidal and toroidal nodes preferably rounded down to powers of two. - Determines the flux surfaces to compute on and resolution of FFTs. - Default grid samples the boundary surface at ρ=1. - X : int - Poloidal Fourier grid resolution to interpolate the angle. - Preferably rounded down to power of 2. - Y : int - Toroidal Chebyshev grid resolution over a single field period - to interpolate the angle. - Preferably rounded down to power of 2. - Y_B : int - Desired resolution for algorithm to compute bounce points. - The bounce points are found with 8th order accuracy in this parameter. - A reference value is 100. - - An error of ε in a bounce point manifests - 𝒪(ε¹ᐧ⁵) error in bounce integrals with (v_∥)¹ and - 𝒪(ε⁰ᐧ⁵) error in bounce integrals with (v_∥)⁻¹. - alpha : jnp.ndarray - Shape (num alpha, ). - Starting field line poloidal labels. - Default is single field line. To compute a surface average - on a rational surface, it is necessary to average over multiple - field lines until the surface is covered sufficiently. - num_transit : int - Number of toroidal transits to follow field line. - In an axisymmetric device, field line integration over a single poloidal - transit is sufficient to capture a surface average. For a 3D - configuration, more transits will approximate surface averages on an - irrational magnetic surface better, with diminishing returns. - num_well : int - Maximum number of wells to detect for each pitch and field line. - Giving ``-1`` will detect all wells but due to current limitations in - JAX this will have worse performance. - Specifying a number that tightly upper bounds the number of wells will - increase performance. In general, an upper bound on the number of wells - per toroidal transit is ``Aι+C`` where ``A``, ``C`` are the poloidal and - toroidal Fourier resolution of B, respectively, in straight-field line - PEST coordinates, and ι is the rotational transform normalized by 2π. - A tighter upper bound than ``num_well=(Aι+C)*num_transit`` is preferable. - The ``check_points`` or ``plot`` methods in ``desc.integrals.Bounce2D`` - are useful to select a reasonable value. - - This is the most important parameter to specify for performance. - num_quad : int - Resolution for quadrature of bounce integrals. Default is 32. - num_pitch : int - Resolution for quadrature over velocity coordinate. Default is 51. - pitch_batch_size : int - Number of pitch values with which to compute simultaneously. - If given ``None``, then ``pitch_batch_size`` is ``num_pitch``. - Default is ``num_pitch``. - surf_batch_size : int - Number of flux surfaces with which to compute simultaneously. - If given ``None``, then ``surf_batch_size`` is ``grid.num_rho``. - Default is ``1``. Only consider increasing if ``pitch_batch_size`` is ``None``. - nufft_eps : float - Precision requested for interpolation with non-uniform fast Fourier - transform (NUFFT). If less than ``1e-14`` then NUFFT will not be used. - use_bounce1d : bool - Set to ``True`` to use ``Bounce1D`` instead of ``Bounce2D``, - basically replacing some pseudo-spectral methods with local splines. - This can be efficient if ``num_transit`` and ``alpha.size`` are small, - depending on hardware and hardware features used by the JIT compiler. - If ``True``, then parameters ``X``, ``Y``, ``nufft_eps`` are ignored. - """ - __doc__ = __doc__.rstrip() + collect_docs( - target_default="``target=0``.", - bounds_default="``target=0``.", - normalize_detail=" Note: Has no effect for this objective.", - normalize_target_detail=" Note: Has no effect for this objective.", - overwrite=_bounce_overwrite, + __doc__ = ( + __doc__.rstrip() + + doc_bounce + + collect_docs( + target_default="``target=0``.", + bounds_default="``target=0``.", + normalize_detail=" Note: Has no effect for this objective.", + normalize_target_detail=" Note: Has no effect for this objective.", + ) ) _static_attrs = _Objective._static_attrs + [ @@ -160,7 +70,7 @@ def __init__( jac_chunk_size=None, name="Effective ripple", grid=None, - X=16, + X=32, Y=32, Y_B=None, alpha=jnp.array([0.0]), @@ -171,6 +81,7 @@ def __init__( pitch_batch_size=None, surf_batch_size=1, nufft_eps=1e-6, + spline=True, use_bounce1d=False, **kwargs, ): @@ -188,9 +99,7 @@ def __init__( if target is None and bounds is None: target = 0.0 - self._use_bounce1d = parse_argname_change( - use_bounce1d, kwargs, "spline", "use_bounce1d" - ) + self._use_bounce1d = use_bounce1d self._grid = grid self._constants = {"quad_weights": 1.0, "alpha": alpha} self._hyperparam = { @@ -204,12 +113,14 @@ def __init__( "pitch_batch_size": pitch_batch_size, "surf_batch_size": surf_batch_size, "nufft_eps": nufft_eps, + "spline": spline, } if use_bounce1d: self._hyperparam.pop("X") self._hyperparam.pop("Y") self._hyperparam.pop("pitch_batch_size") self._hyperparam.pop("nufft_eps") + self._hyperparam.pop("spline") super().__init__( things=eq, @@ -238,7 +149,7 @@ def build(self, use_jit=True, verbose=1): if self._use_bounce1d: return self._build_bounce1d(use_jit, verbose) - Bounce2D._objective_build(self, names="effective ripple", eta=1) + Bounce2D._build(self, names="effective ripple", eta=1) super().build(use_jit=use_jit, verbose=verbose) def compute(self, params, constants=None): diff --git a/desc/objectives/objective_funs.py b/desc/objectives/objective_funs.py index 4e4d0549ce..d35efaad91 100644 --- a/desc/objectives/objective_funs.py +++ b/desc/objectives/objective_funs.py @@ -131,6 +131,95 @@ "jac_chunk_size": doc_jac_chunk_size, } +doc_bounce = """ + Notes + ----- + Consider using an optimizer that uses a scalar output loss function + to improve performance before reducing ``jac_chunk_size``. + + Developer notes: Performance will improve significantly by resolving GitHub issues: + * ``1206`` Upsample data above midplane to full grid assuming stellarator symmetry + * ``1034`` Optimizers/objectives with auxiliary output + * ``2168`` Sparse cotangent pullbacks + + Parameters + ---------- + eq : Equilibrium + ``Equilibrium`` to be optimized. + grid : Grid + Tensor-product grid in (ρ, θ, ζ) with uniformly spaced nodes + (θ, ζ) ∈ [0, 2π) × [0, 2π/NFP). + Number of poloidal and toroidal nodes preferably rounded down to powers of two. + Determines the flux surfaces to compute on and resolution of FFTs. + Default grid samples the boundary surface at ρ=1. + X : int + Poloidal Fourier grid resolution to interpolate the angle. + Preferably rounded down to power of 2. + Default is 32. + Y : int + Toroidal Chebyshev grid resolution over a single field period + to interpolate the angle. + Preferably rounded down to power of 2. + Default is 32. + Y_B : int + Desired resolution for algorithm to compute bounce points. + If the option ``spline`` is ``True``, the bounce points are found with + 8th order accuracy in this parameter. If the option ``spline`` is ``False``, + then the bounce points are found with spectral accuracy in this parameter. + A reference value for the ``spline=True`` option is + ``grid.NFP*(grid.num_theta+grid.num_zeta)//2``. + A reference value for the ``spline=False`` option is + ``(grid.num_theta+grid.num_zeta)//2``. + + An error of ε in a bounce point manifests + 𝒪(ε¹ᐧ⁵) error in bounce integrals with (v_∥)¹ and + 𝒪(ε⁰ᐧ⁵) error in bounce integrals with (v_∥)⁻¹. + alpha : jnp.ndarray + Shape (num alpha, ). + Starting field line poloidal labels. + Default is single field line. To compute a surface average + on a rational surface, it is necessary to average over multiple + field lines until the surface is covered sufficiently. + num_transit : int + Number of toroidal transits to follow field line. + In an axisymmetric device, field line integration over a single poloidal + transit is sufficient to capture a surface average. For a 3D + configuration, more transits will approximate surface averages on an + irrational magnetic surface better, with diminishing returns. + num_well : int + Maximum number of wells to detect for each pitch and field line. + Giving ``-1`` will detect all wells but due to current limitations in + JAX this will have worse performance. + Specifying a number that tightly upper bounds the number of wells will + increase performance. In general, an upper bound on the number of wells + per toroidal transit is ``Aι+C`` where ``A``, ``C`` are the poloidal and + toroidal Fourier resolution of B, respectively, in straight-field line + PEST coordinates, and ι is the rotational transform normalized by 2π. + A tighter upper bound than ``num_well=(Aι+C)*num_transit`` is preferable. + The ``check_points`` or ``plot`` methods in ``desc.integrals.Bounce2D`` + are useful to select a reasonable value. + + This is the most important parameter to specify for performance. + num_quad : int + Resolution for quadrature of bounce integrals. Default is 32. + num_pitch : int + Resolution for quadrature over velocity coordinate. + pitch_batch_size : int + Number of pitch values with which to compute simultaneously. + If given ``None``, then ``pitch_batch_size`` is ``num_pitch``. + Default is ``num_pitch``. + surf_batch_size : int + Number of flux surfaces with which to compute simultaneously. + If given ``None``, then ``surf_batch_size`` is ``grid.num_rho``. + Default is ``1``. Only consider increasing if ``pitch_batch_size`` is ``None``. + nufft_eps : float + Precision requested for interpolation with non-uniform fast Fourier + transform (NUFFT). If less than ``1e-14`` then NUFFT will not be used. + spline : bool + Whether to use cubic splines to compute initial guess for bounce points + instead of Chebyshev series. Default is ``True``. + """.rstrip() + # Note: If we ever switch to Python 3.13 for building the docs, there will probably # be some errors since 3.13 changed how tabs are handled in docstrings. This can be diff --git a/publications/unalmis2025/spectral_reverse_bounce.pdf b/publications/unalmis2025/spectral_reverse_bounce.pdf index 7296ac7e52..5b95e28136 100644 Binary files a/publications/unalmis2025/spectral_reverse_bounce.pdf and b/publications/unalmis2025/spectral_reverse_bounce.pdf differ diff --git a/tests/baseline/test_Gamma_c_Nemov_1D.png b/tests/baseline/test_Gamma_c_Nemov_1D.png deleted file mode 100644 index 8c6e705152..0000000000 Binary files a/tests/baseline/test_Gamma_c_Nemov_1D.png and /dev/null differ diff --git a/tests/baseline/test_Gamma_c_Nemov_2D_0.png b/tests/baseline/test_Gamma_c_Nemov_2D_0.png deleted file mode 100644 index 06d6cb0853..0000000000 Binary files a/tests/baseline/test_Gamma_c_Nemov_2D_0.png and /dev/null differ diff --git a/tests/baseline/test_Gamma_c_Nemov_2D_1e-07.png b/tests/baseline/test_Gamma_c_Nemov_2D_1e-07.png deleted file mode 100644 index 474abc8a31..0000000000 Binary files a/tests/baseline/test_Gamma_c_Nemov_2D_1e-07.png and /dev/null differ diff --git a/tests/baseline/test_Gamma_c_Velasco_1D.png b/tests/baseline/test_Gamma_c_Velasco_1D.png deleted file mode 100644 index 144d3b6b73..0000000000 Binary files a/tests/baseline/test_Gamma_c_Velasco_1D.png and /dev/null differ diff --git a/tests/baseline/test_Gamma_c_Velasco_2D_0.png b/tests/baseline/test_Gamma_c_Velasco_2D_0.png deleted file mode 100644 index 9e760e260d..0000000000 Binary files a/tests/baseline/test_Gamma_c_Velasco_2D_0.png and /dev/null differ diff --git a/tests/baseline/test_Gamma_c_Velasco_2D_1e-07.png b/tests/baseline/test_Gamma_c_Velasco_2D_1e-07.png deleted file mode 100644 index 57c6dc8c18..0000000000 Binary files a/tests/baseline/test_Gamma_c_Velasco_2D_1e-07.png and /dev/null differ diff --git a/tests/baseline/test_effective_ripple_1D.png b/tests/baseline/test_effective_ripple_1D.png deleted file mode 100644 index af4651bb1b..0000000000 Binary files a/tests/baseline/test_effective_ripple_1D.png and /dev/null differ diff --git a/tests/baseline/test_effective_ripple_2D_0.png b/tests/baseline/test_effective_ripple_2D_0.png deleted file mode 100644 index 44fab0cbca..0000000000 Binary files a/tests/baseline/test_effective_ripple_2D_0.png and /dev/null differ diff --git a/tests/baseline/test_effective_ripple_2D_1e-06.png b/tests/baseline/test_effective_ripple_2D_1e-06.png deleted file mode 100644 index 318f720693..0000000000 Binary files a/tests/baseline/test_effective_ripple_2D_1e-06.png and /dev/null differ diff --git a/tests/inputs/master_compute_data_rpz.pkl b/tests/inputs/master_compute_data_rpz.pkl index 0eb15360df..ebdb276f09 100644 Binary files a/tests/inputs/master_compute_data_rpz.pkl and b/tests/inputs/master_compute_data_rpz.pkl differ diff --git a/tests/test_compute_everything.py b/tests/test_compute_everything.py index 5429257ff7..9bc8900e0b 100644 --- a/tests/test_compute_everything.py +++ b/tests/test_compute_everything.py @@ -24,13 +24,14 @@ FourierXYZCurve, ZernikeRZToroidalSection, ) -from desc.grid import LinearGrid +from desc.grid import Grid, LinearGrid +from desc.integrals import Bounce2D from desc.magnetic_fields import ( CurrentPotentialField, FourierCurrentPotentialField, OmnigenousField, ) -from desc.utils import ResolutionWarning, errorif, xyz2rpz, xyz2rpz_vec +from desc.utils import ResolutionWarning, apply, errorif, xyz2rpz, xyz2rpz_vec def _compare_against_master( @@ -259,6 +260,11 @@ def need_special(name): f"Parameterization: {p}. Can't compute " + f"{names - this_branch_data_rpz[p].keys()}." ) + + # now we get the special grid data + this_branch_data_rpz[p].update(fft_grid_data(p)) + this_branch_data_rpz[p].update(raz_grid_data(p)) + # compare data against master branch error_rpz, update_master_data_rpz = _compare_against_master( p, @@ -297,3 +303,75 @@ def need_special(name): pickle.dump(this_branch_data_rpz, file) assert not error_rpz + + +def fft_grid_data(p): + """Compute fft grid quantities.""" + if p != "desc.equilibrium.equilibrium.Equilibrium": + return {} + + # TODO: can automate this later to add omngeneity, boozer transform, etc. + fft_names = ["effective ripple", "Gamma_c", "Gamma_c Velasco"] + + eq = get("W7-X") + # ci and my laptop differ a bunch at rho = 0, so skip that + rho = np.linspace(1e-2, 1, 10) + grid = LinearGrid(rho=rho, M=eq.M_grid, N=eq.N_grid, NFP=eq.NFP, sym=False) + + kwargs = dict( + angle=Bounce2D.angle(eq, X=32, Y=32, rho=rho, tol=1e-10), + Y_B=grid.num_zeta * grid.NFP, + num_transit=5, + num_well=20 * 5, + nufft_eps=1e-10, + ) + data = eq.compute(fft_names, grid, **kwargs) + + # check vectorization too + d = data.copy() + del d["Gamma_c"] + d = eq.compute("Gamma_c", grid, data=d, surf_batch_size=2, **kwargs) + np.testing.assert_allclose( + d["Gamma_c"], + data["Gamma_c"], + rtol=1e-9, + atol=1e-9, + err_msg="Gamma_c vectorization", + ) + # check no nufft + del d["Gamma_c"] + kwargs["nufft_eps"] = 0.0 + d = eq.compute("Gamma_c", grid, data=d, **kwargs) + np.testing.assert_allclose( + d["Gamma_c"], + data["Gamma_c"], + # This is large since no nuffts + spline for bounce points + # are innaccurate due to lack of Newton step after finding bounce point + # approximation with splines. + rtol=0.2, + err_msg="Gamma_c no nufft", + ) + + data = apply(data, grid.compress, fft_names) + return data + + +def raz_grid_data(p): + """Compute field line grid quantities.""" + if p != "desc.equilibrium.equilibrium.Equilibrium": + return {} + + # TODO: can automate this later to add ballooning etc. + raz_names = ["old effective ripple", "old Gamma_c", "old Gamma_c Velasco"] + + eq = get("W7-X") + num_transit = 2 + Y_B = eq.N_grid * 2 * eq.NFP + rho = np.linspace(1e-2, 1, 10) + alpha = np.array([0]) + zeta = np.linspace(0, num_transit * 2 * np.pi, num_transit * Y_B) + grid = Grid.create_meshgrid([rho, alpha, zeta], coordinates="raz") + + data = eq.compute(raz_names, grid, num_well=20 * num_transit, tol=1e-10) + data = apply(data, grid.compress, raz_names) + return data diff --git a/tests/test_compute_funs.py b/tests/test_compute_funs.py index 44d4b79929..39c6c59753 100644 --- a/tests/test_compute_funs.py +++ b/tests/test_compute_funs.py @@ -2049,6 +2049,45 @@ def test_parallel_grad_fd(DummyStellarator): ) +@pytest.mark.unit +@pytest.mark.slow +def test_fieldline_average(): + """Test that fieldline average converges to surface average.""" + rho = np.array([1]) + alpha = np.array([0]) + eq = get("DSHAPE") + iota_grid = LinearGrid(rho=rho, M=eq.M_grid, N=eq.N_grid, NFP=eq.NFP, sym=eq.sym) + iota = iota_grid.compress(eq.compute("iota", grid=iota_grid)["iota"]).item() + # For axisymmetric devices, one poloidal transit must be exact. + zeta = np.linspace(0, 2 * np.pi / iota, 25) + grid = Grid.create_meshgrid([rho, alpha, zeta], coordinates="raz") + data = eq.compute( + ["fieldline length", "fieldline length/volume", "V_r(r)"], grid=grid + ) + np.testing.assert_allclose( + data["fieldline length"] / data["fieldline length/volume"], + data["V_r(r)"] / (4 * np.pi**2), + rtol=1e-3, + ) + assert np.all(data["fieldline length"] > 0) + assert np.all(data["fieldline length/volume"] > 0) + + # Otherwise, many toroidal transits are necessary to sample surface. + eq = get("W7-X") + zeta = np.linspace(0, 40 * np.pi, 300) + grid = Grid.create_meshgrid([rho, alpha, zeta], coordinates="raz") + data = eq.compute( + ["fieldline length", "fieldline length/volume", "V_r(r)"], grid=grid + ) + np.testing.assert_allclose( + data["fieldline length"] / data["fieldline length/volume"], + data["V_r(r)"] / (4 * np.pi**2), + rtol=2e-3, + ) + assert np.all(data["fieldline length"] > 0) + assert np.all(data["fieldline length/volume"] > 0) + + @pytest.mark.unit def test_compute_deprecation_warning(): """Test DeprecationWarning for deprecated compute names.""" diff --git a/tests/test_fast_ion.py b/tests/test_fast_ion.py deleted file mode 100644 index 8835b8dc5d..0000000000 --- a/tests/test_fast_ion.py +++ /dev/null @@ -1,100 +0,0 @@ -"""Test fast ion compute functions.""" - -import matplotlib.pyplot as plt -import numpy as np -import pytest -from tests.test_plotting import tol_1d - -from desc.examples import get -from desc.grid import Grid, LinearGrid -from desc.integrals import Bounce2D - - -@pytest.mark.unit -@pytest.mark.mpl_image_compare(remove_text=True, tolerance=tol_1d) -@pytest.mark.parametrize("nufft_eps", [0, 1e-7]) -def test_Gamma_c_Nemov_2D(nufft_eps): - """Test Γ_c Nemov with W7-X.""" - eq = get("W7-X") - rho = np.linspace(1e-12, 1, 10) - grid = LinearGrid(rho=rho, M=eq.M_grid, N=eq.N_grid, NFP=eq.NFP, sym=False) - num_transit = 10 - data = eq.compute( - "Gamma_c", - grid=grid, - angle=Bounce2D.angle(eq, X=32, Y=32, rho=rho), - Y_B=128, - num_transit=num_transit, - num_well=20 * num_transit, - nufft_eps=nufft_eps, - ) - assert data["Gamma_c"].ndim == 1 - assert np.isfinite(data["Gamma_c"]).all() - fig, ax = plt.subplots() - ax.plot(rho, grid.compress(data["Gamma_c"]), marker="o") - return fig - - -@pytest.mark.unit -@pytest.mark.mpl_image_compare(remove_text=True, tolerance=tol_1d) -@pytest.mark.parametrize("nufft_eps", [0, 1e-7]) -def test_Gamma_c_Velasco_2D(nufft_eps): - """Test Γ_c Velasco with W7-X.""" - eq = get("W7-X") - rho = np.linspace(1e-12, 1, 10) - grid = LinearGrid(rho=rho, M=eq.M_grid, N=eq.N_grid, NFP=eq.NFP, sym=False) - num_transit = 10 - data = eq.compute( - "Gamma_c Velasco", - grid=grid, - angle=Bounce2D.angle(eq, X=32, Y=32, rho=rho), - Y_B=128, - num_transit=num_transit, - num_well=20 * num_transit, - nufft_eps=nufft_eps, - ) - assert data["Gamma_c Velasco"].ndim == 1 - assert np.isfinite(data["Gamma_c Velasco"]).all() - fig, ax = plt.subplots() - ax.plot(rho, grid.compress(data["Gamma_c Velasco"]), marker="o") - return fig - - -@pytest.mark.unit -@pytest.mark.mpl_image_compare(remove_text=True, tolerance=tol_1d) -def test_Gamma_c_Nemov_1D(): - """Test Γ_c Nemov 1D with W7-X.""" - eq = get("W7-X") - Y_B = 100 - num_transit = 10 - num_well = 20 * num_transit - rho = np.linspace(1e-12, 1, 10) - alpha = np.array([0]) - zeta = np.linspace(0, num_transit * 2 * np.pi, num_transit * Y_B) - grid = Grid.create_meshgrid([rho, alpha, zeta], coordinates="raz") - data = eq.compute("old Gamma_c", grid=grid, num_well=num_well, surf_batch_size=2) - - assert np.isfinite(data["old Gamma_c"]).all() - fig, ax = plt.subplots() - ax.plot(rho, grid.compress(data["old Gamma_c"]), marker="o") - return fig - - -@pytest.mark.unit -@pytest.mark.mpl_image_compare(remove_text=True, tolerance=tol_1d) -def test_Gamma_c_Velasco_1D(): - """Test Γ_c Velasco 1D with W7-X.""" - eq = get("W7-X") - Y_B = 100 - num_transit = 10 - num_well = 20 * num_transit - rho = np.linspace(1e-12, 1, 10) - alpha = np.array([0]) - zeta = np.linspace(0, num_transit * 2 * np.pi, num_transit * Y_B) - grid = Grid.create_meshgrid([rho, alpha, zeta], coordinates="raz") - data = eq.compute("old Gamma_c Velasco", grid=grid, num_well=num_well) - - assert np.isfinite(data["old Gamma_c Velasco"]).all() - fig, ax = plt.subplots() - ax.plot(rho, grid.compress(data["old Gamma_c Velasco"]), marker="o") - return fig diff --git a/tests/test_integrals.py b/tests/test_integrals.py index bf11c28853..d26761ec19 100644 --- a/tests/test_integrals.py +++ b/tests/test_integrals.py @@ -732,7 +732,7 @@ def test_z1_first(self): B = CubicHermiteSpline(k, np.cos(k), -np.sin(k)) pitch_inv = 0.5 intersect = B.solve(pitch_inv, extrapolate=False) - z1, z2, _ = bounce_points(pitch_inv, k, B.c.T) + z1, z2 = bounce_points(pitch_inv, k, B.c.T) check_bounce_points(z1, z2, pitch_inv, k, B.c.T, plot=True, include_knots=True) z1, z2 = TestBouncePoints.filter(z1, z2) assert z1.size and z2.size @@ -748,7 +748,7 @@ def test_z2_first(self): B = CubicHermiteSpline(k, np.cos(k), -np.sin(k)) pitch_inv = 0.5 intersect = B.solve(pitch_inv, extrapolate=False) - z1, z2, _ = bounce_points(pitch_inv, k, B.c.T) + z1, z2 = bounce_points(pitch_inv, k, B.c.T) check_bounce_points(z1, z2, pitch_inv, k, B.c.T, plot=True, include_knots=True) z1, z2 = TestBouncePoints.filter(z1, z2) assert z1.size and z2.size @@ -767,7 +767,7 @@ def test_z1_before_extrema(self): k, np.cos(k) + 2 * np.sin(-2 * k), -np.sin(k) - 4 * np.cos(-2 * k) ) pitch_inv = B(B.derivative().roots(extrapolate=False))[3] - 1e-13 - z1, z2, _ = bounce_points(pitch_inv, k, B.c.T) + z1, z2 = bounce_points(pitch_inv, k, B.c.T) check_bounce_points(z1, z2, pitch_inv, k, B.c.T, plot=True, include_knots=True) z1, z2 = TestBouncePoints.filter(z1, z2) assert z1.size and z2.size @@ -792,7 +792,7 @@ def test_z2_before_extrema(self): -np.sin(k) - 4 * np.cos(-2 * k) + 1 / 4, ) pitch_inv = B(B.derivative().roots(extrapolate=False))[2] - z1, z2, _ = bounce_points(pitch_inv, k, B.c.T) + z1, z2 = bounce_points(pitch_inv, k, B.c.T) check_bounce_points(z1, z2, pitch_inv, k, B.c.T, plot=True, include_knots=True) z1, z2 = TestBouncePoints.filter(z1, z2) assert z1.size and z2.size @@ -813,7 +813,7 @@ def test_extrema_first_and_before_z1(self): -np.sin(k) - 4 * np.cos(-2 * k) + 1 / 20, ) pitch_inv = B(B.derivative().roots(extrapolate=False))[2] + 1e-13 - z1, z2, _ = bounce_points(pitch_inv, k[2:], B.c[:, 2:].T) + z1, z2 = bounce_points(pitch_inv, k[2:], B.c[:, 2:].T) check_bounce_points( z1, z2, @@ -845,7 +845,7 @@ def test_extrema_first_and_before_z2(self): -np.sin(k) - 4 * np.cos(-2 * k) + 1 / 10, ) pitch_inv = B(B.derivative().roots(extrapolate=False))[1] - 1e-13 - z1, z2, _ = bounce_points(pitch_inv, k, B.c.T) + z1, z2 = bounce_points(pitch_inv, k, B.c.T) check_bounce_points(z1, z2, pitch_inv, k, B.c.T, plot=True, include_knots=True) z1, z2 = TestBouncePoints.filter(z1, z2) assert z1.size and z2.size @@ -1602,12 +1602,14 @@ def _not_part_of_tutorial_test( np.testing.assert_array_equal(near_zero_nufft, near_zero) np.testing.assert_allclose(num_nufft[near_zero_nufft], num[near_zero]) np.testing.assert_allclose( - num_nufft[~near_zero_nufft], num[~near_zero], rtol=8e-3 + num_nufft[~near_zero_nufft], num[~near_zero], rtol=3e-2 ) bounce = Bounce2D(grid, data, angle, alpha=alpha, num_transit=2, check=True) points = bounce.points(pitch_inv) - z1, z2 = _newton(bounce, pitch_inv[:, None], *points, points[0] < points[1]) + z1, z2 = _newton( + bounce, pitch_inv[:, None], *points, points[0] < points[1], 1e-10 + ) np.testing.assert_allclose(points[0], z1, rtol=5e-6) np.testing.assert_allclose(points[1], z2, rtol=5e-6) diff --git a/tests/test_neoclassical.py b/tests/test_neoclassical.py deleted file mode 100644 index 6c4d688a73..0000000000 --- a/tests/test_neoclassical.py +++ /dev/null @@ -1,115 +0,0 @@ -"""Test neoclassical transport compute functions.""" - -import matplotlib.pyplot as plt -import numpy as np -import pytest -from tests.test_plotting import tol_1d - -from desc.examples import get -from desc.external.neo import NeoIO -from desc.grid import Grid, LinearGrid -from desc.integrals import Bounce2D - - -@pytest.mark.unit -@pytest.mark.mpl_image_compare(remove_text=True, tolerance=tol_1d) -@pytest.mark.parametrize("nufft_eps", [0, 1e-6]) -def test_effective_ripple_2D(nufft_eps): - """Test effective ripple with W7-X against NEO.""" - eq = get("W7-X") - rho = np.linspace(0, 1, 10) - grid = LinearGrid(rho=rho, M=eq.M_grid, N=eq.N_grid, NFP=eq.NFP, sym=False) - num_transit = 10 - data = eq.compute( - "effective ripple 3/2", - grid=grid, - angle=Bounce2D.angle(eq, X=32, Y=32, rho=rho), - Y_B=128, - num_transit=num_transit, - num_well=20 * num_transit, - surf_batch_size=1 if (nufft_eps == 0) else 2, - nufft_eps=nufft_eps, - ) - - assert data["effective ripple 3/2"].ndim == 1 - assert np.isfinite(data["effective ripple 3/2"]).all() - eps_32 = grid.compress(data["effective ripple 3/2"]) - # NEO file generated from DESC equlibrium on 2025-10-23 17:47:07.280264. - neo_rho, neo_eps_32 = NeoIO.read("tests/inputs/neo_out.W7-X") - np.testing.assert_allclose(eps_32, np.interp(rho, neo_rho, neo_eps_32), rtol=0.11) - - fig, ax = plt.subplots() - ax.plot(rho, eps_32, marker="o") - ax.plot(neo_rho, neo_eps_32) - return fig - - -@pytest.mark.unit -@pytest.mark.mpl_image_compare(remove_text=True, tolerance=tol_1d) -def test_effective_ripple_1D(): - """Test effective ripple 1D with W7-X against NEO.""" - eq = get("W7-X") - Y_B = 100 - num_transit = 10 - num_well = 20 * num_transit - rho = np.linspace(0, 1, 10) - alpha = np.array([0]) - zeta = np.linspace(0, num_transit * 2 * np.pi, num_transit * Y_B) - grid = Grid.create_meshgrid([rho, alpha, zeta], coordinates="raz") - data = eq.compute( - "old effective ripple", grid=grid, num_well=num_well, surf_batch_size=2 - ) - - assert np.isfinite(data["old effective ripple"]).all() - np.testing.assert_allclose( - data["old effective ripple 3/2"] ** (2 / 3), - data["old effective ripple"], - err_msg="Bug in source grid logic in eq.compute.", - ) - eps_32 = grid.compress(data["old effective ripple 3/2"]) - neo_rho, neo_eps_32 = NeoIO.read("tests/inputs/neo_out.W7-X") - np.testing.assert_allclose(eps_32, np.interp(rho, neo_rho, neo_eps_32), rtol=0.1) - - fig, ax = plt.subplots() - ax.plot(rho, eps_32, marker="o") - ax.plot(neo_rho, neo_eps_32) - return fig - - -@pytest.mark.unit -@pytest.mark.slow -def test_fieldline_average(): - """Test that fieldline average converges to surface average.""" - rho = np.array([1]) - alpha = np.array([0]) - eq = get("DSHAPE") - iota_grid = LinearGrid(rho=rho, M=eq.M_grid, N=eq.N_grid, NFP=eq.NFP, sym=eq.sym) - iota = iota_grid.compress(eq.compute("iota", grid=iota_grid)["iota"]).item() - # For axisymmetric devices, one poloidal transit must be exact. - zeta = np.linspace(0, 2 * np.pi / iota, 25) - grid = Grid.create_meshgrid([rho, alpha, zeta], coordinates="raz") - data = eq.compute( - ["fieldline length", "fieldline length/volume", "V_r(r)"], grid=grid - ) - np.testing.assert_allclose( - data["fieldline length"] / data["fieldline length/volume"], - data["V_r(r)"] / (4 * np.pi**2), - rtol=1e-3, - ) - assert np.all(data["fieldline length"] > 0) - assert np.all(data["fieldline length/volume"] > 0) - - # Otherwise, many toroidal transits are necessary to sample surface. - eq = get("W7-X") - zeta = np.linspace(0, 40 * np.pi, 300) - grid = Grid.create_meshgrid([rho, alpha, zeta], coordinates="raz") - data = eq.compute( - ["fieldline length", "fieldline length/volume", "V_r(r)"], grid=grid - ) - np.testing.assert_allclose( - data["fieldline length"] / data["fieldline length/volume"], - data["V_r(r)"] / (4 * np.pi**2), - rtol=2e-3, - ) - assert np.all(data["fieldline length"] > 0) - assert np.all(data["fieldline length/volume"] > 0) diff --git a/tests/test_objective_funs.py b/tests/test_objective_funs.py index 9c0cb01b24..b63d8ef556 100644 --- a/tests/test_objective_funs.py +++ b/tests/test_objective_funs.py @@ -4153,7 +4153,15 @@ def test_objective_no_nangrad_effective_ripple(self): obj.build(verbose=0) g = obj.grad(obj.x()) assert not np.any(np.isnan(g)) - np.testing.assert_allclose(g, g_0, atol=1e-6) + # This test needs high tolerance because the no nuffts + spline + # method for bounce points doesn't do a Newton step. Recall + # an O(ε) error in the spline approximation of bounce point + # yields O(ε¹ᐧ⁵) error in integrals with v_||. For the + # gradient it is probably O(ε) in general, but you'd need to work this out + # from the supplementary information. + # TODO: Reduce tolerance after someone implements the Newton step. + # (When we used to do the Newton step the atol could be 1e-6). + np.testing.assert_allclose(g, g_0, atol=0.0025) obj = ObjectiveFunction( _reduced_resolution_objective(eq, EffectiveRipple, use_bounce1d=True) @@ -4175,18 +4183,19 @@ def test_objective_no_nangrad_Gamma_c(self): g_0 = obj_0.grad(obj_0.x()) assert not np.any(np.isnan(g_0)) - # this needs 5e-11 for eps to pass when jax_finufft==1.3.0 - obj = ObjectiveFunction( - _reduced_resolution_objective(eq, GammaC, nufft_eps=5e-11) - ) + obj = ObjectiveFunction(_reduced_resolution_objective(eq, GammaC)) obj.build(verbose=0) g = obj.grad(obj.x()) assert not np.any(np.isnan(g)) - # these are generally sensitive to nufft_eps because - # we are not using enough resolution in other parameters - # in this test to nullify the singularities - # TODO: Do we want to keep this test then if it is so sensitive? - np.testing.assert_allclose(g, g_0, atol=2e-6, rtol=3e-4) + # This test needs high tolerance because the no nuffts + spline + # method for bounce points doesn't do a Newton step. Recall + # an O(ε) error in the spline approximation of bounce point + # yields O(ε⁰ᐧ⁵) error in integrals with 1/v_||. For the gradient + # it is probably O(ε⁰ᐧ³³) in general, but you'd need to work this out + # from the supplementary information. + # TODO: Reduce tolerance after someone implements the Newton step. + # (When we used to do the Newton step the atol could be 1e-6). + np.testing.assert_allclose(g, g_0, atol=0.042) obj = ObjectiveFunction( _reduced_resolution_objective(eq, GammaC, use_bounce1d=True)