TrappedResonance - #2257
Open
ejpaul wants to merge 260 commits into
Open
Conversation
Author
This is a great suggestion. Right now the sum (corresponding to the phase-space average) is performed before returning the objective, so the entire sum is squared. It seems like it might make sense to pass the point-wise values (removing a square), then perform the sum-of-squares through the least-squares objective. I think one square may still be needed within the objective, since |f_r| = sqrt(...). I will take a look at this. |
The least-squares objective already squares whatever compute() returns, so raising the per-point residual to the 4th power (Delta_s^4) before that made the net penalty (island width)^8 instead of the intended ^4. Return Delta_s^2 instead and let the outer least-squares squaring supply the other factor. This reintroduces one unavoidable sqrt (|f_r| = sqrt(f_q_c^2 + f_q_s^2) is now used directly rather than squared away), which needs its own NaN-safe guard since f_q_c == f_q_s == 0 exactly at every invalid grid point. Also fixes the resolution-convergence test, which was exercising this objective on a shared HELIOTRON equilibrium with N=0, M=1 (QA) defaults that do not match HELIOTRON's actual (non-QA) symmetry -- mismatched helicity makes tiny, already-converged differences in iota/|B| flip which resonance crossings get caught, producing spurious resolution sensitivity unrelated to the fix above. Testing on a genuinely QA equilibrium instead brings it back in line with other objectives. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
6 tasks
The clamp guarded a plain division by sqrt|1 - pitch*B|, which vanishes at the bounce points. safediv already guards that case, and is what the other two integrands in this path (_v_tau, _radial_drift) use, so this makes all three consistent. Verified the clamp was not load-bearing: values are unchanged, and the NaN-gradient test still passes, so it was not providing gradient safety either -- the sqrt's derivative is never evaluated at an exact zero, since quadrature nodes fall strictly inside the wells. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
get_alphas assumed a single 1D alpha shared across all flux surfaces. That is wrong for objectives whose field line labels are parameterized by a quantity depending on the rotational transform -- e.g. the omnigenity angle eta, where alpha = eta*(N*nfp - iota*M)/nfp varies per surface. The existing arithmetic already produces the correct result for that case; it only needed the right axis inserted. Accept alpha of shape (num alpha, num rho) alongside the existing (num alpha,), and take num_alpha from shape[0] rather than size, which is identical in 1D. Backward compatible: for 1D alpha the original branch is taken unchanged, and a 2D alpha tiled across rho reproduces the 1D result bit-for-bit. Verified against per-surface reference builds on precise_QA to 2.7e-15. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
Bounce1D consumes data already evaluated on field lines, so the eta and phase-space grids scale as num_eta * knots_per_transit * num_transit per surface and are rebuilt every call through a Newton coordinate map. Bounce2D instead interpolates onto field lines from a 2D Fourier series on a fixed (theta, zeta) grid, which is what removes that cost. Adds the Bounce2D path as the default, keeping Bounce1D available via use_bounce1d=True. This mirrors EffectiveRipple, and matters more here than for a typical migration because the radial drift is badly conditioned (see below), so the two backends need to stay comparable. - _compute2D: the Bounce2D analogue of _compute1D. Broadcasts the global B_crit across surfaces so the shared-pitch semantics are preserved, and carries alpha with rho leading so batch_map slices it consistently. - get_alphas now accepts per-surface field line labels (previous commit), so the eta parameterization no longer needs _build_eta_grid at all; the objective evaluates everything once on the (theta, zeta) grid. - Field line length comes from V_psi/2pi rather than an explicit simpson integral over B^zeta; the two agree to 1.5%, the difference being single field line versus true surface average. nufft_eps defaults to 1e-10, not Bounce2D's usual 1e-6. The radial drift bounce integral is a near-total cancellation -- omnigenity optimization drives it toward zero, so int|f|/|int f| reaches 1e2 to 1e7 -- and at 1e-6 the result has no correct digits and does not converge under refinement. At 1e-10 it agrees with Bounce1D to ~5e-4. Quantities with no cancellation (the poloidal drift, int|f|/|int f| = 1 exactly) are unaffected either way. On precise_QA at num_rho=20 the two backends select the same surfaces and agree to 0.3-1% on the entries that carry the objective. The residual is concentrated in near-zero entries (~1e-6) where the conditioning above makes any two methods disagree. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
The previous comment described measurements taken on the isolated radial drift integral as if they applied to the objective. They do not: the ratio and the phase-space average cancel much of the error, so at the objective level 1e-6 roughly doubles the disagreement with Bounce1D rather than destroying it, and 1e-8 was already converged on precise_QA. Keeps the 1e-10 default as margin for equilibria with worse cancellation, but no longer overstates what was measured. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
Asserts the two backends select the same resonant surfaces and agree on the summed objective. The comparison is on the sum rather than per-surface values because the pointwise relative error is dominated by the near-zero entries, where the radial drift bounce integral is a near-total cancellation and any two methods disagree by tens of percent; the sum is what the optimizer sees and agrees to ~2%. The 10% tolerance sits between that ~2% and the ~700% error produced by reducing the phase space average over a flattened (rho, alpha) axis, so it discriminates rather than rubber-stamps. Guards the settings too: below num_transit=4 no crossings are detected at all, which would leave the test comparing two zero arrays and passing for the wrong reason. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
Removed unnecessary explanation about flux surfaces in docstring.
Removed comments about grid requirements and interpolation methods to streamline the code.
Removed commented-out explanations for various parameters in the test cases to improve code clarity.
Neither belongs in this PR. The .DS_Store entry came in with 6fb110b while clearing scratch artifacts, and jac_chunk_size=1 was added to _test_objective_ripple by the benchmark commits (66fc063, 741e044); upstream master has neither. Both files now match master exactly, so they drop out of the diff. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
Completes the metric change discussed on the PR. The exponent fix removed one square; this moves the phase space sum out of the objective so the least squares machinery performs it, as suggested. Writing the phase space average as sum_j W_j f_j, compute now returns sqrt(W_j) f_j per point, so the objective forms sum_j W_j f_j^2 rather than (sum_j W_j f_j)^2. Note this is not only a matter of cross terms: the former is the phase space mean of the squares and the latter the square of the mean, which differ by the variance of f over phase space. On precise_QA that is a factor of 7.5, since f ranges over orders of magnitude. The new form weights the worst phase space points more heavily, which is the intent of a resonance penalty, but it is a change in what is optimized and not a refactor. W is non-negative but exactly zero at masked points, where d sqrt(W)/dW diverges, so the root is guarded the same way as f_q_abs. Costs: dim_f becomes the number of phase space points rather than flux surfaces (20 -> 1800 at test resolution), enlarging the Jacobian by the same factor. Set pointwise=False for the previous behaviour. A cheaper variant returning sqrt(sum_j W_j f_j^2) per surface would give the same sum without growing dim_f, at the cost of another guarded root. Backend agreement improves, since cross terms had amplified the difference: 2.1% -> 0.6% between Bounce1D and Bounce2D on the summed square. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
Author
I took an additional look at this strategy. If the sum is performed by the least-squares objective itself, it seems like each iteration is about 2.5 x more expensive. It also seems to converge to a pretty similar result. I don't see a clear benefit to this approach. |
This reverts commit 18f533e. Returning one residual per phase space point rather than per flux surface made dim_f 20 -> 1800 at test resolution, enlarging the Jacobian by the same factor. In practice the optimization takes considerably longer and reaches very similar results, so the cross term removal does not pay for its cost. The exponent fix is unaffected and remains: compute returns the island width squared, and the least squares objective supplies the remaining square. If this is worth revisiting, returning sqrt(sum_j W_j f_j^2) per flux surface would give the same cross term free sum without growing dim_f, at the cost of another guarded root. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
Removed comments explaining the handling of zero values in calculations to improve code clarity.
dpanici
reviewed
Aug 10, 2026
| ---------- | ||
| eq : Equilibrium | ||
| Equilibrium that will be optimized to satisfy the Objective. | ||
| num_rho : int, optional |
Collaborator
There was a problem hiding this comment.
We could also allow a user to pass in a custom array here. Reason for this is that some equilibria have pressure profiles that are not well-defined at rho=1 (#1745), in this current setup such an equilibrium gets NaN values for f_tr at the edge and at the node next to the edge due to p_r being nan at the edge, and thus "cvdrift (periodic)" which relies on p_r also is NaN, which trickles down to the trapped metric computation.
Something like rho : int or array, mention in the docstring that it needs to be linearly spaced, and if is a number will be what the current default is, and we can add a check like this in the objective build like
rho = np.linspace(0, 1, self._rho + 1)[1:] if isinstance(self._rho, int) else rho
errorif(not np.allclose(np.diff(rho)), ValueError, msg="rho array must be linearly spaced!")
errorif(np.any(np.isclose(rho,0.0)), ValueError, msg="rho array must not include the axis!")Rename num_rho to rho and accept either an int, which keeps the current np.linspace(0, 1, rho + 1)[1:] behaviour, or an explicit array of surfaces. Some equilibria have pressure profiles that are not well-defined at rho=1 (PlasmaControl#1745). Those give a NaN p_r at the edge, so "cvdrift (periodic)" is NaN there too, which trickles down into the trapped metric and leaves f_tr NaN on the last two surfaces. Passing an array that stops short of the edge avoids this. An array must be increasing, linearly spaced, and must exclude the axis: rho_res is a single scalar step used for the radial finite differences of Omega, and a decreasing array would silently flip its sign. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Move the third-party `interpax_fft` import out of the first-party `desc.*` block and into the third-party group, resolving the I001 flake8 failure introduced by the Bounce2D backend work. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
Omega'(s) sets the island width, and finite differencing it across the radial grid is only as accurate as the grid spacing allows. On ESTELL the finite difference does not converge to the true derivative at any practical num_rho, because dOmega/drho swings over orders of magnitude between adjacent surfaces. Omega_prime_method="analytic" instead pushes the radial derivatives of the eta grid field line data through the bounce stage in forward mode. The bounce integral is a fixed-node quadrature on [-1, 1] with all rho dependence in the integrand and the bounce points, so its derivative is exact: near a bounce point both 1 - lambda*B and its rho derivative vanish linearly, leaving the differentiated integrand with the same integrable inverse square root singularity the sin automorphism already absorbs. The derivative is taken at fixed lambda and fixed eta, matching what the finite difference across neighbouring surfaces measures. Verified by Taylor remainder tests: the first order remainder falls as h^2 for the field line data tangents, for dOmega/drho, and for the objective gradient with respect to the equilibrium degrees of freedom. Two notes on this rebase onto the master merge: - Only the Bounce1D path is supported. The analytic derivative needs the radial derivative of whatever data the bounce stage consumes; on the Bounce2D path that is the FFT grid data, whose rho derivative needs the base grid rho column to be differentiable. Constructing analytic + use_bounce1d=False now raises rather than silently falling back to fd. - Note that with stab_sacrifice=True the objective is f_res = Delta_s_4_sum * Omega_prime_s**2 where Delta_s_4 goes as 1/Omega_prime_s**2, so Omega' cancels algebraically and the two methods are bitwise identical. The option only changes anything at stab_sacrifice=False. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
jnp.where evaluates both branches, and in _root_cubic each branch is singular exactly where the other is selected: arccos leaves [-1, 1] when R^2 >= Q^3, and Q / A divides by zero when Q = R = 0. Discarding the value is not sufficient -- the NaN derivative of the discarded branch survives the where in reverse mode and poisons the gradient. Apply the standard double-where guard, substituting a harmless argument into each branch where it is not selected. Selected values are bitwise unchanged; only the derivative of the discarded branch is affected. This surfaced when differentiating the trapped-EP resonance objective through Bounce2D, but the fix is shared with EffectiveRipple and GammaC. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
Omega'(s) was estimated by finite differences across the radial grid,
which needs no extra bounce integrals but converges only as the grid is
refined. Measured against the analytic value, it does not converge at any
practical num_rho: median relative difference stays near 0.8-1.0 at
num_rho of 10, 20, 40 and 80, and it is biased ~22% low even on the
finest grid tested.
Replace it with the exact radial derivative, obtained by propagating a
radial tangent through the bounce integrals with forward-mode AD
(jax.jvp). Objective gradients then nest reverse-over-forward. Two stages
carry tangents: the field data on the eta grid (_eta_data) and, for
Bounce2D, the FFT stage (_fft_stage). Verified by Taylor remainder tests
-- stage-1 tangents converge at rate 1.964-2.011 across 14 quantities,
dOmega/drho at 2.05 (Bounce1D) and 2.24 (Bounce2D, Y_B=68), and the
objective gradient at 2.00 (Bounce1D). No NaN gradients on either backend.
Both backends are supported. Bounce2D needs low_ram=True to route around
an arccos in interpax_fft.eval1d that is not NaN-safe under AD, and its
gradients carry ~0.5% directional error from the regularized custom_jvp
(regular_points); nufft_eps < 1e-14 bypasses that and restores rate 2.00.
Bounce2D is CPU-only for now -- jax-finufft 1.2.0 has no CUDA lowering
for nufft2. Island widths agree across backends to 0.4%
(8.171e-03 vs 8.207e-03).
The Omega_prime_method kwarg is removed rather than defaulted, since
keeping a known-unconverged method invites silent misuse; the class
docstring records why. Note this changes nothing for stab_sacrifice=True,
where |Omega'| is constant across the summed resonance axis and cancels
exactly -- fd and analytic were bitwise identical there.
Removing fd made several guards unreachable, so they go too. Omega' is
now defined wherever Omega is, which collapses the valid_prime mask onto
valid; the output key is renamed "valid" to match. Rho is validated
off-axis by the objective, so the divide-by-zero guard on 2*rho and the
mask intersection it fed were both dead. Also dropped a broadcast_to
ahead of an op that broadcasts anyway (in _frequencies and the extracted
_barely_trapped_filter), a no-op float() cast, a dead default in
fun_data.pop("iota"), a duplicate "iota" in _FFT_BOUNCE_KEYS, and the
redundant None defaults on kwargs.get.
Tests build radial tangents via _tangent_keys/_seed_tangents and jax.jvp
rather than requesting a finite-difference stencil, and gain a Taylor
test for the analytic derivative.
Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
A mechanical pass (AST scans for unused parameters and for expressions evaluated more than once in a scope) turned up four redundancies: _v_tau was byte-identical to the definition in _fast_ion, confirmed by comparing ASTs rather than by eye. This module already imported _radial_drift from there, so import _v_tau alongside it. safediv(f_q_abs, denom) was evaluated twice to build the island width and its square. Hoist the shared quotient, which also makes it visible that the two profiles are the same quantity with and without the square root. The reshape restoring the flux surface and field line axes after batching appeared verbatim in both phase-space-average branches, differing only in how num_rho_psa is obtained. Hoist it below the branch. This also restores a comment, dropped when _field_line_length was extracted, that records why the reshape matters: _phase_space_average reduces over axis 1, so skipping it corrupts the result silently with only the magnitude wrong. _seed_tangent_key was called twice per key in a dict comprehension; bind it once with a walrus. Left alone deliberately: eta_vals and alpha_psa evaluate the same linspace, but they are the eta grid and the phase-space-average alpha grid, which coincide only because num_alpha_psa == num_eta. Collapsing them would assert an identity that is incidental. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
The linear res_weight branch computed the same bracketing test and the same interpolation weight twice, once toward each neighbouring surface, from four near-identical concatenates. Iterate over the two neighbours instead, so the logic appears once. The tuple is ordered prev then next because the later iteration overwrites the earlier through jnp.where, which is what gives the next surface precedence where both bracket a resonance; a comment records that, since the ordering now carries meaning. Verified against the previous implementation rather than by inspection, because this sits in the differentiated path: values are bitwise identical over 300 randomized cases spanning shapes, Omega scales over four orders of magnitude, and validity fractions from none to all; gradients agree to 1.8e-15 with no NaN introduced on either side. register_compute_fun declared "Psi" as a dependency, but nothing in the module reads it and nothing ever did in tracked history. Drop it. The objective still lists Psi in _keys_1dr, which is the separate mechanism that seeds it onto the field line grids. The resonance table appended +p/q and then repeated four near-identical lines for -p/q, recomputing p/q three times behind an np.logical_and doing scalar work. Loop over the sign instead. Checked equivalent across 392 combinations of p_max, q_max and range window, comparing all three output arrays including ordering. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
test_trapped_resonance_analytic_omega_prime carried no marker, so both CI lanes deselected it and the test validating the analytic derivative had never run. Mark it unit. A static scan of all 36 test files confirms it was the only unmarked test in the suite; every other test carries unit or regression. The TrappedResonance tests cost ~9 min, against ~2 min for the closest peer objective, Gamma_c. Two changes bring the CI share to ~3.4 min. Retune two tests. The no-NaN-gradient test is ~80% XLA compile time for the nested reverse-over-forward gradient graph, so resolution buys little: cutting arithmetic 4x saves 14%, and use_jit=False is slower still. Take the ~20% that is available. The compute-scalar-resolution test has no gradient graph and is compute bound, so halving rho and num_eta halves it. Both carry comments recording the trap that bounds them, because each fails differently and neither is evident from the numbers: rho=5 or num_transit=2 send the linear weighting's gradient to exactly zero, which satisfies a no-NaN assertion while testing nothing, and cutting knots_per_transit or num_quad moves the objective ~60x, which would verify resolution independence in a regime where the quadrature is unconverged. Move the two most expensive tests out of CI behind a new manual marker: backend agreement between Bounce1D and Bounce2D (~155 s, and Bounce2D is not the default backend) and compute-scalar resolution (~170 s). Nothing in this repo previously sat outside CI, so manual is a new convention rather than an existing one, and it is registered in setup.cfg with the rule that it must not be combined with unit or regression. An unmarked test would have worked mechanically but reads as the oversight fixed above. The compute-scalar case was a parametrization of a shared test whose regression marker applies to every parameter, so it could not be excluded per parameter. Extract it into its own method and add TrappedResonance to specials; without the second step it would have kept running through the generic path against HELIOTRON rather than precise_QA. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
…e wrappers
Two independent changes. They land together because they interleave within
_trapped_EP_resonance, and splitting them by hunk would mean committing an
intermediate state that was never run.
Restore Omega_prime_method={"analytic", "fd"}, defaulting to analytic. This
partially reverses 03eb5ca, which removed the finite difference entirely.
Both methods now work on both bounce backends; the earlier code restricted
analytic to Bounce1D.
Restoring the branch is not just re-adding it. Removing it had made two
things dead, and both are load bearing again:
- The sentinel has to survive the divide by 2*rho. A finite difference
writes fill_value on a surface with no valid neighbour, and dividing
that yields a finite, plausible looking number instead of a sentinel.
- valid_prime is a distinct mask again. It had been collapsed onto valid
because the analytic derivative needs no neighbour and the two
coincide. Under finite differences Omega'(s) is undefined on the
surfaces at either end, and four gates read the mask: res_weight in
both weighting methods, the final res_weight, and f_q_abs.
The "valid" output key keeps its name and is now computed from valid_prime,
so downstream diagnostics are unaffected and are identical under analytic.
Replace the docstring's unverified claim that analytic "roughly doubles"
the work with measured figures: 1.27x the objective evaluation and 1.64x
the gradient, timed against fd in one process at dc3273a, which carried
both paths. It is below 2x because forward mode doubles only the bounce
integrals, while the grids, resonance physics and phase space average run
once either way.
Separately, inline three thin wrappers: _barely_trapped_filter and
_field_line_length in the compute module, _seed_tangent_key in the
objective. Each wrapped a few lines behind a signature listing arrays that
were already locals at the single call site. _field_line_length folds into
the if/else on use_bounce1d that already existed, so no branch is added.
The key naming rule was checked against all 1053 compute keys before
inlining, since "append _r unless it already ends in one" is easy to get
subtly wrong.
Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
Author
|
@dario-panici I've now added analytic derivatives for Omega'. |
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
5 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.
Development of TrappedResonance objective function for EP optimization.