outward superbanana

[unalmis]

Resolves #1748 .

• these measure the outward superbanana orbits
• adds batch utility to public api

Summary

This branch adds the Velasco prompt-loss metrics Gamma_delta and
Gamma_alpha, updates the Velasco Gamma_c implementation to the normalized
\check{\Gamma}_c quantity, and introduces a reusable GammaLoss objective for
the new metrics.

It also factors several support utilities used by bounce-integral objectives:

• Bounce1D.batch and Bounce2D.batch now own the common surface-batching,
  reshaping, FFT, and sparse-pullback plumbing.
• Options._compute_objective(...) now owns the repeated objective path that
  maps field-line labels to angle and calls compute functions.
• check_nufft(...) centralizes the jax-finufft availability fallback.
• _LossCone in quad_utils.py implements the periodic interval logic used by
  Gamma_alpha.

This branch does not include the later alpha-folding work #2223 for code simplicity. Here,
Gamma_delta and Gamma_alpha reduce over the opts.alpha samples directly.

1. Fast-Ion Compute Path

The branch adds a shared private compute driver:

_Gamma(reduction, params, transforms, profiles, data, **kwargs)

The driver computes the common bounce-integral quantities

v_tau,
radial drift,
poloidal secular drift,

then delegates the final pitch/alpha/well reduction to one of:

_reduction_gamma_c
_reduction_gamma_delta
_reduction_gamma_alpha

This keeps the expensive bounce computation shared across the Velasco metrics
and makes each metric easier to review as a small mathematical reduction.

The branch also removes per-function theta to angle compatibility logic from
these compute functions. That argument migration is now handled once in
desc.compute.utils.compute(...).

2. Gamma_c Velasco Normalization Change

The branch changes the registered "Gamma_c Velasco" compute quantity from the
older equation-16 normalization to the normalized
\check{\Gamma}_c quantity referenced as equation 20 in Velasco et al.

The reduction itself is still

sum_well mean_alpha [ v_tau * gamma_c^2 ],

with

gamma_c = (2 / pi) * arctan(radial / poloidal).

The scalar normalization in _Gamma(...) is

scalar = pi**3 / 16 * grid.NFP / opts.num_field_periods

3. Gamma_delta

Gamma_delta implements Velasco model I: a pitch/well branch is counted as lost
if there exists an outward superbanana somewhere in alpha.

The code uses monotonicity of arctan to bypass the nonlinearities of an arctan as well as safediv by using the tangent-space threshold stored in opts.thresh:

opts.thresh = tan(pi / 2 * gamma_threshold).

The branch condition is

radial > opts.thresh * abs(poloidal)

This is equivalent to asking whether the signed drift ratio exceeds the
threshold, while preserving the direction of the radial drift. The reduction is

sum_well mean_alpha(v_tau) * H(any_alpha outward_candidate).

4. Gamma_alpha

Gamma_alpha implements Velasco model II: only the alpha interval between
inward and outward superbanana branches is counted as lost.

The code constructs two signed branch scores:

outward_score = radial - thresh
inward_score = -radial - thresh

where

thresh = opts.thresh * abs(poloidal)

Thus:

outward_score > 0  means radial drift is sufficiently outward,
inward_score  > 0  means radial drift is sufficiently inward.

The interval orientation depends on the sign of the poloidal drift:

loss_cone = where(
    poloidal >= 0,
    indicator(inward_score, outward_score),
    indicator(outward_score, inward_score),
)

The edge-case rule is:

if alpha_out and alpha_in both exist:
    use the interval returned by _LossCone
if alpha_out exists and alpha_in does not:
    count the full alpha domain
if alpha_out does not exist:
    count nothing

The final reduction is

sum_well mean_alpha(v_tau * loss_cone).

5. _LossCone Periodic Interval Logic

_LossCone.indicator(...) turns signed branch scores into a periodic loss-cone
indicator over the sampled alpha grid.

For order=0, it returns the sampled boolean interval mask.

For order=1, it improves the interval boundary placement by linearly
interpolating zero crossings of the signed scores. If the previous score is
p <= 0, the current score is s > 0, and the alpha spacing is dx, then the
root is located by

offset from current sample = dx * s / (s - p).

Each start crossing is paired to the first downstream stop crossing using the
forward periodic distance matrix

dist[i, j] = (alpha_j - alpha_i) mod 2*pi.

The order=1 output is a fractional generalized indicator in [0, 1], giving
the fraction of each uniform alpha cell covered by the loss interval.

Review point: this improves the discontinuous interval boundary without
interpolating v_tau; the sampled v_tau values are still used as the cell
values.

NOTE: Forming a distance matrix is not the most efficient strategy for this. However, I chose this implementation because the code is simpler.

For a more detailed derivation of the loss-cone interval construction,
linear crossing formula, and fractional-width quadrature, see the end of this file.

6. Objective API

The branch adds:

GammaLoss(kind, eq, ...)

where

kind = "delta" -> Gamma_delta
kind = "alpha" -> Gamma_alpha

Defaults:

alpha = GammaLoss._default_alpha(eq)
num_field_periods = eq.NFP + 2
gamma_threshold = 0.2

GammaLoss._default_alpha(eq) returns 16 uniformly spaced field-line labels on

[0, (1 + eq.sym) * pi)

with endpoint=False.

The objective uses the shared

Options._build_objective(...)
Options._compute_objective(...)

flow that is also adopted by GammaC, EffectiveRipple, and
AvailableEnergy.

Review point: GammaLoss is one objective class with a required kind string,
not two separate objective classes.

7. Shared Infrastructure Changes

Bounce*D.batch

Bounce1D.batch and Bounce2D.batch now accept:

data,
grid,
names=(),
custom_data=None,
rho_data=None,
batch_size=1,
sparse=True,

Bounce2D.batch additionally requires angle=....

This moves repeated reshape/FFT/surface-batching logic out of individual compute
functions. The old pattern passed a partially prepared fun_data dictionary and
mutated it. The new pattern builds the batched data dictionary internally from
data, names, custom_data, and rho_data.

Options._compute_objective

The shared objective compute path now:

1. Computes iota.
2. Builds the mapped angle through _map_poloidal_coordinates.
3. Calls the selected compute key with angle, alpha, quadrature constants,
   and the objective hyperparameters.
4. Compresses the requested compute output back to the objective grid.

This removes duplicated logic from fast-ion, neoclassical, and turbulence
objectives.

check_nufft

The NUFFT availability fallback now lives in _interp_utils.py and is reused by
the relevant objectives. It preserves the prior behavior: if jax-finufft is
unavailable and nufft_eps requests NUFFT use, warn and set nufft_eps = 0.0.

Math Note: Loss-Cone Indicator For Gamma_alpha

Goal

This note explains the math behind the loss-cone indicator used in
Gamma_alpha, especially why the implementation uses linearly interpolated
crossings and fractional cell widths, but does not interpolate v_tau.

For this branch, fix one flux surface, one pitch value, and one well branch.
All quantities below are functions of the field-line label alpha.

Branch Scores

Let

$$R(\alpha)$$

denote the bounce-integrated radial drift contribution, and let

$$P(\alpha)$$

denote the corresponding poloidal drift contribution. The threshold parameter
is stored in tangent space:

$$\tau_\gamma = \tan(\frac{\pi}{2}\gamma_{\mathrm{threshold}}).$$

The outward and inward superbanana candidate scores are

$$s_{\mathrm{out}}(\alpha) = R(\alpha) - \tau_\gamma |P(\alpha)|,$$

and

$$s_{\mathrm{in}}(\alpha) = -R(\alpha) - \tau_\gamma |P(\alpha)|.$$

Thus

$$s_{\mathrm{out}}(\alpha) > 0$$

means the branch is sufficiently outward drifting, while

$$s_{\mathrm{in}}(\alpha) > 0$$

means the branch is sufficiently inward drifting.

Equivalently, the outward condition is

$$R(\alpha) > \tau_\gamma |P(\alpha)|,$$

and the inward condition is

$$R(\alpha) < -\tau_\gamma |P(\alpha)|.$$

The absolute value is on the poloidal drift only. We do not take
abs(radial), because inward-only branches should not count as outward
superbanana branches.

Loss Interval Orientation

Gamma_alpha counts alpha intervals between inward and outward branches. The
orientation depends on the sign of the poloidal drift.

If

$$P(\alpha) \ge 0,$$

the loss interval starts at the inward branch and stops at the next outward
branch:

$$s_{\mathrm{in}} = 0 \quad\longrightarrow\quad s_{\mathrm{out}} = 0.$$

If

$$P(\alpha) < 0,$$

the orientation reverses:

$$s_{\mathrm{out}} = 0 \quad\longrightarrow\quad s_{\mathrm{in}} = 0.$$

This is why the code calls the interval indicator twice with reversed
start/stop scores and then selects by the sign of poloidal.

Indicator As A Set

For a chosen start score s_start and stop score s_stop, define the start
crossing set

$$\mathcal{S} = {\alpha : s_{\mathrm{start}}(\alpha) = 0 \text{ and crosses from } \le 0 \text{ to } > 0},$$

and the stop crossing set

$$\mathcal{T} = {\alpha : s_{\mathrm{stop}}(\alpha) = 0 \text{ and crosses from } \le 0 \text{ to } > 0}.$$

For each start crossing

$$\alpha_s \in \mathcal{S},$$

we pair it with the first downstream stop crossing

$$\alpha_t = \text{argmin}_{\beta \in \mathcal{T}} ((\beta - \alpha_s) \bmod 2\pi).$$

The corresponding loss interval is

$$I_s = [\alpha_s,, \alpha_s + ((\alpha_t - \alpha_s) \bmod 2\pi)].$$

The continuous loss-cone indicator is then

$$\chi_{\mathrm{loss}}(\alpha) = \begin{cases} 1, & \alpha \in \bigcup_{\alpha_s \in \mathcal{S}} I_s, \ 0, & \text{otherwise}. \end{cases}$$

The code constructs a sampled or fractional approximation to this indicator.

Linear Crossing Formula

Suppose a score crosses zero between two adjacent alpha samples:

$$\alpha_{i-1} \quad\text{and}\quad \alpha_i.$$

Let

$$p = s(\alpha_{i-1}) \le 0, \qquad q = s(\alpha_i) > 0, \qquad \Delta\alpha = \alpha_i - \alpha_{i-1}.$$

Assume s is linear on the cell:

$$s(\alpha_{i-1} + t\Delta\alpha) = p + t(q - p), \qquad 0 \le t \le 1.$$

The root satisfies

$$0 = p + t(q - p),$$

so

$$t = \frac{-p}{q - p}.$$

Thus the root is

$$\alpha_* = \alpha_{i-1} + \Delta\alpha \frac{-p}{q-p}.$$

The implementation stores the distance from the current sample backward to the
root:

$$\alpha_i - \alpha_* = \Delta\alpha \frac{q}{q-p}.$$

That is the formula

offset = dx * score / (score - previous)

in _LossCone._root.

order=0: Sampled Indicator

For order=0, no crossing interpolation is used. The code treats positive
sample values as branch events and marks each sampled alpha as either inside or
outside the nearest start-to-stop interval.

This is simple, but the estimated interval length can be wrong by

$$O(\Delta\alpha)$$

because each crossing is snapped to the nearest sample.

order=1: Fractional Cell Width

For order=1, the code first finds the linearly interpolated start and stop
crossings. Then each alpha sample is assigned a cell

$$C_i = [\alpha_i - \frac{\Delta\alpha}{2},, \alpha_i + \frac{\Delta\alpha}{2}].$$

Instead of returning a boolean value, the loss-cone indicator returns the
fraction of that cell covered by the loss interval:

$$w_i = \frac{|C_i \cap I_{\mathrm{loss}}|}{|C_i|}.$$

So

$$0 \le w_i \le 1.$$

This is the "fractional width". A sample whose cell is fully inside the
loss interval has weight 1. A sample whose cell is fully outside has weight
0. A sample whose cell is cut by a crossing gets a value between 0 and 1.

Why We Do Not Interpolate v_tau

The Gamma_alpha alpha integral has the form

$$\int_0^{2\pi} V(\alpha),\chi_{\mathrm{loss}}(\alpha),\frac{d\alpha}{2\pi},$$

where

$$V(\alpha) = v\tau(\alpha)$$

for the fixed pitch and well branch.

The discontinuous part is

$$\chi_{\mathrm{loss}}(\alpha),$$

not usually V(alpha). The method uses the sampled value

$$V_i = V(\alpha_i)$$

as the value on cell C_i, but integrates the discontinuous domain restriction
with the fractional cell coverage:

$$\int_0^{2\pi} V(\alpha),\chi_{\mathrm{loss}}(\alpha),\frac{d\alpha}{2\pi} \approx \sum_i V_i,\frac{|C_i \cap I_{\mathrm{loss}}|}{2\pi}.$$

Equivalently, because

$$w_i = \frac{|C_i \cap I_{\mathrm{loss}}|}{|C_i|}, \qquad \omega_i = \frac{|C_i|}{2\pi},$$

the quadrature contribution is

$$V_i,w_i,\omega_i = V_i \frac{|C_i \cap I_{\mathrm{loss}}|}{2\pi}.$$

So the fractional method improves the geometric measure of the loss interval
without pretending to know new values of V(alpha).

Why Interpolating Samples Is Not The Main Accuracy Gain

Simply upsampling a quadrature integrand by interpolation does not generally
create new information. If we only interpolated

$$V_i \chi_i,$$

then integrated the interpolant, we would still be limited by the sampled
location of the discontinuity in

$$\chi_i.$$

The useful extra information here is different: the signed scores

$$s_{\mathrm{out}}(\alpha), \qquad s_{\mathrm{in}}(\alpha)$$

contain sub-cell information about where the indicator changes state. Linear
root finding uses that signed information to locate the loss interval boundary
inside the cell.

This improves the convergence of the interval measure even though v_tau is
not interpolated.

Code Map

The main implementation points are:

• _reduction_gamma_alpha

  • Builds outward_score and inward_score.
  • Selects the interval orientation using the sign of poloidal.
  • Applies edge-case logic for missing inward/outward branches.
  • Averages v_tau * loss_cone over alpha and sums over wells.

• _LossCone._root

  • Finds negative-to-positive crossings.
  • Uses the linear formula
    offset = dx * score / (score - previous).

• _LossCone.indicator(..., order=0)

  • Builds the sampled boolean loss interval.

• _LossCone.indicator(..., order=1)

  • Builds the fractional cell coverage weights.