feat: add out-of-place form relativistic equations (#203)

* feat: add out-of-place form normalized relativistic equations

* perf: improve performance with SVector

* Add test trace_relativistic_normalized

---------

Co-authored-by: Hongyang Zhou <[email protected]>

.github/workflows/benchmark.yml

@@ -11,12 +11,13 @@ on:
       - 'docs/**'
       - 'examples/**'
 
-permissions:
-  pull-requests: write
-
 jobs:
-    generate_plots:
+    performance-guard:
         runs-on: ubuntu-latest
+        permissions:
+          contents: write
+          pull-requests: write
+          repository-projects: write
         steps:
             - uses: actions/checkout@v4
             - uses: julia-actions/setup-julia@v2

src/TestParticle.jl

@@ -15,8 +15,8 @@ using ChunkSplitters
 using PrecompileTools: @setup_workload, @compile_workload
 
 export prepare, prepare_gc, sample, get_gc
-export trace!, trace_relativistic!, trace_normalized!, trace, trace_relativistic,
-   trace_relativistic_normalized!, trace_gc!, trace_gc_1st!, trace_gc_drifts!
+export trace!, trace_relativistic!, trace_normalized!, trace_relativistic_normalized!, 
+   trace, trace_relativistic, trace_relativistic_normalized, trace_gc!, trace_gc_1st!, trace_gc_drifts!
 export Proton, Electron, Ion, User
 export Maxwellian, BiMaxwellian
 export get_gyrofrequency, get_gyroperiod, get_gyroradius, get_velocity, get_energy,

src/equations.jl

@@ -146,15 +146,12 @@ the extrapolation function provided by Interpolations.jl.
 function trace_normalized!(dy, y, p::TPNormalizedTuple, t)
    _, E, B = p
 
-   vx, vy, vz = @view y[4:6]
-   Ex, Ey, Ez = E(y, t)
-   Bx, By, Bz = B(y, t)
+   v = @views SVector{3}(y[4:6])
+   E = SVector{3}(E(y, t))
+   B = SVector{3}(B(y, t))
 
-   dy[1], dy[2], dy[3] = vx, vy, vz
-   # E + v × B
-   dy[4] = vy*Bz - vz*By + Ex
-   dy[5] = vz*Bx - vx*Bz + Ey
-   dy[6] = vx*By - vy*Bx + Ez
+   dy[1:3] = v
+   dy[4:6] = v × B + E
 
    return
 end
@@ -166,25 +163,36 @@ Normalized ODE equations for relativistic charged particle (x, γv) moving in st
 """
 function trace_relativistic_normalized!(dy, y, p::TPNormalizedTuple, t)
    _, E, B = p
-   Ex, Ey, Ez = E(y, t)
-   Bx, By, Bz = B(y, t)
+   E = SVector{3}(E(y, t))
+   B = SVector{3}(B(y, t))
+   γv = @views SVector{3}(y[4:6])
 
-   γv = @views SVector{3, eltype(dy)}(y[4:6])
    γ²v² = γv[1]^2 + γv[2]^2 + γv[3]^2
-   if γ²v² > eps(eltype(dy))
-      v̂ = normalize(γv)
-   else # no velocity
-      v̂ = SVector{3, eltype(dy)}(0, 0, 0)
-   end
+   v̂ = normalize(γv)
    vmag = √(γ²v² / (1 + γ²v²))
-   vx, vy, vz = vmag * v̂
+   v = vmag * v̂
 
-   dy[1], dy[2], dy[3] = vx, vy, vz
-   dy[4] = vy*Bz - vz*By + Ex
-   dy[5] = vz*Bx - vx*Bz + Ey
-   dy[6] = vx*By - vy*Bx + Ez
+   dy[1:3] = v
+   dy[4:6] = v × B + E
+end
 
-   return
+"""
+    trace_relativistic_normalized(y, p::TPNormalizedTuple, t)
+
+Normalized ODE equations for relativistic charged particle (x, γv) moving in static EM field with out-of-place form.
+"""
+function trace_relativistic_normalized(y, p::TPNormalizedTuple, t)
+   _, E, B = p
+   E = SVector{3}(E(y, t))
+   B = SVector{3}(B(y, t))
+   γv = @views SVector{3}(y[4:6])
+
+   γ²v² = γv[1]^2 + γv[2]^2 + γv[3]^2
+   vmag = √(γ²v² / (1 + γ²v²))
+   v = vmag * normalize(γv)
+   dv = v × B + E
+
+   return vcat(v, dv)
 end
 
 """

test/runtests.jl

@@ -302,6 +302,12 @@ end
       prob = ODEProblem(trace_relativistic_normalized!, stateinit, tspan, param)
       sol = solve(prob, Vern6())
       @test sol[1,end] == 0.0
+
+      # Test trace_relativistic_normalized
+      stateinit = [0.0, 0.0, 0.0, 0.5, 0.0, 0.0]
+      prob = ODEProblem(trace_relativistic_normalized, SA[stateinit...], tspan, param)
+      sol = solve(prob, Vern6())
+      @test sol[1,end] ≈ 0.38992532495827226
    end
 
    @testset "normalized fields" begin