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/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..761acea64 100644 --- a/src/utility/utility.jl +++ b/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 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()