diff --git a/desc/compute/__init__.py b/desc/compute/__init__.py index 904bd65712..4acdf09c45 100644 --- a/desc/compute/__init__.py +++ b/desc/compute/__init__.py @@ -42,6 +42,7 @@ _profiles, _stability, _surface, + _trapped_resonance, ) from .data_index import _topological_order, all_kwargs, allowed_kwargs, data_index from .utils import ( diff --git a/desc/compute/_trapped_resonance.py b/desc/compute/_trapped_resonance.py new file mode 100644 index 0000000000..9e2fd573f9 --- /dev/null +++ b/desc/compute/_trapped_resonance.py @@ -0,0 +1,1343 @@ +"""Compute functions for trapped energetic particle resonance.""" + +from quadax import simpson + +from desc.backend import jax, jnp +from desc.grid import Grid + +from ..batching import batch_map +from ..integrals.bounce_integral import Bounce1D, Bounce2D +from ..utils import safediv +from ._fast_ion import _radial_drift, _v_tau +from ._neoclassical import _bounce_doc +from .data_index import register_compute_fun + + +def _jnpmean_nz(x, axis=0, fill=jnp.nan): + """Mean over an axis, ignoring zero and fill-value entries.""" + mask = (x != 0.0) & _is_valid_value(x, fill) + count = jnp.sum(mask, axis=axis) + return safediv(jnp.sum(jnp.where(mask, x, 0.0), axis=axis), count, fill=fill) + + +def _is_valid_value(x, fill_value): + """Validity mask compatible with finite sentinel fill values.""" + return x != fill_value + + +def _masked_sum(x, mask, axis=None): + """Sum x over axis, excluding entries where mask is False.""" + return jnp.sum(jnp.where(mask, x, jnp.zeros_like(x)), axis=axis) + + +def _build_eta_source_grid(rhos, alpha_per_rho, zeta): + """Build the (rho, alpha, zeta) grid with alpha derived from uniform eta. + + This is the field line following grid that ``Bounce1D`` integrates along. + It is built from array arithmetic alone, unlike the (rho, theta, zeta) + grid of ``_build_eta_grid``, which needs a coordinate map. + """ + num_rho = len(rhos) + num_eta = alpha_per_rho.shape[1] + num_zeta = len(zeta) + + # Build raz nodes in meshgrid order: alpha fastest, rho middle, zeta slowest. + _, rr, zz = jnp.meshgrid(jnp.arange(num_eta), rhos, zeta, indexing="ij") + alpha_arr = jnp.broadcast_to( + alpha_per_rho.T[:, :, jnp.newaxis], (num_eta, num_rho, num_zeta) + ) + raz_nodes = jnp.column_stack( + [ + rr.flatten(order="F"), + alpha_arr.flatten(order="F"), + zz.flatten(order="F"), + ] + ) + + unique_rho_idx = jnp.arange(num_rho) * num_eta + unique_poloidal_idx = jnp.arange(num_eta) + unique_zeta_idx = jnp.arange(num_zeta) * num_rho * num_eta + inverse_rho_idx = jnp.tile(jnp.repeat(jnp.arange(num_rho), num_eta), num_zeta) + inverse_poloidal_idx = jnp.tile(jnp.arange(num_eta), num_rho * num_zeta) + inverse_zeta_idx = jnp.repeat(jnp.arange(num_zeta), num_rho * num_eta) + + return Grid( + nodes=raz_nodes, + coordinates="raz", + period=(jnp.inf, jnp.inf, jnp.inf), + sort=False, + is_meshgrid=True, + jitable=True, + _unique_rho_idx=unique_rho_idx, + _unique_poloidal_idx=unique_poloidal_idx, + _unique_zeta_idx=unique_zeta_idx, + _inverse_rho_idx=inverse_rho_idx, + _inverse_poloidal_idx=inverse_poloidal_idx, + _inverse_zeta_idx=inverse_zeta_idx, + ) + + +def _build_eta_grid(eq, rhos, alpha_per_rho, zeta, iotas, params): + """Build a DESC grid with per-rho alpha values derived from uniform eta.""" + from desc.equilibrium.coords import map_coordinates + + raz_grid = _build_eta_source_grid(rhos, alpha_per_rho, zeta) + + iota_expanded = raz_grid.expand(jnp.atleast_1d(jnp.asarray(iotas))) + rtz_nodes = map_coordinates( + eq, + raz_grid.nodes, + inbasis=["rho", "alpha", "zeta"], + outbasis=("rho", "theta", "zeta"), + period=(jnp.inf, jnp.inf, jnp.inf), + iota=iota_expanded, + params=params, + ) + + return Grid( + nodes=rtz_nodes, + coordinates="rtz", + source_grid=raz_grid, + sort=False, + jitable=True, + _unique_rho_idx=raz_grid.unique_rho_idx, + _inverse_rho_idx=raz_grid.inverse_rho_idx, + ) + + +def _compute2D( + fun, + fun_data, + data, + grid, + angle, + alpha_per_rho, + num_pitch, + surf_batch_size=1, + simp=True, + pitch_invs=None, + pitch_inv_weight=None, +): + """Compute Bounce2D integral quantity with ``fun``. + + The Bounce2D analogue of ``_compute1D``. ``fun_data`` and ``data`` are + given on a tensor-product (θ, ζ) grid rather than on a field line grid, + and are Fourier transformed here for ``Bounce2D`` to interpolate onto + field lines internally. + + Parameters + ---------- + fun : callable + Function to compute. Receives the batched data dictionary, whose + arrays hold Fourier transforms (pass ``is_fourier=True`` to + ``Bounce2D``). + fun_data : dict[str, jnp.ndarray] + Data to Fourier transform and pass to ``fun``. Modified in place. + data : dict[str, jnp.ndarray] + DESC data dict evaluated on ``grid``. + grid : Grid + Tensor-product (ρ, θ, ζ) grid satisfying ``can_fft2``. + angle : jnp.ndarray + Shape (num ρ, X, Y). Angle returned by ``Bounce2D.angle``. + alpha_per_rho : jnp.ndarray + Shape (num ρ, num α). Field line labels, which differ between flux + surfaces because they derive from the omnigenity angle η. Stored + with ρ leading so that ``batch_map`` slices it consistently with the + rest of the data; ``fun`` transposes it back to (num α, num ρ). + num_pitch : int + Resolution for quadrature over velocity coordinate. + surf_batch_size : int + Number of flux surfaces with which to compute simultaneously. + simp : bool + Whether to use an open Simpson rule instead of uniform weights. + pitch_invs : jnp.ndarray + If specified, use these pitch_inv values rather than ``num_pitch``. + pitch_inv_weight : jnp.ndarray + Quadrature weight paired with ``pitch_invs``. + + """ + for name in Bounce2D.required_names: + fun_data[name] = data[name] + # iota is per-surface, not a (θ, ζ) field, so it must stay out of the + # transform below and come back compressed instead. + fun_data.pop("iota") + for name in fun_data: + fun_data[name] = Bounce2D.fourier(Bounce2D.reshape(grid, fun_data[name])) + fun_data["iota"] = grid.compress(data["iota"]) + fun_data["angle"] = angle + fun_data["alpha_per_rho"] = alpha_per_rho + if pitch_invs is None: + # A single B_crit set is shared across flux surfaces, so broadcast the + # global extrema rather than using each surface's own. + num_rho = grid.num_rho + B_min = jnp.min(grid.compress(data["min_tz |B|"])) + B_max = jnp.max(grid.compress(data["max_tz |B|"])) + ( + fun_data["pitch_inv"], + fun_data["pitch_inv weight"], + ) = Bounce2D.get_pitch_inv_quad( + jnp.full(num_rho, B_min), jnp.full(num_rho, B_max), num_pitch, simp=simp + ) + else: + n = len(pitch_invs) + fun_data["pitch_inv"] = jnp.broadcast_to(pitch_invs, (grid.num_rho, n)) + fun_data["pitch_inv weight"] = jnp.broadcast_to( + ( + jnp.ones(n) * (2 * jnp.pi / n) + if pitch_inv_weight is None + else pitch_inv_weight + ), + (grid.num_rho, n), + ) + return batch_map(fun, fun_data, surf_batch_size) + + +def _compute1D( + fun, + fun_data, + data, + grid, + num_pitch, + surf_batch_size=1, + simp=True, + pitch_invs=None, + pitch_inv_weight=None, +): + """Compute Bounce1D integral quantity with ``fun``. + + Parameters + ---------- + fun : callable + Function to compute. + fun_data : dict[str, jnp.ndarray] + Data to provide to ``fun``. This dict will be modified. + data : dict[str, jnp.ndarray] + DESC data dict. + grid : Grid + Grid that can expand and compress. + num_pitch : int + Resolution for quadrature over velocity coordinate. + surf_batch_size : int + Number of flux surfaces with which to compute simultaneously. + Default is ``1``. + simp : bool + Whether to use an open Simpson rule instead of uniform weights. + pitch_invs : jnp.ndarray + If specified, use the given pitch_invs values rather than using num_pitch. + pitch_inv_weight : jnp.ndarray + Quadrature weight paired with ``pitch_invs``. If ``pitch_invs`` is given + without a matching weight, falls back to uniform weighting. + + """ + for name in Bounce1D.required_names: + fun_data[name] = data[name] + for name in fun_data: + fun_data[name] = Bounce1D.reshape(grid, fun_data[name]) + if pitch_invs is None: + ( + fun_data["pitch_inv"], + fun_data["pitch_inv weight"], + ) = Bounce1D.get_pitch_inv_quad( + grid.compress(data["min_tz |B|"]), + grid.compress(data["max_tz |B|"]), + num_pitch, + simp=simp, + ) + else: # Caller-supplied pitch_invs with matching quadrature weight. + n = len(pitch_invs) + if pitch_inv_weight is None: + pitch_inv_weight = jnp.ones(n) / n + fun_data["pitch_inv"] = jnp.broadcast_to(pitch_invs, (grid.num_rho, n)) + fun_data["pitch_inv weight"] = jnp.broadcast_to( + pitch_inv_weight, (grid.num_rho, n) + ) + + out = batch_map(fun, fun_data, surf_batch_size) + + return out + + +def _alpha_drift_integrand(data, B, pitch): + """Cross-field-line drift integrand for bounce integration. + + Used in ``_trapped_EP_resonance``. + """ + return safediv( + 2 + * ( + data["gbdrift (periodic)"] * pitch * B + + 2 * (1 - pitch * B) * data["cvdrift (periodic)"] + ), + jnp.sqrt(jnp.abs(1 - pitch * B)), + ) + + +# Alpha particle constants, used to turn bounce-averaged +# drifts into physical frequencies. +_M_ALPHA = 6.6446573450e-27 # mass, kg +_E_CHARGE = 1.602e-19 # elementary charge, C +_Z_ALPHA = 2 # charge number +# 3.5 MeV, the birth energy of alpha particles from D-T fusion. +_E_BIRTH = 5.6076e-13 # J + +_BOUNCE_INTEGRAND_KEYS = ("cvdrift0", "gbdrift (periodic)", "cvdrift (periodic)") +_ETA_BOUNCE_KEYS = tuple(Bounce1D.required_names) + _BOUNCE_INTEGRAND_KEYS +_FFT_BOUNCE_KEYS = tuple(Bounce2D.required_names) + _BOUNCE_INTEGRAND_KEYS + + +def _global_pitch_quad(base_grid, data, num_pitch, pitch_invs): + """Pitch grid shared by every flux surface. + + Built from the base grid's min/max |B|, which comes from the equilibrium's + full Fourier resolution and so does not move with ``num_transit``. + + Returns + ------- + pitch_inv, pitch_inv_weight : tuple[jnp.ndarray] + ``pitch_inv_weight`` is ``None`` when ``pitch_invs`` was supplied, in + which case ``_compute1D`` falls back to uniform weights. + + """ + if pitch_invs is not None: + return pitch_invs, None + B_min = jnp.min(base_grid.compress(data["min_tz |B|"])) + B_max = jnp.max(base_grid.compress(data["max_tz |B|"])) + pitch_inv, weight = Bounce1D.get_pitch_inv_quad( + jnp.array([B_min]), jnp.array([B_max]), num_pitch, simp=True + ) + return pitch_inv[0], weight[0] + + +def _frequencies( + alpha_drift_out, s_drift_out, vtau_out, iotas, KE_frac, nfp, M, N, fill_value +): + """Turn bounce-averaged drifts into frequencies and precession Omega. + + Parameters + ---------- + alpha_drift_out, s_drift_out : jnp.ndarray + Shape (rho, alpha, Bcrit, well). + Bounce-averaged poloidal and radial drift, before energy scaling. + vtau_out : jnp.ndarray + Shape (rho, alpha, Bcrit, well). + Bounce integral of v·τ. + iotas : jnp.ndarray, shape (rho,) + Rotational transform per surface. + KE_frac : float + Fraction of the 3.5 MeV D-T fusion alpha-particle birth energy to use + for the energetic particle kinetic energy. + nfp : int + Number of field periods. + M, N : int + Generalized omnigenous helicity. + fill_value : float + Value bounce integration outputs take when no well is found. + + Returns + ------- + dict + ``Omega`` plus the intermediates it is built from. ``valid`` is + ``True`` where a trapped particle exists at every alpha and ``Omega`` + is defined. + + """ + KE = KE_frac * _E_BIRTH + v2 = 2 * KE / _M_ALPHA + + # Bounce-averaged drifts → physical frequencies + alpha_drift = alpha_drift_out * KE / (_Z_ALPHA * _E_CHARGE) + eta_drift = safediv( + nfp * alpha_drift, N * nfp - iotas[:, None, None, None] * M, fill=fill_value + ) + + s_drift = s_drift_out * KE / (_Z_ALPHA * _E_CHARGE) + tau_bounce = vtau_out / jnp.sqrt(v2) + omega_bounce = safediv(2 * jnp.pi, tau_bounce, fill=fill_value) + + # Require particle to be trapped at all alpha/eta values for a given + # (rho, pitch, well). + all_alpha_valid = (_is_valid_value(omega_bounce, fill_value)).all( + axis=1 + ) # (rho, pitch, well) + + # Alpha-averaged frequencies → normalized precession Omega + omega_bounce_avg = _jnpmean_nz(omega_bounce, axis=1, fill=fill_value) + eta_drift_avg = _jnpmean_nz(eta_drift, axis=1, fill=fill_value) + Omega = safediv(eta_drift_avg, omega_bounce_avg, fill=fill_value) + valid = ( + _is_valid_value(eta_drift_avg, fill_value) + & _is_valid_value(omega_bounce_avg, fill_value) + & all_alpha_valid + ) + return { + "Omega": jnp.where(valid, Omega, fill_value), + "valid": valid, + "omega_bounce_avg": omega_bounce_avg, + "eta_drift_avg": eta_drift_avg, + "omega_bounce": omega_bounce, + "eta_drift": eta_drift, + "tau_bounce": tau_bounce, + "s_drift": s_drift, + } + + +_bounce1D_doc = { + "num_quad": _bounce_doc["num_quad"], + "num_pitch": _bounce_doc["num_pitch"], + "surf_batch_size": _bounce_doc["surf_batch_size"], + "quad": _bounce_doc["quad"], +} + +_resonance_doc = { + "M": """int : + Generalized omnigenous helicity. Each B contour closes on itself after + traversing the torus M times toroidally and N times poloidally. + """, + "N": """int : + Generalized omnigenous helicity. Each B contour closes on itself after + traversing the torus M times toroidally and N times poloidally. + """, + "nfp": """int : + Number of field periods. + """, + "KE_frac": """jnp.ndarray : + Fraction of the 3.5 MeV D-T fusion alpha-particle birth energy to use + for the energetic particle kinetic energy. + """, + "pitch_invs": """jnp.ndarray or None : + If not ``None``, sets pitch_invs (Bcrits) to the specified value, and + causes this function to skip the phase-space average and return the + raw per-(rho, pitch, well) resonance-physics dictionary instead of the + phase-space-averaged objective. If ``None``, uses a linspace of + num_pitch between Bmin and Bmax of each flux surface. + """, + "rho_res": """float : + Radial grid spacing. + """, + "eta_res": """float : + Grid spacing for eta. + """, + "res_arr": """jnp.ndarray : + Resonance frequency ratios p/q to check Omega_eta against, for all + combinations of p/q up to p_max/q_max within + [res_range_min, res_range_max]. + """, + "p_arr": """jnp.ndarray : + Numerators of the resonance ratios in ``res_arr``. + """, + "q_arr": """jnp.ndarray : + Denominators (toroidal mode numbers) of the resonance ratios in + ``res_arr``. + """, + "weight_method": """str : + ``"linear"`` or ``"bump"`` resonance weighting. + """, + "Delta_Omega": """float or None : + Half-width of the resonance interval for ``weight_method="bump"``. + If ``None``, defaults to wd_blur × the max |Ω[i+1]-Ω[i]| spacing. + Ignored when ``weight_method="linear"``. + """, + "wd_blur": """float : + Multiplicative blur factor used to compute bump half-width from + adjacent-surface Omega spacing when ``Delta_Omega`` is not provided. + """, + "fill_value": """float : + Value to set bounce integration outputs to if no well is found. + Cannot use ``jnp.nan`` to retain optimization abilities. Cannot use 0 + for confusion with other quantities and averages. + """, + "stab_sacrifice": """bool : + If ``True``, multiply the island-width term by ``Omega_prime_s**2`` + in the objective. If ``False``, omit that factor to preserve + numerical stability. + """, + "bt_filter_flag": """bool : + If ``True``, zero out wells whose poloidal bounce width exceeds 2π + (barely-trapped filter) before the resonance physics calculation. + """, + "cropping_DOmega": """bool : + If ``True``, Delta_Omega calculation is clipped by + ``0.01 * max(Omega_eta) < Delta_Omega < 0.10 * max(Omega_eta)``. + Only used with the ``bump`` weighting method and + ``Delta_Omega = None``. Otherwise this quantity is ignored. + """, + "num_transit": """int : + Number of toroidal transits spanned by ``zeta``. Should be long + enough to capture the full trapping well. + """, + "num_eta": """int : + Number of uniformly spaced eta points in [0, 2π). Alpha values are + derived per rho surface via ``alpha = eta * (N*nfp - iota*M) / nfp``. + """, + "zeta": """jnp.ndarray : + Toroidal angle values spanning ``num_transit`` toroidal transits, + used for field-line integration and to compute the field-line length. + """, + "_eta_grid": """Grid : + Field-line-following grid with per-rho alpha values derived from + uniform eta, built by ``TrappedResonance.compute``. This private + parameter is intended to be used only by developers for objectives. + """, + "_psa_grid": """Grid : + Field-line-following grid uniform in alpha, used for the phase-space + average, built by ``TrappedResonance.compute``. This private + parameter is intended to be used only by developers for objectives. + """, + "_data_eta": """dict[str, jnp.ndarray] : + Field data evaluated on ``_eta_grid``, built by + ``TrappedResonance.compute``. This private parameter is intended to + be used only by developers for objectives. + """, + "Omega_prime_method": """str : + ``"analytic"`` differentiates the bounce integrals with respect to rho + to obtain Ω'(s). ``"fd"`` instead finite differences Ω across the + surfaces of the radial grid. + """, + "_data_fft_r": """dict[str, jnp.ndarray] : + Radial derivatives, at fixed theta and zeta, of the field data in + ``_data_fft``. The Bounce2D counterpart of ``_data_eta_r``. Required + on the Bounce2D path. + """, + "_angle_r": """jnp.ndarray : + Radial derivative of ``_angle``, required alongside ``_data_fft_r`` + because the poloidal angle map moves with rho through both iota and + lambda. + """, + "_data_eta_r": """dict[str, jnp.ndarray] : + Radial derivatives, at fixed eta and zeta, of the field data in + ``_data_eta``, built by ``TrappedResonance.compute``. Only the keys in + ``_ETA_BOUNCE_KEYS`` are read. Required on the Bounce1D path: Ω'(s) is + obtained by differentiating the bounce integrals through them. This + private parameter is intended to be used only by developers for + objectives. + """, + "_data_psa": """dict[str, jnp.ndarray] : + Field data evaluated on ``_psa_grid``, built by + ``TrappedResonance.compute``. This private parameter is intended to + be used only by developers for objectives. + """, +} + + +def _phase_space_average( + vtau_out, + f_res, + pitch_inv, + pitch_inv_weight, + fl_length, + num_alpha=None, + fill_value=jnp.nan, +): + """Phase-space average of f_res. + + Computes = Σ_w ∫dα ∫dλ v·τ_b · f / (2 ∫dα ∫dl/B). + Pitch quadrature uses Gauss-Legendre weights from + ``Bounce1D.get_pitch_inv_quad``. + + Parameters + ---------- + vtau_out : jnp.ndarray, shape (rho, alpha, Bcrit, well) + Bounce integral of v·τ. + f_res : jnp.ndarray, shape (rho, Bcrit, well) + Objective function per (rho, pitch, well). + pitch_inv : jnp.ndarray, shape (rho, Bcrit) + Pitch inverse values. + pitch_inv_weight : jnp.ndarray, shape (rho, Bcrit) + Quadrature weights for pitch integration. + fl_length : jnp.ndarray, shape (rho,) + Mean-alpha fieldline length, i.e. mean_α ∫ dl/B. + num_alpha : int or None, optional + If ``None``, number of field lines considered is consistent with bounce + integration in ``_trapped_EP_resonance``. If not ``None``, specifies number + of total field lines to consider. + Defaults to ``None``. + fill_value : float, optional + Value to set bounce integration outputs to if no well is found. Cannot use + ``jnp.nan`` to retain optimization abilities. Cannot use 0 for confusion + with other quantities and averages. + Defaults to 11.0. + + Returns + ------- + f_res_avg : jnp.ndarray, shape (rho,) + """ + if num_alpha is None: + num_alpha = vtau_out.shape[1] + # Zero out BT-filtered (fill_value sentinel) f_res before weighting by vtau. + f_res_clean = jnp.where( + _is_valid_value(f_res, fill_value), f_res, jnp.zeros_like(f_res) + ) + integrand = vtau_out * f_res_clean[:, jnp.newaxis, :, :] + # 1. Integrate over pitch (per α, per well): ∫dλ g(λ) = ∫dp g(1/p)/p² + pitch_inv_4d = pitch_inv[:, jnp.newaxis, :, jnp.newaxis] + pitch_mask = _is_valid_value(pitch_inv_4d, fill_value) + integrand_mask = _is_valid_value(vtau_out, fill_value) & pitch_mask + safe_pitch_inv = jnp.where(pitch_mask, pitch_inv_4d, jnp.ones_like(pitch_inv_4d)) + pitch_integrated = _masked_sum( + integrand + * pitch_inv_weight[:, jnp.newaxis, :, jnp.newaxis] + / safe_pitch_inv**2, + mask=integrand_mask, + axis=2, + ) # (rho, alpha, well) + # 2. Sum over α (discrete ∫dα) + alpha_summed = pitch_integrated.sum(axis=1) # (rho, well) + # 3. Sum over wells + numerator = _masked_sum( + alpha_summed, + mask=_is_valid_value(alpha_summed, fill_value), + axis=-1, + ) # (rho,) + # Denominator: 2 · Σ_α ∫dl/B = 2 · N_α · mean_α(∫dl/B) + return safediv(numerator, 2 * num_alpha * fl_length) + + +def _resonance_physics( + alpha_drift_out, + s_drift_out, + vtau_out, + iotas, + rhos, + rho_res, + KE_frac, + nfp, + M, + N, + res_arr, + q_arr, + eta_vals, + eta_res, + weight_method, + Delta_Omega, + wd_blur, + fill_value, + stab_sacrifice, + dOmega_drho=None, + cropping_DOmega=False, +): + """Compute resonance frequencies, weights, island widths, and f_res. + + Takes bounce-averaged drifts and converts them to physical frequencies, + computes the normalised precession frequency Omega and its radial + derivative Omega'(s), assigns resonance weights, evaluates Fourier + coefficients of the radial drift, and finally computes island widths. + + Parameters + ---------- + alpha_drift_out : jnp.ndarray, shape (rho, alpha, Bcrit, well) + Bounce-averaged poloidal drift (dimensionless, before energy scaling). + s_drift_out : jnp.ndarray, shape (rho, alpha, Bcrit, well) + Bounce-averaged radial drift (dimensionless, before energy scaling). + vtau_out : jnp.ndarray, shape (rho, alpha, Bcrit, well) + Bounce integral of v·τ. + iotas : jnp.ndarray, shape (rho,) + Rotational transform per surface. + rhos : jnp.ndarray, shape (rho,) + Flux surface labels. + rho_res : float + Radial grid spacing. + KE_frac : float + Fraction of the 3.5 MeV D-T fusion alpha-particle birth energy to use + for the energetic particle kinetic energy. + nfp : int + Number of field periods. + M, N : int + Poloidal and toroidal mode numbers for resonance condition. + res_arr : jnp.ndarray, shape (res,) + Resonance frequency ratios p/q. + q_arr : jnp.ndarray, shape (res,) + Toroidal mode numbers of resonances. + eta_vals : jnp.ndarray, shape (num_eta,) + Uniform eta grid on [0, 2π). + eta_res : float + Grid spacing for eta. + weight_method : str + ``"linear"`` or ``"bump"`` resonance weighting. + Delta_Omega : float or None + Half-width for bump weighting. + wd_blur : float + Multiplicative blur factor used to compute bump half-width from + adjacent-surface Omega spacing when ``Delta_Omega`` is not provided. + stab_sacrifice : bool + Whether to sacrifice accuracy for stability in island widths. + dOmega_drho : jnp.ndarray or None + Shape (rho, pitch, well). + ∂Ω/∂ρ at fixed λ and η, obtained by differentiating the bounce + integrals. If ``None``, ∂Ω/∂ρ is estimated by finite differences + across the surfaces of ``rhos`` instead, which needs both neighbours + of a surface to be valid and is only as accurate as ``rho_res`` + allows. + Defaults to ``None``. + cropping_DOmega : bool + If ``True``, Delta_Omega calculation is clipped by + ``0.01 * max(Omega_eta) < Delta_Omega < 0.10 * max(Omega_eta)``. + This must be when using the ``bump`` weighting method and + ``Delta_Omega = None`` case. Otherwise this quantity is ignored. + Defaults to ``False``. + + Returns + ------- + result : dict + Dictionary containing: + + f_res : jnp.ndarray, shape (rho, pitch, well) + Per-(rho, pitch, well) resonance objective contribution: island + width squared, summed over resonances in ``res_arr`` and weighted + by ``res_weight``, optionally scaled by ``Omega_prime_s`` if + ``stab_sacrifice``. Phase-space averaged elsewhere to form the + "trapped EP resonance" objective. The least-squares objective + built on top of this squares its residual again, so the net + penalty is (island width)^4, optionally scaled by + ``Omega_prime_s**2``, matching the pre-existing scaling without + squaring it twice. + Omega : jnp.ndarray, shape (rho, pitch, well) + Normalized precession frequency, i.e. Omega_eta, + ``eta_drift_avg / omega_bounce_avg``, compared against the + rational ratios in ``res_arr`` to locate resonances. + omega_bounce_avg : jnp.ndarray, shape (rho, pitch, well) + Alpha-averaged bounce frequency ``2π / tau_bounce``. + eta_drift_avg : jnp.ndarray, shape (rho, pitch, well) + Alpha-averaged eta precession frequency ω_η (the numerator of + ``Omega``/Omega_eta). + omega_bounce : jnp.ndarray, shape (rho, alpha, pitch, well) + Bounce frequency per field line, before alpha-averaging. + eta_drift : jnp.ndarray, shape (rho, alpha, pitch, well) + Eta precession frequency ω_η per field line, before + alpha-averaging. + Omega_prime_s : jnp.ndarray, shape (rho, pitch, well) + Radial derivative dOmega/ds, where s = rho², from ``dOmega_drho``. + res_weight : jnp.ndarray, shape (rho, pitch, well, res) + Weight assigning each (rho, pitch, well) to each resonance in + ``res_arr``, via 2-point linear interpolation or a smooth bump + function, depending on ``weight_method``. + f_q_abs : jnp.ndarray, shape (rho, pitch, well, res) + Magnitude of the q-th eta-Fourier harmonic of the bounce-averaged + radial (s) drift. + Delta_s : jnp.ndarray, shape (pitch, well, res) + Resonance-weighted island width (s = rho² units), summed over + rho; a diagnostic quantity. + Delta_s_prof : jnp.ndarray, shape (rho, pitch, well, res) + Per-surface island width (s = rho² units) at each resonance. + s_res : jnp.ndarray, shape (pitch, well, res) + Resonance-weighted mean s = rho² location of each resonance. + valid : jnp.ndarray, shape (rho, pitch, well) + Boolean mask, ``True`` where a trapped particle exists at all + alpha/eta and both ``Omega`` and ``Omega_prime_s`` are defined. + """ + freq = _frequencies( + alpha_drift_out, s_drift_out, vtau_out, iotas, KE_frac, nfp, M, N, fill_value + ) + Omega = freq["Omega"] + valid = freq["valid"] + omega_bounce_avg = freq["omega_bounce_avg"] + eta_drift_avg = freq["eta_drift_avg"] + omega_bounce = freq["omega_bounce"] + eta_drift = freq["eta_drift"] + tau_bounce = freq["tau_bounce"] + s_drift = freq["s_drift"] + + if dOmega_drho is None: + # Omega'(s) by finite differences across the surfaces of ``rhos``. + # Double-where: replace invalid Omega with 0 before the arithmetic so + # no nan derivative flows through a discarded jnp.where branch. + Omega_safe = jnp.where(valid, Omega, 0.0) + valid_prev = jnp.concatenate( + [jnp.zeros((1,) + valid.shape[1:], dtype=bool), valid[:-1]], axis=0 + ) + valid_next = jnp.concatenate( + [valid[1:], jnp.zeros((1,) + valid.shape[1:], dtype=bool)], axis=0 + ) + Omega_prev_safe = jnp.concatenate( + [jnp.zeros((1,) + Omega.shape[1:]), Omega_safe[:-1]], axis=0 + ) + Omega_next_safe = jnp.concatenate( + [Omega_safe[1:], jnp.zeros((1,) + Omega.shape[1:])], axis=0 + ) + grad_central = (Omega_next_safe - Omega_prev_safe) / (2 * rho_res) + grad_forward = (Omega_next_safe - Omega_safe) / rho_res + grad_backward = (Omega_safe - Omega_prev_safe) / rho_res + dOmega_drho = jnp.where( + valid & valid_prev & valid_next, + grad_central, + jnp.where( + valid & valid_next & ~valid_prev, + grad_forward, + jnp.where( + valid & valid_prev & ~valid_next, + grad_backward, + fill_value, + ), + ), + ) + else: + # Analytic ∂Ω/∂ρ needs no neighbouring surface, so it is defined + # wherever Omega itself is. + dOmega_drho = jnp.where(valid, dOmega_drho, fill_value) + + # rho > 0 on every surface, so the division needs no guard, but the + # sentinel must survive it: a finite difference leaves fill_value on a + # surface with no valid neighbour, and dividing that would turn it into a + # plausible looking number. + Omega_prime_s = jnp.where( + _is_valid_value(dOmega_drho, fill_value), + dOmega_drho / (2 * rhos[:, None, None]), + fill_value, + ) + + # Resonance weights + Omega_broad = Omega[..., None] + res_broad = res_arr[None, None, None, :] + + # Narrower than ``valid`` under finite differences, which is undefined on + # the surfaces at either end; identical to it under the analytic path. + valid_prime = valid & _is_valid_value(Omega_prime_s, fill_value) + + if weight_method == "bump": + if Delta_Omega is None: + Omega_safe_bump = jnp.where(valid, Omega, 0.0) + Omega_prev_b = Omega_safe_bump[:-1, :, :] + Omega_next_b = Omega_safe_bump[1:, :, :] + valid_pair = jnp.logical_and(valid[:-1, :, :], valid[1:, :, :]) + domega_arr = jnp.where( + valid_pair, + jnp.abs(Omega_next_b - Omega_prev_b), + 0.0, + ) # (rho-1, pitch, well) + from desc.objectives.utils import softmax as _softmax + + Delta_Omega_val = (wd_blur * _softmax(domega_arr, alpha=50, axis=0) / 2.0)[ + None, :, :, None + ] + if cropping_DOmega: + # Delta_Omega_val needs to be cropped if resolution is too low + # or Omega_eta shear is too high + Omega_max = _softmax(Omega_safe_bump, alpha=50, axis=0)[ + None, :, :, None + ] + Delta_Omega_val_max = ( + 0.1 * Omega_max + ) # DeltaOmega < 10% of maximum Omega + Delta_Omega_val_min = ( + 0.01 * Omega_max + ) # DeltaOmega > 1% of maximum Omega + Delta_Omega_val = jnp.where( + Delta_Omega_val > Delta_Omega_val_max, + Delta_Omega_val_max, + Delta_Omega_val, + ) + Delta_Omega_val = jnp.where( + Delta_Omega_val < Delta_Omega_val_min, + Delta_Omega_val_min, + Delta_Omega_val, + ) + else: + Delta_Omega_val = Delta_Omega + a = res_broad + Delta_Omega_val + b = res_broad - Delta_Omega_val + in_interval = (Omega_broad >= b) & (Omega_broad <= a) + denom = (Omega_broad - b) * (Omega_broad - a) + exp_arg = safediv((2.0 * Delta_Omega_val) ** 2, denom, fill=-1e10) + C_norm = safediv(71.12518788738504, Delta_Omega_val, fill=0.0) + w_raw = rho_res * C_norm * jnp.abs(dOmega_drho[..., None]) * jnp.exp(exp_arg) + # Weight is non-zero only if in interval and valid + res_weight = jnp.where(in_interval & valid_prime[..., None], w_raw, 0) + else: + # Double-where: use Omega_safe (0 at invalid entries) so that + # safediv never sees fill_value operands, preventing NaN gradients. + Omega_safe_lin = jnp.where(valid, Omega, 0.0) + Omega_broad_safe = Omega_safe_lin[..., None] + zero_O = jnp.zeros_like(Omega_safe_lin[:1]) + zero_v = jnp.zeros_like(valid[:1]) + # Neighbouring surfaces, zero padded at the ends where there is none. + # A resonance is bracketed when it lies between this surface's Omega + # and the neighbour's, in either direction; the weight is the linear + # interpolant across that bracket. Ordered prev then next, so that + # next overwrites prev where both bracket the same resonance. + res_weight = 0.0 + for Omega_nb, valid_nb in ( + ( + jnp.concatenate([zero_O, Omega_safe_lin[:-1]], axis=0), + jnp.concatenate([zero_v, valid[:-1]], axis=0), + ), + ( + jnp.concatenate([Omega_safe_lin[1:], zero_O], axis=0), + jnp.concatenate([valid[1:], zero_v], axis=0), + ), + ): + Omega_nb = Omega_nb[..., None] + between = valid_nb[..., None] & ( + ((Omega_nb >= res_broad) & (res_broad >= Omega_broad_safe)) + | ((Omega_nb <= res_broad) & (res_broad <= Omega_broad_safe)) + ) + res_weight = jnp.where( + between, + safediv(Omega_nb - res_broad, Omega_nb - Omega_broad_safe, fill=0.0), + res_weight, + ) + + # Set weight to zero for invalid points. + res_weight = jnp.where(valid_prime[..., None], res_weight, 0) + + # Fourier analysis of radial drift. + # Only perform FT if all eta points are valid. + # Mask out fill_value entries in s_drift (set by bt/rt filters) so they + # don't contaminate the alpha sum. + s_drift_valid = jnp.where(_is_valid_value(s_drift_out, fill_value), s_drift, 0.0) + ft_integrand = s_drift_valid * tau_bounce + + phase = q_arr[None, :] * eta_vals[:, None] + cos_phase = jnp.cos(phase) + sin_phase = jnp.sin(phase) + ft_cos = ft_integrand[..., None] * cos_phase[None, :, None, None, :] + ft_sin = ft_integrand[..., None] * sin_phase[None, :, None, None, :] + ft_prefactor = eta_res / jnp.pi + f_q_c = ft_prefactor * jnp.sum(ft_cos, axis=1) + f_q_s = ft_prefactor * jnp.sum(ft_sin, axis=1) + + f_q_r2 = f_q_c**2 + f_q_s**2 + is_zero = f_q_r2 == 0 + f_q_abs = 0.5 * jnp.sqrt(jnp.where(is_zero, 1.0, f_q_r2)) + f_q_abs = jnp.where(is_zero, 0.0, f_q_abs) + + # Filter FT results to valid points. + f_q_abs = jnp.where(valid_prime[..., None], f_q_abs, 0.0) + + # Island widths + q_iw = q_arr[None, None, None, :] + denom = jnp.pi * q_iw * jnp.abs(Omega_prime_s[..., None]) + # (Delta_s / 4)² = f_q / (pi q |Omega'|), so both forms share one quotient. + width_sq = safediv(f_q_abs, denom, fill=0.0) + Delta_s_profile = 4 * jnp.sqrt(width_sq) + Delta_s_sq_profile = 16 * width_sq + Delta_s_sq_sum = (Delta_s_sq_profile * res_weight).sum(axis=-1) + + if stab_sacrifice: + f_res = Delta_s_sq_sum * Omega_prime_s + else: + f_res = Delta_s_sq_sum + + # Sum over radius to get weighted island width and resonance location. + Delta_s = (Delta_s_profile * res_weight).sum(axis=0) + s_vals = rhos**2 + s_res = (res_weight * s_vals[:, None, None, None]).sum(axis=0) + + return { + "f_res": f_res, # (rho, pitch, well) + "Omega": Omega, # (rho, pitch, well) + "omega_bounce_avg": omega_bounce_avg, # (rho, pitch, well) + "eta_drift_avg": eta_drift_avg, # (rho, pitch, well) + "omega_bounce": omega_bounce, # (rho, alpha, pitch, well) + "eta_drift": eta_drift, # (rho, alpha, pitch, well) + "Omega_prime_s": Omega_prime_s, # (rho, pitch, well) + "res_weight": res_weight, # (rho, pitch, well, res) + "f_q_abs": f_q_abs, # (rho, pitch, well, res) + "Delta_s": Delta_s, # (pitch, well, res), rho-weighted diagnostic + "Delta_s_prof": Delta_s_profile, # (rho, pitch, well, res) + "s_res": s_res, # (pitch, well, res), rho-weighted resonance location + "valid": valid_prime, # (rho, pitch, well) + } + + +@register_compute_fun( + name="trapped EP resonance", + label=("Trapped Energetic Particle Resonance Objective Function"), + units="s^-2", + units_long="seconds squared", + description="Trapped Energetic Particle Resonance Minimizer", + dim=1, + params=[], + transforms={"grid": []}, + profiles=[], + coordinates="r", + data=["iota", "iota_r", "min_tz |B|", "max_tz |B|", "V_psi"], + grid_requirement={"is_meshgrid": True}, + public=False, + **_bounce1D_doc, + **_resonance_doc, +) +def _trapped_EP_resonance(params, transforms, profiles, data, **kwargs): + """Trapped particle resonance penalty. + + Three stages: + 1. Bounce integrals (per-surface, via ``_compute1D`` / ``batch_map``) + 2. Resonance physics (cross-surface, via ``_resonance_physics``) + 3. Phase-space average (via ``_phase_space_average``) + + The eta/PSA grids and the field data evaluated on them (``_eta_grid``, + ``_psa_grid``, ``_data_eta``, ``_data_psa``) are built by the caller (see + ``TrappedResonance.compute``) rather than here, since building them + requires the full ``Equilibrium`` object, which compute functions must + stay pure with respect to (only ``params``/``transforms``/``profiles``/ + ``data``) to remain properly differentiable and dispatchable for any + parameterization. + """ + num_pitch = kwargs.get("num_pitch") + num_well = 1 + M = kwargs.get("M", 1) + N = kwargs.get("N", 1) + nfp = kwargs.get("nfp") + KE_frac = kwargs.get("KE_frac") + pitch_invs = kwargs.get("pitch_invs") + rho_res = kwargs.get("rho_res") + eta_res = kwargs.get("eta_res") + res_arr = kwargs.get("res_arr") + p_arr = kwargs.get("p_arr") + q_arr = kwargs.get("q_arr") + quad = kwargs.get("quad") + surf_batch_size = kwargs.get("surf_batch_size", 1) + num_eta = kwargs.get("num_eta") + weight_method = kwargs.get("weight_method", "linear") + Delta_Omega = kwargs.get("Delta_Omega") + wd_blur = kwargs.get("wd_blur", 1.25) + fill_value = kwargs.get("fill_value", 11) + zeta = kwargs.get("zeta") + stab_sacrifice = kwargs.get("stab_sacrifice", False) + bt_filter_flag = kwargs.get("bt_filter_flag", False) + cropping_DOmega = kwargs.get("cropping_DOmega", False) + eta_grid = kwargs.get("_eta_grid") + psa_grid = kwargs.get("_psa_grid") + data_eta = kwargs.get("_data_eta") + data_psa = kwargs.get("_data_psa") + num_transit = kwargs.get("num_transit", 1) + use_bounce1d = kwargs.get("use_bounce1d", False) + # Bounce2D path only. + angle = kwargs.get("_angle") + fft_grid = kwargs.get("_fft_grid") + data_fft = kwargs.get("_data_fft") + Y_B = kwargs.get("Y_B") + spline = kwargs.get("spline", True) + vander = kwargs.get("_vander") + + nufft_eps = kwargs.get("nufft_eps", 1e-10) + + base_grid = transforms["grid"] + iotas = base_grid.compress(data["iota"]) + iotas_r = base_grid.compress(data["iota_r"]) + rhos = base_grid.compress(base_grid.nodes[:, 0]) + eta_vals = jnp.linspace(0, 2 * jnp.pi, num_eta, endpoint=False) + + # --- 1. Bounce integrals on the eta grid --- + pitch_invs_use, pitch_inv_weight_use = _global_pitch_quad( + base_grid, data, num_pitch, pitch_invs + ) + + drift_names = list(_BOUNCE_INTEGRAND_KEYS) + + if use_bounce1d: + + def drifts(data_in): + bounce = Bounce1D(eta_grid, data_in, quad, is_reshaped=True) + points = bounce.points(data_in["pitch_inv"], num_well=num_well) + v_tau, _alpha_drift, _s_drift = bounce.integrate( + [_v_tau, _alpha_drift_integrand, _radial_drift], + data_in["pitch_inv"], + data_in, + drift_names, + num_well=num_well, + ) + _alpha_drift = safediv(_alpha_drift, v_tau) + _s_drift = 4 * safediv(_s_drift, v_tau) + return _alpha_drift, _s_drift, points, v_tau, data_in["pitch_inv"] + + def bounce_and_omega(data_in, iotas_in): + """Bounce integrals and Omega, as a function of the field line data. + + Kept as its own function so that pushing the radial derivatives of + the field line data through it in forward mode gives the exact + ∂Ω/∂ρ. The pitch grid is closed over rather than passed in, so that + derivative is taken at fixed λ, as the resonance condition requires. + """ + _alpha_drift, _s_drift, _points, _v_tau, _pitch_inv = _compute1D( + drifts, + {name: data_in[name] for name in drift_names}, + data_in, + eta_grid, + num_pitch, + surf_batch_size, + pitch_invs=pitch_invs_use, + pitch_inv_weight=pitch_inv_weight_use, + ) + _Omega = _frequencies( + _alpha_drift, + _s_drift, + _v_tau, + iotas_in, + KE_frac, + nfp, + M, + N, + fill_value, + )["Omega"] + return ( + _alpha_drift, + _s_drift, + _points[0], + _points[1], + _v_tau, + _pitch_inv, + _Omega, + ) + + data_eta_r = kwargs.get("_data_eta_r") + eta_bounce_data = {name: data_eta[name] for name in _ETA_BOUNCE_KEYS} + if data_eta_r is None: + bounce_out = bounce_and_omega(eta_bounce_data, iotas) + dOmega_drho = None + else: + # ∂Ω/∂ρ at fixed λ and η, from the radial derivatives of the field + # line data. Forward mode differentiates the bounce points along + # with the quadrature, so endpoint motion is accounted for exactly. + bounce_out, bounce_dot = jax.jvp( + bounce_and_omega, + (eta_bounce_data, iotas), + ({name: data_eta_r[name] for name in _ETA_BOUNCE_KEYS}, iotas_r), + ) + dOmega_drho = bounce_dot[-1] + alpha_drift_out, s_drift_out, z1, z2, vtau_out, pitch_inv_out = bounce_out[:-1] + points = (z1, z2) + else: + + def drifts(data_in): + bounce = Bounce2D( + fft_grid, + data_in, + data_in["angle"], + Y_B, + data_in["alpha_per_rho"].T, + num_transit, + quad, + nufft_eps=nufft_eps, + is_fourier=True, + spline=spline, + vander=vander, + ) + points = bounce.points(data_in["pitch_inv"], num_well=num_well) + v_tau, _alpha_drift, _s_drift = bounce.integrate( + [_v_tau, _alpha_drift_integrand, _radial_drift], + data_in["pitch_inv"], + data_in, + drift_names, + num_well=num_well, + nufft_eps=nufft_eps, + is_fourier=True, + low_ram=True, + ) + _alpha_drift = safediv(_alpha_drift, v_tau) + _s_drift = 4 * safediv(_s_drift, v_tau) + return _alpha_drift, _s_drift, points, v_tau, data_in["pitch_inv"] + + def bounce_and_omega(data_in, angle_in, iotas_in): + """Bounce integrals and Omega, as a function of the field data. + + The Bounce2D counterpart of the Bounce1D closure above, and + differentiated the same way. ``fft_grid`` enters only through + reshapes and compressions, so it stays fixed while the data it + indexes carries the radial tangents. + """ + # eta -> alpha, which differs between flux surfaces because it + # depends on iota, and so moves with rho. Carried with rho leading + # so ``batch_map`` slices it with the rest of the data; + # ``Bounce2D`` wants (num alpha, num rho). + _alpha_eta = eta_vals[None, :] * (N * nfp - iotas_in[:, None] * M) / nfp + _alpha_drift, _s_drift, _points, _v_tau, _pitch_inv = _compute2D( + drifts, + {name: data_in[name] for name in drift_names}, + data_in, + fft_grid, + angle_in, + _alpha_eta, + num_pitch, + surf_batch_size, + pitch_invs=pitch_invs_use, + pitch_inv_weight=pitch_inv_weight_use, + ) + _Omega = _frequencies( + _alpha_drift, + _s_drift, + _v_tau, + iotas_in, + KE_frac, + nfp, + M, + N, + fill_value, + )["Omega"] + return ( + _alpha_drift, + _s_drift, + _points[0], + _points[1], + _v_tau, + _pitch_inv, + _Omega, + ) + + data_fft_r = kwargs.get("_data_fft_r") + angle_r = kwargs.get("_angle_r") + fft_bounce_data = {name: data_fft[name] for name in _FFT_BOUNCE_KEYS} + if data_fft_r is None: + bounce_out = bounce_and_omega(fft_bounce_data, angle, iotas) + dOmega_drho = None + else: + bounce_out, bounce_dot = jax.jvp( + bounce_and_omega, + (fft_bounce_data, angle, iotas), + ( + {name: data_fft_r[name] for name in _FFT_BOUNCE_KEYS}, + angle_r, + iotas_r, + ), + ) + dOmega_drho = bounce_dot[-1] + alpha_drift_out, s_drift_out, z1, z2, vtau_out, pitch_inv_out = bounce_out[:-1] + points = (z1, z2) + + # --- 1b. Barely-trapped filter --- + # Zero out wells whose poloidal bounce width exceeds 2π. Those particles + # sample the whole poloidal angle within one bounce, so the resonance + # analysis does not describe them. + if bt_filter_flag: + z1, z2 = points + delta_chi = jnp.abs( + jnp.abs(z1 - z2) * (M * iotas[:, None, None, None] - N * nfp) + ) + s_drift_out = jnp.where(delta_chi < 2 * jnp.pi, s_drift_out, fill_value) + + # --- 2. Resonance physics (cross-surface) --- + res = _resonance_physics( + alpha_drift_out, + s_drift_out, + vtau_out, + iotas, + rhos, + rho_res, + KE_frac, + nfp, + M, + N, + res_arr, + q_arr, + eta_vals, + eta_res, + weight_method, + Delta_Omega, + wd_blur, + fill_value, + stab_sacrifice, + dOmega_drho, + cropping_DOmega, + ) + + # --- 3. Phase-space average on the PSA grid (uniform in alpha) --- + # Skip PSA when custom pitch_invs are provided. + if pitch_invs is None: + if use_bounce1d: + num_alpha_psa = psa_grid.num_poloidal + + def drifts_vtau(data_local): + bounce = Bounce1D(psa_grid, data_local, quad, is_reshaped=True) + v_tau = bounce.integrate( + [_v_tau], + data_local["pitch_inv"], + data_local, + [], + num_well=num_well, + )[0] + return v_tau, data_local + + vtau_psa, _data_psa = _compute1D( + drifts_vtau, + {}, + data_psa, + psa_grid, + num_pitch, + surf_batch_size, + pitch_invs=pitch_invs_use, + pitch_inv_weight=pitch_inv_weight_use, + ) + num_rho_psa = psa_grid.num_rho + # Normalize by the alpha-averaged field line length of a single + # toroidal transit, ∫dl/B = ∫dζ/B^ζ, along the field lines of the + # phase-space average grid. + Bzeta_psa = Bounce1D.reshape(psa_grid, data_psa["B^zeta"]) + n_1t = len(zeta) // num_transit + fl_length = jnp.abs( + simpson(1 / Bzeta_psa[..., :n_1t], x=zeta[:n_1t], axis=-1).mean(axis=1) + ) + else: + # The phase-space average is taken over field lines uniform in + # alpha, unlike the eta grid above, so the labels are the same on + # every surface. Reuses the same (θ, ζ) data and angle. + num_alpha_psa = num_eta + alpha_psa = jnp.broadcast_to( + jnp.linspace(0, 2 * jnp.pi, num_alpha_psa, endpoint=False), + (rhos.size, num_alpha_psa), + ) + + def drifts_vtau(data_local): + bounce = Bounce2D( + fft_grid, + data_local, + data_local["angle"], + Y_B, + data_local["alpha_per_rho"].T, + num_transit, + quad, + nufft_eps=nufft_eps, + is_fourier=True, + spline=spline, + vander=vander, + ) + v_tau = bounce.integrate( + [_v_tau], + data_local["pitch_inv"], + data_local, + [], + num_well=num_well, + nufft_eps=nufft_eps, + is_fourier=True, + )[0] + return v_tau, data_local + + vtau_psa, _data_psa = _compute2D( + drifts_vtau, + {}, + data_fft, + fft_grid, + angle, + alpha_psa, + num_pitch, + surf_batch_size, + pitch_invs=pitch_invs_use, + pitch_inv_weight=pitch_inv_weight_use, + ) + num_rho_psa = rhos.size + fl_length = base_grid.compress(data["V_psi"]) / (2 * jnp.pi) + + if vtau_psa.ndim == 3 and vtau_psa.shape[0] == num_rho_psa * num_alpha_psa: + vtau_psa = vtau_psa.reshape( + num_rho_psa, num_alpha_psa, vtau_psa.shape[1], vtau_psa.shape[2] + ) + + f_res_avg = _phase_space_average( + vtau_psa, + res["f_res"], + _data_psa["pitch_inv"], + _data_psa["pitch_inv weight"], + fl_length, + num_alpha=num_alpha_psa, + fill_value=fill_value, + ) + data["trapped EP resonance"] = base_grid.expand(f_res_avg) + else: # Custom pitch_invs specified: skip phase-space average, + # just return the raw resonance physics results + data["trapped EP resonance"] = { + **res, + "pitch_inv": pitch_inv_out, + "res_arr": res_arr, + "p_arr": p_arr, + "q_arr": q_arr, + "rhos": rhos, + } + + return data diff --git a/desc/integrals/_bounce_utils.py b/desc/integrals/_bounce_utils.py index fee1609022..9c7c08fe72 100644 --- a/desc/integrals/_bounce_utils.py +++ b/desc/integrals/_bounce_utils.py @@ -687,8 +687,9 @@ def get_alphas(alpha, iota, num_transit, NFP): Parameters ---------- alpha : jnp.ndarray - Shape (num α, ). - Starting field line poloidal labels {αᵢ₀}. + Shape (num α, ) or (num α, num ρ). + Starting field line poloidal labels {αᵢ₀}. If two-dimensional, then + the labels may differ between flux surfaces. iota : jnp.ndarray Shape (num ρ, ). Rotational transform normalized by 2π. @@ -704,7 +705,7 @@ def get_alphas(alpha, iota, num_transit, NFP): Set of field line poloidal coordinates {Aᵢ | Aᵢ = (αᵢ₀, αᵢ₁, ..., αᵢ₍ₘ₋₁₎)}. """ - alpha = alpha[:, None, None] + alpha = alpha[:, None, None] if alpha.ndim == 1 else alpha[..., None] iota = iota[:, None] return alpha + iota * (2 * jnp.pi / NFP) * jnp.arange(num_transit * NFP) @@ -721,8 +722,8 @@ def theta_on_fieldlines(angle, iota, alpha, num_transit, NFP, *, X_min=24): Shape (num ρ, ). Rotational transform normalized by 2π. alpha : jnp.ndarray - Shape (num α, ). - Starting field line poloidal labels {αᵢ₀}. + Shape (num α, ) or (num α, num ρ). + Starting field line poloidal labels {αᵢ₀}. See ``get_alphas``. num_transit : int Number of toroidal transits to follow field line. NFP : int @@ -777,7 +778,7 @@ def theta_on_fieldlines(angle, iota, alpha, num_transit, NFP, *, X_min=24): """ X = angle.shape[-2] Y = truncate_rule(angle.shape[-1]) - num_alpha = alpha.size + num_alpha = alpha.shape[0] domain = (0, 2 * jnp.pi / NFP) # peeling off field lines diff --git a/desc/integrals/_interp_utils.py b/desc/integrals/_interp_utils.py index b77a32e7ea..babf9279fb 100644 --- a/desc/integrals/_interp_utils.py +++ b/desc/integrals/_interp_utils.py @@ -424,10 +424,25 @@ def root(b, c, d): c = c / a Q = (b**2 - 3 * c) / 9 R = (2 * b**3 - 9 * b * c) / 54 + d / (2 * a) + three_real = R**2 < Q**3 + # Double where. jnp.where evaluates both branches, and each is + # singular exactly where the other is selected: arccos leaves [-1, 1] + # when R² ≥ Q³, and Q / A divides by zero when Q = R = 0. Discarding + # the value is not enough, because the nan derivative of the discarded + # branch survives the where in reverse mode. Substituting a harmless + # argument there leaves every selected value untouched. return jnp.where( - R**2 < Q**3, - irreducible(jnp.abs(Q), R, b), - reducible(Q, R, b), + three_real, + irreducible( + jnp.where(three_real, jnp.abs(Q), 1.0), + jnp.where(three_real, R, 0.0), + b, + ), + reducible( + jnp.where(three_real, 1.0, Q), + jnp.where(three_real, 1.0, R), + b, + ), ) return jnp.where( diff --git a/desc/objectives/__init__.py b/desc/objectives/__init__.py index 52ca059e98..33f411dc3c 100644 --- a/desc/objectives/__init__.py +++ b/desc/objectives/__init__.py @@ -57,6 +57,7 @@ from ._power_balance import FusionPower, HeatingPowerISS04 from ._profiles import Pressure, RotationalTransform, Shear, ToroidalCurrent from ._stability import BallooningStability, MagneticWell, MercierStability +from ._trapped_resonance import TrappedResonance from .getters import ( get_equilibrium_objective, get_fixed_axis_constraints, diff --git a/desc/objectives/_trapped_resonance.py b/desc/objectives/_trapped_resonance.py new file mode 100644 index 0000000000..240f4fb66c --- /dev/null +++ b/desc/objectives/_trapped_resonance.py @@ -0,0 +1,843 @@ +"""Objectives for trapped energetic particle resonance.""" + +import numpy as np +from interpax_fft import cheb_pts, fourier_pts +from orthax.legendre import leggauss +from scipy.constants import elementary_charge + +from desc.backend import jax, jnp +from desc.compute import get_profiles, get_transforms +from desc.compute._trapped_resonance import ( + _FFT_BOUNCE_KEYS, + _build_eta_grid, + _build_eta_source_grid, +) +from desc.compute.data_index import data_index +from desc.compute.utils import _compute as compute_fun +from desc.compute.utils import _parse_parameterization, get_data_deps +from desc.grid import Grid, LinearGrid +from desc.integrals._bounce_utils import Y_B_rule, get_vander_spline +from desc.integrals.bounce_integral import Bounce1D, Bounce2D +from desc.utils import Timer, errorif + +from ..integrals.quad_utils import ( + automorphism_sin, + get_quadrature, + grad_automorphism_sin, +) +from .objective_funs import _Objective +from .utils import _parse_callable_target_bounds + + +def _shift_grid_rho(grid, drho): + """``grid`` with rho displaced by ``drho``, traceable in ``drho``. + + A ``LinearGrid`` is built from numpy and so cannot carry a tangent in rho. + This rebuilds it as a jitable ``Grid`` over the same nodes, which lets the + quantities computed on it be differentiated with respect to rho. + + Everything copied over is either untouched by a radial shift or describes + the (theta, zeta) structure, which a radial shift leaves alone: + ``spacing`` and ``weights`` because surface integrals are over theta and + zeta; ``fft_poloidal`` and ``fft_toroidal`` because they record uniformity + in those angles, and ``_partial_sum`` refuses a grid without the former. + + ``M`` and ``N`` must be copied rather than left to the ``Grid`` defaults, + which report 0 instead of the source ``LinearGrid``'s resolution. Dropping + them zeroes the Fourier resolution silently rather than raising. + + Parameters + ---------- + grid : Grid + drho : float or jnp.ndarray + Radial displacement. + + Returns + ------- + Grid + + """ + rho, theta, zeta = grid.nodes.T + out = Grid( + nodes=jnp.column_stack([rho + drho, theta, zeta]), + spacing=grid.spacing, + weights=grid.weights, + coordinates=grid.coordinates, + period=grid.period, + NFP=grid.NFP, + sort=False, + is_meshgrid=grid.is_meshgrid, + jitable=True, + _unique_rho_idx=grid.unique_rho_idx, + _unique_poloidal_idx=grid.unique_poloidal_idx, + _unique_zeta_idx=grid.unique_zeta_idx, + _inverse_rho_idx=grid.inverse_rho_idx, + _inverse_poloidal_idx=grid.inverse_poloidal_idx, + _inverse_zeta_idx=grid.inverse_zeta_idx, + ) + out._fft_poloidal = grid.fft_poloidal + out._fft_toroidal = grid.fft_toroidal + out._M = grid.M + out._N = grid.N + return out + + +def _seeded_keys(keys, parameterization): + """Per-surface keys seeded onto a field line grid, and those it reads. + + ``_compute`` skips any quantity already present in ``data``, so a seeded + quantity shadows its own dependencies: only the seeded keys that some + recomputed quantity depends on directly can influence the result. Those are + the ones whose radial derivative has to be supplied for an analytic Ω'(s). + + Parameters + ---------- + keys : list[str] + Quantities requested on the field line grid. + parameterization : str + Parameterization of the thing being computed, e.g. + ``"desc.equilibrium.equilibrium.Equilibrium"``. + + Returns + ------- + seeded, consumed : tuple[set[str]] + + """ + index = data_index[parameterization] + closure = set(get_data_deps(keys, obj=parameterization)) | set(keys) + seeded = {k for k in closure if index.get(k, {}).get("coordinates", "") == "r"} + + consumed, visited, stack = set(), set(), list(keys) + while stack: + key = stack.pop() + if key in visited: + continue + visited.add(key) + if key in seeded and key not in keys: + continue # seeded, so it is not recomputed and its deps are unused + for dep in index.get(key, {}).get("dependencies", {}).get("data", []): + if dep in seeded: + consumed.add(dep) + stack.append(dep) + return seeded, consumed + + +def _tangent_keys(parameterization, keys): + """Map each seeded per-surface key to the key holding its radial derivative. + + Raises if a key the field line compute actually reads has no radial + derivative registered in ``data_index``, so that a missing tangent can + never quietly turn into a wrong Ω'(s). + """ + seeded, consumed = _seeded_keys(keys, parameterization) + index = data_index[parameterization] + # DESC spells d/dρ by appending "_r", so a key already carrying one takes + # a bare "r" instead: "iota" -> "iota_r", but "iota_r" -> "iota_rr". + tangent = { + key: t_key + for key in seeded + if (t_key := key + "r" if key.endswith(("_r", "_rr")) else key + "_r") in index + } + # "rho" differentiates to 1; "p_r" is covered by _p_rr. + missing = sorted(consumed - set(tangent) - {"rho", "p_r"}) + if missing: + raise NotImplementedError( + f"TrappedResonance needs the radial derivative of {missing} to " + "differentiate the bounce integrals, and DESC has no compute " + "quantity for it." + ) + return tangent + + +def _p_rr(params, profiles, grid, data): + """d²p/dρ², which has no compute quantity of its own to read it from. + + Mirrors how ``p_r`` itself is computed, one derivative higher. + """ + if profiles.get("pressure") is not None: + return profiles["pressure"].compute(grid, params["p_l"], dr=2) + return elementary_charge * ( + data["ne_rr"] * data["Te"] + + 2 * data["ne_r"] * data["Te_r"] + + data["ne"] * data["Te_rr"] + + data["ni_rr"] * data["Ti"] + + 2 * data["ni_r"] * data["Ti_r"] + + data["ni"] * data["Ti_rr"] + ) + + +def _seed_tangents(params, profiles, grid, data, seed_1d, tangent_keys): + """d/dρ of each per-surface quantity seeded onto a field line grid.""" + seed_dot = {} + for key, val in seed_1d.items(): + if key == "rho": + seed_dot[key] = jnp.ones_like(val) + elif key == "p_r": + seed_dot[key] = _p_rr(params, profiles, grid, data) + elif key in tangent_keys: + seed_dot[key] = data[tangent_keys[key]] + else: + # Nothing on the field line grid reads this, as _tangent_keys + # checked. min_tz |B| and max_tz |B| land here on purpose: the + # pitch grid must stay put so the derivative is taken at fixed λ. + seed_dot[key] = jnp.zeros_like(val) + return seed_dot + + +# New resonance objective from John Anthony Labbate +class TrappedResonance(_Objective): + """Trapped energetic particle resonance penalty. + + Penalizes rational crossings of Omega_eta (the ratio between precessional + motion and bounce frequency) to minimize trapped energetic particle radial + motion due to resonances with magnetic field perturbations from omnigenity. + + Parameters + ---------- + eq : Equilibrium + Equilibrium that will be optimized to satisfy the Objective. + rho : int or ndarray, optional + Flux surfaces on which to evaluate the objective. If an int, the + surfaces are constructed as ``np.linspace(0, 1, rho + 1)[1:]``, giving + ``rho`` uniformly spaced surfaces from ``1/rho`` to ``1`` with spacing + ``1/rho``. If an array, it must be increasing, linearly spaced, and + must not include the magnetic axis (rho=0); e.g. pass an array ending + before rho=1 for equilibria whose pressure profile is not well-defined + at the edge. Default is 10. + num_eta : int, optional + Number of uniformly spaced eta points in [0, 2*pi). + Alpha values are derived per rho surface via + ``alpha = eta * (N*nfp - iota*M) / nfp``. + Default is 10. + weight_method : {"linear", "bump"}, optional + How to weight surfaces near resonance. ``"linear"`` uses 2-point linear + interpolation between bracketing surfaces. ``"bump"`` uses a smooth + normalized bump function. Default is ``"linear"``. + Delta_Omega : float, optional + Half-width of the resonance interval for ``weight_method="bump"``. + If ``None``, defaults to wd_blur × the max |Ω[i+1]-Ω[i]| spacing. + Ignored when ``weight_method="linear"``. + wd_blur : float, optional + Factor multiplying Delta_Omega in case where Delta_Omega = ``None`` + (see Delta_Omega). Otherwise is ignored. + Defaults to 1.25. + num_transit : float, optional + 2π * num_transits sets the extent of zeta for bounce integration. + Defaults to 5. + num_quad : int, optional + Number of quadrature points utilized for any integration in this objective. + Defaults to 32. + num_pitch : int, optional + Number of trapped particle pitches/Bcrit to consider, calculated in + evenly-spaced intervals between Bmin,Bmax on each flux surface. + Defaults to 16. + KE_frac : array, optional + Fraction of 3.5 MeV to use for the energetic particle kinetic energy. + Defaults to np.array([1]). + knots_per_transit : int, optional + knots_per_transit * num_transits gives how many points to use in zeta grid. + Defaults to 100. + batch : bool, optional + Whether or not to calculate multiple trapped particles simultaneously, + especially for bounce integration. + Defaults to True. + pitch_invs : array or None, optional + If not None, sets pitch_invs (Bcrits) to specified value. If None, let's + compute specify a linspace of num_pitch between Bmin and Bmax of each + flux surface. Also causes ``compute`` to skip the phase-space average and + return the raw per-(rho, pitch, well) resonance-physics dictionary instead + of the phase-space-averaged objective. + Defaults to None. + N : int, optional + Generalized omnigenous helicity. Each B contour closes on itself after + traversing the torus M times toroidally and N times poloidally. + Defaults to 0, which is a quasi-axisymmetric configuration. + M : int, optional + Generalized omnigenous helicity. Each B contour closes on itself after + traversing the torus M times toroidally and N times poloidally. + Defaults to 1, which is a quasi-axisymmetric configuration. + p_max : int, optional + Maximum numerator of rational Omega_eta considered. Rational Omega_eta + will be considered for all combinations of p/q up to p_max/q_max. + Defaults to 10. + q_max : int, optional + Maximum denominator of rational Omega_eta considered. Rational Omega_eta + will be considered for all combinations of p/q up to p_max/q_max. + Defaults to 10. + res_range_min : float, optional + Minimum value of rational Omega_eta to consider regardless of p and q. + Defaults to -4. + res_range_max : float, optional + Maximum value of rational Omega_eta to consider regardless of p and q. + Defaults to 4. + fill_value : float, optional + Value to set bounce integration outputs to if no well is found. Cannot + use ``jnp.nan`` to retain optimization abilities. Cannot use 0 for + confusion with other quantities and averages. + Defaults to 11.0. + stab_sacrifice : bool, optional + If ``True``, multiply the island-width term by ``Omega_prime_s**2`` in the + objective. If ``False``, omit that factor to preserve numerical stability. + Defaults to ``False``. + cropping_DOmega : bool, optional + If ``True``, Delta_Omega calculation is clipped by + ``0.01 * max(Omega_eta) < Delta_Omega < 0.10 * max(Omega_eta)``. + This must be when using the ``bump`` weighting method and + ``Delta_Omega = None`` case. Otherwise this quantity is ignored. + Defaults to ``False``. + bt_filter_flag : bool, optional + If ``True``, zero out wells whose poloidal bounce width exceeds 2π + (barely-trapped filter) before the resonance physics calculation. + Defaults to ``False``. + Omega_prime_method : {"analytic", "fd"}, optional + How to obtain Ω'(s), the radial derivative of the normalized precession + frequency that sets the island width. ``"analytic"`` differentiates the + bounce integrals with respect to rho at fixed λ and η, which is exact + for any radial grid and defined on every surface. ``"fd"`` finite + differences Ω across neighbouring surfaces, so its accuracy is limited + by the radial grid spacing and it is undefined on the surfaces at + either end. Defaults to ``"analytic"``; see the Notes. + + Notes + ----- + ``"fd"`` does not converge to Ω'(s) at any practical number of surfaces: + on a stellarator dΩ/dρ swings over orders of magnitude between adjacent + surfaces, and a secant cannot resolve |Ω'| below the scale set by the + radial grid spacing, which biases it low by ~22% even on the finest grid + tested. ``"analytic"`` costs 1.27x the objective evaluation and 1.64x the + gradient, measured on ESTELL at rho=10, num_eta=10. The two agree exactly + when ``stab_sacrifice=True``, where Ω' cancels out of the objective. + """ + + _scalar = False + _coordinates = "r" + _units = "~" + _print_value_fmt = "Trapped EP Resonance Penalty: " + + _static_attrs = _Objective._static_attrs + [ + "_hyperparameters", + "_keys_1dr", + "_key", + "_use_bounce1d", + "_X", + "_Y", + ] + + def __init__( + self, + eq, + target=None, + bounds=None, + weight=1, + normalize=True, + normalize_target=True, + name="TrappedResonance", + jac_chunk_size=None, + verbose=False, + pitch_batch_size=1, + surf_batch_size=1, + rho=10, + num_eta=10, + weight_method="linear", + Delta_Omega=None, + wd_blur=1.25, + num_transit=5, + num_quad=32, + num_pitch=16, + KE_frac=np.array([1]), + knots_per_transit=100, + batch=True, + pitch_invs=None, + N=0, + M=1, + p_max=10, + q_max=10, + res_range_min=-4, + res_range_max=4, + fill_value=11, + stab_sacrifice=False, + cropping_DOmega=False, + bt_filter_flag=False, + use_bounce1d=False, + X=32, + Y=32, + Y_B=None, + spline=True, + nufft_eps=1e-10, + Omega_prime_method="analytic", + ): + if Omega_prime_method not in ("fd", "analytic"): + raise ValueError( + 'Omega_prime_method must be "fd" or "analytic", ' + f"got {Omega_prime_method}." + ) + if target is None and bounds is None: + target = 1e-8 + self._use_bounce1d = bool(use_bounce1d) + self._rho = int(rho) if np.isscalar(rho) else np.atleast_1d(np.asarray(rho)) + self._num_eta = int(num_eta) + if self._num_eta < 2: + raise ValueError(f"num_eta must be >= 2, got {self._num_eta}.") + + self._constants = {"quad_weights": 1} + self._constants["zeta"] = np.linspace( + 0, 2 * np.pi * num_transit, knots_per_transit * num_transit + ) + + self._hyperparameters = { + "num_quad": num_quad, + "num_pitch": num_pitch, + "num_eta": self._num_eta, + "batch": batch, + "KE_frac": KE_frac, + "pitch_invs": pitch_invs, + "N": N, + "M": M, + "p_max": p_max, + "q_max": q_max, + "res_range_min": res_range_min, + "res_range_max": res_range_max, + "verbose": verbose, + "pitch_batch_size": pitch_batch_size, + "surf_batch_size": surf_batch_size, + "num_transit": num_transit, + "weight_method": weight_method, + "Delta_Omega": Delta_Omega, + "fill_value": fill_value, + "wd_blur": wd_blur, + "stab_sacrifice": stab_sacrifice, + "cropping_DOmega": cropping_DOmega, + "bt_filter_flag": bt_filter_flag, + "use_bounce1d": self._use_bounce1d, + "Omega_prime_method": Omega_prime_method, + } + if not self._use_bounce1d: + self._hyperparameters["Y_B"] = Y_B + self._hyperparameters["spline"] = spline + self._hyperparameters["nufft_eps"] = nufft_eps + self._X = int(X) + self._Y = int(Y) + self._keys_1dr = ["iota", "iota_r", "min_tz |B|", "max_tz |B|", "Psi"] + self._key = "trapped EP resonance" + + super().__init__( + things=[eq], + target=target, + bounds=bounds, + weight=weight, + normalize=normalize, + normalize_target=normalize_target, + name=name, + jac_chunk_size=jac_chunk_size, + ) + + def build(self, use_jit=True, verbose=1): + """Build constant arrays. + + Parameters + ---------- + use_jit : bool, optional + Whether to just-in-time compile the objective and derivatives. + verbose : int, optional + Level of output. + + """ + eq = self.things[0] + + rho = ( + np.linspace(0, 1, self._rho + 1)[1:] + if isinstance(self._rho, int) + else self._rho + ) + errorif( + rho.size < 2, + ValueError, + msg=f"rho must have >= 2 surfaces, got {rho.size}.", + ) + errorif( + rho[1] <= rho[0] or not np.allclose(np.diff(rho), rho[1] - rho[0]), + ValueError, + msg="rho array must be increasing and linearly spaced!", + ) + errorif( + np.any(np.isclose(rho, 0.0)), + ValueError, + msg="rho array must not include the axis!", + ) + self._constants["rho"] = rho + self._dim_f = rho.size + + self._grid_1dr = LinearGrid( + rho=rho, + M=eq.M_grid, + N=eq.N_grid, + NFP=eq.NFP, + sym=eq.sym if self._use_bounce1d else False, + ) + self._constants["quad"] = get_quadrature( + leggauss(self._hyperparameters["num_quad"]), + (automorphism_sin, grad_automorphism_sin), + ) + if not self._use_bounce1d: + assert self._grid_1dr.can_fft2 + # Nodes at which the poloidal angle map is interpolated, and the + # transform used to solve for it. Mirrors ``Bounce2D._build``. + self._constants["x"] = fourier_pts(self._X) + self._constants["y"] = cheb_pts(self._Y, (0, 2 * np.pi / eq.NFP))[::-1] + self._constants["lambda"] = get_transforms( + "lambda", + eq, + grid=LinearGrid( + rho=rho, + M=eq.L_basis.M, + zeta=self._constants["y"], + NFP=eq.NFP, + ), + )["L"] + spline = self._hyperparameters["spline"] + Y_B = self._hyperparameters["Y_B"] + if Y_B is None: + Y_B = Y_B_rule(self._grid_1dr, spline) + self._hyperparameters["Y_B"] = Y_B + self._constants["_vander"] = ( + get_vander_spline(self._grid_1dr, self._Y, Y_B, eq.NFP) + if spline + else {} + ) + rho_res = rho[1] - rho[0] + eta_res = 2 * np.pi / self._num_eta + self._params2 = { + "rho_res": rho_res, + "eta_res": eta_res, + } + self._target, self._bounds = _parse_callable_target_bounds( + self._target, self._bounds, rho + ) + + timer = Timer() + if verbose > 0: + print("Precomputing transforms") + timer.start("Precomputing transforms") + + self._constants["transforms_1dr"] = get_transforms( + self._keys_1dr, eq, self._grid_1dr + ) + self._constants["profiles"] = get_profiles( + self._keys_1dr + [self._key], eq, self._grid_1dr + ) + + # Setup rational array + p_max = self._hyperparameters["p_max"] + q_max = self._hyperparameters["q_max"] + res_range_min = self._hyperparameters["res_range_min"] + res_range_max = self._hyperparameters["res_range_max"] + + # Preallocate: max resonances = n_max (m=0) + 2*m_max*n_max (m>0) + n_res_max = q_max + 2 * p_max * q_max + res_arr = np.full(n_res_max, np.nan) + q_arr = np.zeros(n_res_max, dtype=int) + p_arr = np.zeros(n_res_max, dtype=int) + res_arr_set = 0 + + for p in range(0, p_max + 1): + for q in range(1, q_max + 1): + ratio = p / q + if not res_range_min <= ratio <= res_range_max: + continue + # +p/q and -p/q are distinct resonances unless p is zero. + for sign in (1, -1) if p else (1,): + res_arr[res_arr_set] = sign * ratio + q_arr[res_arr_set] = q + p_arr[res_arr_set] = sign * p + res_arr_set += 1 + + res_arr = res_arr[:res_arr_set] + q_arr = q_arr[:res_arr_set] + p_arr = p_arr[:res_arr_set] + + self._hyperparameters["q_arr"] = q_arr + self._hyperparameters["res_arr"] = res_arr + self._hyperparameters["p_arr"] = p_arr + timer.stop("Precomputing transforms") + if verbose > 1: + timer.disp("Precomputing transforms") + + super().build(use_jit=use_jit, verbose=verbose) + + def compute(self, params, constants=None): + """Compute TrappedResonance objective. + + Parameters + ---------- + params : dict + Dictionary of equilibrium degrees of freedom, e.g. + ``Equilibrium.params_dict`` + constants : dict + Dictionary of constant data, e.g. transforms, profiles etc. + Defaults to ``self.constants``. + + Returns + ------- + f_res_avg : ndarray + Phase-space-averaged trapped resonance penalty as a function + of the flux surface label. + + """ + if constants is None: + constants = self._constants + eq = self.things[0] + + data = compute_fun( + eq, + self._keys_1dr, + params, + constants["transforms_1dr"], + constants["profiles"], + ) + quad2 = {} + if "quad2" in constants: + quad2["quad2"] = constants["quad2"] + + base_grid = self._grid_1dr + iotas = base_grid.compress(data["iota"]) + iotas_r = base_grid.compress(data["iota_r"]) + rhos = base_grid.compress(base_grid.nodes[:, 0]) + M = self._hyperparameters["M"] + N = self._hyperparameters["N"] + nfp = eq.NFP + zeta = constants.get("zeta") + num_eta = self._hyperparameters["num_eta"] + analytic = self._hyperparameters["Omega_prime_method"] == "analytic" + + eta_vals = jnp.linspace(0, 2 * jnp.pi, num_eta, endpoint=False) + ft_denom = N * nfp - iotas * M + alpha_per_rho = eta_vals[None, :] * ft_denom[:, None] / nfp + + if not self._use_bounce1d: + fft_keys = list(Bounce2D.required_names) + [ + "cvdrift0", + "gbdrift (periodic)", + "cvdrift (periodic)", + "min_tz |B|", + "max_tz |B|", + ] + fft_profiles = get_profiles(fft_keys, eq) + lambda_grid = constants["lambda"].grid + + def _fft_stage(t): + """Field data and angle map on surfaces displaced by ``t``. + + Written as a function of a radial displacement so that its + forward mode derivative at ``t = 0`` is d/dρ at fixed theta + and zeta, which is what Omega'(s) needs. Unlike the Bounce1D + path there is nothing to seed: the grid carries the whole + radial dependence, so every quantity moves with it. + """ + grid_t = _shift_grid_rho(base_grid, t) + data_t = compute_fun( + eq, + self._keys_1dr, + params, + get_transforms(self._keys_1dr, eq, grid_t, jitable=True), + constants["profiles"], + ) + data_fft_t = compute_fun( + eq, + fft_keys, + params, + get_transforms(fft_keys, eq, grid_t, jitable=True), + fft_profiles, + data=data_t, + ) + # Poloidal angle map, rebuilt every call since it moves with + # rho through both iota and lambda. + angle_t = eq._map_poloidal_coordinates( + grid_t.compress(data_t["iota"]), + constants["x"], + constants["y"], + params["L_lmn"], + get_transforms( + "lambda", + eq, + grid=_shift_grid_rho(lambda_grid, t), + jitable=True, + )["L"], + outbasis="delta", + tol=1e-8, + )[..., ::-1] + return {k: data_fft_t[k] for k in _FFT_BOUNCE_KEYS}, angle_t + + if analytic: + (data_fft, angle), (data_fft_r, angle_r) = jax.jvp( + _fft_stage, (0.0,), (1.0,) + ) + else: + data_fft, angle = _fft_stage(0.0) + data_fft_r = angle_r = None + + data = compute_fun( + eq, + self._key, + params, + get_transforms(self._key, eq, base_grid, jitable=True), + constants["profiles"], + data=data, + quad=constants["quad"], + nfp=nfp, + zeta=zeta, + _angle=angle, + _angle_r=angle_r, + _fft_grid=base_grid, + _data_fft=data_fft, + _data_fft_r=data_fft_r, + _vander=constants["_vander"], + **quad2, + **self._hyperparameters, + **self._params2, + ) + if self._hyperparameters.get("pitch_invs") is not None: + return data[self._key] + return base_grid.compress(data[self._key]) + + # The field line following grid the bounce integrals run along. It + # needs no coordinate map, so it is built here rather than pulled off + # the (rho, theta, zeta) grid, which _eta_data rebuilds at shifted rho. + eta_grid = _build_eta_source_grid(rhos, alpha_per_rho, zeta) + + alpha_psa = jnp.linspace(0, 2 * jnp.pi, num_eta, endpoint=False) + psa_desc_grid = eq._get_rtz_grid( + rhos, + alpha_psa, + zeta, + coordinates="raz", + iota=iotas, + params=params, + ) + psa_grid = psa_desc_grid.source_grid + + eta_data_keys = list(Bounce1D.required_names) + [ + "cvdrift0", + "gbdrift (periodic)", + "cvdrift (periodic)", + "iota", + "min_tz |B|", + "max_tz |B|", + ] + psa_bounce_keys = list(Bounce1D.required_names) + [ + "min_tz |B|", + "max_tz |B|", + "|B|", + ] + all_needed_keys = list(set(eta_data_keys + psa_bounce_keys)) + + # An analytic Omega'(s) needs the radial derivative of every + # per-surface quantity seeded onto the eta grid; without it the seed + # would look constant in rho and the derivative would come out wrong. + _p = _parse_parameterization(eq) + tangent_keys = _tangent_keys(_p, eta_data_keys) if analytic else {} + if analytic: + extra = list(tangent_keys.values()) + if eq.pressure is None: + # _p_rr then builds d²p/dρ² out of the kinetic profiles. + extra += ["ne_rr", "ni_rr", "Te_rr", "Ti_rr"] + all_needed_keys = list(set(all_needed_keys + extra)) + + # Pre-compute all transitive dependencies on the base grid (which has + # spacing for surface integrals). This gives us 1D intermediates like + # iota_den, iota_num, Psi, etc. that the 3D grids cannot compute. + internal_profiles = get_profiles(all_needed_keys, eq) + base_data = compute_fun( + eq, + all_needed_keys, + params, + get_transforms(all_needed_keys, eq, base_grid, jitable=True), + internal_profiles, + data=data, + ) + + # Seed only per-surface (coordinates="r") quantities onto the 3D grids. + # 3D quantities will be recomputed with proper angular resolution. + seed_1d = {} + for key, val in base_data.items(): + entry = data_index.get(_p, {}).get(key) + if entry is not None and entry.get("coordinates", "") == "r": + seed_1d[key] = val + + seed_dot = ( + _seed_tangents( + params, internal_profiles, base_grid, base_data, seed_1d, tangent_keys + ) + if analytic + else None + ) + + def _eta_data(t): + """Field line data on the eta grid, with rho displaced by ``t``. + + Written as a function of a radial displacement so that its forward + mode derivative at ``t = 0`` is d/dρ at fixed eta and zeta. It is + eta, not alpha, that is held fixed: alpha follows rho through iota, + exactly as it does when Omega is instead finite differenced across + neighbouring surfaces. + """ + iotas_t = iotas + t * iotas_r + alpha_t = eta_vals[None, :] * (N * nfp - iotas_t[:, None] * M) / nfp + grid_t = _build_eta_grid(eq, rhos + t, alpha_t, zeta, iotas_t, params) + seed_t = ( + seed_1d + if seed_dot is None + else {key: val + t * seed_dot[key] for key, val in seed_1d.items()} + ) + eta_seed = { + key: grid_t.copy_data_from_other(val, base_grid) + for key, val in seed_t.items() + } + return compute_fun( + eq, + eta_data_keys, + params, + get_transforms(eta_data_keys, eq, grid_t, jitable=True), + internal_profiles, + data=eta_seed, + ) + + if analytic: + data_eta, data_eta_r = jax.jvp(_eta_data, (0.0,), (1.0,)) + else: + data_eta, data_eta_r = _eta_data(0.0), None + + psa_seed = { + key: psa_desc_grid.copy_data_from_other(val, base_grid) + for key, val in seed_1d.items() + } + data_psa = compute_fun( + eq, + psa_bounce_keys, + params, + get_transforms(psa_bounce_keys, eq, psa_desc_grid, jitable=True), + internal_profiles, + data=psa_seed, + ) + + data = compute_fun( + eq, + self._key, + params, + get_transforms(self._key, eq, self._grid_1dr, jitable=True), + constants["profiles"], + data=data, + quad=constants["quad"], + nfp=eq.NFP, + zeta=zeta, + _eta_grid=eta_grid, + _psa_grid=psa_grid, + _data_eta=data_eta, + _data_eta_r=data_eta_r, + _data_psa=data_psa, + **quad2, + **self._hyperparameters, + **self._params2, + ) + if self._hyperparameters.get("pitch_invs") is not None: + return data[self._key] + return self._grid_1dr.compress(data[self._key]) diff --git a/setup.cfg b/setup.cfg index 64d6be0c06..13d7e66aaf 100644 --- a/setup.cfg +++ b/setup.cfg @@ -47,6 +47,9 @@ markers= optimize : marks tests that perform an optimization slow: marks tests as slow (deselect with 'pytest -m "not slow"'). fast: mark tests as fast. + manual: marks tests kept in the repo but deliberately excluded from CI, which + selects only 'unit' and 'regression'. Run these by hand with + 'pytest -m manual'. Do not add 'unit' or 'regression' alongside it. memory: marks tests that check memory usage filterwarnings= error diff --git a/tests/test_objective_funs.py b/tests/test_objective_funs.py index f97bc93a27..12352a5e2d 100644 --- a/tests/test_objective_funs.py +++ b/tests/test_objective_funs.py @@ -11,6 +11,7 @@ import numpy as np import pytest +from orthax.legendre import leggauss from packaging.version import Version from qsc import Qsc from scipy.constants import elementary_charge, mu_0 @@ -25,12 +26,21 @@ MixedCoilSet, initialize_modular_coils, ) -from desc.compute import get_transforms +from desc.compute import get_profiles, get_transforms +from desc.compute._trapped_resonance import _build_eta_grid +from desc.compute.data_index import data_index +from desc.compute.utils import _compute as compute_fun +from desc.compute.utils import _parse_parameterization from desc.equilibrium import Equilibrium from desc.examples import get from desc.geometry import FourierPlanarCurve, FourierRZToroidalSurface, FourierXYZCurve from desc.grid import ConcentricGrid, Grid, LinearGrid, QuadratureGrid -from desc.integrals import Bounce2D +from desc.integrals import Bounce1D, Bounce2D +from desc.integrals.quad_utils import ( + automorphism_sin, + get_quadrature, + grad_automorphism_sin, +) from desc.io import load from desc.magnetic_fields import ( CurrentPotentialField, @@ -91,11 +101,13 @@ SurfaceQuadraticFlux, ToroidalCurrent, ToroidalFlux, + TrappedResonance, VacuumBoundaryError, Volume, get_NAE_constraints, ) from desc.objectives._free_boundary import BoundaryErrorNESTOR +from desc.objectives._trapped_resonance import _seed_tangents, _tangent_keys from desc.objectives.nae_utils import ( _calc_1st_order_NAE_coeffs, _calc_2nd_order_NAE_coeffs, @@ -2171,6 +2183,305 @@ def test_objective_against_compute_bounce(self, use_bounce1d): obj.compute(eq.params_dict), grid.compress(data[names[1]]) ) + @pytest.mark.unit + def test_objective_against_compute_trapped_resonance(self): + """Test TrappedResonance objective matches a direct compute call.""" + eq = get("ESTELL") + with pytest.warns(UserWarning, match="Reducing radial"): + eq.change_resolution(2, 2, 2, 4, 4, 4) + + num_rho = 3 + num_eta = 8 + num_transit = 4 + knots_per_transit = 60 + num_quad = 16 + opts = dict( + num_pitch=8, + KE_frac=np.array([1]), + N=0, + M=1, + p_max=0, + q_max=1, + res_range_min=-1, + res_range_max=1, + weight_method="linear", + use_bounce1d=True, + ) + + rho = np.linspace(0, 1, num_rho + 1)[1:] + zeta = np.linspace(0, 2 * np.pi * num_transit, knots_per_transit * num_transit) + grid = LinearGrid(rho=rho, M=eq.M_grid, N=eq.N_grid, NFP=eq.NFP, sym=eq.sym) + quad = get_quadrature( + leggauss(num_quad), (automorphism_sin, grad_automorphism_sin) + ) + keys_1dr = ["iota", "iota_r", "min_tz |B|", "max_tz |B|", "Psi"] + profiles = get_profiles(keys_1dr + ["trapped EP resonance"], eq, grid) + params = eq.params_dict + data = compute_fun( + eq, keys_1dr, params, get_transforms(keys_1dr, eq, grid), profiles + ) + + # Build the eta/PSA grids and evaluate field data on them by hand, + iotas = grid.compress(data["iota"]) + rhos = grid.compress(grid.nodes[:, 0]) + M, N, nfp = opts["M"], opts["N"], eq.NFP + eta_vals = jnp.linspace(0, 2 * jnp.pi, num_eta, endpoint=False) + ft_denom = N * nfp - iotas * M + alpha_per_rho = eta_vals[None, :] * ft_denom[:, None] / nfp + + eta_desc_grid = _build_eta_grid(eq, rhos, alpha_per_rho, zeta, iotas, params) + eta_grid = eta_desc_grid.source_grid + + alpha_psa = jnp.linspace(0, 2 * jnp.pi, num_eta, endpoint=False) + psa_desc_grid = eq._get_rtz_grid( + rhos, alpha_psa, zeta, coordinates="raz", iota=iotas, params=params + ) + psa_grid = psa_desc_grid.source_grid + + eta_data_keys = list(Bounce1D.required_names) + [ + "cvdrift0", + "gbdrift (periodic)", + "cvdrift (periodic)", + "iota", + "min_tz |B|", + "max_tz |B|", + ] + psa_bounce_keys = list(Bounce1D.required_names) + [ + "min_tz |B|", + "max_tz |B|", + "|B|", + ] + all_needed_keys = list(set(eta_data_keys + psa_bounce_keys)) + # the radial tangents of the seeded per-surface quantities are extra + # compute keys, so they have to be requested before base_data is built + _p = _parse_parameterization(eq) + tangent_keys = _tangent_keys(_p, eta_data_keys) + all_needed_keys = list(set(all_needed_keys + list(tangent_keys.values()))) + internal_profiles = get_profiles(all_needed_keys, eq) + base_data = compute_fun( + eq, + all_needed_keys, + params, + get_transforms(all_needed_keys, eq, grid, jitable=True), + internal_profiles, + data=data, + ) + seed_1d = { + key: val + for key, val in base_data.items() + if data_index.get(_p, {}).get(key, {}).get("coordinates", "") == "r" + } + # Omega'(s) comes from differentiating the bounce integrals, so the + # compute function needs d(field data)/drho as well as the field data. + # Build both the way the objective does: as the forward mode derivative + # of the eta grid data with respect to a radial displacement. + seed_dot = _seed_tangents( + params, internal_profiles, grid, base_data, seed_1d, tangent_keys + ) + + def _eta_data(t): + iotas_t = iotas + t * grid.compress(data["iota_r"]) + alpha_t = eta_vals[None, :] * (N * nfp - iotas_t[:, None] * M) / nfp + grid_t = _build_eta_grid(eq, rhos + t, alpha_t, zeta, iotas_t, params) + seed_t = {k: v + t * seed_dot[k] for k, v in seed_1d.items()} + return compute_fun( + eq, + eta_data_keys, + params, + get_transforms(eta_data_keys, eq, grid_t, jitable=True), + internal_profiles, + data={ + k: grid_t.copy_data_from_other(v, grid) for k, v in seed_t.items() + }, + ) + + data_eta, data_eta_r = jax.jvp(_eta_data, (0.0,), (1.0,)) + psa_seed = { + key: psa_desc_grid.copy_data_from_other(val, grid) + for key, val in seed_1d.items() + } + data_psa = compute_fun( + eq, + psa_bounce_keys, + params, + get_transforms(psa_bounce_keys, eq, psa_desc_grid, jitable=True), + internal_profiles, + data=psa_seed, + ) + + data = compute_fun( + eq, + "trapped EP resonance", + params, + get_transforms("trapped EP resonance", eq, grid, jitable=True), + profiles, + data=data, + quad=quad, + nfp=eq.NFP, + zeta=zeta, + _eta_grid=eta_grid, + _psa_grid=psa_grid, + _data_eta=data_eta, + _data_eta_r=data_eta_r, + _data_psa=data_psa, + num_eta=num_eta, + num_transit=num_transit, + rho_res=1.0 / num_rho, + eta_res=2 * np.pi / num_eta, + res_arr=np.array([0.0]), + q_arr=np.array([1]), + p_arr=np.array([0]), + **opts, + ) + expected = grid.compress(data["trapped EP resonance"]) + + obj = TrappedResonance( + eq, + rho=num_rho, + num_eta=num_eta, + num_transit=num_transit, + knots_per_transit=knots_per_transit, + num_quad=num_quad, + **opts, + ) + obj.build(verbose=0) + actual = obj.compute(eq.params_dict) + np.testing.assert_allclose(actual, expected) + + @pytest.mark.manual + def test_trapped_resonance_bounce2d_matches_bounce1d(self): + """Test TrappedResonance agrees between its two bounce backends. + + Kept out of CI: ~155 s, and Bounce2D is not the default backend. Run + with 'pytest -m manual' after touching either bounce path. + """ + eq = get("precise_QA") + opts = dict( + rho=20, + num_eta=10, + num_transit=4, + knots_per_transit=60, + num_pitch=8, + num_quad=16, + p_max=4, + q_max=4, + N=0, + M=1, + ) + f = {} + for use_bounce1d in (True, False): + obj = TrappedResonance(eq, use_bounce1d=use_bounce1d, **opts) + ObjectiveFunction(obj, use_jit=False).build(verbose=0) + f[use_bounce1d] = np.asarray(obj.compute(eq.params_dict)).ravel() + b1d, b2d = f[True], f[False] + + assert np.count_nonzero(b1d) > 3, "no resonance crossings detected" + # Both backends should mark the same surfaces as resonant. + np.testing.assert_array_equal(b1d != 0, b2d != 0) + np.testing.assert_allclose(b2d.sum(), b1d.sum(), rtol=0.1) + + @pytest.mark.unit + def test_trapped_resonance_rho_array(self): + """Custom rho arrays must be increasing, linearly spaced, off-axis.""" + eq = Equilibrium() + + # an int and the equivalent explicit array build identical grids + obj_int = TrappedResonance(eq, rho=4) + obj_arr = TrappedResonance(eq, rho=np.linspace(0, 1, 5)[1:]) + obj_int.build(verbose=0) + obj_arr.build(verbose=0) + np.testing.assert_allclose(obj_int._constants["rho"], obj_arr._constants["rho"]) + assert obj_int._params2["rho_res"] == obj_arr._params2["rho_res"] + + # a custom array avoiding the edge, e.g. for equilibria whose + # pressure profile is not well-defined at rho=1 + rho = np.linspace(0.1, 0.9, 5) + obj = TrappedResonance(eq, rho=rho) + obj.build(verbose=0) + np.testing.assert_allclose(obj._constants["rho"], rho) + np.testing.assert_allclose(obj._params2["rho_res"], 0.2) + + with pytest.raises(ValueError, match="linearly spaced"): + TrappedResonance(eq, rho=np.array([0.1, 0.3, 0.9])).build(verbose=0) + with pytest.raises(ValueError, match="linearly spaced"): + TrappedResonance(eq, rho=np.array([0.9, 0.6, 0.3])).build(verbose=0) + with pytest.raises(ValueError, match="axis"): + TrappedResonance(eq, rho=np.array([0.0, 0.5, 1.0])).build(verbose=0) + with pytest.raises(ValueError, match=">= 2"): + TrappedResonance(eq, rho=np.array([0.5])).build(verbose=0) + with pytest.raises(ValueError, match=">= 2"): + TrappedResonance(eq, rho=1).build(verbose=0) + + @pytest.mark.unit + def test_trapped_resonance_analytic_omega_prime(self): + """Analytic Ω'(s) matches a small step finite difference of Ω.""" + eq = get("ESTELL") + with pytest.warns(UserWarning, match="Reducing radial"): + eq.change_resolution(2, 2, 2, 4, 4, 4) + + num_rho = 3 + rho0 = np.linspace(0, 1, num_rho + 1)[1:] + # Pinning the pitch grid keeps λ the same on every displaced radial + # grid, which is the variable Ω'(s) is defined at fixed value of, and + # makes compute return the raw per-(rho, pitch, well) dict. + B = eq.compute(["min_tz |B|", "max_tz |B|"]) + pitch_invs = np.linspace( + 1.02 * np.min(B["min_tz |B|"]), 0.98 * np.max(B["max_tz |B|"]), 16 + ) + obj = TrappedResonance( + eq, + rho=num_rho, + use_bounce1d=True, + num_eta=8, + num_transit=4, + knots_per_transit=60, + num_quad=16, + N=0, + M=1, + p_max=1, + q_max=1, + res_range_min=-1, + res_range_max=1, + pitch_invs=pitch_invs, + ) + obj.build(verbose=0) + + def compute_at(rho_shift): + """Rerun the objective with its radial grid displaced.""" + obj._grid_1dr = LinearGrid( + rho=rho0 + rho_shift, M=eq.M_grid, N=eq.N_grid, NFP=eq.NFP, sym=eq.sym + ) + obj._constants["transforms_1dr"] = get_transforms( + obj._keys_1dr, eq, obj._grid_1dr + ) + obj._constants["profiles"] = get_profiles( + obj._keys_1dr + [obj._key], eq, obj._grid_1dr + ) + return {k: np.asarray(v) for k, v in obj.compute(eq.params_dict).items()} + + ref = compute_at(0.0) + analytic = ref["Omega_prime_s"] * 2 * np.asarray(ref["rhos"])[:, None, None] + + fd = {} + for h in (2e-3, 6e-4): + up, down = compute_at(h), compute_at(-h) + fd[h] = ( + (up["Omega"] - down["Omega"]) / (2 * h), + ref["valid"] & up["valid"] & down["valid"], + ) + + # Ω is not smooth in ρ everywhere: wells and the set of trapped field + # lines both change from surface to surface, so a difference quotient + # straddling one of those changes measures nothing. Compare only where + # the finite difference has itself converged. + (coarse, ok_c), (fine, ok_f) = fd[2e-3], fd[6e-4] + converged = ok_c & ok_f & (np.abs(coarse - fine) <= 1e-2 * np.abs(fine)) + assert converged.sum() >= 5, ( + f"only {converged.sum()} points to compare; the test resolution no " + "longer resolves enough trapped particles to be meaningful" + ) + np.testing.assert_allclose(analytic[converged], fine[converged], rtol=1e-2) + @pytest.mark.unit def test_objective_against_compute_ballooning(self): """To avoid issues such as #1424.""" @@ -3314,6 +3625,20 @@ def _reduced_resolution_objective(eq, objective, **kwargs): kwargs["num_well"] = 15 * kwargs["num_transit"] kwargs["num_pitch"] = 24 kwargs["num_quad"] = 16 + if objective is TrappedResonance: + # Trimmed to what still exercises the objective: the test is dominated + # by compiling the gradient graph, not by these sizes, so they buy + # ~20%. Do not cut further without checking the gradient is still + # non-zero -- rho=5 or num_transit=2 send the linear weighting to + # exactly zero, which passes a no-NaN assertion while testing nothing. + kwargs["rho"] = 6 + kwargs["num_eta"] = 8 + kwargs["num_transit"] = 4 + kwargs["knots_per_transit"] = 40 + kwargs["num_pitch"] = 8 + kwargs["num_quad"] = 8 + kwargs["p_max"] = 4 + kwargs["q_max"] = 4 return objective(eq=eq, **kwargs) @@ -3363,6 +3688,8 @@ class TestComputeScalarResolution: ExternalObjective, LinearObjectiveFromUser, ObjectiveFromUser, + # has its own test, kept out of CI by the manual marker + TrappedResonance, ] other_objectives = list(set(objectives) - set(specials)) @@ -3769,9 +4096,10 @@ def test_compute_scalar_resolution_omnigenity(self): ) def test_compute_scalar_resolution_others(self, objective): """All other objectives.""" + rtol = 6e-2 f = np.zeros_like(self.res_array, dtype=float) for i, res in enumerate(self.res_array): - # just change eq resolution and let objective pick the right grid type + # just change eq resolution, let objective pick the grid type self.eq.change_resolution( L_grid=int(self.eq.L * res), M_grid=int(self.eq.M * res), @@ -3782,6 +4110,42 @@ def test_compute_scalar_resolution_others(self, objective): ) obj.build(verbose=0) f[i] = obj.compute_scalar(obj.x()) + np.testing.assert_allclose( + f, f[-1], rtol=rtol, atol=1e-4 if np.max(f) < 1e-3 else 0 + ) + + @pytest.mark.manual + def test_compute_scalar_resolution_trapped_resonance(self): + """TrappedResonance. Kept out of CI: ~170 s, see the manual marker.""" + eq = get("precise_QA") + # rho and num_eta halved from 20: this test is compute bound, so that + # halves its runtime while the spread across res_array stays at 0.004, + # well inside rtol. Leave knots_per_transit and num_quad alone -- + # cutting those moves the objective value by ~60x, which would verify + # resolution independence in a regime where the quadrature itself is + # unconverged. + kwargs = dict( + rho=10, + num_eta=10, + num_transit=4, + knots_per_transit=60, + num_pitch=8, + num_quad=16, + p_max=4, + q_max=4, + N=0, + M=1, + ) + f = np.zeros_like(self.res_array, dtype=float) + for i, res in enumerate(self.res_array): + eq.change_resolution( + L_grid=int(eq.L * res), + M_grid=int(eq.M * res), + N_grid=int(eq.N * res), + ) + obj = ObjectiveFunction(TrappedResonance(eq=eq, **kwargs), use_jit=False) + obj.build(verbose=0) + f[i] = obj.compute_scalar(obj.x()) np.testing.assert_allclose( f, f[-1], rtol=6e-2, atol=1e-4 if np.max(f) < 1e-3 else 0 ) @@ -3876,6 +4240,7 @@ class TestObjectiveNaNGrad: SurfaceCurrentRegularization, SurfaceQuadraticFlux, ToroidalFlux, + TrappedResonance, VacuumBoundaryError, # we do not test these since they depend too much on what the user wants ExternalObjective, @@ -4282,6 +4647,27 @@ def test_objective_no_nangrad_Gamma_c(self): g = obj.grad(obj.x()) assert not np.any(np.isnan(g)) + @pytest.mark.unit + def test_objective_no_nangrad_trapped_resonance(self): + """TrappedResonance.""" + eq = get("ESTELL") + with pytest.warns(UserWarning, match="Reducing radial"): + eq.change_resolution(2, 2, 2, 4, 4, 4) + + obj = ObjectiveFunction( + _reduced_resolution_objective(eq, TrappedResonance, weight_method="linear") + ) + obj.build(verbose=0) + g = obj.grad(obj.x()) + assert not np.any(np.isnan(g)), "linear weighting" + + obj = ObjectiveFunction( + _reduced_resolution_objective(eq, TrappedResonance, weight_method="bump") + ) + obj.build(verbose=0) + g = obj.grad(obj.x()) + assert not np.any(np.isnan(g)), "bump weighting" + @pytest.mark.unit def test_objective_no_nangrad_ballooning(self): """BallooningStability."""