(perf): Newton-based flux-surface extrema refinement

Project.toml

@@ -5,6 +5,7 @@ authors = ["Orso Meneghini <orso82@gmail.com>"]
 
 [deps]
 AbstractTrees = "1520ce14-60c1-5f80-bbc7-55ef81b5835c"
+AdaptiveArrayPools = "4f381ef7-9af0-4cbe-99d4-cf36d7b0f233"
 Compat = "34da2185-b29b-5c13-b0c7-acf172513d20"
 Contour = "d38c429a-6771-53c6-b99e-75d170b6e991"
 CoordinateConventions = "7204ce3a-f536-43d2-be4a-fbed74e90d86"
@@ -43,8 +44,15 @@ SimpleNonlinearSolve = "727e6d20-b764-4bd8-a329-72de5adea6c7"
 StaticArrays = "90137ffa-7385-5640-81b9-e52037218182"
 Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2"
 
+[weakdeps]
+Interpolations = "a98d9a8b-a2ab-59e6-89dd-64a1c18fca59"
+
+[extensions]
+IMASInterpolationsExt = "Interpolations"
+
 [compat]
 AbstractTrees = "0.4"
+AdaptiveArrayPools = "0.3.6"
 Compat = "4.10"
 Contour = "0.6"
 CoordinateConventions = "1"
@@ -58,6 +66,7 @@ Graphs = "1"
 HelpPlots = "1.1"
 IMASdd = "8.3"
 IMASutils = "1.5.1"
+Interpolations = "0.13, 0.14, 0.15, 0.16"
 Jedis = "0.3"
 LaTeXStrings = "1"
 Measurements = "2"
@@ -76,7 +85,7 @@ Random = "1.11.0"
 Ratios = "0.4"
 RecipesBase = "1"
 Roots = "2, 3"
-SimpleNonlinearSolve = "2.12.0"
+SimpleNonlinearSolve = "2"
 StaticArrays = "1"
 Statistics = "1.11.1"
 julia = "1"

claudedocs/cubic_tracer_benchmark.jl

@@ -0,0 +1,23 @@
+# PC vs adaptive-RK4 vs Contour.jl: drift, point count, timing on DIII-D closed surfaces.
+# Scratch diagnostic (not a pass/fail test).
+import IMAS
+using Printf
+
+filename = joinpath(pkgdir(IMAS.IMASdd), "sample", "D3D_eq_ods.json")
+dd = IMAS.json2imas(filename; show_warnings=false)
+eqt = dd.equilibrium.time_slice[1]
+eqt2d = IMAS.findfirst(:rectangular, eqt.profiles_2d)
+r, z, itp = IMAS.ψ_interpolant(eqt2d)
+RA, ZA = eqt.global_quantities.magnetic_axis.r, eqt.global_quantities.magnetic_axis.z
+psi_axis = itp(RA, ZA)
+eqt1d = eqt.profiles_1d
+# a mid-radius closed level
+c = psi_axis + 0.5 * (eqt1d.psi[end] - psi_axis)
+seed = (maximum(r) - 1e-3, ZA)   # outboard; project finds the surface
+seed = IMAS._project_to_level(itp, c, seed)[1]
+
+for method in (:pc, :rk4)
+    t = @elapsed (Rs, Zs, closed) = IMAS._trace_surface_cubic(itp, c, seed; method)
+    drift = maximum(abs(IMAS.FI.value_gradient(itp,(Rs[k],Zs[k]))[1]-c) for k in eachindex(Rs))
+    @printf("%-5s closed=%s  N=%4d  drift=%.2e  t=%.2f ms\n", method, closed, length(Rs), drift, t*1e3)
+end

ext/IMASInterpolationsExt.jl

@@ -0,0 +1,18 @@
+module IMASInterpolationsExt
+
+import IMAS
+import Interpolations
+
+# Interpolations.jl backend for the generic interpolant helpers in IMAS (src/physics/fields.jl).
+# Active only when Interpolations is loaded. It uses positional coords and returns SVector/SMatrix,
+# which index like the FastInterpolations NTuple/Matrix the IMAS math expects.
+@inline IMAS._gradient(itp::Interpolations.AbstractInterpolation, r, z) =
+    Interpolations.gradient(itp, r, z)
+
+@inline IMAS._value_gradient(itp::Interpolations.AbstractInterpolation, r, z) =
+    (itp(r, z), Interpolations.gradient(itp, r, z))
+
+@inline IMAS._hessian!(H, itp::Interpolations.AbstractInterpolation, r, z) =
+    (H .= Interpolations.hessian(itp, r, z); H)
+
+end

src/IMAS.jl

@@ -4,6 +4,7 @@ using Printf
 using Compat:@compat
 import OrderedCollections
 import FastInterpolations as FI
+using AdaptiveArrayPools
 const document = OrderedCollections.OrderedDict()
 
 macro import_all(mod)

src/physics.jl

@@ -20,6 +20,7 @@ include(joinpath("physics", "tf.jl"))
 include(joinpath("physics", "currents.jl"))
 include(joinpath("physics", "fields.jl"))
 include(joinpath("physics", "fluxsurfaces.jl"))
+include(joinpath("physics", "fluxsurfaces_cubic.jl"))
 include(joinpath("physics", "rf.jl"))
 include(joinpath("physics", "neoclassical.jl"))
 include(joinpath("physics", "profiles.jl"))

src/physics/fields.jl

@@ -3,10 +3,13 @@ import SimpleNonlinearSolve
 
 document[Symbol("Physics fields")] = Symbol[]
 
-# gradient of a ψ interpolant: FI gets the native fast path,
-# the open version is for users still passing an Interpolations.jl interpolant
-@inline _psi_gradient(itp::FI.AbstractInterpolant, r, z) = FI.gradient(itp, (r, z))
-@inline _psi_gradient(itp, r, z) = parentmodule(typeof(itp)).gradient(itp, r, z)
+# Analytic gradient / value+gradient / Hessian of a ψ interpolant, dispatched on the backend.
+# FastInterpolations gets the native in-place fast path here; other backends (e.g. Interpolations.jl)
+# are added by package extensions, loaded only when that backend is (see ext/IMASInterpolationsExt.jl).
+# Results are indexed g[1]/g[2] and H[i,j], so any backend returning indexable values works.
+@inline _gradient(itp::FI.AbstractInterpolant, r, z) = FI.gradient(itp, (r, z))
+@inline _value_gradient(itp::FI.AbstractInterpolant, r, z) = FI.value_gradient(itp, (r, z))
+@inline _hessian!(H, itp::FI.AbstractInterpolant, r, z) = FI.hessian!(H, itp, (r, z))
 
 """
     Br_Bz(eqt2d::IMAS.equilibrium__time_slice___profiles_2d)
@@ -30,7 +33,7 @@ end
     Br_Bz(PSI_interpolant, r::T, z::T) where {T<:Real}
 """
 function Br_Bz(PSI_interpolant, r::T, z::T) where {T<:Real}
-    grad = _psi_gradient(PSI_interpolant, r, z)
+    grad = _gradient(PSI_interpolant, r, z)
     inv_twopi_r = 1.0 / (2π * r)
     Br = grad[2] * inv_twopi_r
     Bz = -grad[1] * inv_twopi_r