Use SVector in all trace methods (#207)

* Use SVector in all trace methods

* Add memory benchmark output

* Use smaller grid for precompilation

* Recover the zero velocity check

* No line character limit for docstrings

* Increase benchmark running time for reducing noise

* Bump minor version

.github/workflows/benchmark.yml

@@ -47,7 +47,7 @@ jobs:
                 benchpkg ${{ steps.extract-package-name.outputs.package_name }} --rev="${{github.event.repository.default_branch}},${{github.event.pull_request.head.sha}}" --url=${{ github.event.repository.clone_url }} --bench-on="${{github.event.repository.default_branch}}" --output-dir=results/ --tune
             - name: Create markdown table from benchmarks
               run: |
-                benchpkgtable ${{ steps.extract-package-name.outputs.package_name }} --rev="${{github.event.repository.default_branch}},${{github.event.pull_request.head.sha}}" --input-dir=results/ --ratio > table.md
+                benchpkgtable ${{ steps.extract-package-name.outputs.package_name }} --rev="${{github.event.repository.default_branch}},${{github.event.pull_request.head.sha}}" --input-dir=results/ --ratio --mode=time,memory > table.md
                 echo '### Benchmark Results' > body.md
                 echo '' >> body.md
                 echo '' >> body.md

Project.toml

@@ -1,7 +1,7 @@
 name = "TestParticle"
 uuid = "953b605b-f162-4481-8f7f-a191c2bb40e3"
 authors = ["Hongyang Zhou <[email protected]>, and Tiancheng Liu <[email protected]>"]
-version = "0.11.4"
+version = "0.12.0"
 
 [deps]
 ChunkSplitters = "ae650224-84b6-46f8-82ea-d812ca08434e"

benchmark/benchmarks.jl

@@ -36,7 +36,7 @@ mesh = CartesianGrid((length(x)-1, length(y)-1, length(z)-1),
     (x[1], y[1], z[1]),
     (Δx, Δy, Δz))
 
-tspan = (0.0, 1.0)
+tspan = (0.0, 10.0)
 x0 = [0.0, 0.0, 0.0] # initial position, [m]
 u0 = [1.0, 0.0, 0.0] # initial velocity, [m/s]
 stateinit = [x0..., u0...]

src/equations.jl

@@ -6,37 +6,31 @@
 
 ODE equations for charged particle moving in static EM field with in-place form.
 
-ODE equations for charged particle moving in static EM field and external force field with
-in-place form.
+ODE equations for charged particle moving in static EM field and external force field with in-place form.
 """
 function trace!(dy, y, p::TPTuple, t)
-   q2m, E, B = p
+   q2m, Efunc, Bfunc = 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}(Efunc(y, t))
+   B = SVector{3}(Bfunc(y, t))
 
-   dy[1], dy[2], dy[3] = vx, vy, vz
-   # q/m*(E + v × B)
-   dy[4] = q2m*(vy*Bz - vz*By + Ex)
-   dy[5] = q2m*(vz*Bx - vx*Bz + Ey)
-   dy[6] = q2m*(vx*By - vy*Bx + Ez)
+   dy[1:3] = v
+   dy[4:6] = q2m * (v × B + E)
 
    return
 end
 
 function trace!(dy, y, p::FullTPTuple, t)
-   q, m, E, B, F = p
+   q, m, Efunc, Bfunc, Ffunc = p
 
-   vx, vy, vz = @view y[4:6]
-   Ex, Ey, Ez = E(y, t)
-   Bx, By, Bz = B(y, t)
-   Fx, Fy, Fz = F(y, t)
+   v = @views SVector{3}(y[4:6])
+   E = SVector{3}(Efunc(y, t))
+   B = SVector{3}(Bfunc(y, t))
+   F = SVector{3}(Ffunc(y, t))
 
-   dy[1], dy[2], dy[3] = vx, vy, vz
-   dy[4] = (q*(vy*Bz - vz*By + Ex) + Fx) / m
-   dy[5] = (q*(vz*Bx - vx*Bz + Ey) + Fy) / m
-   dy[6] = (q*(vx*By - vy*Bx + Ez) + Fz) / m
+   dy[1:3] = v
+   dy[4:6] = (q * (v × B + E) + F) / m
 
    return
 end
@@ -47,37 +41,30 @@ end
 
 ODE equations for charged particle moving in static EM field with out-of-place form.
 
-ODE equations for charged particle moving in static EM field and external force field with
-out-of-place form.
+ODE equations for charged particle moving in static EM field and external force field with out-of-place form.
 """
 function trace(y, p::TPTuple, t)
-   q2m, E, B = p
-   vx, vy, vz = @view y[4:6]
-   Ex, Ey, Ez = E(y, t)
-   Bx, By, Bz = B(y, t)
-
-   dx, dy, dz = vx, vy, vz
-   # q/m*(E + v × B)
-   dux = q2m*(vy*Bz - vz*By + Ex)
-   duy = q2m*(vz*Bx - vx*Bz + Ey)
-   duz = q2m*(vx*By - vy*Bx + Ez)
-   SVector{6}(dx, dy, dz, dux, duy, duz)
+   q2m, Efunc, Bfunc = p
+   v = @views SVector{3}(y[4:6])
+   E = SVector{3}(Efunc(y, t))
+   B = SVector{3}(Bfunc(y, t))
+
+   dv = q2m * (v × B + E)
+
+   vcat(v, dv)
 end
 
 function trace(y, p::FullTPTuple, t)
-   q, m, E, B, F = p
+   q, m, Efunc, Bfunc, Ffunc = p
 
-   vx, vy, vz = @view y[4:6]
-   Ex, Ey, Ez = E(y, t)
-   Bx, By, Bz = B(y, t)
-   Fx, Fy, Fz = F(y, t)
+   v = @views SVector{3}(y[4:6])
+   E = SVector{3}(Efunc(y, t))
+   B = SVector{3}(Bfunc(y, t))
+   F = SVector{3}(Ffunc(y, t))
 
-   dx, dy, dz = vx, vy, vz
-   dux = (q*(vy*Bz - vz*By + Ex) + Fx) / m
-   duy = (q*(vz*Bx - vx*Bz + Ey) + Fy) / m
-   duz = (q*(vx*By - vy*Bx + Ez) + Fz) / m
+   dv = (q * (v × B + E) + F) / m
 
-   SVector{6}(dx, dy, dz, dux, duy, duz)
+   vcat(v, dv)
 end
 
 """
@@ -86,9 +73,9 @@ end
 ODE equations for relativistic charged particle (x, γv) moving in static EM field with in-place form.
 """
 function trace_relativistic!(dy, y, p::TPTuple, t)
-   q2m, E, B = p
-   Ex, Ey, Ez = E(y, t)
-   Bx, By, Bz = B(y, t)
+   q2m, Efunc, Bfunc = p
+   E = SVector{3}(Efunc(y, t))
+   B = SVector{3}(Bfunc(y, t))
 
    γv = @views SVector{3, eltype(dy)}(y[4:6])
    γ²v² = γv[1]^2 + γv[2]^2 + γv[3]^2
@@ -98,12 +85,10 @@ function trace_relativistic!(dy, y, p::TPTuple, t)
       v̂ = SVector{3, eltype(dy)}(0, 0, 0)
    end
    vmag = √(γ²v² / (1 + γ²v²/c^2))
-   vx, vy, vz = vmag * v̂
-
-   dy[1], dy[2], dy[3] = vx, vy, vz
-   dy[4] = q2m * (vy*Bz - vz*By + Ex)
-   dy[5] = q2m * (vz*Bx - vx*Bz + Ey)
-   dy[6] = q2m * (vx*By - vy*Bx + Ez) 
+   v = vmag * v̂
+
+   dy[1:3] = v
+   dy[4:6] = q2m * (v × B + E)
 
    return
 end
@@ -114,9 +99,9 @@ end
 ODE equations for relativistic charged particle (x, γv) moving in static EM field with out-of-place form.
 """
 function trace_relativistic(y, p::TPTuple, t)
-   q2m, E, B = p
-   Ex, Ey, Ez = E(y, t)
-   Bx, By, Bz = B(y, t)
+   q2m, Efunc, Bfunc = p
+   E = SVector{3}(Efunc(y, t))
+   B = SVector{3}(Bfunc(y, t))
 
    γv = @views SVector{3, eltype(y)}(y[4:6])
    γ²v² = γv[1]^2 + γv[2]^2 + γv[3]^2
@@ -126,14 +111,10 @@ function trace_relativistic(y, p::TPTuple, t)
       v̂ = SVector{3, eltype(y)}(0, 0, 0)
    end
    vmag = √(γ²v² / (1 + γ²v²/c^2))
-   vx, vy, vz = vmag * v̂
-
-   dx, dy, dz = vx, vy, vz
-   dux = q2m * (vy*Bz - vz*By + Ex)
-   duy = q2m * (vz*Bx - vx*Bz + Ey)
-   duz = q2m * (vx*By - vy*Bx + Ez)
+   v = vmag * v̂
+   dv = q2m * (v × B + E)
 
-   SVector{6}(dx, dy, dz, dux, duy, duz)
+   vcat(v, dv)
 end
 
 """
@@ -144,11 +125,11 @@ If the field is in 2D X-Y plane, periodic boundary should be applied for the fie
 the extrapolation function provided by Interpolations.jl.
 """
 function trace_normalized!(dy, y, p::TPNormalizedTuple, t)
-   _, E, B = p
+   _, Efunc, Bfunc = p
 
    v = @views SVector{3}(y[4:6])
-   E = SVector{3}(E(y, t))
-   B = SVector{3}(B(y, t))
+   E = SVector{3}(Efunc(y, t))
+   B = SVector{3}(Bfunc(y, t))
 
    dy[1:3] = v
    dy[4:6] = v × B + E
@@ -162,18 +143,24 @@ end
 Normalized ODE equations for relativistic charged particle (x, γv) moving in static EM field with in-place form.
 """
 function trace_relativistic_normalized!(dy, y, p::TPNormalizedTuple, t)
-   _, E, B = p
-   E = SVector{3}(E(y, t))
-   B = SVector{3}(B(y, t))
+   _, Efunc, Bfunc = p
+   E = SVector{3}(Efunc(y, t))
+   B = SVector{3}(Bfunc(y, t))
    γv = @views SVector{3}(y[4:6])
 
    γ²v² = γv[1]^2 + γv[2]^2 + γv[3]^2
-   v̂ = normalize(γv)
+   if γ²v² > eps(eltype(dy))
+      v̂ = normalize(γv)
+   else # no velocity
+      v̂ = SVector{3, eltype(dy)}(0, 0, 0)
+   end
    vmag = √(γ²v² / (1 + γ²v²))
    v = vmag * v̂
 
    dy[1:3] = v
    dy[4:6] = v × B + E
+
+   return
 end
 
 """
@@ -182,56 +169,61 @@ end
 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))
+   _, Efunc, Bfunc = p
+   E = SVector{3}(Efunc(y, t))
+   B = SVector{3}(Bfunc(y, t))
    γv = @views SVector{3}(y[4:6])
 
    γ²v² = γv[1]^2 + γv[2]^2 + γv[3]^2
+   if γ²v² > eps(eltype(y))
+      v̂ = normalize(γv)
+   else # no velocity
+      v̂ = SVector{3, eltype(y)}(0, 0, 0)
+   end
    vmag = √(γ²v² / (1 + γ²v²))
-   v = vmag * normalize(γv)
+   v = vmag * v̂
    dv = v × B + E
 
-   return vcat(v, dv)
+   vcat(v, dv)
 end
 
 """
     trace_gc_drifts!(dx, x, p, t)
 
-Equations for tracing the guiding center using analytical drifts, including the grad-B
-drift, the curvature drift, the ExB drift. Parallel velocity is also added. This expression
-requires the full particle trajectory `p.sol`.
+Equations for tracing the guiding center using analytical drifts, including the grad-B drift, curvature drift, and ExB drift.
+Parallel velocity is also added. This expression requires the full particle trajectory `p.sol`.
 """
 function trace_gc_drifts!(dx, x, p, t)
-   q2m, E, B, sol = p
+   q2m, Efunc, Bfunc, sol = p
    xu = sol(t)
    v = @views SVector{3, eltype(dx)}(xu[4:6])
-   abs_B(x) = norm(B(x))
-   gradient_B = ForwardDiff.gradient(abs_B, x)
-   Bv = B(x)
-   b = normalize(Bv)
+   E = SVector{3}(Efunc(x))
+   B = SVector{3}(Bfunc(x))
+
+   Bmag(x) = √(Bfunc(x) ⋅ Bfunc(x))
+   ∇B = ForwardDiff.gradient(Bmag, x)
+   b = normalize(B)
    v_par = (v ⋅ b) .* b
    v_perp = v - v_par
-   Ω = q2m*norm(Bv)
-   κ = ForwardDiff.jacobian(B, x)*Bv  # B⋅∇B
+   Ω = q2m*norm(B)
+   κ = ForwardDiff.jacobian(Bfunc, x) * B  # B⋅∇B
    ## v⟂^2*(B×∇|B|)/(2*Ω*B^2) + v∥^2*(B×(B⋅∇B))/(Ω*B^3) + (E×B)/B^2 + v∥
-   dx[1:3] = norm(v_perp)^2*(Bv × gradient_B)/(2*Ω*norm(Bv)^2) +
-      norm(v_par)^2*(Bv × κ)/Ω/norm(Bv)^3 + (E(x) × Bv)/norm(Bv)^2 + v_par
+   dx[1:3] = norm(v_perp)^2*(B × ∇B)/(2*Ω*norm(B)^2) +
+      norm(v_par)^2*(B × κ)/Ω/norm(B)^3 + (E × B)/(B ⋅ B) + v_par
 end
 
 """
     trace_gc!(dy, y, p::TPTuple, t)
 
-Guiding center equations for nonrelativistic charged particle moving in static EM field with
-in-place form. Variable `y = (x, y, z, u)`, where `u` is the velocity along the magnetic
-field at (x,y,z).
+Guiding center equations for nonrelativistic charged particle moving in static EM field with in-place form.
+Variable `y = (x, y, z, u)`, where `u` is the velocity along the magnetic field at (x,y,z).
 """
 function trace_gc!(dy, y, p::GCTuple, t)
    q, m, μ, Efunc, Bfunc = p
    q2m = q / m
-   X = @view y[1:3]
-   E = Efunc(X, t)
-   B = Bfunc(X, t)
+   X = @views SVector{3}(y[1:3])
+   E = SVector{3}(Efunc(X, t))
+   B = SVector{3}(Bfunc(X, t))
    b̂ = normalize(B) # unit B field at X
 
    Bmag(x) = √(Bfunc(x) ⋅ Bfunc(x))

src/precompile.jl

@@ -3,9 +3,9 @@
 @setup_workload begin
    @compile_workload begin
       # numerical field parameters
-      x = range(-10, 10, length=15)
-      y = range(-10, 10, length=20)
-      z = range(-10, 10, length=25)
+      x = range(-10, 10, length=4)
+      y = range(-10, 10, length=6)
+      z = range(-10, 10, length=8)
       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]
 

test/runtests.jl

@@ -301,13 +301,17 @@ end
       stateinit = zeros(6)
       prob = ODEProblem(trace_relativistic_normalized!, stateinit, tspan, param)
       sol = solve(prob, Vern6())
-      @test sol[1,end] == 0.0
+      @test sol[1,end] == 0.0 && length(sol) == 6
 
-      # 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)
+      stateinit = SA[0.0, 0.0, 0.0, 0.5, 0.0, 0.0]
+      prob = ODEProblem(trace_relativistic_normalized, stateinit, tspan, param)
       sol = solve(prob, Vern6())
       @test sol[1,end] ≈ 0.38992532495827226
+
+      stateinit = @SVector zeros(6)
+      prob = ODEProblem(trace_relativistic_normalized, stateinit, tspan, param)
+      sol = solve(prob, Vern6())
+      @test sol[1,end] == 0.0 && length(sol) == 6
    end
 
    @testset "normalized fields" begin