Relativistic momentum (#200)

* Solve momentum in the relativistic case #198 * Add relativistic benchmark * Add utility methods * Add normalized relativistic demo * Update dimensionless demos --------- Co-authored-by: Beforerr <[email protected]>
henry2004y · Oct 28, 2024 · f1f6197 · f1f6197
1 parent c1034ec
commit f1f6197
Show file tree

Hide file tree

Showing 8 changed files with 210 additions and 166 deletions.
diff --git a/benchmark/benchmarks.jl b/benchmark/benchmarks.jl
@@ -43,25 +43,36 @@ stateinit = [x0..., u0...]
 
 param_analytic = prepare(E_analytic, B_analytic)
 prob_ip = ODEProblem(trace!, stateinit, tspan, param_analytic) # in place
+prob_rel_ip = ODEProblem(trace_relativistic!, stateinit, tspan, param_analytic) # in place
 prob_oop = ODEProblem(trace, SA[stateinit...], tspan, param_analytic) # out of place
-SUITE["trace"]["analytic field"]["in place"] = @benchmarkable solve($prob_ip, Tsit5(); save_idxs=[1,2,3])
-SUITE["trace"]["analytic field"]["out of place"] = @benchmarkable solve($prob_oop, Tsit5(); save_idxs=[1,2,3])
+SUITE["trace"]["analytic field"]["in place"] =
+    @benchmarkable solve($prob_ip, Tsit5(); save_idxs=[1,2,3])
+SUITE["trace"]["analytic field"]["out of place"] =
+    @benchmarkable solve($prob_oop, Tsit5(); save_idxs=[1,2,3])
+SUITE["trace"]["analytic field"]["in place relativistic"] =
+    @benchmarkable solve($prob_rel_ip, Tsit5(); save_idxs=[1,2,3])
 
 param_numeric = prepare(mesh, E_numeric, B_numeric)
 prob_ip = ODEProblem(trace!, stateinit, tspan, param_numeric) # in place
 prob_oop = ODEProblem(trace, SA[stateinit...], tspan, param_numeric) # out of place
 prob_boris = TraceProblem(stateinit, tspan, param_numeric)
 
-SUITE["trace"]["numerical field"]["in place"] = @benchmarkable solve($prob_ip, Tsit5(); save_idxs=[1,2,3])
-SUITE["trace"]["numerical field"]["out of place"] = @benchmarkable solve($prob_oop, Tsit5(); save_idxs=[1,2,3])
-SUITE["trace"]["numerical field"]["Boris"] = @benchmarkable TestParticle.solve($prob_boris; dt=1/7, savestepinterval=10)
-SUITE["trace"]["numerical field"]["Boris ensemble"] = @benchmarkable TestParticle.solve($prob_boris; dt=1/7, savestepinterval=10, trajectories=2)
+SUITE["trace"]["numerical field"]["in place"] =
+    @benchmarkable solve($prob_ip, Tsit5(); save_idxs=[1,2,3])
+SUITE["trace"]["numerical field"]["out of place"] =
+    @benchmarkable solve($prob_oop, Tsit5(); save_idxs=[1,2,3])
+SUITE["trace"]["numerical field"]["Boris"] =
+    @benchmarkable TestParticle.solve($prob_boris; dt=1/7, savestepinterval=10)
+SUITE["trace"]["numerical field"]["Boris ensemble"] =
+    @benchmarkable TestParticle.solve($prob_boris; dt=1/7, savestepinterval=10, trajectories=2)
 
 param_td = prepare(E_td, B_td, F_td)
 prob_ip = ODEProblem(trace!, stateinit, tspan, param_td) # in place
 prob_oop = ODEProblem(trace, SA[stateinit...], tspan, param_td) # out of place
-SUITE["trace"]["time-dependent field"]["in place"] = @benchmarkable solve($prob_ip, Tsit5(); save_idxs=[1,2,3])
-SUITE["trace"]["time-dependent field"]["out of place"] = @benchmarkable solve($prob_oop, Tsit5(); save_idxs=[1,2,3])
+SUITE["trace"]["time-dependent field"]["in place"] =
+    @benchmarkable solve($prob_ip, Tsit5(); save_idxs=[1,2,3])
+SUITE["trace"]["time-dependent field"]["out of place"] =
+    @benchmarkable solve($prob_oop, Tsit5(); save_idxs=[1,2,3])
 
 stateinit_gc, param_gc = TestParticle.prepare_gc(stateinit, E_analytic, B_analytic,
    species=Proton, removeExB=true)

diff --git a/docs/examples/advanced/demo_proton_dipole.jl b/docs/examples/advanced/demo_proton_dipole.jl
@@ -62,9 +62,9 @@ f = DisplayAs.PNG(f) #hide
 # In the above we used Verner's “Most Efficient” 9/8 Runge-Kutta method. Let's check other algorithms.
 
 function get_energy_ratio(sol)
-   vx = getindex.(sol.u, 4)
-   vy = getindex.(sol.u, 5)
-   vz = getindex.(sol.u, 6)
+   vx = @view sol[4,:]
+   vy = @view sol[5,:]
+   vz = @view sol[6,:]
 
    Einit = vx[1]^2 + vy[1]^2 + vz[1]^2
    Eend = vx[end]^2 + vy[end]^2 + vz[end]^2

diff --git a/docs/examples/basics/demo_dimensionless.jl b/docs/examples/basics/demo_dimensionless.jl
@@ -1,27 +1,26 @@
 # ---
 # title: Dimensionless Units
 # id: demo_dimensionless
-# date: 2023-12-14
+# date: 2024-10-28
 # author: "[Hongyang Zhou](https://github.com/henry2004y)"
-# julia: 1.9.4
+# julia: 1.11.1
 # description: Tracing charged particle in dimensionless units
 # ---
 
 # This example shows how to trace charged particles in dimensionless units.
 # After normalization, ``q=1, B=1, m=1`` so that the gyroradius ``r_L = mv_\perp/qB = v_\perp``. All the quantities are given in dimensionless units.
 # If the magnetic field is homogeneous and the initial perpendicular velocity is 4, then the gyroradius is 4.
 # To convert them to the original units, ``v_\perp = 4*U_0`` and ``r_L = 4*l_0``.
-# Check [Demo: single tracing with additional diagnostics](@ref demo_savingcallback) for explaining the unit conversion.
+# Check [Demo: single tracing with additional diagnostics](@ref demo_savingcallback) and [Demo: Dimensionless and Dimensional Tracing](@ref demo_dimensionless_dimensional) for explaining the unit conversion.
 
 # Tracing in dimensionless units is beneficial for many scenarios. For example, MHD simulations do not have intrinsic scales. Therefore, we can do dimensionless particle tracing in MHD fields, and then convert to any scale we would like.
 
-# Now let's demonstrate this with `trace_normalized!`.
+# Now let's demonstrate this with `trace_normalized!` and `trace_relativistic_normalized!`.
 
 import DisplayAs #hide
 using TestParticle
 using TestParticle: qᵢ, mᵢ
-using OrdinaryDiffEq
-
+using OrdinaryDiffEq, StaticArrays
 using CairoMakie
 CairoMakie.activate!(type = "png") #hide
 
@@ -67,6 +66,56 @@ ax = Axis(f[1, 1],
    aspect = DataAspect()
 )
 
-lines!(ax, sol, vars=(1,2))
+lines!(ax, sol, idxs=(1,2))
+
+f = DisplayAs.PNG(f) #hide
+
+# In the relativistic case,
+# * Velocity is normalized by speed of light c, $V = V^\prime c$;
+# * Magnetic field is normalized by the characteristic magnetic field, $B = B^\prime B\_0$;
+# * Electric field is normalized by $E_0 = c\,B_0$, $E = E^\prime E_0$;
+# * Location is normalized by the $L = c / \Omega_0$, where $\Omega_0 = q\,B_0 / m$, and
+# * Time is normalized by $\Omega_0^{-1}$, $t = t^\prime * \Omega_0^{-1}$.
+
+# In the small velocity scenario, it should behave similar to the non-relativistic case:
+
+param = prepare(xu -> SA[0.0, 0.0, 0.0], xu -> SA[0.0, 0.0, 1.0]; species=User)
+tspan = (0.0, π) # half period
+stateinit = [0.0, 0.0, 0.0, 0.01, 0.0, 0.0]
+prob = ODEProblem(trace_relativistic_normalized!, stateinit, tspan, param)
+sol = solve(prob, Vern9())
+
+### Visualization
+f = Figure(fontsize = 18)
+ax = Axis(f[1, 1],
+   title = "Relativistic particle trajectory",
+   xlabel = "X",
+   ylabel = "Y",
+   #limits = (-0.6, 0.6, -1.1, 0.1),
+   aspect = DataAspect()
+)
+
+lines!(ax, sol, idxs=(1,2))
+
+f = DisplayAs.PNG(f) #hide
+
+# In the large velocity scenario, relativistic effect takes place:
+
+param = prepare(xu -> SA[0.0, 0.0, 0.0], xu -> SA[0.0, 0.0, 1.0]; species=User)
+tspan = (0.0, π) # half period
+stateinit = [0.0, 0.0, 0.0, 0.9, 0.0, 0.0]
+prob = ODEProblem(trace_relativistic_normalized!, stateinit, tspan, param)
+sol = solve(prob, Vern9())
+
+### Visualization
+f = Figure(fontsize = 18)
+ax = Axis(f[1, 1],
+   title = "Relativistic particle trajectory",
+   xlabel = "X",
+   ylabel = "Y",
+   aspect = DataAspect()
+)
+
+lines!(ax, sol, idxs=(1,2))
 
 f = DisplayAs.PNG(f) #hide
diff --git a/docs/examples/basics/demo_dimensionless_dimensional.jl b/docs/examples/basics/demo_dimensionless_dimensional.jl
@@ -3,7 +3,7 @@
 # id: demo_dimensionless_dimensional
 # date: 2023-12-19
 # author: "[Hongyang Zhou](https://github.com/henry2004y)"
-# julia: 1.9.4
+# julia: 1.11.1
 # description: Tracing charged particle in both dimensional and dimensionless units. 
 # ---
 
@@ -83,8 +83,8 @@ sol2 = solve(prob2, Vern9(); reltol=1e-4, abstol=1e-6)
 ### Visualization
 f = Figure(fontsize=18)
 ax = Axis(f[1, 1],
-    xlabel = "x",
-    ylabel = "y",
+    xlabel = "x [km]",
+    ylabel = "y [km]",
     aspect = DataAspect(),
 )
 
@@ -93,7 +93,9 @@ lines!(ax, sol1, idxs=(1, 2))
 xp, yp = let trange = range(tspan2..., length=40)
    sol2.(trange, idxs=1) .* l₀, sol2.(trange, idxs=2) .* l₀
 end
-lines!(ax, xp, yp, linestyle=:dashdot, linewidth=5)
+lines!(ax, xp, yp, linestyle=:dashdot, linewidth=5, color=Makie.wong_colors()[2])
+invL = inv(1e3)
+scale!(ax.scene, invL, invL)
 
 f = DisplayAs.PNG(f) #hide
 

diff --git a/src/TestParticle.jl b/src/TestParticle.jl
@@ -19,7 +19,8 @@ export trace!, trace_relativistic!, trace_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
+export get_gyrofrequency, get_gyroperiod, get_gyroradius, get_velocity, get_energy,
+   energy2velocity
 export orbit, monitor
 export TraceProblem
 

diff --git a/src/equations.jl b/src/equations.jl
@@ -80,71 +80,58 @@ function trace(y, p::FullTPTuple, t)
    SVector{6}(dx, dy, dz, dux, duy, duz)
 end
 
-const FTLError = """
-Particle faster than the speed of light!
-
-If the initial velocity is slower than light and
-adaptive timestepping of the solver is turned on, it
-is better to set a small initial stepsize (dt) or
-maximum dt for adaptive timestepping (dtmax).
-
-More details about the keywords of initial stepsize
-can be found in this documentation page:
-https://diffeq.sciml.ai/stable/basics/common_solver_opts/#Stepsize-Control
-"""
-
 """
     trace_relativistic!(dy, y, p::TPTuple, t)
 
-ODE equations for relativistic charged particle moving in static EM field with in-place
-form.
+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
-
-   u2 = y[4]^2 + y[5]^2 + y[6]^2
-   c2 = c^2
-   if u2 ≥ c2
-      throw(DomainError(u2, FTLError))
-   end
-
-   γInv3 = √(1.0 - u2/c2)^3
-   vx, vy, vz = @view y[4:6]
    Ex, Ey, Ez = E(y, t)
    Bx, By, Bz = B(y, t)
 
+   γ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
+   vmag = √(γ²v² / (1 + γ²v²/c^2))
+   vx, vy, vz = vmag * v̂
+
    dy[1], dy[2], dy[3] = vx, vy, vz
-   dy[4] = q2m*γInv3*(vy*Bz - vz*By + Ex)
-   dy[5] = q2m*γInv3*(vz*Bx - vx*Bz + Ey)
-   dy[6] = q2m*γInv3*(vx*By - vy*Bx + Ez)
+   dy[4] = q2m * (vy*Bz - vz*By + Ex)
+   dy[5] = q2m * (vz*Bx - vx*Bz + Ey)
+   dy[6] = q2m * (vx*By - vy*Bx + Ez) 
 
    return
 end
 
 """
-    trace_relativistic(y, p::TPTuple, t) -> SVector{6, Float64}
+    trace_relativistic(y, p::TPTuple, t) -> SVector{6}
 
-ODE equations for relativistic charged particle moving in static EM field with out-of-place
-form.
+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
-
-   u2 = y[4]^2 + y[5]^2 + y[6]^2
-   c2 = c^2
-   if u2 ≥ c2
-      throw(DomainError(u2, FTLError))
-   end
-
-   γInv3 = √(1.0 - u2/c2)^3
-   vx, vy, vz = @view y[4:6]
    Ex, Ey, Ez = E(y, t)
    Bx, By, Bz = B(y, t)
 
+   γv = @views SVector{3, eltype(y)}(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²/c^2))
+   vx, vy, vz = vmag * v̂
+
    dx, dy, dz = vx, vy, vz
-   dux = q2m*γInv3*(vy*Bz - vz*By + Ex)
-   duy = q2m*γInv3*(vz*Bx - vx*Bz + Ey)
-   duz = q2m*γInv3*(vx*By - vy*Bx + Ez)
+   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)
 end
@@ -175,26 +162,27 @@ end
 """
     trace_relativistic_normalized!(dy, y, p::TPNormalizedTuple, t)
 
-Normalized ODE equations for relativistic charged particle moving in static EM field with
-in-place form.
+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
-
-   u2 = y[4]^2 + y[5]^2 + y[6]^2
-   if u2 ≥ 1
-      throw(DomainError(u2, FTLError))
-   end
-
-   γInv3 = √(1.0 - u2)^3
-   vx, vy, vz = @view y[4:6]
+   _, E, B = p
    Ex, Ey, Ez = E(y, t)
    Bx, By, Bz = B(y, t)
 
+   γ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
+   vmag = √(γ²v² / (1 + γ²v²))
+   vx, vy, vz = vmag * v̂
+
    dy[1], dy[2], dy[3] = vx, vy, vz
-   dy[4] = Ω*γInv3*(vy*Bz - vz*By + Ex)
-   dy[5] = Ω*γInv3*(vz*Bx - vx*Bz + Ey)
-   dy[6] = Ω*γInv3*(vx*By - vy*Bx + Ez)
+   dy[4] = vy*Bz - vz*By + Ex
+   dy[5] = vz*Bx - vx*Bz + Ey
+   dy[6] = vx*By - vy*Bx + Ez
 
    return
 end
@@ -209,7 +197,7 @@ requires the full particle trajectory `p.sol`.
 function trace_gc_drifts!(dx, x, p, t)
    q2m, E, B, sol = p
    xu = sol(t)
-   v = @view xu[4:6]
+   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)

diff --git a/src/utility/utility.jl b/src/utility/utility.jl
@@ -104,4 +104,46 @@ end
 function get_gyroperiod(B::AbstractFloat=5e-9; q::AbstractFloat=qᵢ, m::AbstractFloat=mᵢ)
    ω = get_gyrofrequency(B; q, m)
    2π / ω
-end
+end
+
+"Return velocity from relativistic γv in `sol`."
+function get_velocity(sol)
+   v = Array{eltype(sol.u[1]), 2}(undef, 3, length(sol))
+   for is in axes(v, 2)
+      γv = @view sol[4:6, is]
+      γ²v² = γv[1]^2 + γv[2]^2 + γv[3]^2
+      v² = γ²v² / (1 + γ²v²/c^2)
+      γ = 1 / √(1 - v²/c^2)
+      for i in axes(v, 1)
+         v[i,is] = γv[i] / γ
+      end
+   end
+
+   v
+end
+
+"Return the energy [eV] from relativistic `sol`."
+function get_energy(sol::AbstractODESolution; m=mᵢ, q=qᵢ)
+   e = Vector{eltype(sol.u[1])}(undef, length(sol))
+   for i in eachindex(e)
+      γv = @view sol[4:6, i]
+      γ²v² = γv[1]^2 + γv[2]^2 + γv[3]^2
+      v² = γ²v² / (1 + γ²v²/c^2)
+      γ = 1 / √(1 - v²/c^2)
+      e[i] = (γ-1)*m*c^2/abs(q)
+   end
+
+   e
+end
+
+"Calculate the energy [eV] of a relativistic particle from γv."
+function get_energy(γv; m=mᵢ, q=qᵢ)
+   γ²v² = γv[1]^2 + γv[2]^2 + γv[3]^2
+   v² = γ²v² / (1 + γ²v²/c^2)
+   γ = 1 / √(1 - v²/c^2)
+
+   (γ-1)*m*c^2/abs(q)
+end
+
+"Return velocity magnitude from energy in [eV]."
+energy2velocity(Ek; m=mᵢ, q=qᵢ) = c*sqrt(1 - 1/(1+Ek*abs(q)/(m*c^2))^2)