available energy for turbulent transport, fast ion, and outward superbanana optimization - #2215
Open
unalmis wants to merge 40 commits into
Open
available energy for turbulent transport, fast ion, and outward superbanana optimization#2215unalmis wants to merge 40 commits into
unalmis wants to merge 40 commits into
Conversation
Contributor
Memory benchmark result| Test Name | %Δ | Master (MB) | PR (MB) | Δ (MB) | Time PR (s) | Time Master (s) |
| -------------------------------------- | ------------ | ------------------ | ------------------ | ------------ | ------------------ | ------------------ |
test_objective_jac_w7x | 3.75 % | 3.933e+03 | 4.081e+03 | 147.37 | 31.44 | 29.94 |
test_proximal_jac_w7x_with_eq_update | 2.10 % | 6.444e+03 | 6.579e+03 | 135.29 | 156.00 | 156.48 |
test_proximal_freeb_jac | 0.55 % | 1.331e+04 | 1.338e+04 | 73.48 | 80.62 | 81.13 |
test_proximal_freeb_jac_blocked | 1.05 % | 7.662e+03 | 7.742e+03 | 80.77 | 70.18 | 69.88 |
test_proximal_freeb_jac_batched | -0.34 % | 7.627e+03 | 7.601e+03 | -25.87 | 69.19 | 70.10 |
test_proximal_jac_ripple | -1.75 % | 3.545e+03 | 3.483e+03 | -61.98 | 56.26 | 55.76 |
test_proximal_jac_ripple_bounce1d | 4.51 % | 2.042e+02 | 2.134e+02 | 9.21 | 4.41 | 4.51 |
test_eq_solve | 0.16 % | 2.053e+03 | 2.056e+03 | 3.18 | 88.54 | 89.03 |For the memory plots, go to the summary of |
Improves description of Available Energy metric and plotting utility.
6 tasks
unalmis
changed the title
5 tasks
- Fixes bugs in `e_rho_rr`, `e_theta_rtt`, `e_zeta_rzz`, `iota_num_rr limit`. and fix so that redl bootrap uses independent ion density gradient Te * d_ne_d_s + Ti * d_ni_d_s.
unalmis
changed the title
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
changed the title
Port the normalization and field-line handling corrections from da4fb3c while retaining the fixed energy quadrature on ku/available_energy.
unalmis
added
stable
unalmis
requested review from
SebereX,
YigitElma,
byoungj,
ekolemen and
gretahibbard
August 3, 2026 16:16
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Can use the implementation on https://github.com/unalmis/DESC for significantly improved performance.