From 1618924695fa9e3f255fd31c9bb3564839e1ad33 Mon Sep 17 00:00:00 2001 From: Hongyang Zhou Date: Tue, 10 Feb 2026 15:45:58 -0500 Subject: [PATCH 1/4] refactor the hybrid solver and add a new demo --- docs/examples/features/demo_hybrid.jl | 286 ++++++++++++++++++++++++++ docs/make.jl | 1 + src/boris.jl | 2 +- src/hybrid.jl | 61 +++--- src/prepare.jl | 2 +- test/test_hybrid.jl | 127 +++++++++--- 6 files changed, 422 insertions(+), 57 deletions(-) create mode 100644 docs/examples/features/demo_hybrid.jl diff --git a/docs/examples/features/demo_hybrid.jl b/docs/examples/features/demo_hybrid.jl new file mode 100644 index 000000000..28cde3493 --- /dev/null +++ b/docs/examples/features/demo_hybrid.jl @@ -0,0 +1,286 @@ +# # Hybrid GC-FO Solver +# +# This example demonstrates the adaptive hybrid solver that dynamically switches +# between full orbit (FO) and guiding center (GC) tracing based on the local +# adiabaticity parameter `ε = ρ_L / R_c`. +# +# We use a magnetic bottle field where the adiabaticity varies spatially: +# near the midplane `ε` is large (non-adiabatic → FO), and near the mirror +# points `ε` is small (adiabatic → GC). + +# import DisplayAs #hide +using TestParticle, OrdinaryDiffEq, StaticArrays +import TestParticle as TP +using LinearAlgebra, Random +using CairoMakie +CairoMakie.activate!(type = "png") #hide + +# ## Magnetic Bottle Configuration +# +# The magnetic field satisfies `∇·B = 0`: +# ```math +# B_z = B_0 (1 + α z^2), \quad +# B_x = -B_0 α x z, \quad +# B_y = -B_0 α y z +# ``` +# A larger `α` produces stronger curvature at the midplane, which +# raises the adiabaticity parameter `ε` there. + +B0 = 1.0e-4 # [T] +α = 1.0e-2 # [m⁻²] + +function bottle_B(x, t) + Bz = B0 * (1 + α * x[3]^2) + Bx = -B0 * α * x[1] * x[3] + By = -B0 * α * x[2] * x[3] + return SA[Bx, By, Bz] +end + +B_field = TP.Field(bottle_B) +E_field = TP.Field((x, t) -> SA[0.0, 0.0, 0.0]) + +m = TP.mᵢ +q = TP.qᵢ +q2m = q / m + +# ## Step 1: Full Orbit Reference Trace +# +# First we trace the proton using the standard ODE solver to confirm +# it is trapped and bounces inside the magnetic bottle. + +x0 = SA[0.0, 0.0, 0.0] +v_perp = 5.0e4 # [m/s] +v_par = 1.0e5 # [m/s] +v0 = SA[v_perp, 0.0, v_par] +u0 = vcat(x0, v0) + +Ω = abs(q2m) * B0 +T_gyro = 2π / Ω + +## Trace long enough to see several bounces +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) + +## Verify trapping: z should oscillate, not diverge +z_fo = [u[3] for u in sol_fo.u] +@assert maximum(abs, z_fo) < 1.0e4 "Particle escaped the bottle!" + +# ## Step 2: Guiding Center Reference Trace +# +# For comparison, we also trace with the guiding center equations. + +bottle_B_static(x) = bottle_B(x, 0.0) +bottle_E_static(x) = SA[0.0, 0.0, 0.0] + +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()) + +# ## Step 3: Hybrid Solver +# +# The adaptive hybrid solver switches between FO and GC based on the +# adiabaticity parameter and a threshold. + +## The classical adiabaticity criterion: ε = ρ_L / R_c < 0.1 +threshold = 0.1 + +p = (q2m, m, E_field, B_field, TP.ZeroField()) + +alg = AdaptiveHybrid(; + threshold, + dtmax = T_gyro, + dtmin = 1.0e-4 * T_gyro, + maxiters = 500_000, +) + +Random.seed!(1234) +## Set verbose = true to see the dynamic switching +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) + +## Common colorbar range across all solvers +ε_clamp_lo = 1.0e-3 +clims = extrema(log10.(clamp.(vcat(ε_fo, ε_gc, ε_hybrid), ε_clamp_lo, Inf))) + +# ## Visualization +# +# We define a helper to plot a 3D trajectory colored by ε, with +# start/end markers and a colorbar. + +## Plot a 3D trajectory colored by ε onto an existing Axis3 +## When `npts` is set, the solution is interpolated for a smoother curve. +function plot_trajectory!(ax, sol, ε_vals; npts = nothing) + if isnothing(npts) + xs = [u[1] for u in sol.u] + ys = [u[2] for u in sol.u] + zs = [u[3] for u in sol.u] + ε_plot = ε_vals + else + t_new = range(sol.t[1], sol.t[end]; length = npts) + xs = Vector{Float64}(undef, npts) + ys = Vector{Float64}(undef, npts) + zs = Vector{Float64}(undef, npts) + ε_plot = Vector{Float64}(undef, npts) + for (i, t) in enumerate(t_new) + u = sol(t) + xs[i], ys[i], zs[i] = u[1], u[2], u[3] + ## Linearly interpolate ε to the new time grid + idx = clamp(searchsortedlast(sol.t, t), 1, length(sol.t) - 1) + frac = (t - sol.t[idx]) / (sol.t[idx + 1] - sol.t[idx]) + ε_plot[i] = ε_vals[idx] + frac * (ε_vals[idx + 1] - ε_vals[idx]) + end + end + ε_log = log10.(clamp.(ε_plot, ε_clamp_lo, Inf)) + lines!( + ax, xs, ys, zs; + color = ε_log, colormap = :turbo, colorrange = clims, linewidth = 1.0 + ) + scatter!(ax, [Point3f(xs[1], ys[1], zs[1])]; color = :purple, markersize = 12) + return scatter!( + ax, [Point3f(xs[end], ys[end], zs[end])]; + color = :red, marker = :rect, markersize = 12 + ) +end + +## Create a standalone figure with a 3D trajectory colored by ε +function plot_trajectory(sol, ε_vals, title_str; figsize = (700, 500), npts = nothing) + f = Figure(; size = figsize, fontsize = 18) + ax = Axis3( + f[1, 1], + xlabel = "x [m]", ylabel = "y [m]", zlabel = "z [m]", + title = title_str, aspect = :data, + ) + plot_trajectory!(ax, sol, ε_vals; npts) + Colorbar(f[1, 2]; colormap = :turbo, limits = clims, label = L"\log_{10}(\epsilon)") + return f +end + +# ## Visualization + +f = Figure(; size = (1400, 900), fontsize = 18) + +## Top row: three 3D trajectories + shared colorbar +ax_fo = Axis3( + f[1, 1], + xlabel = "x [m]", ylabel = "y [m]", zlabel = "z [m]", + title = "Full Orbit", aspect = :data, +) +plot_trajectory!(ax_fo, sol_fo, ε_fo; npts = 5000) + +ax_gc = Axis3( + f[1, 2], + xlabel = "x [m]", ylabel = "y [m]", zlabel = "z [m]", + title = "Guiding Center", aspect = :data, +) +plot_trajectory!(ax_gc, sol_gc, ε_gc; npts = 5000) + +ax_hyb = Axis3( + f[1, 3], + xlabel = "x [m]", ylabel = "y [m]", zlabel = "z [m]", + title = "Hybrid Mode", aspect = :data, +) +plot_trajectory!(ax_hyb, sol, ε_hybrid) + +Colorbar(f[1, 4]; colormap = :turbo, limits = clims, label = L"\log_{10}(\epsilon)") + +## Bottom row: time series of adiabaticity with solver mode +ax_ts = Axis( + f[2, 1:4], + xlabel = L"t / T_\text{gyro}", + ylabel = L"\epsilon = \rho_L / R_c", + yscale = log10, + title = "Adiabaticity & Solver Mode", + limits = (nothing, (ε_clamp_lo, 10^ceil(log10(maximum(ε_hybrid))))), +) + +t_norm = sol.t ./ T_gyro + +## Shade FO and GC regions +is_fo = ε_hybrid .>= threshold +i_region = 1 +while i_region <= length(t_norm) + mode_fo = is_fo[i_region] + j_region = i_region + while j_region < length(t_norm) && + is_fo[j_region + 1] == mode_fo + j_region += 1 + end + t_lo = t_norm[i_region] + t_hi = t_norm[j_region] + if mode_fo + vspan!(ax_ts, t_lo, t_hi; color = (:red, 0.2)) + else + vspan!(ax_ts, t_lo, t_hi; color = (:blue, 0.2)) + end + global i_region = j_region + 1 +end + +lines!(ax_ts, t_norm, ε_hybrid; color = :black, linewidth = 1.0) +hlines!( + ax_ts, [threshold]; + color = :gray50, linestyle = :dash, linewidth = 1.5, + label = "threshold = $(round(threshold; sigdigits = 2))" +) + +## Legend entries for solver modes +poly!( + ax_ts, Point2f[(NaN, NaN)]; + color = (:red, 0.4), strokewidth = 0, label = "Full Orbit" +) +poly!( + ax_ts, Point2f[(NaN, NaN)]; + color = (:blue, 0.4), strokewidth = 0, + label = "Guiding Center" +) +Legend(f[3, 1:4], ax_ts; orientation = :horizontal) + +rowsize!(f.layout, 1, Relative(0.55)) + +f +# f = DisplayAs.PNG(f) #hide diff --git a/docs/make.jl b/docs/make.jl index 669ac9db2..17d74dee5 100644 --- a/docs/make.jl +++ b/docs/make.jl @@ -70,6 +70,7 @@ features_order = [ "demo_extrasaving.jl", "demo_flux.jl", "demo_gc.jl", + "demo_hybrid.jl", "demo_gpu.md", "demo_array.md", "demo_diskarray.jl", diff --git a/src/boris.jl b/src/boris.jl index 95594a7b9..a0548add8 100644 --- a/src/boris.jl +++ b/src/boris.jl @@ -194,7 +194,7 @@ Trace particles using the Boris method with specified `prob`. - `save_fields::Bool`: save the electric and magnetic fields. Default is `false`. - `save_work::Bool`: save the work done by the electric field. Default is `false`. """ -@inline function solve( +function solve( prob::TraceProblem, ensemblealg::EA = EnsembleSerial(); trajectories::Int = 1, savestepinterval::Int = 1, dt::AbstractFloat, isoutofdomain::F = ODE_DEFAULT_ISOUTOFDOMAIN, n::Int = 1, diff --git a/src/hybrid.jl b/src/hybrid.jl index 962e43a65..589d09222 100644 --- a/src/hybrid.jl +++ b/src/hybrid.jl @@ -60,45 +60,52 @@ end function solve( prob::TraceHybridProblem, alg::AdaptiveHybrid, - ensemblealg::BasicEnsembleAlgorithm = EnsembleSerial(); + ensemblealg::EA = EnsembleSerial(); trajectories::Int = 1, savestepinterval::Int = 1, - isoutofdomain::Function = ODE_DEFAULT_ISOUTOFDOMAIN, - save_start::Bool = true, save_end::Bool = true, save_everystep::Bool = true - ) + isoutofdomain::F = ODE_DEFAULT_ISOUTOFDOMAIN, + save_start::Bool = true, save_end::Bool = true, + save_everystep::Bool = true, verbose::Bool = false + ) where {EA <: BasicEnsembleAlgorithm, F} return _solve( - ensemblealg, prob, trajectories, alg, savestepinterval, isoutofdomain, - save_start, save_end, save_everystep + ensemblealg, prob, trajectories, alg, + savestepinterval, isoutofdomain, + save_start, save_end, save_everystep, verbose ) end -function _solve( - ::EnsembleSerial, prob, trajectories, alg::AdaptiveHybrid, savestepinterval, - isoutofdomain, save_start, save_end, save_everystep - ) +@inline function _solve( + ::EnsembleSerial, prob::TraceHybridProblem, + trajectories, alg::AdaptiveHybrid, + savestepinterval, isoutofdomain::F, + save_start, save_end, save_everystep, verbose + ) where {F} sol_type = _get_sol_type(prob, zero(eltype(prob.tspan))) sols = Vector{sol_type}(undef, trajectories) irange = 1:trajectories _hybrid_adaptive!( - sols, prob, irange, alg, savestepinterval, isoutofdomain, - save_start, save_end, save_everystep + sols, prob, irange, alg, savestepinterval, + isoutofdomain, save_start, save_end, save_everystep, verbose ) return sols end -function _solve( - ::EnsembleThreads, prob, trajectories, alg::AdaptiveHybrid, savestepinterval, - isoutofdomain, save_start, save_end, save_everystep - ) +@inline function _solve( + ::EnsembleThreads, prob::TraceHybridProblem, + trajectories, alg::AdaptiveHybrid, + savestepinterval, isoutofdomain::F, + save_start, save_end, save_everystep, verbose + ) where {F} sol_type = _get_sol_type(prob, zero(eltype(prob.tspan))) sols = Vector{sol_type}(undef, trajectories) nchunks = Threads.nthreads() Threads.@threads for irange in index_chunks(1:trajectories; n = nchunks) _hybrid_adaptive!( - sols, prob, irange, alg, savestepinterval, isoutofdomain, - save_start, save_end, save_everystep + sols, prob, irange, alg, savestepinterval, + isoutofdomain, save_start, save_end, + save_everystep, verbose ) end @@ -175,17 +182,16 @@ function _gc_to_full_at_t(state_gc, E_field, B_field, q, m, μ, t, phase = 0) return SVector{6}(x[1], x[2], x[3], v[1], v[2], v[3]) end -function _hybrid_adaptive!( - sols, prob, irange, alg, savestepinterval, isoutofdomain, - save_start, save_end, save_everystep - ) +@inline function _hybrid_adaptive!( + sols, prob::TraceHybridProblem, irange, alg, + savestepinterval, isoutofdomain::F, + save_start, save_end, save_everystep, verbose + ) where {F} (; tspan, p, u0) = prob - # p = (q2m, m, Efunc, Bfunc, Ffunc) q2m, m, Efunc, Bfunc, _ = p q = q2m * m T = eltype(u0) - # Determine if fields are time dependent is_td = is_time_dependent(Efunc) || is_time_dependent(Bfunc) # Safety parameters for RK45 (from gc_solver.jl) @@ -193,7 +199,7 @@ function _hybrid_adaptive!( max_growth = 5.0 min_growth = 0.2 - @fastmath @inbounds for i in irange + @inbounds for i in irange # Initialize solution containers with sizehint! # Estimate initial capacity based on timespan and maximum possible dt estimated_steps = ceil(Int, (tspan[2] - tspan[1]) / alg.dtmax) @@ -225,6 +231,7 @@ function _hybrid_adaptive!( Bmag = norm(B_vec) omega = abs(q2m * Bmag) dt = 0.5 * 2π / omega + verbose && @info "Initial mode: GC" ϵ t else mode = :FO # Initial dt for FO @@ -234,6 +241,7 @@ function _hybrid_adaptive!( dt = alg.safety_fo / omega dt = clamp(dt, alg.dtmin, alg.dtmax) v = update_velocity(v, r, p, -0.5 * dt, t) + verbose && @info "Initial mode: FO" ϵ t end if save_start @@ -250,6 +258,7 @@ function _hybrid_adaptive!( if ϵ >= alg.threshold # Switch to FO (GC -> FO) mode = :FO + verbose && @info "Switch GC → FO" ϵ t r = xv_gc[SVector(1, 2, 3)] xv_fo_vec = _gc_to_full_at_t( xv_gc, Efunc, Bfunc, q, m, μ, t, 2π * rand() ) @@ -313,6 +322,7 @@ function _hybrid_adaptive!( if ϵ < alg.threshold # Switch to GC mode = :GC + verbose && @info "Switch FO → GC" ϵ t r = X_gc xv_gc = SVector{4, T}(X_gc[1], X_gc[2], X_gc[3], vpar) p_gc = (q, q2m, μ, Efunc, Bfunc) @@ -371,7 +381,6 @@ function _hybrid_adaptive!( end # Final Save - if save_end && (isempty(tsave) || tsave[end] != t) if mode == :GC push!(traj, _gc_to_full_at_t(xv_gc, Efunc, Bfunc, q, m, μ, t)) diff --git a/src/prepare.jl b/src/prepare.jl index 302736389..00684ed01 100644 --- a/src/prepare.jl +++ b/src/prepare.jl @@ -6,7 +6,7 @@ Judge whether the field function is time dependent. is_time_dependent(f::Function) = applicable(f, zeros(6), 0.0) ? true : false """ - Field{itd, F} <: AbstractField{itd} + Field{itd, F} <: AbstractField{itd} A representation of a field function `f`, defined by: diff --git a/test/test_hybrid.jl b/test/test_hybrid.jl index 0697823cc..4d6b855d6 100644 --- a/test/test_hybrid.jl +++ b/test/test_hybrid.jl @@ -4,44 +4,113 @@ using LinearAlgebra using Test @testset "Adaptive Hybrid Solver" begin - # 1. Pure Adiabatic Case (Uniform Field) - B_func(x, t) = SA[0.0, 0.0, 0.01] - B_field = TestParticle.Field(B_func) - E_field = TestParticle.Field((x, t) -> SA[0.0, 0.0, 0.0]) - m = TestParticle.mᵢ q = TestParticle.qᵢ q2m = q / m - x0 = SA[0.0, 0.0, 0.0] - v0 = SA[1.0e4, 0.0, 1.0e3] - u0 = vcat(x0, v0) - tspan = (0.0, 1.0e-4) + E_field = TestParticle.Field((x, t) -> SA[0.0, 0.0, 0.0]) - p = (q2m, m, E_field, B_field, TestParticle.ZeroField()) - alg = AdaptiveHybrid(threshold = 0.1, dtmax = 1.0e-6) + # 1. Pure Adiabatic Case (Uniform Field) + let + B_func(x, t) = SA[0.0, 0.0, 0.01] + B_field = TestParticle.Field(B_func) - sols = TestParticle.solve(TraceHybridProblem(u0, tspan, p), alg) - @test sols[1].retcode == TestParticle.ReturnCode.Success + x0 = SA[0.0, 0.0, 0.0] + v0 = SA[1.0e4, 0.0, 1.0e3] + u0 = vcat(x0, v0) + tspan = (0.0, 1.0e-4) - # 2. Switching Case (Curved/Sheared Field) - # B = [B0*cos(kx), B0*sin(kx), 0] - # This field line is curved in the xy plane. - function sheared_B_func(x, t) - B0 = 0.01 - k = 100.0 # High curvature - return SA[B0 * cos(k * x[1]), B0 * sin(k * x[1]), 0.0] + p = (q2m, m, E_field, B_field, TestParticle.ZeroField()) + alg = AdaptiveHybrid(; threshold = 0.1, dtmax = 1.0e-6) + + sols = TestParticle.solve(TraceHybridProblem(u0, tspan, p), alg) + @test sols[1].retcode == TestParticle.ReturnCode.Success end - sheared_B = TestParticle.Field(sheared_B_func) - p2 = (q2m, m, E_field, sheared_B, TestParticle.ZeroField()) - # Threshold is 0.1. ρ ≈ 0.01. Rc ≈ 1/k = 0.01. - # ϵ = ρ / Rc ≈ 1.0 > 0.1 => Should start in FO or switch to FO. + # 2. Non-adiabatic Case (Highly Curved Field) + let + function sheared_B_func(x, t) + B0 = 0.01 + k = 100.0 + return SA[ + B0 * cos(k * x[1]), + B0 * sin(k * x[1]), + 0.0, + ] + end + sheared_B = TestParticle.Field(sheared_B_func) + + x0 = SA[0.0, 0.0, 0.0] + v0 = SA[1.0e4, 0.0, 1.0e3] + u0 = vcat(x0, v0) + tspan = (0.0, 1.0e-5) - tspan2 = (0.0, 1.0e-5) - sols2 = TestParticle.solve(TraceHybridProblem(u0, tspan2, p2), alg) - @test sols2[1].retcode == TestParticle.ReturnCode.Success + p = (q2m, m, E_field, sheared_B, TestParticle.ZeroField()) + alg = AdaptiveHybrid(; threshold = 0.1, dtmax = 1.0e-6) + + sols = TestParticle.solve(TraceHybridProblem(u0, tspan, p), alg) + @test sols[1].retcode == TestParticle.ReturnCode.Success + end - # Let's use a lower k for adiabatic start and higher k for transition if possible. - # Or just verify with logs/print. + # 3. Dynamic GC ↔ FO Switching (Magnetic Bottle) + # + # B_z(z) = B0*(1 + α*z²), with B_r from ∇·B = 0: + # B_x ≈ -B0*α*x*z, B_y ≈ -B0*α*y*z + # + # Near the midplane (z ≈ 0) the field is weak and + # curvature is high → FO. Away from midplane the + # field strengthens and becomes more uniform → GC. + # A bouncing particle triggers repeated GC → FO → GC + # transitions. + let + B0 = 1.0e-4 # [T] background field + α = 1.0e-4 # [m⁻²] mirror ratio parameter + + function bottle_B(x, t) + Bz = B0 * (1 + α * x[3]^2) + Bx = -B0 * α * x[1] * x[3] + By = -B0 * α * x[2] * x[3] + return SA[Bx, By, Bz] + end + B_bottle = TestParticle.Field(bottle_B) + + # Proton with moderate perpendicular velocity + # and parallel velocity to bounce in the bottle + x0 = SA[0.0, 0.0, 0.0] + v_perp = 5.0e4 # [m/s] + v_par = 2.0e5 # [m/s] + v0 = SA[v_perp, 0.0, v_par] + u0 = vcat(x0, v0) + + Ω = abs(q2m) * B0 + T_gyro = 2π / Ω + tspan = (0.0, 50 * T_gyro) + + p = ( + q2m, m, E_field, B_bottle, + TestParticle.ZeroField(), + ) + alg = AdaptiveHybrid(; + threshold = 0.05, + dtmax = T_gyro, + dtmin = 1.0e-4 * T_gyro, + ) + + sols = TestParticle.solve(TraceHybridProblem(u0, tspan, p), alg) + sol = sols[1] + @test sol.retcode == TestParticle.ReturnCode.Success + @test length(sol.t) > 10 + + # Energy conservation: kinetic energy should be + # approximately conserved (no E field) + KE_init = 0.5 * m * sum(abs2, v0) + v_end = sol.u[end][SA[4, 5, 6]] + KE_end = 0.5 * m * sum(abs2, v_end) + @test KE_end ≈ KE_init rtol = 0.1 + + # The particle should remain confined by the + # mirror: z should not diverge + z_vals = [u[3] for u in sol.u] + @test maximum(abs, z_vals) < 1.0e6 + end end From b91f85198125ac505183c2d450cac4ea99b7f464 Mon Sep 17 00:00:00 2001 From: Hongyang Zhou Date: Tue, 10 Feb 2026 18:50:11 -0500 Subject: [PATCH 2/4] Add back @inline --- src/boris.jl | 2 +- src/hybrid.jl | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/src/boris.jl b/src/boris.jl index a0548add8..95594a7b9 100644 --- a/src/boris.jl +++ b/src/boris.jl @@ -194,7 +194,7 @@ Trace particles using the Boris method with specified `prob`. - `save_fields::Bool`: save the electric and magnetic fields. Default is `false`. - `save_work::Bool`: save the work done by the electric field. Default is `false`. """ -function solve( +@inline function solve( prob::TraceProblem, ensemblealg::EA = EnsembleSerial(); trajectories::Int = 1, savestepinterval::Int = 1, dt::AbstractFloat, isoutofdomain::F = ODE_DEFAULT_ISOUTOFDOMAIN, n::Int = 1, diff --git a/src/hybrid.jl b/src/hybrid.jl index 589d09222..90a497dc4 100644 --- a/src/hybrid.jl +++ b/src/hybrid.jl @@ -58,7 +58,7 @@ function TraceHybridProblem(u0, tspan, p; prob_func = DEFAULT_PROB_FUNC) ) end -function solve( +@inline function solve( prob::TraceHybridProblem, alg::AdaptiveHybrid, ensemblealg::EA = EnsembleSerial(); trajectories::Int = 1, savestepinterval::Int = 1, From 48e059371c174619adb3a8a63acd8fd08559843d Mon Sep 17 00:00:00 2001 From: Hongyang Zhou Date: Tue, 10 Feb 2026 18:50:35 -0500 Subject: [PATCH 3/4] doc: improve the hybrid demo --- docs/examples/features/demo_hybrid.jl | 35 +++++++++++++-------------- 1 file changed, 17 insertions(+), 18 deletions(-) diff --git a/docs/examples/features/demo_hybrid.jl b/docs/examples/features/demo_hybrid.jl index 28cde3493..8a545b9f1 100644 --- a/docs/examples/features/demo_hybrid.jl +++ b/docs/examples/features/demo_hybrid.jl @@ -8,7 +8,7 @@ # near the midplane `ε` is large (non-adiabatic → FO), and near the mirror # points `ε` is small (adiabatic → GC). -# import DisplayAs #hide +import DisplayAs #hide using TestParticle, OrdinaryDiffEq, StaticArrays import TestParticle as TP using LinearAlgebra, Random @@ -243,22 +243,22 @@ t_norm = sol.t ./ T_gyro ## Shade FO and GC regions is_fo = ε_hybrid .>= threshold -i_region = 1 -while i_region <= length(t_norm) - mode_fo = is_fo[i_region] - j_region = i_region - while j_region < length(t_norm) && - is_fo[j_region + 1] == mode_fo - j_region += 1 - end - t_lo = t_norm[i_region] - t_hi = t_norm[j_region] - if mode_fo - vspan!(ax_ts, t_lo, t_hi; color = (:red, 0.2)) - else - vspan!(ax_ts, t_lo, t_hi; color = (:blue, 0.2)) +let i_region = 1 + while i_region <= length(t_norm) + mode_fo = is_fo[i_region] + j_region = i_region + while j_region < length(t_norm) && is_fo[j_region + 1] == mode_fo + j_region += 1 + end + t_lo = t_norm[i_region] + t_hi = t_norm[j_region] + if mode_fo + vspan!(ax_ts, t_lo, t_hi; color = (:red, 0.2)) + else + vspan!(ax_ts, t_lo, t_hi; color = (:blue, 0.2)) + end + i_region = j_region + 1 end - global i_region = j_region + 1 end lines!(ax_ts, t_norm, ε_hybrid; color = :black, linewidth = 1.0) @@ -282,5 +282,4 @@ Legend(f[3, 1:4], ax_ts; orientation = :horizontal) rowsize!(f.layout, 1, Relative(0.55)) -f -# f = DisplayAs.PNG(f) #hide +f = DisplayAs.PNG(f) #hide From 0959d055131c58035bb0579bc92ab39bc026cf35 Mon Sep 17 00:00:00 2001 From: Hongyang Zhou Date: Tue, 10 Feb 2026 20:52:51 -0500 Subject: [PATCH 4/4] Add hybrid solver threading test --- test/test_hybrid.jl | 51 ++++++++++++++++++++++++++------------------- 1 file changed, 30 insertions(+), 21 deletions(-) diff --git a/test/test_hybrid.jl b/test/test_hybrid.jl index 4d6b855d6..4dbbdf513 100644 --- a/test/test_hybrid.jl +++ b/test/test_hybrid.jl @@ -62,6 +62,8 @@ using Test # field strengthens and becomes more uniform → GC. # A bouncing particle triggers repeated GC → FO → GC # transitions. + # + # Also tests EnsembleThreads against EnsembleSerial. let B0 = 1.0e-4 # [T] background field α = 1.0e-4 # [m⁻²] mirror ratio parameter @@ -74,8 +76,6 @@ using Test end B_bottle = TestParticle.Field(bottle_B) - # Proton with moderate perpendicular velocity - # and parallel velocity to bounce in the bottle x0 = SA[0.0, 0.0, 0.0] v_perp = 5.0e4 # [m/s] v_par = 2.0e5 # [m/s] @@ -86,31 +86,40 @@ using Test T_gyro = 2π / Ω tspan = (0.0, 50 * T_gyro) - p = ( - q2m, m, E_field, B_bottle, - TestParticle.ZeroField(), - ) + p = (q2m, m, E_field, B_bottle, TestParticle.ZeroField()) alg = AdaptiveHybrid(; threshold = 0.05, dtmax = T_gyro, dtmin = 1.0e-4 * T_gyro, ) - sols = TestParticle.solve(TraceHybridProblem(u0, tspan, p), alg) - sol = sols[1] - @test sol.retcode == TestParticle.ReturnCode.Success - @test length(sol.t) > 10 - - # Energy conservation: kinetic energy should be - # approximately conserved (no E field) + ntraj = 4 KE_init = 0.5 * m * sum(abs2, v0) - v_end = sol.u[end][SA[4, 5, 6]] - KE_end = 0.5 * m * sum(abs2, v_end) - @test KE_end ≈ KE_init rtol = 0.1 - - # The particle should remain confined by the - # mirror: z should not diverge - z_vals = [u[3] for u in sol.u] - @test maximum(abs, z_vals) < 1.0e6 + + # Serial + sols_serial = TestParticle.solve( + TraceHybridProblem(u0, tspan, p), alg, EnsembleSerial(); + trajectories = ntraj, + ) + + @test all(s -> s.retcode == TestParticle.ReturnCode.Success, sols_serial) + @test all(s -> length(s.t) > 10, sols_serial) + @test all(sols_serial) do s + isapprox(0.5 * m * sum(abs2, s.u[end][SA[4, 5, 6]]), KE_init; rtol = 0.1) + end + @test all(s -> maximum(u -> abs(u[3]), s.u) < 1.0e6, sols_serial) + + # Threaded — should match serial + sols_threads = TestParticle.solve( + TraceHybridProblem(u0, tspan, p), alg, EnsembleThreads(); + trajectories = ntraj, + ) + + @test all(s -> s.retcode == TestParticle.ReturnCode.Success, sols_threads) + @test all(i -> length(sols_threads[i].t) == length(sols_serial[i].t), 1:ntraj) + @test all(sols_threads) do s + isapprox(0.5 * m * sum(abs2, s.u[end][SA[4, 5, 6]]), KE_init; rtol = 0.1) + end + @test all(s -> maximum(u -> abs(u[3]), s.u) < 1.0e6, sols_threads) end end