TrappedResonance - #2257
Open
ejpaul wants to merge 247 commits into
Open
Conversation
…sics Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
The threshold in code is 2*pi (float(2 * jnp.pi)), not 2.5*pi. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
The old Returns section was just a bare name list, and had gone stale: it listed f_res twice and included f_q_c/f_q_s (which are locals, not returned), while omitting Delta_s_prof and valid_prime (which are). Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
Reverts to plain 2-point linear interpolation for res_weight; the 4-point Catmull-Rom Hermite refinement (_shift_rho, _hermite_crossing) was WIP that didn't pan out. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
Both denominators are already floored away from zero via jnp.maximum(..., 1e-30), so a plain division is safe; matches the convention elsewhere in the codebase of using one zero-avoidance mechanism, not both. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
safediv already handles the pitch*B == 1 turning point (returns 0 by default there, same as _radial_drift and friends); flooring the value under the sqrt on top of that was redundant double protection. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
…ption f_q_abs and Delta_s_profile aren't ancestors of f_res (which is built from the independent f_q_abs_sq/denom**2 path), so their jnp.maximum floors weren't protecting anything reachable by autodiff in the standard objective; both quantities are always >= 0 regardless. num_well was accepted as a TrappedResonance constructor argument but never stored in _hyperparameters, so it was silently ignored; the compute function already hardcodes num_well=1. Remove the dead parameter and its doc entry, and drop it from _bounce1D_doc so it's no longer advertised as a configurable kwarg. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
ejpaul
marked this pull request as ready for review
July 21, 2026 19:37
_trapped_EP_resonance read 26 kwargs via kwargs.get(...), but only 4 were declared through register_compute_fun (the rest reached the function only because both call sites bypass validation by calling the private _compute directly instead of the public eq.compute(...), which checks kwargs against allowed_kwargs). Add _resonance_doc with an entry for each of the remaining kwargs, adapting descriptions already written for TrappedResonance's __init__ docstring where available. Rename the four kwargs that only exist to pass grids/data built by TrappedResonance.compute (eta_grid, psa_grid, data_eta, data_psa) to a leading-underscore form (matching the existing _vander convention for parameters not meant to be set by external callers), and document them as such. Verified via allowed_kwargs that all needed names are now declared, and confirmed no test regressions. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
Previously excluded because resonance detection converges slower with grid resolution than the shared 6% tolerance allows. Give it its own tolerance (10%) instead, and reduce p_max/q_max in the test's reduced-resolution settings so the sum is over fewer resonances, markedly cutting sensitivity to which crossings a coarse rho grid happens to catch (empirically ~16% -> ~8% spread across the sweep). Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
The top-level import already covers this; the local re-import (which shadowed it) was the only user of simpson, so both flake8 warnings (unused top-level import, redefinition) go away with one deletion. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
Collaborator
|
Another thought on the metric: does |
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>
5 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 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.