Energy conservation test

[henry2004y]

The classical way of performing an energy conserved PIC simulation is to adopt the Boris method (see here and here). Mathmatically, this belongs to the family of symplectic integrators, where energy is conserved.

In OrdinaryDiffEq.jl, most of the symplectic integrators are defined only for DynamicalODEs, as can be found in this page. These integrators require the definition of DynamicalODEProblem only, so we need to split the ODE equation for test particle into two parts. In the general integrator list, only a few symplectic choices are listed: ImplicitMidpoint(), Trapezoid(), and SSPSDIRK2(). As explained by Chris in this thread, most symplectic integrators require a fixed timestep. We cannot use the symplectic integrators designed for DynamicalODEProblem for the general ODEProblem.

Let us test this using a single proton in a uniform magnetic field without the electric field:

Static B field

```julia # Tracing charged particle in a static EM field. # TestParticle.jl v0.8.2 using TestParticle using OrdinaryDiffEq using StaticArrays using LinearAlgebra: × const B₀ = 1e-8 # [T] "f2" function location!(dx, v, x, p::TestParticle.TPTuple, t) dx .= v end "f1" function lorentz!(dv, v, x, p::TestParticle.TPTuple, t) q2m, E, B = p dv .= q2m*(E(x, t) + v × (B(x, t))) end ## Initialize field function uniform_B(x) return SA[0, 0, B₀] end function uniform_E(x) return SA[0.0, 0.0, 0.0] end "Check energy conservation." E(dx, dy, dz) = 1 // 2 * (dx^2 + dy^2 + dz^2) ## Initialize particles x0 = [0.0, 0, 0] v0 = [0.0, 1e2, 0.0] stateinit = [x0..., v0...] tspan_proton = (0.0, 2000.0) param_proton = prepare(uniform_E, uniform_B, species=Proton) ## Solve for the trajectories prob_p = ODEProblem(trace!, stateinit, tspan_proton, param_proton) #prob_p = DynamicalODEProblem(lorentz!, location!, v0, x0, tspan_proton, param_proton) Ωᵢ = TestParticle.qᵢ * B₀ / TestParticle.mᵢ Tᵢ = 2π / Ωᵢ println("Number of gyrations: ", tspan_proton[2] / Tᵢ) function main(prob_p, Tᵢ) schemes = [ (ImplicitMidpoint(), true), (SSPSDIRK2(), true), (RosenbrockW6S4OS(), true), (Vern9(), false), (Vern8(), false), (Trapezoid(), false), (Vern6(), false), (Tsit5(), false), (BS3(), false), #(Nystrom4(), true), # only can be used for DynamicalODEProblem ] for i in eachindex(schemes) @info schemes[i][1] if schemes[i][2] # fixed timestep @time sol = solve(prob_p, schemes[i][1], dt = Tᵢ/15) else @time sol = solve(prob_p, schemes[i][1]) end #energy = map(x -> E(x[1:3]...), sol.u) # for DynamicalODEProblem defined above energy = map(x -> E(x[4:6]...), sol.u) # for ODEProblem energy_diff = (energy[end] - energy[1]) / energy[1] println("Energy diff ratio: ", round(energy_diff, digits=4)) end @info "Boris:" prob = TraceProblem(stateinit, tspan_proton, Tᵢ/15, prob_p.p) @time sol = TestParticle.solve(prob; savestepinterval=10); energy_diff = (E(sol[1].u[4:6,end]...) - E(sol[1].u[4:6,1]...)) / E(sol[1].u[4:6,1]...) println("Energy diff ratio: ", round(energy_diff, digits=4)) end main(prob_p, Tᵢ) println("-------------------------------------------------------") main(prob_p, Tᵢ) ```

First time execution:

Number of gyrations: 304.7332664344071
[ Info: ImplicitMidpoint(; linsolve = nothing, nlsolve = NLNewton{Rational{Int64}, Rational{Int64}, Rational{Int64}, Rational{Int64}}(1//100, 10, 1//5, 1//5, false, true, 0//1), precs = DEFAULT_PRECS, extrapolant = linear,)
  1.782182 seconds (1.79 M allocations: 122.504 MiB, 4.74% gc time, 99.56% compilation time)
Energy diff ratio: -0.0
[ Info: SSPSDIRK2(; linsolve = nothing, nlsolve = NLNewton{Rational{Int64}, Rational{Int64}, Rational{Int64}, Rational{Int64}}(1//100, 10, 1//5, 1//5, false, true, 0//1), precs = DEFAULT_PRECS, smooth_est = true, extrapolant = constant, controller = PI,)
  1.097879 seconds (1.17 M allocations: 80.464 MiB, 98.75% compilation time)
Energy diff ratio: 0.0
[ Info: RosenbrockW6S4OS(; linsolve = nothing, precs = DEFAULT_PRECS,)
  1.782408 seconds (1.41 M allocations: 93.581 MiB, 18.06% gc time, 99.29% compilation time)
Energy diff ratio: 0.0023
[ Info: Vern9(; stage_limiter! = trivial_limiter!, step_limiter! = trivial_limiter!, thread = static(false), lazy = true,)
  2.891703 seconds (3.39 M allocations: 219.685 MiB, 11.00% gc time, 99.96% compilation time)
Energy diff ratio: -0.0059
[ Info: Vern8(; stage_limiter! = trivial_limiter!, step_limiter! = trivial_limiter!, thread = static(false), lazy = true,)
  3.370465 seconds (1.41 M allocations: 88.495 MiB, 99.97% compilation time)
Energy diff ratio: 0.0017
[ Info: Trapezoid(; linsolve = nothing, nlsolve = NLNewton{Rational{Int64}, Rational{Int64}, Rational{Int64}, Rational{Int64}}(1//100, 10, 1//5, 1//5, false, true, 0//1), precs = DEFAULT_PRECS, extrapolant = linear, controller = PI,)
  1.345352 seconds (1.94 M allocations: 132.975 MiB, 4.05% gc time, 98.50% compilation time)
Energy diff ratio: 0.0114
[ Info: Vern6(; stage_limiter! = trivial_limiter!, step_limiter! = trivial_limiter!, thread = static(false), lazy = true,)
  1.420187 seconds (1.01 M allocations: 67.282 MiB, 1.54% gc time, 99.88% compilation time)
Energy diff ratio: -0.0536
[ Info: Tsit5(; stage_limiter! = trivial_limiter!, step_limiter! = trivial_limiter!, thread = static(false),)
  0.749923 seconds (961.48 k allocations: 64.967 MiB, 99.82% compilation time)
Energy diff ratio: 0.2059
[ Info: BS3(; stage_limiter! = trivial_limiter!, step_limiter! = trivial_limiter!, thread = static(false),)
  0.690156 seconds (779.21 k allocations: 53.929 MiB, 99.47% compilation time)
Energy diff ratio: -0.8973
[ Info: Boris:
  0.240403 seconds (299.67 k allocations: 20.174 MiB, 11.22% gc time, 99.89% compilation time)
Energy diff ratio: -0.0

Second time:

[ Info: ImplicitMidpoint(; linsolve = nothing, nlsolve = NLNewton{Rational{Int64}, Rational{Int64}, Rational{Int64}, Rational{Int64}}(1//100, 10, 1//5, 1//5, false, true, 0//1), precs = DEFAULT_PRECS, extrapolant = linear,)
  0.007811 seconds (32.11 k allocations: 2.974 MiB)
Energy diff ratio: -0.0
[ Info: SSPSDIRK2(; linsolve = nothing, nlsolve = NLNewton{Rational{Int64}, Rational{Int64}, Rational{Int64}, Rational{Int64}}(1//100, 10, 1//5, 1//5, false, true, 0//1), precs = DEFAULT_PRECS, smooth_est = true, extrapolant = constant, controller = PI,)
  0.013953 seconds (45.83 k allocations: 4.090 MiB)
Energy diff ratio: 0.0
[ Info: RosenbrockW6S4OS(; linsolve = nothing, precs = DEFAULT_PRECS,)
  0.014184 seconds (45.84 k allocations: 3.464 MiB)
Energy diff ratio: 0.0023
[ Info: Vern9(; stage_limiter! = trivial_limiter!, step_limiter! = trivial_limiter!, thread = static(false), lazy = true,)
  0.001035 seconds (14.80 k allocations: 1.662 MiB)
Energy diff ratio: -0.0059
[ Info: Vern8(; stage_limiter! = trivial_limiter!, step_limiter! = trivial_limiter!, thread = static(false), lazy = true,)
  0.000923 seconds (14.91 k allocations: 1.707 MiB)
Energy diff ratio: 0.0017
[ Info: Trapezoid(; linsolve = nothing, nlsolve = NLNewton{Rational{Int64}, Rational{Int64}, Rational{Int64}, Rational{Int64}}(1//100, 10, 1//5, 1//5, false, true, 0//1), precs = DEFAULT_PRECS, extrapolant = linear, controller = PI,)
  0.022075 seconds (51.05 k allocations: 5.088 MiB)
Energy diff ratio: 0.0114
[ Info: Vern6(; stage_limiter! = trivial_limiter!, step_limiter! = trivial_limiter!, thread = static(false), lazy = true,)
  0.001432 seconds (18.22 k allocations: 2.036 MiB)
Energy diff ratio: -0.0536
[ Info: Tsit5(; stage_limiter! = trivial_limiter!, step_limiter! = trivial_limiter!, thread = static(false),)
  0.001400 seconds (20.16 k allocations: 2.309 MiB)
Energy diff ratio: 0.2059
[ Info: BS3(; stage_limiter! = trivial_limiter!, step_limiter! = trivial_limiter!, thread = static(false),)
  0.002885 seconds (32.01 k allocations: 3.503 MiB)
Energy diff ratio: -0.8973
[ Info: Boris:
  0.000148 seconds (15 allocations: 47.219 KiB)
Energy diff ratio: -0.0

1. Many schemes (including a lot of symplectic integrators) are unstable with a timestep dt = Tᵢ/15: KahanLi8(), DPRKN6(), DPRKN12(), IRKN4(), VelocityVerlet(), McAte5(), McAte2(), SymplecticEuler(), VerletLeapfrog(), Ruth3(), etc. Based on the description in here, my f1 does not satisfy the independency condition.
2. Vern9 has much longer compilation time than the others, probably due to the lack of precompilation. However, it is actually pretty fast.
3. Note the order of the two functions, as well as the order of x and v.
4. Vern8 or Vern9 is probably the best here. This is also consistent with the SciML Benchmarks for nonstiff problems.

Scheme	Energy difference ratio	Time [ms]	Memory [MiB]	fixed time step
Boris	0%	0.1	0.047	true
ImplicitMidpoint	0%	6.1	2.974	true
SSPSDIRK2	0%	11.7	4.090	true
Vern8	0.17%	0.9	1.707	false
RosenbrockW6S4OS	0.23%	11.5	3.464	true
Vern9	-0.59%	1.1	1.662	false
Trapezoid	1.14%	20.2	5.088	false
Vern6	-5.36%	1.1	2.036	false
Tsit5	20.59%	1.2	2.309	false
BS3	-89.73%	2.6	3.503	false

@TCLiuu

[henry2004y]

Another test involves tracing protons in a magnetic mirror:

Magnetic mirror

```julia # Test for energy conservation between different schemes in a magnetic mirror. using TestParticle using TestParticle: getB_mirror using OrdinaryDiffEq using StaticArrays using TestParticleMakie using GLMakie using FieldTracer using Random using Statistics: mean using LinearAlgebra: × using RecursiveArrayTools function location!(dx, x, v, p::TestParticle.TPTuple, t) dx .= v end function lorentz!(dv, x, v, p::TestParticle.TPTuple, t) q, m, E, B = p dv .= q/m*(E(x, t) + v × (B(x, t))) end ## Obtain field # Magnetic bottle parameters in SI units const I1 = 1. # current in the solenoid const N1 = 1 # number of windings const distance = 10. # distance between solenoids const a = 2.5 # radius of each coil # Unfortunately there is no analytical differentiation to that special function! function getB(xu) SVector{3}(getB_mirror(xu[1], xu[2], xu[3], distance, a, I1*N1)) #SVector{3}(0.0, 0.0, 5.88e-8) end function getE(xu) SA[0.0, 0.0, 0.0] end function isoutofdomain(u, p, t) if abs(u[3]) > 6.0 return true else return false end end struct Trajectory{T<:AbstractFloat} x::Vector{T} y::Vector{T} z::Vector{T} vx::Vector{T} vy::Vector{T} vz::Vector{T} end "Uniform sampling on a sphere." function sample_unit_sphere(n::Int=1000, method::Symbol=:standard) x = zeros(n) y = zeros(n) z = zeros(n) if method == :spherical for i in 1:n theta = 2π * rand() phi = acos(1 - 2 * rand()) x[i] = sin(phi) * cos(theta) y[i] = sin(phi) * sin(theta) z[i] = cos(phi) end elseif method == :standard i = 0 while i < n x1, x2, x3 = rand(3) .* 2 .- 1 λ = hypot(x1, x2, x3) if λ ≤ 1 x[i+1] = x1 / λ y[i+1] = x2 / λ z[i+1] = x3 / λ i += 1 end end end x, y, z end function sample_equator(n::Int, rmin::Float64=0.1, rmax::Float64=1.0) x = zeros(n) y = zeros(n) i = 0 while i < n x1, x2 = (rand(2) .* 2 .- 1) .* rmax λ = hypot(x1, x2) if rmin ≤ λ ≤ rmax x[i+1] = x1 y[i+1] = x2 i += 1 end end z = zeros(n) x, y, z end function sample_particles(n::Int=10, energy::Float64=1e-7, rmin::Float64=0.0, rmax::Float64=0.1) x, y, z = sample_equator(n, rmin, rmax) vx, vy, vz = sample_unit_sphere(n) kB = 1.38064852e-23 T = energy * TestParticle.qᵢ / kB vth = √(2kB*T/TestParticle.mᵢ) vx .*= vth vy .*= vth vz .*= vth x, y, z, vx, vy, vz end function trace_particle(; species::TestParticle.Species=Proton, tspan=(0.0, 1.0), trajectories::Int=1, saveat::Float64=0.0) #stateinit = SA[0.0, 0.0, 0.0, 0.0, 0.0, 0.0] stateinit = [0.0, 0.0, 0.0, 0.0, 0.0, 0.0] param = prepare(getE, getB; species) #prob = ODEProblem(TestParticle.trace, stateinit, tspan, param) prob = DynamicalODEProblem(location!, lorentz!, stateinit[1:3], stateinit[4:6], tspan, param) sols = Vector{Trajectory}(undef, trajectories) particles = sample_particles(trajectories) for i in 1:trajectories #u0 = SVector{6}([p[i] for p in particles]) #prob = remake(prob; u0) u0 = [p[i] for p in particles] prob = remake(prob; u0=ArrayPartition(u0[1:3], u0[4:6])) if saveat == 0.0 #sol = solve(prob, Tsit5(); save_idxs=[1,2,3,4,5,6], isoutofdomain, verbose=false) #sol = solve(prob, KahanLi6(); dt=0.1, save_idxs=[1,2,3,4,5,6], isoutofdomain, verbose=false) #sol = solve(prob, Trapezoid(autodiff=false); save_idxs=[1,2,3,4,5,6], isoutofdomain, verbose=false) #sol = solve(prob, Vern6(); save_idxs=[1,2,3,4,5,6], isoutofdomain, verbose=true) sol = solve(prob, Vern9(); save_idxs=[1,2,3,4,5,6], isoutofdomain, verbose=true) else sol = solve(prob, Tsit5(); save_idxs=[1,2,3,4,5,6], isoutofdomain, saveat, verbose=false) end x = getindex.(sol.u, 1) y = getindex.(sol.u, 2) z = getindex.(sol.u, 3) vx = getindex.(sol.u, 4) vy = getindex.(sol.u, 5) vz = getindex.(sol.u, 6) sols[i] = Trajectory(x, y, z, vx, vy, vz) end sols end ##### xrange = range(-2, 2, length=20) yrange = range(-2, 2, length=20) zrange = range(-5, 5, length=20) Bx, By, Bz = let x=xrange, y=yrange, z=zrange Bx = zeros(length(x), length(y), length(z)) By = zeros(length(x), length(y), length(z)) Bz = zeros(length(x), length(y), length(z)) for k in eachindex(z), j in eachindex(y), i in eachindex(x) Bx[i,j,k], By[i,j,k], Bz[i,j,k] = getB([x[i], y[j], z[k]]) end Bx, By, Bz end f = Figure() ax = Axis3(f[1, 1]; aspect = :data ) for i in 0:8 if i == 0 xs, ys, zs = 0.0, 0.0, 0.0 else xs, ys, zs = 1.5*cos(2π*(i-1)/8), 1.5*sin(2π*(i-1)/8), 0.0 end x1, y1, z1 = FieldTracer.trace(Bx, By, Bz, xs, ys, zs, xrange, yrange, zrange; ds=0.1, maxstep=1000) lines!(ax, x1, y1, z1, color=:black) end f tspan = (0.0, 1000.) trajectories = 1 Random.seed!(1234) sols = trace_particle(; species=Proton, tspan, trajectories) KE1 = let vx = [s.vx[1] for s in sols] vy = [s.vy[1] for s in sols] vz = [s.vz[1] for s in sols] KE1 = mean(@. √(vx^2 + vy^2 + vz^2)) end KE2 = let vx = [s.vx[end] for s in sols] vy = [s.vy[end] for s in sols] vz = [s.vz[end] for s in sols] KE2 = mean(@. √(vx^2 + vy^2 + vz^2)) end println("Number of particles: ", trajectories) # energy conservation println("Avg energy: ", KE1, " to ", KE2) println("Change ratio: ", abs(KE2-KE1)/KE1) for s in sols lines!(ax, s.x, s.y, s.z) end f ```

I use DynamicalODEProblem for all to include the test for symplectic integrators like KahanLi6. The parameters chosen here correspond to a gyroperiod of approximately 1 s, so in a total simulation time of 1000 s each particle gyrates for about 1000 rounds.

Scheme	Tsit5	Tsit5(dtmax=0.1)	Trapezoid	Vern6	Vern9	KahanLi6(dt=0.1)	SERK2	ESERK5
Initial energy	4.375	4.375	4.375	4.375	4.375	4.375	4.375	4.375
Final energy	3.743	4.366	4.006	4.334	4.403	7.300	4.400	5.677
change ratio	14.5%	0.2%	8.4%	0.9%	0.6%	66.8%	0.5%	29.7%

---

The numbers in the table above have been updated after merging #76

[henry2004y]

A side finding from performing these tests is that auto-differentiation may fail for complicated functions like ellipke that is used when generating the magnetic field from a coil. In this situation there are three options:

1. Turn off AD and switch to numerical differentiation by setting autodiff=false.
2. Use dual numbers. (I have not found a working solution yet.)
3. Define an analytic Jacobian. (The analytic form may not exist!)

[TCLiuu]

Sorry, I am very interested in this issue, and the results you have presented here seem worthy of further exploration. But recently I need to write a very important article. Therefore, it may take one to two weeks before I can conduct further testing on this topic. Before that, we may not be able to have an in-depth discussion. Please forgive me.

[henry2004y]

No worry, we are not in a rush @TCLiuu!

[henry2004y]

Now I have some better understandings of the issue. For some algorithms like Tsit5 and Vern6, much better energy conservation can be achieved if we set dtmax=T/15, which means the largest time step for an adaptive scheme should not exceed 1/15 of a gyroperiod. However, for many if not all fixed timestep methods, setting the time step to T/15 is still not enough to conserve energy to a decent level. So far, Tsit5, Vern6/7/8/9, and SERK2 are among the best options together with dtmax; if no dtmax is provided, SERK2 is still decent.

The gyroperiod should be estimated from the smallest value, i.e. largest B field area. These are inspired from the reply of Lubos here.

---

When we wish to generalize the above restriction, we need to modify the timesteps. This can be achieved by the DiscreteCallback methods provided in SciMLBase, or indirectly via the DiffEqCallbacks.jl pkg:

using DiffEqCallbacks

# p = (q, m, E, B)
dtFE(u, p, t) = @views p[2] / (2π * p[1] * hypot(p[4](u, t)...))

cb = StepsizeLimiter(dtFE; safety_factor=1 // 15, max_step=true)
sol = solve(prob, Tsit5(); callback=cb, adaptive=false, dt=0.1, isoutofdomain)

This should be included in TestParticle.jl, at least partly.

---

Now the most important remaining question is: why do all symplectic integrators for the DynamicalODEProblem not working properly?