From 2b07c7b3d79eb0d13239fb277b77f33184639551 Mon Sep 17 00:00:00 2001 From: Hongyang Zhou Date: Tue, 10 Feb 2026 23:16:37 -0500 Subject: [PATCH 1/3] Accept trace solution in get_adiabaticity directly --- docs/examples/features/demo_hybrid.jl | 46 ++++----------------------- src/utility/utility.jl | 46 +++++++++++++++++++++++++++ 2 files changed, 53 insertions(+), 39 deletions(-) diff --git a/docs/examples/features/demo_hybrid.jl b/docs/examples/features/demo_hybrid.jl index 8a545b9f1..104df29aa 100644 --- a/docs/examples/features/demo_hybrid.jl +++ b/docs/examples/features/demo_hybrid.jl @@ -41,7 +41,7 @@ E_field = TP.Field((x, t) -> SA[0.0, 0.0, 0.0]) m = TP.mᵢ q = TP.qᵢ -q2m = q / m +q2m = q / m; # ## Step 1: Full Orbit Reference Trace # @@ -62,7 +62,7 @@ tspan = (0.0, 30 * T_gyro) param_fo = prepare(E_field, B_field; species = Proton) prob_fo = ODEProblem(trace, u0, tspan, param_fo) -sol_fo = solve(prob_fo, Vern9(); save_everystep = true) +sol_fo = solve(prob_fo, Vern9()) ## Verify trapping: z should oscillate, not diverge z_fo = [u[3] for u in sol_fo.u] @@ -79,7 +79,7 @@ stateinit_gc, param_gc = TP.prepare_gc( u0, bottle_E_static, bottle_B_static; species = Proton ) prob_gc = ODEProblem(trace_gc!, stateinit_gc, tspan, param_gc) -sol_gc = solve(prob_gc, Vern9()) +sol_gc = solve(prob_gc, Vern9()); # ## Step 3: Hybrid Solver # @@ -100,7 +100,7 @@ alg = AdaptiveHybrid(; Random.seed!(1234) ## Set verbose = true to see the dynamic switching -sol = TP.solve(TraceHybridProblem(u0, tspan, p), alg; verbose = false)[1] +sol = TP.solve(TraceHybridProblem(u0, tspan, p), alg; verbose = false)[1]; # ## Step 4: Compute Adiabaticity # @@ -108,41 +108,9 @@ sol = TP.solve(TraceHybridProblem(u0, tspan, p), alg; verbose = false)[1] # guiding center approximation holds. When `ε < threshold`, GC is # valid; when `ε ≥ threshold`, we need full orbit tracing. -function compute_adiabaticity(sol, Bfunc, q, m) - n = length(sol.t) - ε = Vector{Float64}(undef, n) - for i in 1:n - xv = sol.u[i] - x = xv[SA[1, 2, 3]] - v = xv[SA[4, 5, 6]] - t = sol.t[i] - - B = Bfunc(x, t) - Bmag = norm(B) - b̂ = B / Bmag - vpar_val = v ⋅ b̂ - vperp_vec = v - vpar_val * b̂ - μ = m * (vperp_vec ⋅ vperp_vec) / (2 * Bmag) - - ε[i] = get_adiabaticity(x, Bfunc, q, m, μ, t) - end - return ε -end - -function compute_adiabaticity_gc(sol_gc, Bfunc, q, m, μ) - n = length(sol_gc.t) - ε = Vector{Float64}(undef, n) - for i in 1:n - x = sol_gc.u[i][SA[1, 2, 3]] - t = sol_gc.t[i] - ε[i] = get_adiabaticity(x, Bfunc, q, m, μ, t) - end - return ε -end - -ε_fo = compute_adiabaticity(sol_fo, B_field, q, m) -ε_gc = compute_adiabaticity_gc(sol_gc, B_field, q, m, param_gc[3]) -ε_hybrid = compute_adiabaticity(sol, B_field, q, m) +ε_fo = get_adiabaticity(sol_fo) +ε_gc = get_adiabaticity(sol_gc) +ε_hybrid = get_adiabaticity(sol) ## Common colorbar range across all solvers ε_clamp_lo = 1.0e-3 diff --git a/src/utility/utility.jl b/src/utility/utility.jl index 449ee9da1..81b12772c 100644 --- a/src/utility/utility.jl +++ b/src/utility/utility.jl @@ -652,3 +652,49 @@ end function get_adiabaticity(r, Bfunc, μ, t::Number = 0.0; species = Proton) return get_adiabaticity(r, Bfunc, species.q, species.m, μ, t) end + +""" + get_adiabaticity(sol::AbstractODESolution) + +Calculate the adiabaticity parameter `ϵ` from a solution `sol`. +The parameters `q`, `m`, `Bfunc` are extracted from `sol.prob.p`. +""" +function get_adiabaticity(sol::AbstractODESolution) + n = length(sol.t) + ε = Vector{Float64}(undef, n) + state_len = length(sol.u[1]) + + if state_len == 6 # FO or Hybrid + q2m, m, Efunc, Bfunc, _ = sol.prob.p + q = q2m * m + + for i in 1:n + xv = sol.u[i] + x = xv[SA[1, 2, 3]] + v = xv[SA[4, 5, 6]] + t = sol.t[i] + + B = Bfunc(x, t) + Bmag = norm(B) + b̂ = B / Bmag + vpar_val = v ⋅ b̂ + vperp_vec = v - vpar_val * b̂ + μ = m * (vperp_vec ⋅ vperp_vec) / (2 * Bmag) + + ε[i] = get_adiabaticity(x, Bfunc, q, m, μ, t) + end + elseif state_len == 4 # GC + q, q2m, μ, Efunc, Bfunc = sol.prob.p + m = q / q2m + + for i in 1:n + x = sol.u[i][SA[1, 2, 3]] + t = sol.t[i] + ε[i] = get_adiabaticity(x, Bfunc, q, m, μ, t) + end + else + error("Unsupported solution state dimension: $state_len") + end + + return ε +end From c66f95551b6cab09e5d5bb000ffc1b3c24402a89 Mon Sep 17 00:00:00 2001 From: Hongyang Zhou Date: Wed, 11 Feb 2026 00:04:37 -0500 Subject: [PATCH 2/3] Dispatch on get_adiabaticity --- docs/examples/analytic/demo_proton_dipole.jl | 4 +- src/utility/utility.jl | 38 ++++++++++++++++++ test/test_utility.jl | 42 ++++++++++++++++++++ 3 files changed, 81 insertions(+), 3 deletions(-) diff --git a/docs/examples/analytic/demo_proton_dipole.jl b/docs/examples/analytic/demo_proton_dipole.jl index 7ea72cc45..edf3b459f 100644 --- a/docs/examples/analytic/demo_proton_dipole.jl +++ b/docs/examples/analytic/demo_proton_dipole.jl @@ -78,9 +78,7 @@ end # Calculate adiabaticity: ρ/Rc # The 1st adiabaticity stands for the ratio between the curvature radius and the gyroradius. Here we show the inverse of adiabaticity for plotting. ratio = let - q, q2m, μ, _, Bfunc = param_gc - m = q / q2m - adiabaticity = [get_adiabaticity(sol_gc(t)[1:3], Bfunc, q, m, μ, t) for t in tsample] + adiabaticity = [get_adiabaticity(sol_gc, t) for t in tsample] 1 ./ adiabaticity end diff --git a/src/utility/utility.jl b/src/utility/utility.jl index 81b12772c..c1a8c636f 100644 --- a/src/utility/utility.jl +++ b/src/utility/utility.jl @@ -698,3 +698,41 @@ function get_adiabaticity(sol::AbstractODESolution) return ε end + +""" + get_adiabaticity(sol::AbstractODESolution, t) + +Calculate the adiabaticity parameter `ϵ` from a solution `sol` at time `t`. +`t` must be a scalar. +""" +function get_adiabaticity(sol::AbstractODESolution, t) + state_len = length(sol.u[1]) + + if state_len == 6 # FO or Hybrid + q2m, m, Efunc, Bfunc, _ = sol.prob.p + q = q2m * m + + xv = sol(t) + x = xv[SA[1, 2, 3]] + v = xv[SA[4, 5, 6]] + + B = Bfunc(x, t) + Bmag = norm(B) + b̂ = B / Bmag + vpar_val = v ⋅ b̂ + vperp_vec = v - vpar_val * b̂ + μ = m * (vperp_vec ⋅ vperp_vec) / (2 * Bmag) + + return get_adiabaticity(x, Bfunc, q, m, μ, t) + + elseif state_len == 4 # GC + q, q2m, μ, Efunc, Bfunc = sol.prob.p + m = q / q2m + + xv = sol(t) + x = xv[SA[1, 2, 3]] + return get_adiabaticity(x, Bfunc, q, m, μ, t) + else + error("Unsupported solution state dimension: $state_len") + end +end diff --git a/test/test_utility.jl b/test/test_utility.jl index 2353a7da4..1cd62692d 100644 --- a/test/test_utility.jl +++ b/test/test_utility.jl @@ -581,8 +581,50 @@ import TestParticle as TP # Should return Inf B_zero_test(x, t) = SA[0.0, 0.0, 0.0] @test get_adiabaticity(r, B_zero_test, q, m, μ, 0.0) == Inf + # Simple setup for Solution Dispatch + B0 = 1.0e-8 + # Define fields as closures to capture B0, or just use constant + B_func_disp(x, t) = SA[0.0, 0.0, 1.0e-8] + B_func_disp(x) = B_func_disp(x, 0.0) + E_func_disp(x, t) = SA[0.0, 0.0, 0.0] + E_func_disp(x) = E_func_disp(x, 0.0) + + # FO Solution + x0 = SA[1.0, 0.0, 0.0] + v0 = SA[0.0, 1.0e5, 0.0] + stateinit = [x0..., v0...] + tspan = (0.0, 1.0e-4) # Short duration + param_fo = prepare(E_func_disp, B_func_disp; species = Ion(1, 1)) + prob_fo = ODEProblem(trace!, stateinit, tspan, param_fo) + sol_fo = solve(prob_fo, Tsit5()) + + # Test vector output + ε_fo = get_adiabaticity(sol_fo) + @test length(ε_fo) == length(sol_fo.t) + @test all(isfinite, ε_fo) + # Uniform field -> 0 + @test all(ε_fo .== 0.0) + + # Test scalar output + ε_fo_scalar = get_adiabaticity(sol_fo, 0.5e-4) + @test ε_fo_scalar == 0.0 + + # GC Solution + stateinit_gc, param_gc = prepare_gc(stateinit, E_func_disp, B_func_disp; species = Ion(1, 1)) + prob_gc = ODEProblem(trace_gc!, stateinit_gc, tspan, param_gc) + sol_gc = solve(prob_gc, Tsit5()) + + # Test vector output + ε_gc = get_adiabaticity(sol_gc) + @test length(ε_gc) == length(sol_gc.t) + @test all(ε_gc .== 0.0) + + # Test scalar output + ε_gc_scalar = get_adiabaticity(sol_gc, 0.5e-4) + @test ε_gc_scalar == 0.0 end end + @testset "ZeroField Unitful" begin # Check that ZeroField returns ZeroVector for Unitful types zf = TP.ZeroField() From 0c31599e4924634361dff8ae2ff8c412a5f110d2 Mon Sep 17 00:00:00 2001 From: Hongyang Zhou Date: Wed, 11 Feb 2026 10:15:32 -0500 Subject: [PATCH 3/3] Refactor get_adiabaticity --- src/utility/utility.jl | 91 +++++++++++++++--------------------------- 1 file changed, 32 insertions(+), 59 deletions(-) diff --git a/src/utility/utility.jl b/src/utility/utility.jl index c1a8c636f..761acea64 100644 --- a/src/utility/utility.jl +++ b/src/utility/utility.jl @@ -653,66 +653,13 @@ function get_adiabaticity(r, Bfunc, μ, t::Number = 0.0; species = Proton) return get_adiabaticity(r, Bfunc, species.q, species.m, μ, t) end -""" - get_adiabaticity(sol::AbstractODESolution) - -Calculate the adiabaticity parameter `ϵ` from a solution `sol`. -The parameters `q`, `m`, `Bfunc` are extracted from `sol.prob.p`. -""" -function get_adiabaticity(sol::AbstractODESolution) - n = length(sol.t) - ε = Vector{Float64}(undef, n) - state_len = length(sol.u[1]) - - if state_len == 6 # FO or Hybrid - q2m, m, Efunc, Bfunc, _ = sol.prob.p - q = q2m * m - - for i in 1:n - xv = sol.u[i] - x = xv[SA[1, 2, 3]] - v = xv[SA[4, 5, 6]] - t = sol.t[i] - - B = Bfunc(x, t) - Bmag = norm(B) - b̂ = B / Bmag - vpar_val = v ⋅ b̂ - vperp_vec = v - vpar_val * b̂ - μ = m * (vperp_vec ⋅ vperp_vec) / (2 * Bmag) - - ε[i] = get_adiabaticity(x, Bfunc, q, m, μ, t) - end - elseif state_len == 4 # GC - q, q2m, μ, Efunc, Bfunc = sol.prob.p - m = q / q2m - - for i in 1:n - x = sol.u[i][SA[1, 2, 3]] - t = sol.t[i] - ε[i] = get_adiabaticity(x, Bfunc, q, m, μ, t) - end - else - error("Unsupported solution state dimension: $state_len") - end - - return ε -end - -""" - get_adiabaticity(sol::AbstractODESolution, t) - -Calculate the adiabaticity parameter `ϵ` from a solution `sol` at time `t`. -`t` must be a scalar. -""" -function get_adiabaticity(sol::AbstractODESolution, t) - state_len = length(sol.u[1]) +function get_adiabaticity(xv, p::Tuple, t) + state_len = length(xv) if state_len == 6 # FO or Hybrid - q2m, m, Efunc, Bfunc, _ = sol.prob.p + q2m, m, Efunc, Bfunc, _ = p q = q2m * m - xv = sol(t) x = xv[SA[1, 2, 3]] v = xv[SA[4, 5, 6]] @@ -724,15 +671,41 @@ function get_adiabaticity(sol::AbstractODESolution, t) μ = m * (vperp_vec ⋅ vperp_vec) / (2 * Bmag) return get_adiabaticity(x, Bfunc, q, m, μ, t) - elseif state_len == 4 # GC - q, q2m, μ, Efunc, Bfunc = sol.prob.p + q, q2m, μ, Efunc, Bfunc = p m = q / q2m - xv = sol(t) x = xv[SA[1, 2, 3]] return get_adiabaticity(x, Bfunc, q, m, μ, t) else error("Unsupported solution state dimension: $state_len") end end + +""" + get_adiabaticity(sol::AbstractODESolution) + +Calculate the adiabaticity parameter `ϵ` from a solution `sol` at all time steps. +The parameters `q`, `m`, `Bfunc` are extracted from `sol.prob.p`. +""" +function get_adiabaticity(sol::AbstractODESolution) + n = length(sol.t) + ε = Vector{Float64}(undef, n) + p = sol.prob.p + + for i in 1:n + ε[i] = get_adiabaticity(sol.u[i], p, sol.t[i]) + end + + return ε +end + +""" + get_adiabaticity(sol::AbstractODESolution, t) + +Calculate the adiabaticity parameter `ϵ` from a solution `sol` at time `t`. +`t` must be a scalar. +""" +function get_adiabaticity(sol::AbstractODESolution, t) + return get_adiabaticity(sol(t), sol.prob.p, t) +end