Energy Conservation

This example demonstrates the energy conservation of a single particle motion in three cases.

1. Constant B field, Zero E field.

2. Constant E field, Zero B field.

3. Magnetic Mirror.

The tests are performed in dimensionless units with q=1, m=1. We compare multiple solvers: Tsit5, Vern7, Vern9, BS3, ImplicitMidpoint, Boris, and Boris Multistep.

using TestParticle
using TestParticle: ZeroField
using OrdinaryDiffEq
using StaticArrays
using LinearAlgebra: ×, norm
using CairoMakie

const q = 1.0
const m = 1.0
const B₀ = 1.0
const E₀ = 1.0
const Ω = q * B₀ / m
const T = 2π / Ω

# Helper function to run tests
function run_test(case_name, param, x0, v0, tspan, expected_energy_func;
      uselog = true, dt = 0.1, ymin = nothing, ymax = nothing)
   u0 = [x0..., v0...]
   prob_ode = ODEProblem(trace_normalized!, u0, tspan, param)
   prob_tp = TraceProblem(u0, tspan, param)

   f = Figure(size = (1000, 600), fontsize = 18)
   if uselog
      yscale = log10
   else
      yscale = identity
   end

   ax = Axis(f[1, 1],
      title = "$case_name: Energy Error",
      xlabel = "Time",
      ylabel = "Rel. Energy Error |(E - E_ref)/E_ref|",
      yscale = yscale
   )

   if !isnothing(ymin) && !isnothing(ymax)
      ylims!(ax, ymin, ymax)
   end
   function plot_energy_error!(sol, label, marker)
      # Calculate energy
      v_mag = [norm(u[4:6]) for u in sol.u]
      E = 0.5 .* m .* v_mag .^ 2

      # Expected energy
      t = sol.t
      x = @views [u[1:3] for u in sol.u]
      # Pass velocity to expected_energy_func just in case
      E_ref = @views [expected_energy_func(ti, xi, u[4:6])
                      for (ti, xi, u) in zip(t, x, sol.u)]

      # Error (Avoid division by zero if E_ref is 0)
      error = abs.(E .- E_ref) ./ (abs.(E_ref) .+ 1e-16)

      scatter!(ax, t, error, label = label, marker = marker)
   end

   # Run ODE solvers
   for (name, alg, marker) in ode_solvers
      sol = solve(prob_ode, alg; adaptive = false, dt, dense = false)
      plot_energy_error!(sol, name, marker)
   end

   # Run native solvers
   for (name, marker, kwargs) in native_solvers
      sol = TestParticle.solve(prob_tp; dt, kwargs...)[1]
      plot_energy_error!(sol, name, marker)
   end

   f[1, 2] = Legend(f, ax, "Solvers", framevisible = false, labelsize = 24)

   return f
end

# Solvers to test
const ode_solvers = [
   ("Tsit5", Tsit5(), :circle),
   ("Vern7", Vern7(), :rect),
   ("Vern9", Vern9(), :diamond),
   ("BS3", BS3(), :utriangle),
   ("ImplicitMidpoint", ImplicitMidpoint(), :pentagon)
]

const native_solvers = [
   ("Boris", :star5, Dict{Symbol, Any}()),
   ("Boris Multistep (n=2)", :star8, Dict{Symbol, Any}(:n => 2))
];

Case 1: Constant B, Zero E

Energy should be conserved.

uniform_B(x) = SA[0, 0, B₀]

param1 = prepare(ZeroField(), uniform_B; q = q, m = m)
x0_1 = [0.0, 0.0, 0.0]
v0_1 = [1.0, 0.0, 0.0]
tspan1 = (0.0, 50.0)
E_func1(t, x, v) = 0.5 * m * norm(v0_1)^2 # Constant energy

f = run_test(
   "Constant B", param1, x0_1, v0_1, tspan1, E_func1; dt = T / 20, ymin = 1e-16, ymax = 1.0)

Case 2: Constant E, Zero B

Energy increases due to work done by the electric field. For a particle starting from rest in constant E field:

$$\begin{array}{rl}\mathbf{a}& =\frac{q}{m}\mathbf{E}\\ \mathbf{v}(t)& =\mathbf{a}t(\mathrm{if}{v}_0=0)\end{array}$$$$E_{kin}=\frac{1}{2}m{v}^2=\frac{1}{2}m(\frac{q{E}_0}{m}t{)}^2$$

constant_E(x, t) = SA[E₀, 0.0, 0.0]

param2 = prepare(constant_E, ZeroField(); q = q, m = m)
x0_2 = [0.0, 0.0, 0.0]
v0_2 = [0.0, 0.0, 0.0] # Start from rest
tspan2 = (0.0, 40.0)

function E_func2(t, x, v)
   v_theo = (q * E₀ / m) * t # analytical energy
   return 0.5 * m * v_theo^2
end

f = run_test("Constant E", param2, x0_2, v0_2, tspan2,
   E_func2; dt = T / 20, ymin = 1e-16, ymax = 1e-13)

Case 3: Magnetic Mirror

Energy should be conserved (E=0). The particle bounces back and forth between regions of high magnetic field. We set a divergence-free B field in cylindrical symmetry

$$\begin{array}{rl}{B}_x& =-\alpha B_{0}xz\\ B_y& =-\alpha B_{0}yz\\ B_z& =B_0(1+\alpha z^2)\end{array}$$

function mirror_B(x)
   α = 0.1
   Bz = B₀ * (1 + α * x[3]^2)
   Bx = -B₀ * α * x[1] * x[3]
   By = -B₀ * α * x[2] * x[3]
   return SA[Bx, By, Bz]
end

param3 = prepare(ZeroField(), mirror_B; q = q, m = m)
x0_3 = [0.1, 0.0, 0.0]
v0_3 = [0.5, 0.5, 1.0]
tspan3 = (0.0, 200.0)
E_init_3 = 0.5 * m * norm(v0_3)^2
E_func3(t, x, v) = E_init_3

f = run_test("Magnetic Mirror", param3, x0_3, v0_3, tspan3, E_func3;
   dt = T / 20, ymin = 1e-16, ymax = 2.0)