(perf): Newton-based flux-surface extrema refinement - #308
Merged
Conversation
Standalone utility that refines a flux-surface geometric extremum by solving {psi = target_psi, dpsi/dn = 0} from a rough seed, using the cubic interpolant's analytic gradient and Hessian. extremum_of=:R finds max_r/min_r (dpsi/dZ=0); extremum_of=:Z finds max_z/min_z (dpsi/dR=0); the seed selects the basin. Falls back to the seed on non-convergence or a degenerate Jacobian, and rejects an invalid extremum_of with ArgumentError.
Building block for replacing the Contour-based extrema refinement in refine_extrema (integration deferred). Adds runtests_refine_extremum.jl (four-extrema recovery to ~1e-15 on an analytic ellipse, fallback, invalid-arg) and registers it in the default suite.
…nto fluxsurfaces_cubic.jl Move the cubic-interpolant + Newton extremum-refinement helpers (_newton2d, _extremum_residual, _critical_residual, _refine_extremum, _refine_extremum_bounded) out of the 2549-line fluxsurfaces.jl into a dedicated fluxsurfaces_cubic.jl. Pure relocation: these helpers are not yet wired into trace_surfaces! (the Contour + Optim.Brent path remains the active extremum step), so there is no behavior change. Establishes a dedicated home for the cubic-interpolant flux-surface engine, separate from the Contour-based one, ready for the planned general cubic tracer. runtests_refine_extremum.jl passes.
…rameter for genuine extremum recovery
…e+resample+reorder)
…2.1, 12.2) _resample_contour: handle m==1 (fill), guard L>0, detect closed polyline and sample half-open [0,L) so n outputs are distinct with uniform wrap segment. Call reorder_flux_surface! with force_close=false to preserve the half-open representation (n points, no appended duplicate). _find_xpoint: add domain guard — reject Newton criticals outside the interpolant grid extents, preventing spurious extrapolated saddles from being returned as genuine X-points. Tests: tighten _resample_contour assertions (distinctness, wrap parity, 1e-5 on-surface, m==1 case); add _find_xpoint :extremum branch (O-point) and out-of-domain negative case. 56/56 pass, no regressions.
…9.1,10.1,10.2,11b.1,13.1)
…_cubic docstring (audit 10.2)
…ompatible, unbiased arclength)
… extrema refinement - trace_surfaces_cubic(eqt, wall_r, wall_z) high-level overload mirrors trace_surfaces for 1:1 comparison - refine_extrema option wires _refine_extremum (cubic analogue of the Contour Optim refinement) - graceful inner-proxy fallback + warning when a level (separatrix) cannot close
…e Newton
Swap trace_surfaces' geometric-extremum refinement from the
Contour.lines(Br=0)/(Bz=0) + Optim.Brent search to the X-point-aware 2-D
Newton _refine_extremum! on the psi interpolant (analytic grad/Hessian).
The same extremum conditions ({psi=c, dpsi/dZ=0} for R-extrema,
{psi=c, dpsi/dR=0} for Z-extrema) are solved point-locally per surface.
- _refine_extremum! (in-place, caller-owned 2x2 Hessian buffer) + non-bang
_refine_extremum wrapper; both batch tracers share one buffer per call
- all hessian -> hessian! (value_gradient/gradient with tuple args are
already 0-heap), eliminating per-call 2x2 matrix allocations
- remove experimental trace_surfaces2 twin (now folded into trace_surfaces)
- BR/BZ args retained for API compatibility (now unused)
Refine ~17x faster (~278us -> ~16us refine-only), ~7000 -> ~2 allocations;
extrema agree with the prior Optim result to <0.5mm. Full suite passes.
After the Newton extrema-refine swap, the precomputed BR/BZ field grids are no longer used by trace_surfaces. Remove them from the low-level signature and drop the now-dead Br_Bz(eqt2d) calls that fed them (high-level eqt overload, flux_surfaces! pipeline, 3 test sites). The Br_Bz routine itself is unchanged and still used elsewhere. The high-level trace_surfaces(eqt, wall_r, wall_z) signature is unchanged. No external caller of the low-level form exists in the dev ecosystem (FUSE does not call trace_surfaces; IMASdd uses only the high-level form). Dropping the two dead full-grid Br_Bz(eqt2d) computations also slightly speeds up the high-level overload and flux_surfaces! pipeline. BREAKING CHANGE: trace_surfaces low-level form no longer takes BR/BZ matrices; callers must remove those two positional args.
…dmax/findmin Replace the four separate findmax/findmin calls (each scanning a coordinate array, value discarded then re-loaded by index) with a single sweep that tracks all four geometric extrema and their indices at once. An order of magnitude faster on a typical flux-surface polyline, zero allocations, bit-identical results (tie-breaking preserved: first extreme index, matching findmax/findmin). findmax/findmin carry NaN-aware ordering plus a (value,index) reduction that defeats vectorization; plain >/< is safe here because traced flux-surface coordinates are finite. Adds a length guard for the @inbounds sweep. Adds a fluxsurface_extrema unit test (explicit hand-checked polyline incl. a tie, equivalence with findmax/findmin on a D-shaped polyline, the length guard).
The Newton refine wired into trace_surfaces used FastInterpolations-only APIs (FI.gradient/value_gradient/hessian!), so passing a non-FI ψ interpolant — e.g. an Interpolations.jl one, as FRESCO builds — raised a MethodError on the refine. Route the refine's interpolant access through _psi_* adapters (extending the existing _psi_gradient FI-fast / open-fallback split to value_gradient and hessian!), and drop the FI.AbstractInterpolant restriction from _refine_extremum!/_extremum_residual!/_critical_residual! and the wrapper. The FI path is unchanged (adapters are @inline passthroughs to the same in-place calls); other backends use their gradient/hessian via parentmodule dispatch. Add Interpolations as a test dependency and tests exercising the non-FI path for real: _refine_extremum on an Interpolations cubic of an analytic ellipse, and trace_surfaces driven by an Interpolations interpolant matching the FastInterpolations result on DIII-D.
…paths x^2.0 (Float exponent) dispatches to the general libm pow() (exp/log based); x^2 (Int literal) compiles to x*x via Base.literal_pow — about an order of magnitude faster per element, numerically identical. Replace .^2.0 / ^2.0 with .^2 / ^2 in the flux-surface (Bp2, fluxexpansion), SOL, particle, and wall-flux-plot paths.
…nsion Replace the parentmodule(typeof(itp)) duck-typing for non-FastInterpolations interpolants with a proper weakdep package extension (ext/IMASInterpolationsExt.jl), active only when Interpolations.jl is loaded. Rename the backend helpers _psi_gradient/_psi_value_gradient/_psi_hessian! to the generic _gradient/_value_gradient/_hessian! (they work on any interpolant). FastInterpolations keeps the built-in in-place fast path; the FI path and behavior are unchanged. Interpolations becomes a [weakdeps]/[extensions] entry with compat. Also tidy the refine block (compress the comment, drop redundant broadcast dots from the scalar k=1 frac-scaling) in trace_surfaces.
…ls for memory management
These two helpers were used only by the Contour.lines(Br=0)/(Bz=0) + Optim.Brent extrema-refine block, which was replaced by the Newton _refine_extremum! backend. They now have zero callers (src and tests) — delete (~67 lines).
The pure-Newton + mirror extremum refine produced off-surface garbage for
outer flux surfaces near an X-point: the curvature-sign genuineness test
false-rejects a correct outboard point (grid-edge ψ_ZZ artifact), the
critical-point Newton then diverges, and the mirror point lands off-grid —
corrupting r_outboard/r_inboard. On KDEMO this crashed FUSE.warmup via
nuestar's `@assert a_eq .> 0` (FUSE CI run 28135611069).
Add `_robust_refine_extremum!` and switch trace_surfaces to it:
- solve {ψ=c, ∂ψ/∂n=0} with a globalized damped (backtracking) Newton
(`_damped_newton2d`) — no divergence where det(J)=ψ_R·ψ_ZZ→0 near the
separatrix blows up the plain Newton step
- classify the result by physical region instead of curvature sign:
axis-relative direction + confined side of every X-point
(`(p−xp)·(axis−xp)>0`), using the precomputed O-point (axis) and X-points
(`eqt.boundary.x_point`, threaded through as `xpoints`) — no per-refine
critical-point search, generic over 0/1/2 X-points
- recover a bad seed by re-seeding toward the (confined) magnetic axis and
re-solving, instead of mirroring across a possibly-divergent critical point
The previous `_refine_extremum!` is kept (unused in production) for A/B
comparison. Verified: KDEMO/D3D a_eq all positive & monotonic, stress test
recovers correct extrema from deliberately bad seeds (ψ_N>1 SOL and private
region), golden runtests_interpolations green, FUSE warmup(:KDEMO) no longer
crashes. Refine cost ~25 µs / 0 allocations on omas_sample (vs ~294 µs on the
old Optim/Contour path).
Add a regression test for _robust_refine_extremum! covering the original requirement: from any seed on the correct side of the axis — including one outside the separatrix (ψ_N>1) or in the private flux region across an X-point — it must still reach the correct CONFINED extremum. - analytic ellipse: far outside / above seeds still find max_r / max_z - DIII-D: outside-separatrix and private-region (X-point-mirrored) seeds recover the same extrema as the good seed, and land confined (ψ_N≤1, below the upper X-point) This guards both observed failures — KDEMO (a_eq<0 from corrupted r_outboard/r_inboard) and MANTA (elongation≈0 → sqrt DomainError on aarch64). Adds a non-bang `_robust_refine_extremum` convenience wrapper.
The damped Newton had no domain bound: from interior seeds, intermediate iterates were measured leaving the grid by up to ~3.6 m (into the cubic extrapolation region) before coming back. The final result was still gated, but the search wandering off-grid is fragile (extrapolated ψ/Hessian). Clamp every iterate (and the initial seed) to the ψ grid box [r1,rN]×[z1,zN], threaded as `lo`/`hi` from trace_surfaces. The iterate now physically cannot escape the domain. Confirmed: an active boundary crossing (result≠seed AND ψ_N>1) never occurs over a 3196-seed LCFS sweep. Also extend the regression tests: the production-faithful (clamped) outside/ private DIII-D cases, and the hardest case — at ψ_N=0.999 the private region across the upper X-point has its own ψ_N=0.999 ∂ψ/∂R=0 solution; seeds at/above the X-point must still return the confined max_z below it (the MANTA geometry).
…e X-point Strengthen the robust refine with the user's "stay on the confined side of the X-point" constraint, encoded as a box bound on the extremized coordinate (a square Newton system cannot take a 3rd equation, but it can take bounds). For a Z-extremum, tighten the search box in Z to [Z_axis-side .. relevant X-point Z]: the iterate then physically cannot cross into the private flux region above/below the X-point, so that wrong solution is unreachable by construction (a seed already in the private lobe is clamped back into the confined band and converges to the genuine extremum without needing the axis-ward reseed). R-extrema keep the grid box — the top/bottom X-points sit near R≈R_axis and a Z-style bound would wrongly clip the outboard/inboard. The region classification + reseed remain as a general safety net. Tests (near-separatrix private-lobe, extreme seeds) and golden all green.
The de1653a Z-bound clamped a Z-extremum's coordinate to the relevant X-point's Z. That is a 1-D cutoff: for an inboard X-point (small R) it wrongly clips a confined max_z that sits at larger R and higher Z than the X-point. The region classification already uses the full 2-D dot test `(p−xp)·(axis−xp) > 0` (same side of each X-point as the axis), which handles those cases correctly — the confined surface is closed inside the separatrix and never wraps past an X-point, so every confined point is on the axis side. The R-component keeps the dot positive for inboard X-points where a Z-only bound would over-clip. Verified on constructed weird equilibria (inboard single-X and inboard double-null): the dot-based refine finds the correct extrema from production-like seeds; golden and refine_extremum tests stay green.
_damped_newton2d now probes the line search with a residual-only eval (value_gradient, which also yields the ψ=target Jacobian row for free) and computes the Hessian row only on accepted steps — instead of a combined residual+Jacobian eval that recomputed and then discarded the probe's Hessian. This halves the interpolant work per Newton iterate. New _extremum_eqs builds the split (residual, hessrow) system; the old combined _extremum_residual!/_newton2d/_refine_extremum! path is untouched. Bit-identical extrema (320/320 refine results across omas+D3D), 0 allocations preserved; refine step -20% (omas) / -22% (D3D).
…on everywhere Remove the fragile pure-Newton + mirror refine routine (_refine_extremum! / _refine_extremum) that caused the KDEMO (a_eq<0) and MANTA (elongation≈0) failures, along with its now-unused helpers (_newton2d, _extremum_residual!, _critical_residual!). Production trace_surfaces already refines with the robust _robust_refine_extremum!; this makes it the only refine path. The two genuine consumers of the plain Newton are repointed to the globalized _damped_newton2d so they converge safely too: - _find_xpoint solves ∇ψ=0 via _damped_newton2d + the new split-form _critical_eqs (unclamped, so the domain guard still rejects out-of-grid seeds); - the standalone cubic tracer (trace_surfaces_cubic) refines extrema with the grid-clamped _robust_refine_extremum! instead of the mirror routine. Tests: drop the _refine_extremum testsets (behaviors already covered by the _robust_refine_extremum testsets) and convert the backend-agnostic Interpolations.jl test to the robust path. Full suite green.
Member
Author
jmcclena
approved these changes
Jul 3, 2026
Member
Author
|
@jmcclena Thank you for merging this! By the way, if you plan to update the IMAS version, including this PR, |
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.
Summary
trace_surfacesnow refines geometric extrema with an X-point-aware 2-D Newton solver on the ψ interpolant (analytic ∇ψ / Hessian) instead ofContour.lines(Br/Bz=0)+Optim.Brent:trace_surfacedirectly)Also lands a standalone cubic-interpolant flux-surface tracer (which is not yet wired to the production path)
What changed
trace_surfacesrefine → Newton (_refine_extremum!on the ψ interpolant), with in-place Hessian + pooled scratch so the refine allocates almost nothing.[weakdeps]package extension (ext/IMASInterpolationsExt.jl), loaded only when that package is.src/physics/fluxsurfaces_cubic.jl(new): the Newton refine + a standalone predictor-corrector cubic tracer (trace_surface_cubic/trace_surfaces_cubic), A/B-validated against the Contour path and kept separate pending review.fluxsurface_extrema(was 4×findmax/findmin), integer exponents in hot loops (^2not^2.0), removed dead_extrema_index/_extrema_cost.trace_surfaces(psi, f, r, z, PSI, …)no longer takes theBR/BZmatrices (unused after the swap). The high-leveltrace_surfaces(eqt, wall_r, wall_z)is unchanged.Performance
@btime, high-leveltrace_surfaces(eqt, first_wall(wall)...)onomas_sample.json, master vs branch:The Newton refine now adds 0 allocations and ~19 µs (refine-on and refine-off allocate identically), versus ~294 µs / ~7,200 allocs on master — the refine step itself is ~16× faster. Extrema match the previous Optim result within < 0.5 mm. Full test suite green.
Note
FUSE CI will be tested with this PR, and will be updated if it passes