Extend Boris to ensemble problems (#123)

* Allow ensemble prob for Boris
* Add Boris to precompile workflow

docs/examples/basics/demo_boris.jl

@@ -46,7 +46,7 @@ sol2 = solve(prob, Tsit5());
 
 f = Figure(size=(700, 600))
 ax = Axis(f[1, 1], aspect=1)
-@views lines!(ax, traj[1,:], traj[2,:], label="Boris")
+@views lines!(ax, traj[1][1,:], traj[1][2,:], label="Boris")
 lines!(ax, sol1, linestyle=:dashdot, label="Tsit5 fixed")
 lines!(ax, sol2, linestyle=:dot, label="Tsit5 adaptive")
 axislegend(position=:lt, framevisible=false)

docs/examples/basics/demo_ensemble.jl

@@ -1,9 +1,9 @@
 # ---
 # title: Ensemble tracing
 # id: demo_ensemble
-# date: 2023-04-20
+# date: 2024-01-23
 # author: "[Hongyang Zhou](https://github.com/henry2004y)"
-# julia: 1.9.0
+# julia: 1.10.0
 # description: Tracing multiple charged particles in a static EM field
 # ---
 
@@ -64,4 +64,33 @@ for i in eachindex(sols)
    lines!(ax, sols[i], label="$i")
 end
 
+f = DisplayAs.PNG(f) #hide
+
+# We can also solve this problem with the native [Boris pusher](@ref demo_boris).
+# Note that the Boris pusher requires a additional parameters: a fixed timestep, and an output save interval.
+
+dt = 0.1
+savestepinterval = 1
+
+## Solve for the trajectories
+
+paramBoris = BorisMethod(param)
+prob = TraceProblem(stateinit, tspan, dt, paramBoris; prob_func)
+trajs = trace_trajectory(prob; trajectories, savestepinterval)
+
+## Visualization
+
+f = Figure(fontsize = 18)
+ax = Axis3(f[1, 1],
+   title = "Electron trajectories",
+   xlabel = "X",
+   ylabel = "Y",
+   zlabel = "Z",
+   aspect = :data,
+)
+
+for i in eachindex(trajs)
+   @views lines!(ax, trajs[i][1,:], trajs[i][2,:], trajs[i][3,:], label="$i")
+end
+
 f = DisplayAs.PNG(f) #hide

docs/examples/basics/demo_proton_dipole.jl

@@ -118,7 +118,7 @@ dt = 1e-4
 paramBoris = BorisMethod(param)
 prob = TraceProblem(stateinit, tspan, dt, paramBoris)
 traj = trace_trajectory(prob)
-get_energy_ratio(traj)
+get_energy_ratio(traj[1])
 
 # The Boris method requires a fixed time step. It takes about 0.05s and consumes 53 MiB memory. In this specific case, the time step is determined empirically. If we increase the time step to `1e-2` seconds, the trajectory becomes completely off (but the energy is still conserved).
 # Therefore, as a rule of thumb, we should not use the default `Tsit5()` scheme without decreasing `reltol`. Use adaptive `Vern9()` for an unfamiliar field configuration, then switch to more accurate schemes if needed. A more thorough test can be found [here](https://github.com/henry2004y/TestParticle.jl/issues/73).

src/precompile.jl

@@ -1,31 +1,40 @@
 # Precompiling workloads
 
 @setup_workload begin
-   # numerical field parameters
-   x = range(-10, 10, length=15)
-   y = range(-10, 10, length=20)
-   z = range(-10, 10, length=25)
-   B = fill(0.0, 3, length(x), length(y), length(z)) # [T]
-   E = fill(0.0, 3, length(x), length(y), length(z)) # [V/m]
+   @compile_workload begin
+      # numerical field parameters
+      x = range(-10, 10, length=15)
+      y = range(-10, 10, length=20)
+      z = range(-10, 10, length=25)
+      B = fill(0.0, 3, length(x), length(y), length(z)) # [T]
+      E = fill(0.0, 3, length(x), length(y), length(z)) # [V/m]
 
-   B[3,:,:,:] .= 10e-9
-   E[3,:,:,:] .= 5e-10
+      B[3,:,:,:] .= 10e-9
+      E[3,:,:,:] .= 5e-10
 
-   Δx = x[2] - x[1]
-   Δy = y[2] - y[1]
-   Δz = z[2] - z[1]
+      Δx = x[2] - x[1]
+      Δy = y[2] - y[1]
+      Δz = z[2] - z[1]
 
-   mesh = CartesianGrid((length(x)-1, length(y)-1, length(z)-1),
-      (x[1], y[1], z[1]),
-      (Δx, Δy, Δz))
+      mesh = CartesianGrid((length(x)-1, length(y)-1, length(z)-1),
+         (x[1], y[1], z[1]),
+         (Δx, Δy, Δz))
 
-   @compile_workload begin
       vdf = Maxwellian([0.0, 0.0, 0.0], 1.0)
       v = sample(vdf, 2)
       # numerical field
       param = prepare(x, y, z, E, B)
       param = prepare(mesh, E, B)
       # analytical field
       param = prepare(getE_dipole, getB_dipole)
+      # Boris pusher
+      stateinit = let x0 = [0.0, 0.0, 0.0], v0 = [0.0, 1e5, 0.0]
+         [x0..., v0...]
+      end
+      tspan = (0.0, 1.0)
+      dt = 0.5
+      paramBoris = BorisMethod(param)
+      prob = TraceProblem(stateinit, tspan, dt, paramBoris)
+      traj = trace_trajectory(prob; savestepinterval=100)
    end
 end

src/pusher.jl

@@ -1,10 +1,17 @@
 # Native particle pusher
 
-struct TraceProblem{TS, T<:Real, TP}
-   stateinit::TS
+struct TraceProblem{TS, T<:Real, TP, PF}
+   u0::TS
    tspan::Tuple{T, T}
    dt::T
    param::TP
+   prob_func::PF
+end
+
+DEFAULT_PROB_FUNC(prob, i, repeat) = prob
+
+function TraceProblem(u0, tspan, dt, param; prob_func=DEFAULT_PROB_FUNC)
+   TraceProblem(u0, tspan, dt, param, prob_func)
 end
 
 struct BorisMethod{T, TV}
@@ -95,42 +102,49 @@ function cross!(v1, v2, vout)
    return
 end
 
-function trace_trajectory(prob::TraceProblem;
+function trace_trajectory(prob::TraceProblem; trajectories::Int=1, 
    savestepinterval::Int=1, isoutofdomain::Function=ODE_DEFAULT_ISOUTOFDOMAIN)
-   (; stateinit, tspan, dt, param) = prob
-   xv = copy(stateinit)
-   # prepare advancing
-   ttotal = tspan[2] - tspan[1]
-   nt = Int(ttotal ÷ dt)
-   iout, nout = 1, nt ÷ savestepinterval + 1
-   traj = zeros(eltype(stateinit), 6, nout)
-
-   traj[:,1] = xv
-
-   # push velocity back in time by 1/2 dt
-   update_velocity!(xv, param, -0.5*dt)
-
-   for it in 1:nt
-      update_velocity!(xv, param, dt)
-      update_location!(xv, dt)
-      if it % savestepinterval == 0
-      	iout += 1
-      	traj[:,iout] .= xv
+
+   trajs = Vector{Array{eltype(prob.u0),2}}(undef, trajectories)
+
+   for i in 1:trajectories
+      new_prob = prob.prob_func(prob, i, false)
+      (; u0, tspan, dt, param) = new_prob
+      xv = copy(u0)
+      # prepare advancing
+      ttotal = tspan[2] - tspan[1]
+      nt = Int(ttotal ÷ dt)
+      iout, nout = 1, nt ÷ savestepinterval + 1
+      traj = zeros(eltype(u0), 6, nout)
+
+      traj[:,1] = xv
+
+      # push velocity back in time by 1/2 dt
+      update_velocity!(xv, param, -0.5*dt)
+
+      for it in 1:nt
+         update_velocity!(xv, param, dt)
+         update_location!(xv, dt)
+         if it % savestepinterval == 0
+         	iout += 1
+         	traj[:,iout] .= xv
+         end
+         isoutofdomain(xv) && break
       end
-      isoutofdomain(xv) && break
-   end
 
-   if iout == nout # regular termination
-      # final step if needed
-      dtfinal = ttotal - nt*dt
-      if dtfinal > 1e-3
-      	update_velocity!(xv, param, dtfinal)
-      	update_location!(xv, dtfinal)
-      	traj = hcat(traj, xv)
+      if iout == nout # regular termination
+         # final step if needed
+         dtfinal = ttotal - nt*dt
+         if dtfinal > 1e-3
+         	update_velocity!(xv, param, dtfinal)
+         	update_location!(xv, dtfinal)
+         	traj = hcat(traj, xv)
+         end
+      else # early termination
+         traj = traj[:, 1:iout]
       end
-   else # early termination
-      traj = traj[:, 1:iout]
+      trajs[i] = traj
    end
 
-   traj
+   trajs
 end

test/runtests.jl

@@ -383,7 +383,7 @@ end
 
       traj = trace_trajectory(prob; savestepinterval=10)
 
-      @test traj[:, end] == [-7.84237771267459e-5, 5.263661571564935e-5, 0.0,
+      @test traj[1][:, end] == [-7.84237771267459e-5, 5.263661571564935e-5, 0.0,
          -93512.6374393526, -35431.43574759836, 0.0]
    end
 end