fold alpha to improve convergence by unalmis · Pull Request #2223 · PlasmaControl/DESC

unalmis · 2026-05-21T03:37:17Z

Summary

This PR takes the simple implementation of #2220 and adds folded-alpha handling for Gamma_delta and Gamma_alpha.
The core idea is that bounce wells are found on long field lines spanning
multiple field periods. The new code
maps long-field-line wells back to effective single-period alpha samples, then
applies the Gamma_delta and Gamma_alpha reductions on that folded alpha
grid.

The implementation keeps the existing Gamma_c Velasco path unchanged.
Only Gamma_delta and Gamma_alpha call the folded-alpha path.

1. Folding Long-Field-Line Wells Into Alpha Samples

For each bounce well with endpoints (zeta_1, zeta_2), the code assigns the
well to a field-period segment using the midpoint approximation

zeta_m = (zeta_1 + zeta_2) / 2,
k = floor(NFP * zeta_m / (2*pi)).

This is implemented in _well_field_period(...).

For a base field-line label alpha_j, a well in field-period copy k
corresponds to the effective single-period label

alpha_eff = alpha_j + iota * k * (2*pi / NFP) mod 2*pi.

This follows the field-line map logic: moving along the long field line by one
field period shifts the equivalent single-period alpha label by
iota * 2*pi / NFP.

The folded representation is indexed conceptually as

(alpha_eff, pitch, local_well)

with optional surface or batch axes in front. A boolean mask records which
folded entries correspond to real wells.

2. Why Local Well Ranking Is Needed

The original well axis is sorted along the extended field line. After folding,
that is not the right branch index. We need the well index local to each
single-field-period copy.

Example:

global well axis:     w = 0    1    2    3    4    5
field period:         k = 0    0    1    1    2    2
local well in period: r = 0    1    0    1    0    1

_local_well_rank(...) computes this local index r. This keeps one physical
local branch coherent across the folded alpha grid. Without it, a single local
well branch could be split across different global long-line well slots, which
would make the alpha scans in Gamma_alpha inconsistent.

3. Gamma_delta Reduction

The threshold is stored in tangent space, so the outward candidate condition is
implemented without forming arctan:

radial > thresh * abs(poloidal).

For folded alpha data, the reduction computes

sum_well [ integral d alpha v_tau(alpha, well) ]
         * H(there exists an outward alpha candidate for this well).

In code:

weight = _alpha_weights(alpha, mask)
v_tau = (v_tau * weight).sum(-3)
has_alpha_out_candidate = alpha_out_candidate.any(-3)
return (v_tau * has_alpha_out_candidate).sum(-1)

The mask remains essential for alpha quadrature weights, but it is not needed
inside has_alpha_out_candidate: dummy wells have zero radial and zero
poloidal, so the strict candidate condition cannot be true for them.

4. Gamma_alpha Reduction

Gamma_alpha uses two signed branch scores:

alpha_out_candidate = radial - thresh * abs(poloidal),
alpha_in_candidate  = -radial - thresh * abs(poloidal).

Positive alpha_out_candidate means the radial drift is sufficiently outward.
Positive alpha_in_candidate means the radial drift is sufficiently inward.

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

if poloidal >= 0:
    interval starts at alpha_in and stops at alpha_out
else:
    interval starts at alpha_out and stops at alpha_in

The edge-case logic is:

if both alpha_out and alpha_in branches exist:
    use the interval returned by _LossCone
if alpha_out exists but alpha_in does not:
    count the whole alpha domain as lost
if alpha_out does not exist:
    count nothing

The folded path integrates

sum_well integral d alpha v_tau(alpha, well) * loss_cone(alpha, well).

5. Loss-Cone Indicator And Fractional Cells

_LossCone builds periodic intervals from two signed scores:

S_start(alpha), S_stop(alpha).

A start crossing occurs when S_start crosses from nonpositive to positive.
For adjacent samples with previous score p <= 0, current score s > 0, and
spacing h, linear interpolation gives

S(alpha_prev + t h) = p + t (s - p),
t = -p / (s - p).

Equivalently, the distance backward from the current sample to the root is

offset = h * s / (s - p).

That is the formula used by _root(...) and _root_nonuniform(...).

After roots are found, each start root is paired with the nearest downstream
stop root using forward periodic distance:

d(alpha_start, alpha_stop) = (alpha_stop - alpha_start) mod 2*pi.

For order=1, the output is a fractional generalized indicator, not just a
boolean mask. Each alpha sample receives

measure(cell_i intersect loss_interval) / measure(cell_i).

For folded nonuniform alpha samples, cell_i is a periodic Voronoi cell. The
neighbor distances and cell widths are computed by _periodic_voronoi_widths.

The final alpha contribution is

v_tau(alpha_i) * loss_fraction_i * alpha_weight_i.

This improves the placement of the discontinuous loss-cone boundary without
interpolating v_tau.

Co-authored-by: Kaya Unalmis <kayaunalmis@proton.me>

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: ```python _Gamma(reduction, params, transforms, profiles, data, **kwargs) ``` The driver computes the common bounce-integral quantities ```text v_tau, radial drift, poloidal secular drift, ``` then delegates the final pitch/alpha/well reduction to one of: ```python _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 ```text sum_well mean_alpha [ v_tau * gamma_c^2 ], ``` with ```text gamma_c = (2 / pi) * arctan(radial / poloidal). ``` The scalar normalization in `_Gamma(...)` is ```python 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`: ```text opts.thresh = tan(pi / 2 * gamma_threshold). ``` The branch condition is ```python 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 ```text 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: ```python outward_score = radial - thresh inward_score = -radial - thresh ``` where ```python thresh = opts.thresh * abs(poloidal) ``` Thus: ```text 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: ```python loss_cone = where( poloidal >= 0, indicator(inward_score, outward_score), indicator(outward_score, inward_score), ) ``` The edge-case rule is: ```text 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 ```text 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 ```text 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 ```text 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: ```python GammaLoss(kind, eq, ...) ``` where ```text kind = "delta" -> Gamma_delta kind = "alpha" -> Gamma_alpha ``` Defaults: ```text 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 ```text [0, (1 + eq.sym) * pi) ``` with `endpoint=False`. The objective uses the shared ```python 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: ```python 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 ```python 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.

unalmis added 17 commits May 18, 2026 17:49

Fix bugs in compute quantities

428ac62

Apply suggestions from code review

345384d

Co-authored-by: Kaya Unalmis <kayaunalmis@proton.me>

Apply suggestions from code review

508663c

Co-authored-by: Kaya Unalmis <kayaunalmis@proton.me>

Merge branch 'ku/available_energy' into ku/compute_bugs

d5cc10b

Merge branch 'ku/available_energy' into ku/compute_bugs

f6253bf

Fixing flunked merges hopefully

e187315

Merge branch 'ku/available_energy' into ku/compute_bugs

f81d3d2

Merge branch 'ku/compute_bugs' into ku/gamma

9f9df69

fixing docs

88b1be9

think this fixes everything

4d9f7bb

fix reference links

8577da4

clarify docstring

1d64f6e

fix failing doc

399a280

make logic clearer for reviewer

26d59ec

old

9b5a5da

This commit adds alpha folding.

6205eca

Clean up.

afcceca

unalmis self-assigned this May 21, 2026

update comptue data

6864aff

unalmis requested review from a team, YigitElma, ddudt, dpanici, f0uriest and rahulgaur104 and removed request for a team May 21, 2026 04:35

unalmis linked an issue May 21, 2026 that may be closed by this pull request

Energetic particle confinement proxies Gamma_delta, Gamma_alpha #1748

Open

unalmis added the stable Besides merging master, other updates require a child PR that should be merged to master later. label May 21, 2026

unalmis mentioned this pull request May 21, 2026

outward superbanana #2220

Merged

clean

d16ee52

unalmis mentioned this pull request Jun 22, 2026

Gamma unalmis/DESC#6

Merged

unalmis added 21 commits July 19, 2026 18:57

Merge branch 'ku/available_energy' into ku/compute_bugs

bb566ed

Merge branch 'ku/compute_bugs' into ku/gamma

c530f15

Merge branch 'ku/gamma' into ku/alpha_fold

53ada21

Merge branch 'ku/available_energy' into ku/compute_bugs

3fd37df

Merge branch 'ku/compute_bugs' into ku/gamma

c17d522

Merge branch 'ku/gamma' into ku/alpha_fold

79bd4e1

Merge branch 'ku/available_energy' into ku/compute_bugs

04e9f31

Merge branch 'ku/compute_bugs' into ku/gamma

deff24e

Merge branch 'ku/gamma' into ku/alpha_fold

0c18bfa

Merge branch 'ku/available_energy' into ku/compute_bugs

3b9738f

Merge branch 'ku/compute_bugs' into ku/gamma

1ffe586

Merge branch 'ku/gamma' into ku/alpha_fold

a353ed0

Merge branch 'ku/available_energy' into ku/compute_bugs

8f1b1d9

Merge branch 'ku/compute_bugs' into ku/gamma

ae3de60

Merge branch 'ku/gamma' into ku/alpha_fold

f01f3bd

Merge branch 'ku/available_energy' into ku/compute_bugs

954d8fb

Merge branch 'ku/compute_bugs' into ku/gamma

4bac33b

Merge branch 'ku/gamma' into ku/alpha_fold

2c088e7

Merge branch 'ku/available_energy' into ku/compute_bugs

f0f86b3

Merge branch 'ku/compute_bugs' into ku/gamma

866558e

Merge branch 'ku/gamma' into ku/alpha_fold

6828091

Base automatically changed from ku/gamma to ku/available_energy August 1, 2026 00:12

unalmis added 6 commits July 31, 2026 17:17

Merge branch 'ku/available_energy' into ku/alpha_fold

5014ff4

Merge branch 'ku/available_energy' into ku/alpha_fold

9cf71fd

Merge branch 'ku/available_energy' into ku/alpha_fold

7a59a9c

Merge branch 'ku/available_energy' into ku/alpha_fold

53ab60f

Merge branch 'ku/available_energy' into ku/alpha_fold

27bc71b

Merge branch 'ku/available_energy' into ku/alpha_fold

8c5693c

Navigation Menu

Uh oh!

fold alpha to improve convergence - #2223

fold alpha to improve convergence#2223
unalmis wants to merge 83 commits into
ku/available_energyfrom
ku/alpha_fold

unalmis commented May 21, 2026 •

edited

Loading

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

1 participant

Uh oh!

Conversation

unalmis commented May 21, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Summary

1. Folding Long-Field-Line Wells Into Alpha Samples

2. Why Local Well Ranking Is Needed

3. Gamma_delta Reduction

4. Gamma_alpha Reduction

5. Loss-Cone Indicator And Fractional Cells

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

1 participant

unalmis commented May 21, 2026 •

edited

Loading