From 75ee267cb962f37b8710ad56acd5567321363831 Mon Sep 17 00:00:00 2001 From: Hongyang Zhou Date: Thu, 31 Oct 2024 12:49:17 -0400 Subject: [PATCH] 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 | 2 +- Project.toml | 2 +- benchmark/benchmarks.jl | 2 +- src/equations.jl | 176 +++++++++++++++----------------- src/precompile.jl | 6 +- test/runtests.jl | 12 ++- 6 files changed, 98 insertions(+), 102 deletions(-) diff --git a/.github/workflows/benchmark.yml b/.github/workflows/benchmark.yml index eef1d2b3..9995dcda 100644 --- a/.github/workflows/benchmark.yml +++ b/.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 diff --git a/Project.toml b/Project.toml index 1af84ef5..d7179ac0 100644 --- a/Project.toml +++ b/Project.toml @@ -1,7 +1,7 @@ name = "TestParticle" uuid = "953b605b-f162-4481-8f7f-a191c2bb40e3" authors = ["Hongyang Zhou , and Tiancheng Liu "] -version = "0.11.4" +version = "0.12.0" [deps] ChunkSplitters = "ae650224-84b6-46f8-82ea-d812ca08434e" diff --git a/benchmark/benchmarks.jl b/benchmark/benchmarks.jl index 7e984f62..d8a83ea2 100644 --- a/benchmark/benchmarks.jl +++ b/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...] diff --git a/src/equations.jl b/src/equations.jl index d8cbc1da..ec3f5729 100644 --- a/src/equations.jl +++ b/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)) diff --git a/src/precompile.jl b/src/precompile.jl index 54ede7a9..9fc8a30a 100644 --- a/src/precompile.jl +++ b/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] diff --git a/test/runtests.jl b/test/runtests.jl index 63c816cc..d92e88dc 100644 --- a/test/runtests.jl +++ b/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