diff --git a/README.md b/README.md index 36434fb25..274dada8d 100644 --- a/README.md +++ b/README.md @@ -56,7 +56,7 @@ sol = TestParticle.solve(prob; dt, savestepinterval)[1] Besides the standard Boris method, we also support various versions of Boris solvers including: - **Multistep Boris**: fixed time step with `n > 1` substeps. -- **Adaptive Boris**: uses `AdaptiveBoris()` for automatic time step selection based on local gyroperiod. +- **Adaptive Boris**: uses `Boris(safety = 0.1)` to automatically adjust time step based on local gyroperiod. For plotting with Makie, diff --git a/benchmark/benchmarks.jl b/benchmark/benchmarks.jl index 106519682..1eceae722 100644 --- a/benchmark/benchmarks.jl +++ b/benchmark/benchmarks.jl @@ -168,7 +168,7 @@ SUITE["trace"]["numerical field"]["Hyper Boris (n=2, N=4)"] = @benchmarkable TP. SUITE["trace"]["numerical field"]["Hyper Boris (n=2, N=6)"] = @benchmarkable TP.solve( $prob_boris; dt = 1 / 7, savestepinterval = 10, n = 2, N = 6 ) -alg_adaptive = AdaptiveBoris(safety = 0.1) +alg_adaptive = Boris(safety = 0.1) SUITE["trace"]["numerical field"]["Adaptive Boris"] = @benchmarkable TP.solve( $prob_boris, $alg_adaptive ) diff --git a/docs/examples/features/demo_adaptive_accuracy.jl b/docs/examples/features/demo_adaptive_accuracy.jl index cea0b4418..7b21f5523 100644 --- a/docs/examples/features/demo_adaptive_accuracy.jl +++ b/docs/examples/features/demo_adaptive_accuracy.jl @@ -1,7 +1,7 @@ # # Adaptive Solver Accuracy # # This example compares the accuracy and energy conservation of solvers that -# use **changing time steps**: the native `AdaptiveBoris` solver and the +# use **changing time steps**: the native `Boris(safety = η)` solver and the # adaptive ODE solvers `Tsit5` and `Vern6`. # # Unlike the fixed-step tests in [Solver Accuracy Analysis](@ref) and @@ -91,16 +91,16 @@ prob_ode1 = ODEProblem(trace_normalized!, u0_1, tspan1, param1) adaptive_solvers_1 = [ ( - "AdaptiveBoris (η=1/16)", - AdaptiveBoris(; safety = 1 / 16), + "Boris (η=1/16)", + Boris(; safety = 1 / 16), ), ( - "AdaptiveBoris (η=1/8)", - AdaptiveBoris(; safety = 1 / 8), + "Boris (η=1/8)", + Boris(; safety = 1 / 8), ), ( - "AdaptiveBoris (η=1/4)", - AdaptiveBoris(; safety = 1 / 4), + "Boris (η=1/4)", + Boris(; safety = 1 / 4), ), ] @@ -134,7 +134,7 @@ ax1b = Axis( f1[1, 2]; xlabel = "Time [Gyroperiod]", ylabel = L"\Delta t / T_\mathrm{gyro}", - title = "AdaptiveBoris: Time Step History", + title = "Adaptive Boris: Time Step History", ) for (i, (name, alg)) in enumerate(adaptive_solvers_1) @@ -201,7 +201,7 @@ plot_table(results1) #hide # # With E = [0, 0.5, 0.1] and B = [0, 0, 1], an exact velocity # solution exists (Section 6, Zenitani & Kato 2025). We sweep the -# `safety` parameter η (representing dt / T_gyro) for `AdaptiveBoris` +# `safety` parameter η (representing dt / T_gyro) for `Boris(safety = η)` # and the tolerance for the ODE solvers, then plot maximum velocity # error vs. number of time steps (computational cost). @@ -245,7 +245,7 @@ boris_steps = Int[] boris_errors = Float64[] for s in safety_values - alg = AdaptiveBoris(; safety = s) + alg = Boris(; safety = s) sol = TP.solve(prob_tp2, alg)[1] push!(boris_steps, length(sol.t)) push!(boris_errors, max_velocity_error(sol)) @@ -288,7 +288,7 @@ ax2 = Axis( scatterlines!( ax2, boris_steps, boris_errors; - label = "AdaptiveBoris", + label = "Adaptive Boris", marker = :circle, linewidth = 2, ) @@ -335,12 +335,12 @@ prob_ode3 = ODEProblem(trace_normalized!, u0_2, tspan3, param3) adaptive_solvers_3 = [ ( - "AdaptiveBoris (η=0.1)", - AdaptiveBoris(; safety = 0.1), + "Boris (η=0.1)", + Boris(; safety = 0.1), ), ( - "AdaptiveBoris (η=0.2)", - AdaptiveBoris(; safety = 0.2), + "Boris (η=0.2)", + Boris(; safety = 0.2), ), ] diff --git a/docs/examples/features/demo_boris.jl b/docs/examples/features/demo_boris.jl index 6afc13454..3ff5b806e 100644 --- a/docs/examples/features/demo_boris.jl +++ b/docs/examples/features/demo_boris.jl @@ -99,7 +99,7 @@ sol_boris_2 = TP.solve(prob; dt, n = 2)[1]; sol_boris_4 = TP.solve(prob; dt, n = 4)[1]; sol_boris_hyper = TP.solve(prob; dt, n = 2, N = 4)[1]; -alg_adaptive = AdaptiveBoris(safety = 0.1) +alg_adaptive = Boris(safety = 0.1) sol_boris_adaptive = TP.solve(prob, alg_adaptive)[1]; # Let's compare against the default ODE solver `Tsit5` from DifferentialEquations.jl, in both fixed time step mode and adaptive mode: @@ -148,7 +148,7 @@ sol_boris_4 = TP.solve(prob_boris; dt, savestepinterval = 36, n = 4)[1]; sol_boris_hyper = TP.solve(prob_boris; dt, savestepinterval = 36, n = 2, N = 4)[1]; sol_boris_adaptive = TP.solve( prob_boris, - AdaptiveBoris(safety = 0.1) + Boris(safety = 0.1) )[1] sol1 = solve(prob, Tsit5(); adaptive = false, dt, dense = false, saveat = dt); sol2 = solve(prob, Tsit5()); diff --git a/docs/src/tutorial/advanced_boris.md b/docs/src/tutorial/advanced_boris.md index 7fffc373d..48adf1084 100644 --- a/docs/src/tutorial/advanced_boris.md +++ b/docs/src/tutorial/advanced_boris.md @@ -74,7 +74,7 @@ Combining both $n > 1$ and $N > 2$ ensures ultra-high stability tracking over dr ## 3. Adaptive Boris Method -The `AdaptiveBoris` solver adjusts the time step $\Delta t$ dynamically based on the local cyclotron frequency $\Omega_c = |q B / m|$. This is particularly useful in systems with strong magnetic field gradients, such as magnetic mirrors or planetary magnetospheres, where the required resolution varies significantly along the particle's trajectory. +The `Boris(safety = η)` solver adjusts the time step $\Delta t$ dynamically based on the local cyclotron frequency $\Omega_c = |q B / m|$. This is particularly useful in systems with strong magnetic field gradients, such as magnetic mirrors or planetary magnetospheres, where the required resolution varies significantly along the particle's trajectory. The time step is determined by: ```math @@ -113,10 +113,10 @@ sol = TestParticle.solve(prob; dt, n=2, N=4) ### Adaptive Boris -You can use the adaptive solver by passing an `AdaptiveBoris` object as the second argument to `solve`. +You can use the adaptive solver by passing `Boris(safety = η)` as the second argument to `solve`. ```julia # Adaptive Boris with safety factor 0.05 (20 steps per period) -alg = AdaptiveBoris(safety=0.05) +alg = Boris(safety=0.05) sol = TestParticle.solve(prob, alg)[1] ``` diff --git a/lib/OrdinaryDiffEqBoris/Project.toml b/lib/OrdinaryDiffEqBoris/Project.toml new file mode 100644 index 000000000..486b7d568 --- /dev/null +++ b/lib/OrdinaryDiffEqBoris/Project.toml @@ -0,0 +1,27 @@ +name = "OrdinaryDiffEqBoris" +uuid = "1e8b2ed2-ea1c-4b6e-b3f5-67c29d6bd197" +version = "0.1.0" + +[deps] +LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e" +MuladdMacro = "46d2c3a1-f734-5fdb-9937-b9b9aeba4221" +OrdinaryDiffEqCore = "bbf590c4-e513-4bbe-9b18-05decba2e5d8" +RecursiveArrayTools = "731186ca-8d62-57ce-b412-fbd966d074cd" +Reexport = "189a3867-3050-52da-a836-e630ba90ab69" +SciMLBase = "0bca4576-84f4-4d90-8ffe-ffa030f20462" +StaticArrays = "90137ffa-7385-5640-81b9-e52037218182" + +[compat] +MuladdMacro = "0.2" +OrdinaryDiffEqCore = "3" +RecursiveArrayTools = "3.24, 4" +Reexport = "1.2.2" +SciMLBase = "2, 3" +StaticArrays = "1" +julia = "1.10" + +[extras] +Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40" + +[targets] +test = ["Test"] diff --git a/lib/OrdinaryDiffEqBoris/src/OrdinaryDiffEqBoris.jl b/lib/OrdinaryDiffEqBoris/src/OrdinaryDiffEqBoris.jl new file mode 100644 index 000000000..94bd1e1c8 --- /dev/null +++ b/lib/OrdinaryDiffEqBoris/src/OrdinaryDiffEqBoris.jl @@ -0,0 +1,21 @@ +module OrdinaryDiffEqBoris + +using Reexport +@reexport using SciMLBase +import OrdinaryDiffEqCore: OrdinaryDiffEqAlgorithm, OrdinaryDiffEqAdaptiveAlgorithm, + OrdinaryDiffEqMutableCache, OrdinaryDiffEqConstantCache, + alg_order, alg_cache, isfsal, initialize!, perform_step! +using RecursiveArrayTools +using StaticArrays +using MuladdMacro +using LinearAlgebra + +include("algorithms.jl") +include("alg_utils.jl") +include("boris_caches.jl") +include("boris_perform_step.jl") + +export Boris +export MultistepBoris2, MultistepBoris4, MultistepBoris6 + +end diff --git a/lib/OrdinaryDiffEqBoris/src/alg_utils.jl b/lib/OrdinaryDiffEqBoris/src/alg_utils.jl new file mode 100644 index 000000000..932c26962 --- /dev/null +++ b/lib/OrdinaryDiffEqBoris/src/alg_utils.jl @@ -0,0 +1,5 @@ +alg_order(alg::Boris) = 2 +isfsal(alg::Boris) = false + +alg_order(alg::MultistepBoris{N}) where {N} = 2 +isfsal(alg::MultistepBoris{N}) where {N} = false diff --git a/lib/OrdinaryDiffEqBoris/src/algorithms.jl b/lib/OrdinaryDiffEqBoris/src/algorithms.jl new file mode 100644 index 000000000..c6c4c4138 --- /dev/null +++ b/lib/OrdinaryDiffEqBoris/src/algorithms.jl @@ -0,0 +1,45 @@ +""" + Boris() + +The standard Boris method for particle pushing in electric and magnetic fields. + +This solver expects a problem where `p` is structured as `(q2m, m, E, B, ...)`, which matches the signature used by TestParticle.jl's `TraceProblem`. +""" +struct Boris{T} <: OrdinaryDiffEqAlgorithm + safety::T +end +Boris(; safety = 0.0) = Boris(safety) + +""" + MultistepBoris{N}(; n=1) + +The Multistep/Hyper Boris method of order `N`. +`n` specifies the number of subcycles. +`N` specifies the gyrophase correction order. `N=2` corresponds to the Multicycle solver, while `N=4` or `N=6` are the Hyper Boris solvers. +""" +struct MultistepBoris{N, T} <: OrdinaryDiffEqAlgorithm + n::Int + safety::T +end +MultistepBoris{N}(; n::Int = 1, safety = 0.0) where {N} = MultistepBoris{N, typeof(safety)}(n, safety) + +""" + MultistepBoris2(; n=1) + +The Multicycle Boris method (MultistepBoris with N=2). +""" +const MultistepBoris2 = MultistepBoris{2} + +""" + MultistepBoris4(; n=1) + +The 4th order Hyper Boris method (MultistepBoris with N=4). +""" +const MultistepBoris4 = MultistepBoris{4} + +""" + MultistepBoris6(; n=1) + +The 6th order Hyper Boris method (MultistepBoris with N=6). +""" +const MultistepBoris6 = MultistepBoris{6} diff --git a/lib/OrdinaryDiffEqBoris/src/boris_caches.jl b/lib/OrdinaryDiffEqBoris/src/boris_caches.jl new file mode 100644 index 000000000..a28d40aee --- /dev/null +++ b/lib/OrdinaryDiffEqBoris/src/boris_caches.jl @@ -0,0 +1,49 @@ +struct BorisConstantCache <: OrdinaryDiffEqConstantCache end + +struct BorisCache{uType, rateType} <: OrdinaryDiffEqMutableCache + u::uType + uprev::uType + tmp::uType + k::rateType +end + +function alg_cache( + alg::Boris, u, rate_prototype, ::Type{uEltypeNoUnits}, + ::Type{uBottomEltypeNoUnits}, ::Type{tTypeNoUnits}, uprev, uprev2, f, t, dt, reltol, p, calck, + ::Val{false}, args...; kwargs... + ) where {uEltypeNoUnits, uBottomEltypeNoUnits, tTypeNoUnits} + return BorisConstantCache() +end + +function alg_cache( + alg::Boris, u, rate_prototype, ::Type{uEltypeNoUnits}, + ::Type{uBottomEltypeNoUnits}, ::Type{tTypeNoUnits}, uprev, uprev2, f, t, dt, reltol, p, calck, + ::Val{true}, args...; kwargs... + ) where {uEltypeNoUnits, uBottomEltypeNoUnits, tTypeNoUnits} + return BorisCache(u, uprev, similar(u), similar(rate_prototype)) +end + +struct MultistepBorisConstantCache <: OrdinaryDiffEqConstantCache end + +struct MultistepBorisCache{uType, rateType} <: OrdinaryDiffEqMutableCache + u::uType + uprev::uType + tmp::uType + k::rateType +end + +function alg_cache( + alg::MultistepBoris{N}, u, rate_prototype, ::Type{uEltypeNoUnits}, + ::Type{uBottomEltypeNoUnits}, ::Type{tTypeNoUnits}, uprev, uprev2, f, t, dt, reltol, p, calck, + ::Val{false}, args...; kwargs... + ) where {uEltypeNoUnits, uBottomEltypeNoUnits, tTypeNoUnits, N} + return MultistepBorisConstantCache() +end + +function alg_cache( + alg::MultistepBoris{N}, u, rate_prototype, ::Type{uEltypeNoUnits}, + ::Type{uBottomEltypeNoUnits}, ::Type{tTypeNoUnits}, uprev, uprev2, f, t, dt, reltol, p, calck, + ::Val{true}, args...; kwargs... + ) where {uEltypeNoUnits, uBottomEltypeNoUnits, tTypeNoUnits, N} + return MultistepBorisCache(u, uprev, similar(u), similar(rate_prototype)) +end diff --git a/lib/OrdinaryDiffEqBoris/src/boris_perform_step.jl b/lib/OrdinaryDiffEqBoris/src/boris_perform_step.jl new file mode 100644 index 000000000..c0ad1a7e5 --- /dev/null +++ b/lib/OrdinaryDiffEqBoris/src/boris_perform_step.jl @@ -0,0 +1,252 @@ +const TN_MAG_THRESHOLD = 1.0e-4 + +@inline @muladd function boris_velocity_update(v, E, B, qdt_2m) + t_rotate = qdt_2m * B + t_mag2 = sum(abs2, t_rotate) + s_rotate = 2 * t_rotate / (1 + t_mag2) + + v⁻ = v + qdt_2m * E + v′ = v⁻ + (v⁻ × t_rotate) + v⁺ = v⁻ + (v′ × s_rotate) + + v_new = v⁺ + qdt_2m * E + + return v_new +end + +function initialize!(integrator, cache::BorisConstantCache) + integrator.fsalfirst = integrator.f(integrator.uprev, integrator.p, integrator.t) + integrator.stats.nf += 1 + integrator.kshortsize = 0 + return integrator.k = typeof(integrator.k)(undef, integrator.kshortsize) +end + +@muladd function perform_step!(integrator, cache::BorisConstantCache, repeat_step = false) + t = integrator.t + dt = integrator.dt + uprev = integrator.uprev + p = integrator.p + + r = uprev[SVector(1, 2, 3)] + v = uprev[SVector(4, 5, 6)] + + q2m = p[1] + Efunc = p[3] + Bfunc = p[4] + + r_half = r + v * (dt / 2) + t_half = t + dt / 2 + E = Efunc(r_half, t_half) + B = Bfunc(r_half, t_half) + + qdt_2m = q2m * 0.5 * dt + v_new = boris_velocity_update(v, E, B, qdt_2m) + + r_new = r_half + v_new * (dt / 2) + + integrator.u = vcat(r_new, v_new) + + alg = integrator.alg + if alg.safety > 0.0 + Bmag = norm(Bfunc(r_new, t + dt)) + dt_new = integrator.tdir * (2π * alg.safety) / (abs(q2m) * Bmag) + set_proposed_dt!(integrator, dt_new) + end + return integrator.u +end + +function initialize!(integrator, cache::BorisCache) + integrator.fsalfirst = integrator.f(integrator.uprev, integrator.p, integrator.t) + integrator.stats.nf += 1 + integrator.kshortsize = 0 + return integrator.k = typeof(integrator.k)(undef, integrator.kshortsize) +end + +@muladd function perform_step!(integrator, cache::BorisCache, repeat_step = false) + t = integrator.t + dt = integrator.dt + uprev = integrator.uprev + p = integrator.p + + r = SVector(uprev[1], uprev[2], uprev[3]) + v = SVector(uprev[4], uprev[5], uprev[6]) + + q2m = p[1] + Efunc = p[3] + Bfunc = p[4] + + r_half = r + v * (dt / 2) + t_half = t + dt / 2 + E = Efunc(r_half, t_half) + B = Bfunc(r_half, t_half) + + qdt_2m = q2m * 0.5 * dt + v_new = boris_velocity_update(v, E, B, qdt_2m) + + r_new = r_half + v_new * (dt / 2) + + integrator.u[1] = r_new[1] + integrator.u[2] = r_new[2] + integrator.u[3] = r_new[3] + integrator.u[4] = v_new[1] + integrator.u[5] = v_new[2] + integrator.u[6] = v_new[3] + + alg = integrator.alg + if alg.safety > 0.0 + Bmag = norm(Bfunc(r_new, t + dt)) + dt_new = integrator.tdir * (2π * alg.safety) / (abs(q2m) * Bmag) + set_proposed_dt!(integrator, dt_new) + end + return +end + +_get_val_N(::MultistepBoris{N}) where {N} = Val{N}() + +@muladd function update_velocity_multistep(v, r, dt, t, n::Int, ::Val{N}, param) where {N} + q2m, Efunc, Bfunc = param[1], param[3], param[4] + + E = Efunc(r, t) + B = Bfunc(r, t) + + factor = q2m * dt / (2 * n) + + t_n = factor * B + e_n = factor * E + + if N != 2 + t_mag2 = sum(abs2, t_n) + if N == 4 + f_N = 1 + t_mag2 / 3 + e_corr_factor = -1 / 3 + else # N == 6 + f_N = 1 + t_mag2 / 3 + 2 * t_mag2 * t_mag2 / 15 + e_corr_factor = -1 / 3 - 2 * t_mag2 / 15 + end + + e_dot_t = e_n ⋅ t_n + e_n = f_N * e_n + (e_corr_factor * e_dot_t) * t_n + t_n = f_N * t_n + end + + t_n_mag2 = sum(abs2, t_n) + t_n_mag = sqrt(t_n_mag2) + + if t_n_mag < TN_MAG_THRESHOLD + c_n1 = 1 - 2 * n * n * t_n_mag2 + + n_term1 = 2 * n + n_term3 = 4 * n * n * n + + c_n2 = n_term1 - (n_term1 + n_term3) / 3 * t_n_mag2 + c_n3 = 2 * n * n - (4 * n * n + 2 * n * n * n * n) / 3 * t_n_mag2 + c_n6 = (n_term1 + n_term3) / 3 + else + alpha_n = atan(t_n_mag) + n_alpha_n = n * alpha_n + sin_n_alpha, cos_n_alpha = sincos(n_alpha_n) + sin_2n_alpha = 2 * sin_n_alpha * cos_n_alpha + cos_2n_alpha = cos_n_alpha * cos_n_alpha - sin_n_alpha * sin_n_alpha + + c_n1 = cos_2n_alpha + c_n2 = sin_2n_alpha / t_n_mag + c_n3 = 2 * sin_n_alpha * sin_n_alpha / t_n_mag2 + c_n6 = (2 * n - c_n2) / t_n_mag2 + end + + c_n4 = c_n2 + c_n5 = c_n3 + + v_dot_t = v ⋅ t_n + e_dot_t = e_n ⋅ t_n + + v_cross_t = v × t_n + e_cross_t = e_n × t_n + + v_new = c_n1 * v + + c_n2 * v_cross_t + + c_n3 * v_dot_t * t_n + + c_n4 * e_n + + c_n5 * e_cross_t + + c_n6 * e_dot_t * t_n + + return v_new +end + +function initialize!(integrator, cache::MultistepBorisConstantCache) + integrator.fsalfirst = integrator.f(integrator.uprev, integrator.p, integrator.t) + integrator.stats.nf += 1 + integrator.kshortsize = 0 + return integrator.k = typeof(integrator.k)(undef, integrator.kshortsize) +end + +@muladd function perform_step!(integrator, cache::MultistepBorisConstantCache, repeat_step = false) + t = integrator.t + dt = integrator.dt + uprev = integrator.uprev + p = integrator.p + alg = integrator.alg + + r = uprev[SVector(1, 2, 3)] + v = uprev[SVector(4, 5, 6)] + + r_half = r + v * (dt / 2) + t_half = t + dt / 2 + + v_new = update_velocity_multistep(v, r_half, dt, t_half, alg.n, _get_val_N(alg), p) + + r_new = r_half + v_new * (dt / 2) + + integrator.u = vcat(r_new, v_new) + + if alg.safety > 0.0 + q2m = p[1] + Bfunc = p[4] + Bmag = norm(Bfunc(r_new, t + dt)) + dt_new = integrator.tdir * (2π * alg.safety) / (abs(q2m) * Bmag) + set_proposed_dt!(integrator, dt_new) + end + return integrator.u +end + +function initialize!(integrator, cache::MultistepBorisCache) + integrator.fsalfirst = integrator.f(integrator.uprev, integrator.p, integrator.t) + integrator.stats.nf += 1 + integrator.kshortsize = 0 + integrator.k = typeof(integrator.k)(undef, integrator.kshortsize) + return +end + +@muladd function perform_step!(integrator, cache::MultistepBorisCache, repeat_step = false) + t = integrator.t + dt = integrator.dt + uprev = integrator.uprev + p = integrator.p + alg = integrator.alg + + r = SVector(uprev[1], uprev[2], uprev[3]) + v = SVector(uprev[4], uprev[5], uprev[6]) + + r_half = r + v * (dt / 2) + t_half = t + dt / 2 + + v_new = update_velocity_multistep(v, r_half, dt, t_half, alg.n, _get_val_N(alg), p) + + r_new = r_half + v_new * (dt / 2) + + integrator.u[1] = r_new[1] + integrator.u[2] = r_new[2] + integrator.u[3] = r_new[3] + integrator.u[4] = v_new[1] + integrator.u[5] = v_new[2] + integrator.u[6] = v_new[3] + + if alg.safety > 0.0 + q2m = p[1] + Bfunc = p[4] + Bmag = norm(Bfunc(r_new, t + dt)) + dt_new = integrator.tdir * (2π * alg.safety) / (abs(q2m) * Bmag) + set_proposed_dt!(integrator, dt_new) + end + return +end diff --git a/lib/OrdinaryDiffEqBoris/test/runtests.jl b/lib/OrdinaryDiffEqBoris/test/runtests.jl new file mode 100644 index 000000000..5f439f664 --- /dev/null +++ b/lib/OrdinaryDiffEqBoris/test/runtests.jl @@ -0,0 +1,143 @@ +using Test +using OrdinaryDiffEqBoris +using StaticArrays +using LinearAlgebra: norm + +@testset "OrdinaryDiffEqBoris.jl" begin + # Definitions + constant_Ey(x, t) = SA[0.0, 1.0, 0.0] + constant_Bz(x, t) = SA[0.0, 0.0, 1.0] + ZeroField() = (x, t) -> SA[0.0, 0.0, 0.0] + + dummy_f(u, p, t) = u # ODEProblem requires an f, but Boris uses p directly + + @testset "E cross B drift - Standard Boris" begin + # E = (0, 1, 0), B = (0, 0, 1) + # Analytic solution: particle moves at constant velocity v_drift = (1, 0, 0) + param = (1.0, 1.0, constant_Ey, constant_Bz, ZeroField()) + u0 = SA[0.0, 0.0, 0.0, 1.0, 0.0, 0.0] + tspan = (0.0, 10.0) + prob = ODEProblem(dummy_f, u0, tspan, param) + + dt = 0.1 + # Test standard Boris + sol_boris = solve(prob, Boris(); dt) + @test sol_boris.u[end][1] ≈ 10.0 atol = 1.0e-6 + @test sol_boris.u[end][4] ≈ 1.0 atol = 1.0e-6 + end + + @testset "E cross B drift - Multistep Boris" begin + param = (1.0, 1.0, constant_Ey, constant_Bz, ZeroField()) + u0 = SA[0.0, 0.0, 0.0, 1.0, 0.0, 0.0] + tspan = (0.0, 10.0) + prob = ODEProblem(dummy_f, u0, tspan, param) + + dt = 0.1 + sol_multi_2 = solve(prob, MultistepBoris2(n = 2); dt) + @test sol_multi_2.u[end][1] ≈ 10.0 atol = 1.0e-6 + @test sol_multi_2.u[end][4] ≈ 1.0 atol = 1.0e-6 + + # Test Hyper Boris N=4 + sol_hyper_4 = solve(prob, MultistepBoris4(n = 2); dt) + @test sol_hyper_4.u[end][1] ≈ 10.0 atol = 1.0e-6 + end + + @testset "Gyrating particle" begin + param_gyro = (1.0, 1.0, ZeroField(), constant_Bz, ZeroField()) + u0_gyro = SA[0.0, 0.0, 0.0, 1.0, 0.0, 0.0] + prob_gyro = ODEProblem(dummy_f, u0_gyro, (0.0, 2π), param_gyro) + + # After one period, should return to origin + sol_1step_gyro = solve(prob_gyro, MultistepBoris2(n = 1); dt = 0.1) + @test hypot(@views sol_1step_gyro.u[end][1:3]...) < 0.02 + + sol_hyper_6_gyro = solve(prob_gyro, MultistepBoris6(n = 4); dt = 0.1) + @test hypot(@views sol_hyper_6_gyro.u[end][1:3]...) < 0.02 + end + + @testset "Adaptive Boris" begin + constant_E(x, t) = SA[0.0, 1.0e2, 0.0] + gradient_B(x, t) = SA[0.0, 0.0, 1.0 * (1.0 + x[1])] + param = (-1.0, 1.0, constant_E, gradient_B, ZeroField()) # q2m = -1 + + u0 = SA[0.0, 0.0, 0.0, 10.0, 0.0, 0.0] + tspan = (0.0, 10.0) + prob = ODEProblem(dummy_f, u0, tspan, param) + + safety = 0.1 + init_Bmag = norm(gradient_B(u0[1:3], 0.0)) + dt_init = safety * 2π / (abs(param[1]) * init_Bmag) + + sol = solve(prob, Boris(safety = safety); dt = dt_init) + + # Check if step sizes are not uniform (hence adapted) + dts = diff(sol.t) + @test !all(y -> isapprox(y, dts[1], rtol = 1.0e-5), dts) + @test sol.t[end] ≈ tspan[2] + + # Energy conservation check + total_energy(u) = 0.5 * param[2] * norm(u[4:6])^2 + (param[1] * param[2]) * (-1.0e2 * u[2]) + E_start = total_energy(sol.u[1]) + E_end = total_energy(sol.u[end]) + @test isapprox(E_end, E_start, rtol = 1.0e-3) + end + + @testset "Nonzero tspan[1]" begin + # B field that is only present when t > 5 + B_field(r, t) = t > 5 ? SA[0.0, 0.0, 0.01] : SA[0.0, 0.0, 0.0] + E_field(r, t) = SA[0.0, 0.0, 0.0] + + # param = (q2m, m, E, B, F) + param = (-1.0e11, 1.0e-30, E_field, B_field, ZeroField()) + + # Start at t=10. If absolute time is used, B should be 0.01. + tspan = (10.0, 10.0 + 1.0e-7) + dt = 1.0e-9 + u0 = SA[0.0, 0.0, 0.0, 1.0e5, 0.0, 0.0] + prob = ODEProblem(dummy_f, u0, tspan, param) + + # Standard Boris + sol_boris = solve(prob, Boris(); dt = dt) + # If B was 0.01, vx should have deviated from initial 1.0e5 + @test abs(sol_boris.u[end][4] - 1.0e5) > 1.0e-4 + + # Adaptive Boris + sol_adaptive = solve(prob, Boris(safety = 0.1); dt = dt) + @test abs(sol_adaptive.u[end][4] - 1.0e5) > 1.0e-4 + end + + @testset "SciML Output saving flags" begin + # Setup + x0 = [0.0, 0.0, 0.0] + v0 = [0.0, 1.0e5, 0.0] + stateinit = [x0..., v0...] + tspan = (0.0, 3.0e-8) + dt = 3.0e-11 + + zero_E(x, t) = SA[0.0, 0.0, 0.0] + uniform_B(x, t) = SA[0.0, 0.0, 0.01] + param = (-1.0e11, 1.0e-30, zero_E, uniform_B, zero_E) + prob = ODEProblem(dummy_f, stateinit, tspan, param) + + # Baseline: save_everystep=true (default) + sol = solve(prob, Boris(); dt) + @test length(sol.t) == 1001 # 3e-8 / 3e-11 = 1000 steps + start + + # Scenario 2: Only final state + sol = solve(prob, Boris(); dt, save_everystep = false, save_start = false, save_on = false) + @test length(sol.t) == 1 + @test sol.t[1] ≈ tspan[2] + + # Scenario 3: Start and End + sol = solve(prob, Boris(); dt, save_everystep = false, save_start = true, save_on = false) + @test length(sol.t) == 2 + @test sol.t[1] == tspan[1] + @test sol.t[end] ≈ tspan[2] + + # Multistep Boris test flags + sol_ms = solve(prob, MultistepBoris2(n = 2); dt, save_everystep = false, save_start = true, save_on = false) + @test length(sol_ms.t) == 2 + @test sol_ms.t[1] == tspan[1] + @test sol_ms.t[end] ≈ tspan[2] + end +end diff --git a/src/TestParticle.jl b/src/TestParticle.jl index 0b034be71..8b2f57d67 100644 --- a/src/TestParticle.jl +++ b/src/TestParticle.jl @@ -37,7 +37,7 @@ export trace!, trace_relativistic!, trace_normalized!, trace_relativistic_normal get_gc_velocity, full_to_gc, gc_to_full export Proton, Electron, Ion export Maxwellian, BiMaxwellian, Kappa, BiKappa -export AdaptiveBoris, AdaptiveHybrid +export AdaptiveHybrid, AdaptiveBoris, Boris, MultistepBoris, MultistepBoris2, MultistepBoris4, MultistepBoris6 export get_gyrofrequency, get_gyroperiod, get_gyroradius, get_velocity, get_energy, get_mean_magnitude, energy2velocity, get_curvature_radius, get_adiabaticity, @@ -63,7 +63,6 @@ include("gc_solver.jl") include("equations.jl") include("boris.jl") include("boris_kernel.jl") -include("adaptive_boris.jl") include("hybrid.jl") include("fieldline.jl") diff --git a/src/adaptive_boris.jl b/src/adaptive_boris.jl deleted file mode 100644 index 8be6319c3..000000000 --- a/src/adaptive_boris.jl +++ /dev/null @@ -1,303 +0,0 @@ -# Adaptive Boris method - -struct AdaptiveBoris{T} - safety::T -end - -""" - AdaptiveBoris(; safety=0.1) - -Adaptive Boris method with adaptive time stepping based on local gyroperiod. -The time step is determined by `dt = safety * T_gyro = safety * 2π / |qB/m|`. -""" -function AdaptiveBoris(; safety = 0.1) - T = typeof(safety) - return AdaptiveBoris{T}(T(safety)) -end -""" - solve(prob::TraceProblem, alg::AdaptiveBoris, - ensemblealg::BasicEnsembleAlgorithm=EnsembleSerial(); - trajectories::Int=1, savestepinterval::Int=1, - isoutside::Function=ODE_DEFAULT_ISOUTOFDOMAIN, - save_start::Bool=true, save_end::Bool=true, save_everystep::Bool=true, - save_fields::Bool=false, save_work::Bool=false, - batch_size::Int = max(1, trajectories ÷ Threads.nthreads())) - -Trace particles using the Adaptive Boris method with specified `prob` and `alg`. - -# keywords - - - `trajectories::Int`: number of trajectories to trace. - - `savestepinterval::Int`: saving output interval. - - `isoutside::Function`: a function with input of position and velocity vector `xv` that determines whether to stop tracing. - - `save_start::Bool`: save the initial condition. Default is `true`. - - `save_end::Bool`: save the final condition. Default is `true`. - - `save_everystep::Bool`: save the state at every `savestepinterval`. Default is `true`. - - `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`. - - `batch_size::Int`: the number of trajectories to process per worker in `EnsembleDistributed` and `EnsembleSplitThreads`. Default is `max(1, trajectories ÷ nworkers())` for distributed and 1 for others. - - -""" -function solve( - prob::TraceProblem, alg::AdaptiveBoris, - ensemblealg::BasicEnsembleAlgorithm = EnsembleSerial(); - trajectories::Int = 1, savestepinterval::Int = 1, - isoutside::F = ODE_DEFAULT_ISOUTOFDOMAIN, - save_start::Bool = true, save_end::Bool = true, save_everystep::Bool = true, - save_fields::Bool = false, save_work::Bool = false, - batch_size::Int = (ensemblealg isa EnsembleDistributed || ensemblealg isa EnsembleSplitThreads) ? - max(1, trajectories ÷ nworkers()) : 1 - ) where {F} - - return if save_fields - if save_work - return _solve( - ensemblealg, prob, trajectories, alg, savestepinterval, isoutside, - save_start, save_end, save_everystep, Val(true), Val(true), batch_size - ) - else - return _solve( - ensemblealg, prob, trajectories, alg, savestepinterval, isoutside, - save_start, save_end, save_everystep, Val(true), Val(false), batch_size - ) - end - else - if save_work - return _solve( - ensemblealg, prob, trajectories, alg, savestepinterval, isoutside, - save_start, save_end, save_everystep, Val(false), Val(true), batch_size - ) - else - return _solve( - ensemblealg, prob, trajectories, alg, savestepinterval, isoutside, - save_start, save_end, save_everystep, Val(false), Val(false), batch_size - ) - end - end -end - -function _solve( - ::EnsembleSerial, prob, trajectories, alg::AdaptiveBoris, savestepinterval, - isoutside, save_start, save_end, save_everystep, - ::Val{SaveFields}, ::Val{SaveWork}, batch_size - ) where {SaveFields, SaveWork} - # We cannot precalculate nt for adaptive steps - sol_type = _get_sol_type(prob, zero(eltype(prob.tspan)), Val(SaveFields), Val(SaveWork)) - sols = Vector{sol_type}(undef, trajectories) - irange = 1:trajectories - - _adaptive_boris!( - sols, prob, irange, alg, savestepinterval, isoutside, - save_start, save_end, save_everystep, Val(SaveFields), Val(SaveWork) - ) - - return sols -end - -function _solve( - ::EnsembleThreads, prob, trajectories, alg::AdaptiveBoris, savestepinterval, - isoutside, save_start, save_end, save_everystep, ::Val{SaveFields}, ::Val{SaveWork}, - batch_size - ) where {SaveFields, SaveWork} - sol_type = _get_sol_type(prob, zero(eltype(prob.tspan)), Val(SaveFields), Val(SaveWork)) - sols = Vector{sol_type}(undef, trajectories) - - nchunks = Threads.nthreads() - Threads.@threads for irange in index_chunks(1:trajectories; n = nchunks) - _adaptive_boris!( - sols, prob, irange, alg, savestepinterval, isoutside, - save_start, save_end, save_everystep, Val(SaveFields), Val(SaveWork) - ) - end - - return sols -end - -"See `_solve_single_boris` for the rationale behind the `single_prob` construction." -function _solve_single_adaptive_boris( - prob, i, alg::AdaptiveBoris, savestepinterval, isoutside, - save_start, save_end, save_everystep, ::Val{SaveFields}, ::Val{SaveWork} - ) where {SaveFields, SaveWork} - new_prob = prob.prob_func(prob, i, false) - single_prob = TraceProblem(new_prob.u0, new_prob.tspan, new_prob.p) - sol_type = _get_sol_type( - single_prob, zero(eltype(single_prob.tspan)), Val(SaveFields), Val(SaveWork) - ) - local_sols = Vector{sol_type}(undef, 1) - _adaptive_boris!( - local_sols, single_prob, 1:1, alg, savestepinterval, isoutside, - save_start, save_end, save_everystep, Val(SaveFields), Val(SaveWork) - ) - return local_sols[1] -end - -function _solve( - ::EnsembleDistributed, prob, trajectories, alg::AdaptiveBoris, savestepinterval, - isoutside, save_start, save_end, save_everystep, ::Val{SaveFields}, ::Val{SaveWork}, - batch_size - ) where {SaveFields, SaveWork} - return pmap(1:trajectories; batch_size = batch_size) do i - _solve_single_adaptive_boris( - prob, i, alg, savestepinterval, isoutside, - save_start, save_end, save_everystep, Val(SaveFields), Val(SaveWork) - ) - end -end - -function _solve( - ::EnsembleSplitThreads, prob, trajectories, alg::AdaptiveBoris, savestepinterval, - isoutside, save_start, save_end, save_everystep, ::Val{SaveFields}, ::Val{SaveWork}, - batch_size - ) where {SaveFields, SaveWork} - # _solve_single_adaptive_boris wraps each trajectory in a fresh TraceProblem(u0, tspan, p) - # with DEFAULT_PROB_FUNC. We get a sample problem from prob_func to ensure the - # u0 type (uType) matches. - sample_prob = prob.prob_func(prob, 1, false) - dummy_prob = TraceProblem(sample_prob.u0, sample_prob.tspan, sample_prob.p) - sol_type = _get_sol_type( - dummy_prob, zero(eltype(dummy_prob.tspan)), Val(SaveFields), Val(SaveWork) - ) - ichunks = index_chunks(1:trajectories; size = batch_size) - results = pmap(ichunks) do irange - local_sols = Vector{sol_type}(undef, length(irange)) - Threads.@threads for k in eachindex(irange) - i = irange[k] - local_sols[k] = _solve_single_adaptive_boris( - prob, i, alg, savestepinterval, isoutside, - save_start, save_end, save_everystep, Val(SaveFields), Val(SaveWork) - ) - end - local_sols - end - return reduce(vcat, results) -end - - -@muladd function _adaptive_boris!( - sols, prob, irange, alg, savestepinterval, isoutside, - save_start, save_end, save_everystep, ::Val{SaveFields}, ::Val{SaveWork} - ) where {SaveFields, SaveWork} - (; tspan, p, u0) = prob - q2m, _, _, Bfunc, _ = p - T = eltype(u0) - vars_dim = 6 - if SaveFields - vars_dim += 6 - end - if SaveWork - vars_dim += 4 - end - - # Pre-calculate common factors for time step calculation - # dt = safety * 2π / (abs(q2m) * Bmag) - C = (2π * alg.safety * sign(tspan[2] - tspan[1])) / abs(q2m) - - @fastmath @inbounds for i in irange - # Initialize solution containers - initial_capacity = 1000 - traj = Vector{SVector{vars_dim, T}}(undef, 0) - tsave = Vector{typeof(tspan[1])}(undef, 0) - sizehint!(traj, initial_capacity) - sizehint!(tsave, initial_capacity) - - new_prob = prob.prob_func(prob, i, false) - # Load independent r and v SVector from u0 - u0_i = SVector{6, T}(new_prob.u0) - r = u0_i[SVector(1, 2, 3)] - v = u0_i[SVector(4, 5, 6)] - t = tspan[1] - ttotal = tspan[2] - tspan[1] - - if save_start - data = _prepare_saved_data(u0_i, p, t, Val(SaveFields), Val(SaveWork)) - push!(traj, data) - push!(tsave, t) - end - - # Initial dt calculation - Bmag = norm(Bfunc(r, t)) - dt = C / Bmag - - # Backstep velocity: v(0) -> v(-1/2) using dt - v = update_velocity(v, r, -0.5 * dt, t, p) - it = 1 - should_save_final = save_end - retcode = ReturnCode.Success - - while abs(t - tspan[1]) < abs(ttotal) - # Check if next step exceeds tspan[2] - if abs(t + dt - tspan[1]) > abs(ttotal) - dt_step = tspan[2] - t - # Resync v from `t - 0.5*dt` to `t - 0.5*dt_step` - v = update_velocity(v, r, 0.5 * dt, t, p) - v = update_velocity(v, r, -0.5 * dt_step, t, p) - dt = dt_step - end - - # Saving logic (start of step) - if save_everystep && (it - 1) > 0 && (it - 1) % savestepinterval == 0 - # Advance to t to get v_n - v_save = update_velocity(v, r, 0.5 * dt, t, p) - - xv_s = vcat(r, v_save) - data = _prepare_saved_data(xv_s, p, t, Val(SaveFields), Val(SaveWork)) - push!(traj, data) - push!(tsave, t) - end - - # Update velocity to v_{n+1/2} - t_mid = t + 0.5 * dt - v_new = update_velocity(v, r, dt, t_mid, p) - - # Update location x_{n} -> x_{n+1} - r_next = r + v_new * dt - t_next = t + dt - - xv_new = vcat(r_next, v_new) - if isoutside(xv_new, p, t_next) - should_save_final = true - retcode = ReturnCode.Terminated - break - end - - r = r_next - t = t_next - v = v_new - - # New dt - Bmag = norm(Bfunc(r, t)) - dt_new = C / Bmag - - # Resync v_{n+1/2}(dt) to v_{n+1/2}(dt_new) - # v is at t_{new} - 0.5 * dt_old (relative to t_{new}) - # i.e. it is v_{n+1/2} from step we just took. - v = update_velocity(v, r, 0.5 * dt, t, p) - v = update_velocity(v, r, -0.5 * dt_new, t, p) - dt = dt_new - it += 1 - if save_everystep && (it - 1) % savestepinterval == 0 - should_save_final = true - end - end - - if should_save_final && (isempty(tsave) || tsave[end] != t) - # v is at t - 0.5*dt. To get v at t, advance by 0.5*dt - v_final = update_velocity(v, r, 0.5 * dt, t, p) - - xv_f = vcat(r, v_final) - data = _prepare_saved_data(xv_f, p, t, Val(SaveFields), Val(SaveWork)) - push!(traj, data) - push!(tsave, t) - end - - # Construct solution - sol_alg = :adaptive_boris - interp = LinearInterpolation(tsave, traj) - stats = nothing - - sols[i] = build_solution(prob, sol_alg, tsave, traj; interp, retcode, stats) - end - - return -end diff --git a/src/boris.jl b/src/boris.jl index c0c4f5079..f6b978104 100644 --- a/src/boris.jl +++ b/src/boris.jl @@ -1,5 +1,56 @@ # Native particle pusher +abstract type AbstractBoris end + +""" + Boris(; safety=0.0) + +The standard Boris method for particle pushing in electric and magnetic fields. +When `safety > 0.0`, it uses adaptive time stepping based on local gyroperiod. +""" +struct Boris{T} <: AbstractBoris + safety::T +end +Boris(; safety = 0.0) = Boris(safety) + +""" + AdaptiveBoris(; safety=0.1) + +The adaptive Boris method with adaptive time stepping based on local gyroperiod. +Returns a `Boris` instance with the specified `safety` factor. +""" +struct AdaptiveBoris{T} <: AbstractBoris + safety::T +end +function AdaptiveBoris(; safety = 0.1) + @warn "AdaptiveBoris is deprecated. Use Boris(safety=$safety) instead." maxlog = 1 + return AdaptiveBoris(safety) +end + +""" + MultistepBoris{N}(; n=1, safety=0.0) + +The Multistep/Hyper Boris method of order `N`. +`n` specifies the number of subcycles. +`N` specifies the gyrophase correction order: 2 (standard), 4, or 6 (Hyper-Boris). +When `safety > 0.0`, it uses adaptive time stepping. +""" +struct MultistepBoris{N, T} <: AbstractBoris + n::Int + safety::T +end +@inline function MultistepBoris{N}(; n::Int = 1, safety = 0.0) where {N} + if N ∉ (2, 4, 6) + throw(ArgumentError("Multistep Boris order N must be 2, 4, or 6.")) + end + return MultistepBoris{N, typeof(safety)}(n, safety) +end + +const MultistepBoris2 = MultistepBoris{2} +const MultistepBoris4 = MultistepBoris{4} +const MultistepBoris6 = MultistepBoris{6} + + struct TraceProblem{uType, tType, isinplace, P, F <: AbstractODEFunction, PF} <: AbstractODEProblem{uType, tType, isinplace} f::F @@ -96,6 +147,252 @@ Update location in one timestep `dt`. return end +@inline _get_val_N(::MultistepBoris{N}) where {N} = Val(N) + +""" + update_velocity_multistep(v, r, dt, t, n, ::Val{N}, param) where {N} + +Update velocity using the Multistep/Hyper Boris method. +""" +@muladd function update_velocity_multistep(v, r, dt, t, n::Int, ::Val{N}, param) where {N} + q2m, _, Efunc, Bfunc, _ = param + E = Efunc(r, t) + B = Bfunc(r, t) + + # t_n and e_n vectors + factor = q2m * dt / (2 * n) + + t_n = factor * B # (q/m * dt/(2n)) * B + e_n = factor * E # (q/m * dt/(2n)) * E + + # Hyper Boris N-th order gyrophase correction + if N != 2 + t_mag2 = sum(abs2, t_n) + if N == 4 + f_N = 1 + t_mag2 / 3 + e_corr_factor = -1 / 3 + else # N == 6 + f_N = 1 + t_mag2 / 3 + 2 * t_mag2 * t_mag2 / 15 + e_corr_factor = -1 / 3 - 2 * t_mag2 / 15 + end + + e_dot_t = e_n ⋅ t_n + e_n = f_N * e_n + (e_corr_factor * e_dot_t) * t_n + t_n = f_N * t_n + end + + t_n_mag2 = sum(abs2, t_n) + t_n_mag = sqrt(t_n_mag2) + + # Calculate coefficients + if t_n_mag < TN_MAG_THRESHOLD + c_n1 = 1 - 2 * n * n * t_n_mag2 + n_term1 = 2 * n + n_term3 = 4 * n * n * n + c_n2 = n_term1 - (n_term1 + n_term3) / 3 * t_n_mag2 + c_n3 = 2 * n * n - (4 * n * n + 2 * n * n * n * n) / 3 * t_n_mag2 + c_n6 = (n_term1 + n_term3) / 3 + else + alpha_n = atan(t_n_mag) + n_alpha_n = n * alpha_n + sin_n_alpha, cos_n_alpha = sincos(n_alpha_n) + sin_2n_alpha = 2 * sin_n_alpha * cos_n_alpha + cos_2n_alpha = cos_n_alpha * cos_n_alpha - sin_n_alpha * sin_n_alpha + + c_n1 = cos_2n_alpha + c_n2 = sin_2n_alpha / t_n_mag + c_n3 = 2 * sin_n_alpha * sin_n_alpha / t_n_mag2 + c_n6 = (2 * n - c_n2) / t_n_mag2 + end + + c_n4 = c_n2 + c_n5 = c_n3 + + v_dot_t = v ⋅ t_n + e_dot_t = e_n ⋅ t_n + v_cross_t = v × t_n + e_cross_t = e_n × t_n + + v_new = c_n1 * v + + c_n2 * v_cross_t + + c_n3 * v_dot_t * t_n + + c_n4 * e_n + + c_n5 * e_cross_t + + c_n6 * e_dot_t * t_n + + return v_new +end + +@inline _update_velocity(v, r, dt, t, p, ::Boris) = update_velocity(v, r, dt, t, p) +@inline function _update_velocity(v, r, dt, t, p, alg::MultistepBoris) + return update_velocity_multistep(v, r, dt, t, alg.n, _get_val_N(alg), p) +end + +@inline @muladd function _boris_loop!( + traj, tsave, r, v, p, dt::T, tspan, + savestepinterval::Int, save_everystep::Bool, isoutside::F1, + alg, ::Val{SaveFields}, ::Val{SaveWork}, maxiters::Int + ) where {T, F1, SaveFields, SaveWork} + t = tspan[1] + ttotal = tspan[2] - tspan[1] + it = 1 + retcode = ReturnCode.Success + + C = zero(eltype(v)) + if alg.safety > 0.0 + q2m, _, _, Bfunc, _ = p + C = (2π * alg.safety * sign(ttotal)) / abs(q2m) + Bmag = norm(Bfunc(r, t)) + dt = C / Bmag + end + + # push velocity back in time by 1/2 dt + v = _update_velocity(v, r, -0.5 * dt, t, p, alg) + + while abs(t - tspan[1]) < abs(ttotal) + if abs(t + dt - tspan[1]) >= abs(ttotal) - 100 * eps(eltype(v)(abs(ttotal))) + dt_step = tspan[2] - t + if dt_step != dt + # Resync v from `t - 0.5*dt` to `t - 0.5*dt_step` + v = _update_velocity(v, r, 0.5 * dt, t, p, alg) + v = _update_velocity(v, r, -0.5 * dt_step, t, p, alg) + dt = dt_step + end + end + + if save_everystep && (it - 1) > 0 && (it - 1) % savestepinterval == 0 + # Advance to t to get v_n + v_save = _update_velocity(v, r, 0.5 * dt, t, p, alg) + data = vcat(r, v_save) + push!(traj, _prepare_saved_data(data, p, t, Val(SaveFields), Val(SaveWork))) + push!(tsave, t) + end + + v_new = _update_velocity(v, r, dt, t + 0.5 * dt, p, alg) + r_next = r + v_new * dt + t_next = t + dt + + # NaN check + if any(isnan, r_next) || any(isnan, v_new) + retcode = ReturnCode.Unstable + break + end + + if isoutside(vcat(r_next, v_new), p, t_next) + retcode = ReturnCode.Terminated + break + end + + r, v, t = r_next, v_new, t_next + + if alg.safety > 0.0 && abs(t - tspan[1]) < abs(ttotal) + q2m, _, _, Bfunc, _ = p + Bmag = norm(Bfunc(r, t)) + dt_new = C / Bmag + # Resync v_{n+1/2}(dt) to v_{n+1/2}(dt_new) + v = _update_velocity(v, r, 0.5 * dt, t, p, alg) + v = _update_velocity(v, r, -0.5 * dt_new, t, p, alg) + dt = dt_new + end + if it > maxiters + retcode = ReturnCode.MaxIters + break + end + it += 1 + end + + return r, v, t, dt, retcode +end + +@inline @muladd function boris_generic_loop!( + sols, prob::TraceProblem, irange, alg, dt, + savestepinterval::Int, isoutside::F1, save_start, save_end, + save_everystep, ::Val{SaveFields}, ::Val{SaveWork}, maxiters::Int + ) where {SaveFields, SaveWork, F1} + (; tspan, p, u0) = prob + T = eltype(u0) + + vars_dim = 6 + if SaveFields + vars_dim += 6 + end + if SaveWork + vars_dim += 4 + end + + algorithm_name = alg isa Boris ? :boris : :multistep_boris + + # Calculate exact nout for fixed-step to avoid memory regression + ttotal = tspan[2] - tspan[1] + nout_fixed = if !iszero(dt) + nt = round(Int, abs(ttotal / dt)) + nsteps = 0 + if save_start + nsteps += 1 + end + if save_everystep + steps = nt ÷ savestepinterval + last_is_step = (nt > 0) && (nt % savestepinterval == 0) + nsteps += steps + if !save_end && last_is_step + nsteps -= 1 + end + if save_end && !last_is_step + nsteps += 1 + end + elseif save_end + nsteps += 1 + end + nsteps + else + 1000 # Default capacity for adaptive Boris + end + + @inbounds for i in irange + traj = Vector{SVector{vars_dim, T}}(undef, 0) + tsave = Vector{typeof(tspan[1] + (iszero(dt) ? 0.0 : dt))}(undef, 0) + sizehint!(traj, nout_fixed) + sizehint!(tsave, nout_fixed) + + # set initial conditions for each trajectory i + new_prob = prob.prob_func(prob, i, false) + u0_i = SVector{6, T}(new_prob.u0) + r = u0_i[SVector(1, 2, 3)] + v = u0_i[SVector(4, 5, 6)] + + # If dt is not provided, it must be adaptive (safety > 0) + _dt = iszero(dt) ? zero(T) : T(dt) + + if save_start + push!(traj, _prepare_saved_data(u0_i, p, tspan[1], Val(SaveFields), Val(SaveWork))) + push!(tsave, tspan[1]) + end + + r, v, t, _dt_final, retcode = _boris_loop!( + traj, tsave, r, v, p, _dt, tspan, + savestepinterval, save_everystep, isoutside, alg, + Val(SaveFields), Val(SaveWork), maxiters + ) + + should_save_final = save_end + if should_save_final && (isempty(tsave) || tsave[end] != t) + # v is at t - 0.5*_dt_final. To get v at t, advance by 0.5*_dt_final + v_save = _update_velocity(v, r, 0.5 * _dt_final, t, p, alg) + data = vcat(r, v_save) + push!(traj, _prepare_saved_data(data, p, t, Val(SaveFields), Val(SaveWork))) + push!(tsave, t) + end + + alg_sol = algorithm_name + interp = LinearInterpolation(tsave, traj) + stats = nothing + + sols[i] = build_solution(prob, alg_sol, tsave, traj; interp, retcode, stats) + end + + return +end + """ In-place cross product. """ @@ -108,182 +405,168 @@ In-place cross product. end """ - solve(prob::TraceProblem; trajectories::Int=1, dt::AbstractFloat, + solve(prob::TraceProblem, alg::Union{Boris, MultistepBoris}=Boris(); + trajectories::Int=1, dt=nothing, savestepinterval::Int=1, isoutside::Function=ODE_DEFAULT_ISOUTOFDOMAIN, - n::Int=1, save_start::Bool=true, save_end::Bool=true, save_everystep::Bool=true, + save_start::Bool=true, save_end::Bool=true, save_everystep::Bool=true, save_fields::Bool=false, save_work::Bool=false) Trace particles using the Boris method with specified `prob`. - -# keywords - - - `trajectories::Int`: number of trajectories to trace. - - `dt::AbstractFloat`: time step. - - `savestepinterval::Int`: saving output interval. - - `isoutside::Function`: pinpointing impact or checking boundaries. - - `n::Int=1`: number of substeps for the Multistep Boris method. 1 is standard Boris. - - `N::Int=2`: order of the Hyper Boris gyrophase correction (2, 4, or 6). 2 is uncorrected. - - `save_start::Bool=true`: save the initial condition. - - `save_end::Bool=true`: save the final condition. - - `save_everystep::Bool=true`: save the state at every `savestepinterval`. - - `save_fields::Bool=false`: save the electric and magnetic fields. - - `save_work::Bool=false`: save the work done by the electric field. - - `batch_size::Int=max(1, trajectories ÷ nworkers())`: the number of trajectories to process per worker in `EnsembleDistributed` and `EnsembleSplitThreads`. - """ @inline function solve( - prob::TraceProblem, ensemblealg::EA = EnsembleSerial(); - trajectories::Int = 1, savestepinterval::Int = 1, dt::AbstractFloat, - isoutside::F = ODE_DEFAULT_ISOUTOFDOMAIN, n::Int = 1, N::Int = 2, + prob::TraceProblem, alg::Union{Boris, MultistepBoris, AdaptiveBoris} = Boris(), + ensemblealg::EA = EnsembleSerial(); + trajectories::Int = 1, savestepinterval::Int = 1, dt = 0.0, + isoutside::F = ODE_DEFAULT_ISOUTOFDOMAIN, save_start::Bool = true, save_end::Bool = true, save_everystep::Bool = true, save_fields::Bool = false, save_work::Bool = false, maxiters::Int = 1_000_000, batch_size::Int = (ensemblealg isa EnsembleDistributed || ensemblealg isa EnsembleSplitThreads) ? - max(1, trajectories ÷ nworkers()) : 1 + max(1, trajectories ÷ nworkers()) : 1, + n::Int = 1, N::Int = 2 ) where {EA <: BasicEnsembleAlgorithm, F} - if N ∉ (2, 4, 6) - throw(ArgumentError("N must be 2, 4, or 6")) + # Backward compatibility: promote AdaptiveBoris to Boris + if alg isa AdaptiveBoris + alg = Boris(alg.safety) + end + + # Backward compatibility: promote Boris to MultistepBoris if n > 1 or N != 2 + if alg isa Boris && (n > 1 || N != 2) + if N ∉ (2, 4, 6) + throw(ArgumentError("Multistep Boris order N must be 2, 4, or 6.")) + end + alg = MultistepBoris{N}(n = n, safety = alg.safety) + end + + if iszero(dt) && alg.safety <= 0.0 + throw(ArgumentError("Time step dt must be provided for fixed-step Boris solver.")) + end + + if !iszero(dt) + if abs(dt) < eps(eltype(dt)) * 100 + throw(ArgumentError("Time step dt is too small.")) + end + ttotal = prob.tspan[2] - prob.tspan[1] + if abs(ttotal / dt) > maxiters + throw(ArgumentError("Total steps exceed maxiters. Increase maxiters or dt.")) + end end return _solve( - ensemblealg, prob, trajectories, dt, savestepinterval, isoutside, n, N, + ensemblealg, prob, trajectories, alg, dt, savestepinterval, isoutside, save_start, save_end, save_everystep, Val(save_fields), Val(save_work), maxiters, batch_size ) end +function solve(prob::TraceProblem, ensemblealg::BasicEnsembleAlgorithm; kwargs...) + return solve(prob, Boris(), ensemblealg; kwargs...) +end + function _dispatch_boris!( - sols, prob::TraceProblem, irange, savestepinterval, dt, nt, nout, isoutside::F, - n, N, save_start, save_end, save_everystep, ::Val{SaveFields}, ::Val{SaveWork} + sols, prob::TraceProblem, irange, alg, dt, savestepinterval::Int, + isoutside::F, save_start, save_end, save_everystep, ::Val{SaveFields}, + ::Val{SaveWork}, maxiters::Int ) where {SaveFields, SaveWork, F} - return if n == 1 && N == 2 - _boris!( - sols, prob, irange, savestepinterval, dt, nt, nout, isoutside, - save_start, save_end, save_everystep, Val(SaveFields), Val(SaveWork) - ) - else - _multistep_boris!( - sols, prob, irange, savestepinterval, dt, nt, nout, isoutside, n, N, - save_start, save_end, save_everystep, Val(SaveFields), Val(SaveWork) - ) - end + return boris_generic_loop!( + sols, prob, irange, alg, dt, savestepinterval, isoutside, + save_start, save_end, save_everystep, Val(SaveFields), Val(SaveWork), + maxiters + ) end @inline function _solve( - ::EnsembleSerial, prob::TraceProblem, trajectories, dt, savestepinterval, - isoutside::F, n, N, save_start, save_end, save_everystep, - ::Val{SaveFields}, ::Val{SaveWork}, maxiters, batch_size + ::EnsembleSerial, prob::TraceProblem, trajectories, alg, dt, + savestepinterval::Int, isoutside::F, save_start, save_end, save_everystep, + ::Val{SaveFields}, ::Val{SaveWork}, maxiters::Int, batch_size ) where {SaveFields, SaveWork, F} - sols, nt, - nout = _prepare( - prob, trajectories, dt, savestepinterval, - save_start, save_end, save_everystep, Val(SaveFields), Val(SaveWork), maxiters + sols = _prepare( + prob, trajectories, alg, dt, save_start, save_end, save_everystep, + Val(SaveFields), Val(SaveWork) ) irange = 1:trajectories _dispatch_boris!( - sols, prob, irange, savestepinterval, dt, nt, nout, isoutside, n, N, - save_start, save_end, save_everystep, Val(SaveFields), Val(SaveWork) + sols, prob, irange, alg, dt, savestepinterval, isoutside, + save_start, save_end, save_everystep, Val(SaveFields), Val(SaveWork), + maxiters ) return sols end @inline function _solve( - ::EnsembleThreads, prob::TraceProblem, trajectories, dt, savestepinterval, - isoutside::F, n, N, save_start, save_end, save_everystep, - ::Val{SaveFields}, ::Val{SaveWork}, maxiters, batch_size + ::EnsembleThreads, prob::TraceProblem, trajectories, alg, dt, + savestepinterval::Int, isoutside::F, save_start, save_end, save_everystep, + ::Val{SaveFields}, ::Val{SaveWork}, maxiters::Int, batch_size ) where {SaveFields, SaveWork, F} - sols, nt, - nout = _prepare( - prob, trajectories, dt, savestepinterval, - save_start, save_end, save_everystep, Val(SaveFields), Val(SaveWork), maxiters + sols = _prepare( + prob, trajectories, alg, dt, save_start, save_end, save_everystep, + Val(SaveFields), Val(SaveWork) ) nchunks = Threads.nthreads() Threads.@threads for irange in index_chunks(1:trajectories; n = nchunks) _dispatch_boris!( - sols, prob, irange, savestepinterval, dt, nt, nout, isoutside, n, N, - save_start, save_end, save_everystep, Val(SaveFields), Val(SaveWork) + sols, prob, irange, alg, dt, savestepinterval, isoutside, + save_start, save_end, save_everystep, Val(SaveFields), Val(SaveWork), + maxiters ) end return sols end -""" - _solve_single_boris(prob, i, ...) - -Solve a single trajectory `i` of `prob` for use with `EnsembleDistributed`. - -`_generic_boris!` uses the loop index `i` for two purposes simultaneously: -applying `prob_func(prob, i, false)` to select per-particle initial conditions, -and storing the result at `sols[i]`. For a distributed worker handling only one -trajectory at a time, a 1-element `local_sols` would be out of bounds if `i > 1` -were passed directly. Decoupling the two uses would require threading a storage -offset through the entire `_dispatch_boris!` → `_generic_boris!` call chain. - -Instead, we pre-apply `prob_func` to get the correct IC for trajectory `i` and -wrap the result in a fresh `TraceProblem` with the default (identity) `prob_func`, -so `_generic_boris!` can safely iterate `1:1` without applying `prob_func` again. -The `TraceProblem` construction is a negligible struct copy relative to the -simulation cost and the serialization overhead inherent in `pmap`. -""" function _solve_single_boris( - prob::TraceProblem, i, savestepinterval, dt, nt, nout, isoutside::F, n, N, - save_start, save_end, save_everystep, ::Val{SaveFields}, ::Val{SaveWork} + prob::TraceProblem, i, alg, dt, savestepinterval, isoutside::F, + save_start, save_end, save_everystep, ::Val{SaveFields}, ::Val{SaveWork}, + maxiters ) where {SaveFields, SaveWork, F} new_prob = prob.prob_func(prob, i, false) single_prob = TraceProblem(new_prob.u0, new_prob.tspan, new_prob.p) - sol_type = _get_sol_type(single_prob, dt, Val(SaveFields), Val(SaveWork)) + sol_type = _get_sol_type( + single_prob, iszero(dt) ? zero(eltype(single_prob.tspan)) : dt, + Val(SaveFields), Val(SaveWork) + ) local_sols = Vector{sol_type}(undef, 1) _dispatch_boris!( - local_sols, single_prob, 1:1, savestepinterval, dt, nt, nout, - isoutside, n, N, save_start, save_end, save_everystep, - Val(SaveFields), Val(SaveWork) + local_sols, single_prob, 1:1, alg, dt, savestepinterval, + isoutside, save_start, save_end, save_everystep, + Val(SaveFields), Val(SaveWork), maxiters ) return local_sols[1] end @inline function _solve( - ::EnsembleDistributed, prob::TraceProblem, trajectories, dt, savestepinterval, - isoutside::F, n, N, save_start, save_end, save_everystep, + ::EnsembleDistributed, prob::TraceProblem, trajectories, alg, dt, savestepinterval, + isoutside::F, save_start, save_end, save_everystep, ::Val{SaveFields}, ::Val{SaveWork}, maxiters, batch_size ) where {SaveFields, SaveWork, F} - _, nt, nout = _prepare( - prob, trajectories, dt, savestepinterval, - save_start, save_end, save_everystep, Val(SaveFields), Val(SaveWork), maxiters - ) return pmap(1:trajectories; batch_size = batch_size) do i _solve_single_boris( - prob, i, savestepinterval, dt, nt, nout, isoutside, n, N, - save_start, save_end, save_everystep, Val(SaveFields), Val(SaveWork) + prob, i, alg, dt, savestepinterval, isoutside, + save_start, save_end, save_everystep, Val(SaveFields), Val(SaveWork), + maxiters ) end end @inline function _solve( - ::EnsembleSplitThreads, prob::TraceProblem, trajectories, dt, savestepinterval, - isoutside::F, n, N, save_start, save_end, save_everystep, + ::EnsembleSplitThreads, prob::TraceProblem, trajectories, alg, dt, savestepinterval, + isoutside::F, save_start, save_end, save_everystep, ::Val{SaveFields}, ::Val{SaveWork}, maxiters, batch_size ) where {SaveFields, SaveWork, F} - _, nt, nout = _prepare( - prob, trajectories, dt, savestepinterval, - save_start, save_end, save_everystep, Val(SaveFields), Val(SaveWork), maxiters - ) - # _solve_single_boris wraps each trajectory in a fresh TraceProblem(u0, tspan, p) - # with DEFAULT_PROB_FUNC. We get a sample problem from prob_func to ensure the - # u0 type (uType) matches. sample_prob = prob.prob_func(prob, 1, false) dummy_prob = TraceProblem(sample_prob.u0, sample_prob.tspan, sample_prob.p) - sol_type = _get_sol_type(dummy_prob, dt, Val(SaveFields), Val(SaveWork)) + sol_type = _get_sol_type(dummy_prob, zero(eltype(dummy_prob.tspan)), Val(SaveFields), Val(SaveWork)) ichunks = index_chunks(1:trajectories; size = batch_size) results = pmap(ichunks) do irange local_sols = Vector{sol_type}(undef, length(irange)) Threads.@threads for k in eachindex(irange) i = irange[k] local_sols[k] = _solve_single_boris( - prob, i, savestepinterval, dt, nt, nout, isoutside, n, N, - save_start, save_end, save_everystep, Val(SaveFields), Val(SaveWork) + prob, i, alg, dt, savestepinterval, isoutside, + save_start, save_end, save_everystep, Val(SaveFields), Val(SaveWork), + maxiters ) end local_sols @@ -296,7 +579,6 @@ function _get_sol_type(prob, dt, ::Val{SaveFields}, ::Val{SaveWork}) where {Save tspan = prob.tspan T_t = typeof(tspan[1] + dt) t = Vector{T_t}(undef, 0) - # Force u to be Vector{SVector{6, T}} as used in _boris! T = eltype(u0) n_vars = 6 @@ -315,331 +597,45 @@ function _get_sol_type(prob, dt, ::Val{SaveFields}, ::Val{SaveWork}) where {Save return typeof(sol) end -""" -Prepare for advancing. -""" function _prepare( - prob::TraceProblem, trajectories, dt, savestepinterval, - save_start, save_end, save_everystep, ::Val{SaveFields}, ::Val{SaveWork}, maxiters + prob::TraceProblem, trajectories, alg, dt, + save_start, save_end, save_everystep, ::Val{SaveFields}, ::Val{SaveWork} ) where {SaveFields, SaveWork} - if abs(dt) < 10 * eps(typeof(dt)) - throw(ArgumentError("time step dt is too small, violating min_dt = 10 * eps(typeof(dt))")) - end - ttotal = prob.tspan[2] - prob.tspan[1] - nt = round(Int, ttotal / dt) |> abs - if nt > maxiters - throw(ArgumentError("number of iterations nt ($nt) exceeds maxiters ($maxiters)")) - end - - nout = 0 - if save_start - nout += 1 - end - - if save_everystep - steps = nt ÷ savestepinterval - last_is_step = (nt > 0) && (nt % savestepinterval == 0) - nout += steps - if !save_end && last_is_step - nout -= 1 - end - if save_end && !last_is_step - nout += 1 - end - elseif save_end - nout += 1 - end - - sol_type = _get_sol_type(prob, dt, Val(SaveFields), Val(SaveWork)) + sol_type = _get_sol_type( + prob, iszero(dt) ? zero(eltype(prob.tspan)) : dt, + Val(SaveFields), Val(SaveWork) + ) sols = Vector{sol_type}(undef, trajectories) - return sols, nt, nout + return sols end @inline function _prepare_saved_data(xv, p, t, ::Val{SaveFields}, ::Val{SaveWork}) where {SaveFields, SaveWork} data = xv - - # Pre-declare variables to share between blocks if both are true - local E_field, magnetic_props - - if SaveFields + if SaveFields || SaveWork r = get_x(xv) T = eltype(xv) + Bfunc = get_BField(p) E_field = SVector{3, T}(get_EField(p)(r, t)) - # We need magnetic properties for work, so if SaveWork is also true, compute them now if SaveWork - q2m, m, Efunc, Bfunc, _ = p # get_magnetic_properties returns (B, ∇B, κ, b̂, Bmag) magnetic_props = get_magnetic_properties(r, t, Bfunc) - B_vec = SVector{3, T}(magnetic_props[1]) - data = vcat(data, E_field, B_vec) - else - B_vec = SVector{3, T}(get_BField(p)(r, t)) - data = vcat(data, E_field, B_vec) - end - end - - if SaveWork - # If SaveFields was true, we already computed E_field and magnetic_props - if SaveFields + if SaveFields + B_vec = SVector{3, T}(magnetic_props[1]) + data = vcat(data, E_field, B_vec) + end work = get_work_rates(xv, p, t, magnetic_props, E_field) - else - work = get_work_rates(xv, p, t) + data = vcat(data, work) + elseif SaveFields + B_vec = SVector{3, T}(Bfunc(r, t)) + data = vcat(data, E_field, B_vec) end - data = vcat(data, work) end return data end -""" -Apply Boris method for particles with index in `irange`. -""" -@inline @muladd function _boris_loop!( - traj, tsave, iout, r, v, p, dt, nt, tspan, - savestepinterval, save_everystep, isoutside::F1, velocity_updater::F2, - ::Val{SaveFields}, ::Val{SaveWork} - ) where {F1, F2, SaveFields, SaveWork} - it = 1 - t = tspan[1] - 0.5 * dt - while it <= nt - v_prev = v - t += dt - v = velocity_updater(v, r, dt, t, p) - - r_next = r + v * dt - t_next = t + 0.5 * dt - if isoutside(vcat(r_next, v), p, t_next) - return it - 1, iout, r, v_prev - end - - if save_everystep && (it - 1) > 0 && (it - 1) % savestepinterval == 0 - iout += 1 - if iout <= length(traj) - t_current = t - 0.5 * dt - v_save = velocity_updater(v_prev, r, 0.5 * dt, t_current, p) - data = vcat(r, v_save) - traj[iout] = _prepare_saved_data(data, p, t_current, Val(SaveFields), Val(SaveWork)) - tsave[iout] = t_current - end - end - - r = r_next - it += 1 - end - return it - 1, iout, r, v -end - -@inline @muladd function _generic_boris!( - sols, prob::TraceProblem, irange, savestepinterval, dt, nt, nout, isoutside::F1, - save_start, save_end, save_everystep, ::Val{SaveFields}, ::Val{SaveWork}, - velocity_updater::F2, alg_name - ) where {SaveFields, SaveWork, F1, F2} - (; tspan, p, u0) = prob - T = eltype(u0) - - vars_dim = 6 - if SaveFields - vars_dim += 6 - end - if SaveWork - vars_dim += 4 - end - - @inbounds for i in irange - traj = Vector{SVector{vars_dim, T}}(undef, nout) - tsave = Vector{typeof(tspan[1] + dt)}(undef, nout) - - # set initial conditions for each trajectory i - iout = 0 - new_prob = prob.prob_func(prob, i, false) - u0_i = SVector{6, T}(new_prob.u0) - r = u0_i[SVector(1, 2, 3)] - v = u0_i[SVector(4, 5, 6)] - - if save_start - iout += 1 - traj[iout] = _prepare_saved_data(u0_i, p, tspan[1], Val(SaveFields), Val(SaveWork)) - tsave[iout] = tspan[1] - end - - # push velocity back in time by 1/2 dt - v = velocity_updater(v, r, -0.5 * dt, tspan[1], p) - - it, iout, r, v = _boris_loop!( - traj, tsave, iout, r, v, p, dt, nt, tspan, - savestepinterval, save_everystep, isoutside, velocity_updater, - Val(SaveFields), Val(SaveWork) - ) - - final_step = it - should_save_final = false - if save_end - should_save_final = true - elseif save_everystep && (final_step > 0) && (final_step % savestepinterval == 0) - should_save_final = true - end - - if iout < nout && should_save_final - t_final = final_step == nt ? tspan[2] : tspan[1] + final_step * dt - if iout == 0 || tsave[iout] < t_final - iout += 1 - dt_final = t_final - (tspan[1] + (final_step - 0.5) * dt) - v_final = velocity_updater(v, r, dt_final, t_final, p) - - data = vcat(r, v_final) - traj[iout] = _prepare_saved_data( - data, p, t_final, Val(SaveFields), Val(SaveWork) - ) - tsave[iout] = t_final - end - end - - if iout < nout - resize!(traj, iout) - resize!(tsave, iout) - retcode = ReturnCode.Terminated - else - retcode = ReturnCode.Success - end - traj_save = traj - t = tsave - - alg = alg_name - interp = LinearInterpolation(t, traj_save) - stats = nothing - - sols[i] = build_solution(prob, alg, t, traj_save; interp, retcode, stats) - end - - return -end - -""" -Apply Boris method for particles with index in `irange`. -""" -@inline @muladd function _boris!( - sols, prob::TraceProblem, irange, savestepinterval, dt, nt, nout, isoutside::F, - save_start, save_end, save_everystep, ::Val{SaveFields}, ::Val{SaveWork} - ) where {SaveFields, SaveWork, F} - - _generic_boris!( - sols, prob, irange, savestepinterval, dt, nt, nout, isoutside, - save_start, save_end, save_everystep, Val(SaveFields), Val(SaveWork), - update_velocity, :boris - ) - - return -end - -""" - update_velocity_multistep(v, r, dt, t, n, N, param) - -Update velocity using the Multistep/Hyper Boris method, returning the new velocity as an SVector. -`n` specifies the number of subcycles. -`N` specifies the gyrophase correction order. When N=2, it corresponds to the Multicycle solver. When N=4 or N=6, it is the Hyper Boris solver. -Reference: [Zenitani & Kato 2025](https://arxiv.org/abs/2505.02270) -""" -@muladd function update_velocity_multistep(v, r, dt, t, n::Int, N::Int, param) - q2m, _, Efunc, Bfunc, _ = param - E = Efunc(r, t) - B = Bfunc(r, t) - - # t_n and e_n vectors - factor = q2m * dt / (2 * n) - - t_n = factor * B # (q/m * dt/(2n)) * B - e_n = factor * E # (q/m * dt/(2n)) * E - - # Hyper Boris N-th order gyrophase correction - if N != 2 - t_mag2 = sum(abs2, t_n) - if N == 4 - f_N = 1 + t_mag2 / 3 - e_corr_factor = -1 / 3 - else # N == 6 - f_N = 1 + t_mag2 / 3 + 2 * t_mag2 * t_mag2 / 15 - e_corr_factor = -1 / 3 - 2 * t_mag2 / 15 - end - - e_dot_t = e_n ⋅ t_n - e_n = f_N * e_n + (e_corr_factor * e_dot_t) * t_n - t_n = f_N * t_n - end - - t_n_mag2 = sum(abs2, t_n) - t_n_mag = sqrt(t_n_mag2) - - # Calculate coefficients - # Check for small t_n to avoid division by zero or precision loss - if t_n_mag < TN_MAG_THRESHOLD - # Taylor expansion limits as t_n -> 0 - c_n1 = 1 - 2 * n * n * t_n_mag2 - - n_term1 = 2 * n - n_term3 = 4 * n * n * n - - c_n2 = n_term1 - (n_term1 + n_term3) / 3 * t_n_mag2 - c_n3 = 2 * n * n - (4 * n * n + 2 * n * n * n * n) / 3 * t_n_mag2 - c_n6 = (n_term1 + n_term3) / 3 - else - alpha_n = atan(t_n_mag) - n_alpha_n = n * alpha_n - sin_n_alpha, cos_n_alpha = sincos(n_alpha_n) - sin_2n_alpha = 2 * sin_n_alpha * cos_n_alpha - cos_2n_alpha = cos_n_alpha * cos_n_alpha - sin_n_alpha * sin_n_alpha - - c_n1 = cos_2n_alpha - c_n2 = sin_2n_alpha / t_n_mag - c_n3 = 2 * sin_n_alpha * sin_n_alpha / t_n_mag2 - c_n6 = (2 * n - c_n2) / t_n_mag2 - end - - c_n4 = c_n2 - c_n5 = c_n3 - - v_dot_t = v ⋅ t_n - e_dot_t = e_n ⋅ t_n - - v_cross_t = v × t_n - e_cross_t = e_n × t_n - - # Update velocity - # Equation 39: - # v_new = c_n1*v + c_n2*(v x t_n) + c_n3*(v . t_n)t_n + c_n4*e_n + c_n5*(e_n x t_n) + c_n6*(e_n . t_n)t_n - v_new = c_n1 * v + - c_n2 * v_cross_t + - c_n3 * v_dot_t * t_n + - c_n4 * e_n + - c_n5 * e_cross_t + - c_n6 * e_dot_t * t_n - - return v_new -end - -@inline @muladd function _multistep_boris!( - sols, prob::TraceProblem, irange, savestepinterval, dt, nt, nout, isoutside::F, - n_steps::Int, N_order::Int, save_start, save_end, save_everystep, ::Val{SaveFields}, ::Val{SaveWork} - ) where {SaveFields, SaveWork, F} - - velocity_updater = (v, r, dt, t, p) -> - update_velocity_multistep(v, r, dt, t, n_steps, N_order, p) - - _generic_boris!( - sols, prob, irange, savestepinterval, dt, nt, nout, isoutside, - save_start, save_end, save_everystep, Val(SaveFields), Val(SaveWork), - velocity_updater, :multistep_boris - ) - - return -end - -""" - get_fields(sol::AbstractODESolution) - -Return the electric and magnetic fields from the solution `sol`. -""" function get_fields(sol::AbstractODESolution) Efunc, Bfunc = _get_field_funcs(sol.prob) @@ -650,7 +646,6 @@ function get_fields(sol::AbstractODESolution) end function _get_field_funcs(prob::TraceGCProblem) - # p = (q, q2m, μ, Efunc, Bfunc) p = prob.p return p[4], p[5] end @@ -660,11 +655,6 @@ function _get_field_funcs(prob) return get_EField(p), get_BField(p) end -""" - get_work(sol::AbstractODESolution) - -Return the work done by the electric field from the solution `sol`. -""" function get_work(sol::AbstractODESolution) return _get_work(sol, sol.prob) end diff --git a/src/precompile.jl b/src/precompile.jl index 70dace791..af5ff1899 100644 --- a/src/precompile.jl +++ b/src/precompile.jl @@ -37,7 +37,7 @@ sol_kernel = solve(prob, CPU(); dt, savestepinterval = 100) # Adaptive Boris - alg_adaptive = AdaptiveBoris(safety = 0.1) + alg_adaptive = Boris(safety = 0.1) sol_adaptive = solve(prob, alg_adaptive)[1] sol_adaptive_sf = solve(prob, alg_adaptive; save_fields = true, save_work = true) diff --git a/test/test_boris.jl b/test/test_boris.jl index 613018af8..7a700f392 100644 --- a/test/test_boris.jl +++ b/test/test_boris.jl @@ -188,7 +188,7 @@ using Distributed @testset "Adaptive Boris" begin # Check constructor default - alg1 = AdaptiveBoris(safety = 0.1) + alg1 = Boris(safety = 0.1) @test alg1.safety == 0.1 x0 = [0.0, 0.0, 0.0] @@ -198,7 +198,7 @@ using Distributed tperiod = abs(TP.get_gyroperiod(0.01; q = TP.qₑ, m = TP.mₑ)) tspan = (0.0, 200 * tperiod) - alg_adaptive = AdaptiveBoris(safety = 0.1) + alg_adaptive = Boris(safety = 0.1) param = prepare(constant_E, gradient_B, species = Electron) prob = TraceProblem(stateinit, tspan, param) @@ -243,11 +243,31 @@ using Distributed @test abs(sol_boris.u[end][4]) < 1.0e5 - 100 # Adaptive Boris - alg_adaptive = AdaptiveBoris(safety = 0.1) + alg_adaptive = Boris(safety = 0.1) sol_adaptive = TP.solve(prob, alg_adaptive)[1] @test abs(sol_adaptive.u[end][4]) < 1.0e5 - 100 end + @testset "Backward tracing" begin + # B = [0, 0, 0.01], E = 0 + param = prepare(ZeroField(), uniform_B2, species = Proton) + u0 = SA[0.0, 0.0, 0.0, 1.0e5, 0.0, 0.0] + # Forward tspan + tspan_fw = (0.0, 1.0e-5) + dt = 1.0e-7 + sol_fw = TP.solve(TraceProblem(u0, tspan_fw, param); dt)[1] + + # Backward tspan + u_end = sol_fw.u[end] + tspan_bw = (1.0e-5, 0.0) + prob_bw = TraceProblem(u_end, tspan_bw, param) + sol_bw = TP.solve(prob_bw; dt = -dt)[1] + + # Should return close to original IC + @test sol_bw.u[end][1:3] ≈ u0[1:3] atol = 1.0e-4 + @test sol_bw.u[end][4:6] ≈ u0[4:6] atol = 1.0e-4 + end + @testset "Output saving flags" begin # Setup x0 = [0.0, 0.0, 0.0] @@ -407,8 +427,7 @@ using Distributed @test abs(work_ms[3]) > 0.0 # Test Adaptive Boris with save_work - # Use simple AdaptiveBoris - alg_adaptive = AdaptiveBoris(safety = 0.1) + alg_adaptive = Boris(safety = 0.1) sol_adaptive = TP.solve(prob, alg_adaptive; save_work = true, save_everystep = true)[1] @test length(sol_adaptive.u[1]) == 10 work_adaptive = sol_adaptive.u[1][7:10] @@ -523,9 +542,9 @@ using Distributed end end - @testset "AdaptiveBoris" begin + @testset "BorisAdaptive" begin tperiod = abs(TP.get_gyroperiod(0.01; q = TP.qₑ, m = TP.mₑ)) - alg_adaptive = AdaptiveBoris(; safety = 0.1) + alg_adaptive = Boris(; safety = 0.1) sols_serial = TP.solve( prob_dist, alg_adaptive, EnsembleSerial(); trajectories diff --git a/test/test_boundary.jl b/test/test_boundary.jl index b320c60af..9255a14f8 100644 --- a/test/test_boundary.jl +++ b/test/test_boundary.jl @@ -45,7 +45,7 @@ using OrdinaryDiffEq: ReturnCode @testset "Adaptive Boris" begin prob = TraceProblem(u0, tspan, p) sol = TP.solve( - prob, AdaptiveBoris(safety = 0.05); + prob, Boris(safety = 0.05); isoutside = callback.condition )[1] @@ -82,13 +82,13 @@ using OrdinaryDiffEq: ReturnCode E_field_fixed(r, t = 0.0) = SA[0.0, 0.0, 0.0] param = prepare(E_field_fixed, B_field_fixed; species = Proton) - @testset "Spatial rejection (AdaptiveBoris)" begin + @testset "Spatial rejection (Boris)" begin isoutside(u, p, t) = u[1] > 0.5 u0 = [0.0, 0.0, 0.0, 1.0e5, 0.0, 0.0] tspan = (0.0, 1.0e-5) prob = TraceProblem(u0, tspan, param) - alg = AdaptiveBoris(safety = 0.1) + alg = Boris(safety = 0.1) sol = TP.solve(prob, alg; isoutside)[1] @test sol.u[end][1] <= 0.5 @test sol.t[end] < tspan[2]