doc: improve hybrid demo

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
 

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,49 +100,17 @@ 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
 #
 # The adiabaticity parameter `ε = ρ_L / R_c` measures how well the
 # 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

src/utility/utility.jl

@@ -652,3 +652,60 @@ end
 function get_adiabaticity(r, Bfunc, μ, t::Number = 0.0; species = Proton)
     return get_adiabaticity(r, Bfunc, species.q, species.m, μ, t)
 end
+
+function get_adiabaticity(xv, p::Tuple, t)
+    state_len = length(xv)
+
+    if state_len == 6 # FO or Hybrid
+        q2m, m, Efunc, Bfunc, _ = p
+        q = q2m * m
+
+        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 = p
+        m = q / q2m
+
+        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

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()