From 6a08c50fa0ed76d9dfd3df995f3ad10b07dcaae2 Mon Sep 17 00:00:00 2001
From: Beforerr <58776897+Beforerr@users.noreply.github.com>
Date: Wed, 30 Oct 2024 19:02:27 -0700
Subject: [PATCH] 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 <hyzhou@umich.edu>
---
 .github/workflows/benchmark.yml |  9 +++---
 src/TestParticle.jl             |  4 +--
 src/equations.jl                | 52 +++++++++++++++++++--------------
 test/runtests.jl                |  6 ++++
 4 files changed, 43 insertions(+), 28 deletions(-)

diff --git a/.github/workflows/benchmark.yml b/.github/workflows/benchmark.yml
index f52515fe..cd821e88 100644
--- a/.github/workflows/benchmark.yml
+++ b/.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
diff --git a/src/TestParticle.jl b/src/TestParticle.jl
index d5015d28..6ef4c749 100644
--- a/src/TestParticle.jl
+++ b/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,
diff --git a/src/equations.jl b/src/equations.jl
index 7a15a160..d8cbc1da 100644
--- a/src/equations.jl
+++ b/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
 
 """
diff --git a/test/runtests.jl b/test/runtests.jl
index df0351e3..63c816cc 100644
--- a/test/runtests.jl
+++ b/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