From 8ca91dd30a4d443e5ac4029e0dd6ca9c0ef59594 Mon Sep 17 00:00:00 2001 From: Hongyang Zhou Date: Wed, 21 Feb 2024 10:40:54 -0500 Subject: [PATCH] Reorganize package --- src/TestParticle.jl | 635 +---------------------------------- src/equations.jl | 200 +++++++++++ src/prepare.jl | 212 ++++++++++++ src/sampler.jl | 91 +++++ src/utility/interpolation.jl | 95 ++++++ src/utility/utility.jl | 40 +++ 6 files changed, 644 insertions(+), 629 deletions(-) create mode 100644 src/equations.jl create mode 100644 src/prepare.jl create mode 100644 src/sampler.jl create mode 100644 src/utility/interpolation.jl diff --git a/src/TestParticle.jl b/src/TestParticle.jl index e2aa304b..01073260 100644 --- a/src/TestParticle.jl +++ b/src/TestParticle.jl @@ -21,8 +21,6 @@ export Maxwellian, BiMaxwellian export orbit, monitor export TraceProblem -include("utility/utility.jl") - """ Type for the particles, `Proton`, `Electron`, `Ion`, or `User`. """ @@ -33,638 +31,17 @@ Abstract type for tracing solutions. """ abstract type AbstractTraceSolution{T, N, S} <: AbstractODESolution{T, N, S} end -""" -Abstract type for velocity distribution functions. -""" -abstract type VDF end - -""" -Type for Maxwellian velocity distributions. -""" -struct Maxwellian{T<:AbstractFloat} <: VDF - "Bulk velocity" - u0::Vector{T} - "Thermal speed" - uth::T - - function Maxwellian(u0::Vector{T}, uth::T) where T - @assert length(u0) == 3 "Bulk velocity must have length 3!" - new{T}(u0, uth) - end -end - -""" -Type for BiMaxwellian velocity distributions. -""" -struct BiMaxwellian{T<:AbstractFloat, U} <: VDF - "Bulk velocity" - u0::Vector{T} - "Parallel thermal speed" - uthpar::T - "Perpendicular thermal speed" - uthperp::T - "Unit magnetic field" - b0::Vector{U} - - function BiMaxwellian(u0::Vector{T}, upar::T, uperp::T, B::Vector{U}) where - {T <: AbstractFloat, U <: AbstractFloat} - @assert length(u0) == 3 && length(B) == 3 "The field vector must have length 3!" - b0 = B ./ hypot(B...) - new{T, U}(u0, upar, uperp, b0) - end -end - -function getchargemass(species::Species, q::AbstractFloat, m::AbstractFloat) - if species == Proton - q = qᵢ - m = mᵢ - elseif species == Electron - q = qₑ - m = mₑ - end - q, m -end - -function makegrid(grid::CartesianGrid{3, T}) where T - gridmin = coordinates(minimum(grid)) - gridmax = coordinates(maximum(grid)) - Δx = spacing(grid) - - gridx = range(gridmin[1], gridmax[1], step=Δx[1]) - gridy = range(gridmin[2], gridmax[2], step=Δx[2]) - gridz = range(gridmin[3], gridmax[3], step=Δx[3]) - - gridx, gridy, gridz -end - -function makegrid(grid::CartesianGrid{2, T}) where T - gridmin = coordinates(minimum(grid)) - gridmax = coordinates(maximum(grid)) - Δx = spacing(grid) - - gridx = range(gridmin[1], gridmax[1], step=Δx[1]) - gridy = range(gridmin[2], gridmax[2], step=Δx[2]) - - gridx, gridy -end - -""" - getinterp(A, gridx, gridy, gridz, order::Int=1, bc::Int=1) - getinterp(A, gridx, gridy, order::Int=1, bc::Int=1) - -Return a function for interpolating array `A` on the grid given by `gridx`, `gridy`, and -`gridz`. - -# Arguments -- `order::Int=1`: order of interpolation in [1,2,3]. -- `bc::Int=1`: type of boundary conditions, 1 -> NaN, 2 -> periodic. -""" -function getinterp(A, gridx, gridy, gridz, order::Int=1, bc::Int=1) - @assert size(A,1) == 3 && ndims(A) == 4 "Inconsistent 3D force field and grid!" - - Ax = @view A[1,:,:,:] - Ay = @view A[2,:,:,:] - Az = @view A[3,:,:,:] - - itpx, itpy, itpz = _getinterp(Ax, Ay, Az, order, bc) - - interpx = scale(itpx, gridx, gridy, gridz) - interpy = scale(itpy, gridx, gridy, gridz) - interpz = scale(itpz, gridx, gridy, gridz) - - # Return field value at a given location. - function get_field(xu) - r = @view xu[1:3] - - return SA[interpx(r...), interpy(r...), interpz(r...)] - end - - return get_field -end - -function getinterp(A, gridx, gridy, order::Int=1, bc::Int=2) - @assert size(A,1) == 3 && ndims(A) == 3 "Inconsistent 2D force field and grid!" - - Ax = @view A[1,:,:] - Ay = @view A[2,:,:] - Az = @view A[3,:,:] - - itpx, itpy, itpz = _getinterp(Ax, Ay, Az, order, bc) - - interpx = scale(itpx, gridx, gridy) - interpy = scale(itpy, gridx, gridy) - interpz = scale(itpz, gridx, gridy) - - # Return field value at a given location. - function get_field(xu) - r = @view xu[1:2] - - return SA[interpx(r...), interpy(r...), interpz(r...)] - end - - return get_field -end - -function _getinterp(Ax, Ay, Az, order::Int, bc::Int) - gtype = OnCell() - - if bc == 1 - if order == 1 - itpx = extrapolate(interpolate(Ax, BSpline(Linear())), NaN) - itpy = extrapolate(interpolate(Ay, BSpline(Linear())), NaN) - itpz = extrapolate(interpolate(Az, BSpline(Linear())), NaN) - elseif order == 2 - itpx = extrapolate(interpolate(Ax, BSpline(Quadratic())), NaN) - itpy = extrapolate(interpolate(Ay, BSpline(Quadratic())), NaN) - itpz = extrapolate(interpolate(Az, BSpline(Quadratic())), NaN) - elseif order == 3 - itpx = extrapolate(interpolate(Ax, BSpline(Cubic(Line(gtype)))), NaN) - itpy = extrapolate(interpolate(Ay, BSpline(Cubic(Line(gtype)))), NaN) - itpz = extrapolate(interpolate(Az, BSpline(Cubic(Line(gtype)))), NaN) - end - else - bctype = Periodic() - if order == 1 - itpx = extrapolate(interpolate(Ax, BSpline(Linear(bctype))), bctype) - itpy = extrapolate(interpolate(Ay, BSpline(Linear(bctype))), bctype) - itpz = extrapolate(interpolate(Az, BSpline(Linear(bctype))), bctype) - elseif order == 2 - itpx = extrapolate(interpolate(Ax, BSpline(Quadratic(Periodic(gtype)))), bctype) - itpy = extrapolate(interpolate(Ay, BSpline(Quadratic(Periodic(gtype)))), bctype) - itpz = extrapolate(interpolate(Az, BSpline(Quadratic(Periodic(gtype)))), bctype) - elseif order == 3 - itpx = extrapolate(interpolate(Ax, BSpline(Cubic(Periodic(gtype)))), bctype) - itpy = extrapolate(interpolate(Ay, BSpline(Cubic(Periodic(gtype)))), bctype) - itpz = extrapolate(interpolate(Az, BSpline(Cubic(Periodic(gtype)))), bctype) - end - end - - itpx, itpy, itpz -end - -"Judge whether the field function is time dependent." -function is_time_dependent(f::Function) - applicable(f, zeros(6), 0.0) ? true : false -end - -abstract type AbstractField{itd} end - -""" - Field{itd, F} <: AbstractField{itd} - -A representation of a field function `f`, defined by: - -time-independent field -```math -\\mathbf{F} = F(\\mathbf{x}), -``` - -time-dependent field -```math -\\mathbf{F} = F(\\mathbf{x}, t). -``` - -# Arguments -- `field_function::Function`: the function of field. -- `itd::Bool`: whether the field function is time dependent. -- `F`: the type of `field_function`. -""" -struct Field{itd, F} <: AbstractField{itd} - field_function::F - Field{itd, F}(field_function::F) where {itd, F} = - isa(itd, Bool) ? new(field_function) : throw(ArgumentError("itd must be a boolean.")) -end - -Field(f::Function) = Field{is_time_dependent(f), typeof(f)}(f) - -(f::AbstractField{true})(xu, t) = f.field_function(xu, t) -(f::AbstractField{true})(xu) = throw(ArgumentError("Time-dependent field function must have a time argument.")) -(f::AbstractField{false})(xu, t) = f.field_function(xu) -(f::AbstractField{false})(xu) = f.field_function(xu) - -"The type of parameter tuple for full test particle problem." -FullTPTuple = Tuple{Float64, Float64, AbstractField, AbstractField, AbstractField} - -"The type of parameter tuple for normal test particle problem." -TPTuple = Tuple{Float64, AbstractField, AbstractField} - -"The type of parameter tuple for normalized test particle problem." -TPNormalizedTuple = Tuple{AbstractFloat, AbstractField, AbstractField} - -""" - sample(vdf::Maxwellian, nparticles::Int) - -Sample velocities from a [`Maxwellian`](@ref) distribution `vdf` with `npoints`. - - sample(vdf::BiMaxwellian, nparticles::Int) - -Sample velocities from a [`BiMaxwellian`](@ref) distribution `vdf` with `npoints`. -""" -function sample(vdf::Maxwellian, nparticles::Int) - sqr2 = typeof(vdf.uth)(√2) - # Convert from thermal speed to std - σ = vdf.uth / sqr2 - v = σ .* randn(typeof(vdf.uth), 3, nparticles) .+ vdf.u0 -end - -function sample(vdf::BiMaxwellian{T, U}, nparticles::Int) where {T, U} - sqr2 = T(√2) - # Convert from thermal speed to std - σpar = vdf.uthpar / sqr2 - σperp = vdf.uthperp / sqr2 - # Transform to Cartesian grid - v = fill(vdf.u0, (3, nparticles)) - vrand = σpar .* randn(T, nparticles) - vrand = reshape(vrand, 1, nparticles) - vpar = repeat(vrand, outer=3) - @inbounds for i in 1:3, ip in 1:nparticles - vpar[i,ip] = vpar[i,ip]*vdf.b0[i] + vdf.u0[i] - end - # Sample vectors on a 2D plane - μ = zeros(SVector{2,T}) - σ = SA[σperp 0; 0 σperp] - d = MvNormal(μ, σ) - vrand = rand(d, nparticles) - - vperp = zeros(T, (3, nparticles)) - # Rotate vectors to be perpendicular to b̂ - k = SVector{3, T}(0, 0, 1) - axis = vdf.b0 × k::SVector{3, T} - θ = acos(vdf.b0 ⋅ k) - R = get_rotation_matrix(axis, θ) - @inbounds for ip in 1:nparticles - vperp[:,ip] = R * SA[vrand[1,ip], vrand[2,ip], 0] - end - - v = vpar .+ vperp -end - -""" - prepare(grid::CartesianGrid, E, B; species=Proton) -> (q2m, E, B) - -Return a tuple consists of particle charge-mass ratio for a prescribed `species` and -interpolated EM field functions. - -# keywords -- `order::Int=1`: order of interpolation in [1,2,3]. -- `bc::Int=1`: type of boundary conditions, 1 -> NaN, 2 -> periodic. - - prepare(grid::CartesianGrid, E, B, F; species=Proton, q=1.0, m=1.0) -> (q, m, E, B, F) - -Return a tuple consists of particle charge, mass for a prescribed `species` of charge `q` -and mass `m`, interpolated EM field functions, and external force `F`. - - prepare(x::AbstractRange, y::AbstractRange, z::AbstractRange, E, B) -> (q2m, E, B) - prepare(x, y, E, B) -> (q2m, E, B) - -Direct range input for uniform grid in 2/3D is also accepted. - - prepare(E, B; species=Proton, q=1.0, m=1.0) -> (q2m, E, B) - -Return a tuple consists of particle charge-mass ratio for a prescribed `species` of charge -`q` and mass `m` and analytic EM field functions. Prescribed `species` are `Electron` and -`Proton`; other species can be manually specified with `species=Ion/User`, `q` and `m`. - - prepare(E, B, F; species=Proton, q=1.0, m=1.0) -> (q, m, E, B, F) - -Return a tuple consists of particle charge, mass for a prescribed `species` of charge `q` -and mass `m`, analytic EM field functions, and external force `F`. -""" -function prepare(grid::CartesianGrid, E::TE, B::TB; species::Species=Proton, - q::AbstractFloat=1.0, m::AbstractFloat=1.0, order::Int=1, bc::Int=1) where {TE, TB} - - q, m = getchargemass(species, q, m) - - if embeddim(grid) == 3 - gridx, gridy, gridz = makegrid(grid) - E = TE <: AbstractArray ? getinterp(E, gridx, gridy, gridz, order, bc) : E - B = TB <: AbstractArray ? getinterp(B, gridx, gridy, gridz, order, bc) : B - elseif embeddim(grid) == 2 - gridx, gridy = makegrid(grid) - E = TE <: AbstractArray ? getinterp(E, gridx, gridy, order, bc) : E - B = TB <: AbstractArray ? getinterp(B, gridx, gridy, order, bc) : B - end - - q/m, Field(E), Field(B) -end - -function prepare(grid::CartesianGrid, E::TE, B::TB, F::TF; species::Species=Proton, - q::AbstractFloat=1.0, m::AbstractFloat=1.0, order::Int=1, bc::Int=1) where {TE, TB, TF} - - q, m = getchargemass(species, q, m) - - gridx, gridy, gridz = makegrid(grid) - - E = TE <: AbstractArray ? getinterp(E, gridx, gridy, gridz, order, bc) : E - B = TB <: AbstractArray ? getinterp(B, gridx, gridy, gridz, order, bc) : B - F = TF <: AbstractArray ? getinterp(F, gridx, gridy, gridz, order, bc) : F - - q, m, Field(E), Field(B), Field(F) -end - -function prepare(x::T, y::T, E::TE, B::TB; species::Species=Proton, q::AbstractFloat=1.0, - m::AbstractFloat=1.0, order::Int=1, bc::Int=1) where {T<:AbstractRange, TE, TB} - - q, m = getchargemass(species, q, m) - - E = TE <: AbstractArray ? getinterp(E, x, y, order, bc) : E - B = TB <: AbstractArray ? getinterp(B, x, y, order, bc) : B - - q/m, Field(E), Field(B) -end - -function prepare(x::T, y::T, z::T, E::TE, B::TB; - species::Species=Proton, q::AbstractFloat=1.0, m::AbstractFloat=1.0, order::Int=1, - bc::Int=1) where {T<:AbstractRange, TE, TB} - - q, m = getchargemass(species, q, m) - - E = TE <: AbstractArray ? getinterp(E, x, y, z, order, bc) : E - B = TB <: AbstractArray ? getinterp(B, x, y, z, order, bc) : B - - q/m, Field(E), Field(B) -end - -function prepare(E, B; species::Species=Proton, q::AbstractFloat=1.0, m::AbstractFloat=1.0) - q, m = getchargemass(species, q, m) - - q/m, Field(E), Field(B) -end - -function prepare(E, B, F; species::Species=Proton, q::AbstractFloat=1.0, - m::AbstractFloat=1.0) - q, m = getchargemass(species, q, m) - - q, m, Field(E), Field(B), Field(F) -end - -""" - trace!(dy, y, p::TPTuple, t) - trace!(dy, y, p::FullTPTuple, t) - -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. -""" -function trace!(dy, 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) - - 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) - - return -end - -function trace!(dy, y, p::FullTPTuple, t) - q, m, E, B, F = 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) - - 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 - - return -end - -""" - trace(y, p::TPTuple, t) -> SVector{6, Float64} - trace(y, p::FullTPTuple, t) -> SVector{6, Float64} - -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. -""" -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) -end - -function trace(y, p::FullTPTuple, t) - q, m, E, B, F = 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) - - 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 - - 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. -""" -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) - - 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) - - return -end - -""" - trace_relativistic(y, p::TPTuple, t) -> SVector{6, Float64} - -ODE equations for relativistic charged particle 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) - - 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) - - SVector{6}(dx, dy, dz, dux, duy, duz) -end - -""" - trace_normalized!(dy, y, p::TPNormalizedTuple, t) - -Normalized ODE equations for charged particle moving in static EM field with in-place form. -If the field is in 2D X-Y plane, periodic boundary should be applied for the field in z via -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) - - 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 - - return -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. -""" -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] - Ex, Ey, Ez = E(y, t) - Bx, By, Bz = B(y, t) - - 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) - - return -end - - -""" - guiding_center(xu, param::Union{TPTuple, FullTPTuple}) - -Calculate the coordinates of the guiding center according to the phase space coordinates of a particle. -Reference: https://en.wikipedia.org/wiki/Guiding_center - -A simple definition: -```math -\\mathbf{X}=\\mathbf{x}-m\\frac{\\mathbf{v}\\times\\mathbf{B}}{qB} -``` -""" -function guiding_center(xu, param::TPTuple) - q2m, _, B_field = param - t = xu[end] - v = @view xu[4:6] - Bv = B_field(xu, t) - B = hypot(Bv...) - # unit vector along B - b = Bv./B - # the vector of Larmor radius - ρ = (b×v)./(q2m*B) - X = xu[1:3] - ρ - return X -end - -function guiding_center(xu, param::FullTPTuple) - q, m, _, B_field = param - t = xu[end] - v = @view xu[4:6] - Bv = B_field(xu, t) - B = hypot(Bv...) - # unit vector along B - b = Bv./B - # the vector of Larmor radius - ρ = (b×v)./(q*B/m) - X = xu[1:3] - ρ - return X -end - -""" - get_gc(param::Union{TPTuple, FullTPTuple}) - -Get three functions for plotting the orbit of guiding center. - -For example: -```julia -param = prepare(E, B; species=Proton) -gc = get_gc(params) -# The definitions of stateinit, tspan, E and B are ignored. -prob = ODEProblem(trace!, stateinit, tspan, param) -sol = solve(prob, Vern7(); dt=2e-11) - -f = Figure(fontsize=18) -ax = Axis3(f[1, 1], aspect = :data) -gc_plot(x,y,z,vx,vy,vz) = (gc([x,y,z,vx,vy,vz])...,) -lines!(ax, sol, idxs=(gc_plot, 1, 2, 3, 4, 5, 6)) -``` -""" -function get_gc(param::Union{TPTuple, FullTPTuple}) - gc(xu) = guiding_center(xu, param) -end +include("utility/utility.jl") +include("utility/interpolation.jl") +include("sampler.jl") +include("prepare.jl") +include("equations.jl") +include("pusher.jl") function orbit end function monitor end -include("pusher.jl") include("precompile.jl") end \ No newline at end of file diff --git a/src/equations.jl b/src/equations.jl new file mode 100644 index 00000000..b38e61fb --- /dev/null +++ b/src/equations.jl @@ -0,0 +1,200 @@ +# Tracing equations. + +""" + trace!(dy, y, p::TPTuple, t) + trace!(dy, y, p::FullTPTuple, t) + +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. +""" +function trace!(dy, 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) + + 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) + + return +end + +function trace!(dy, y, p::FullTPTuple, t) + q, m, E, B, F = 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) + + 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 + + return +end + +""" + trace(y, p::TPTuple, t) -> SVector{6, Float64} + trace(y, p::FullTPTuple, t) -> SVector{6, Float64} + +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. +""" +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) +end + +function trace(y, p::FullTPTuple, t) + q, m, E, B, F = 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) + + 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 + + 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. +""" +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) + + 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) + + return +end + +""" + trace_relativistic(y, p::TPTuple, t) -> SVector{6, Float64} + +ODE equations for relativistic charged particle 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) + + 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) + + SVector{6}(dx, dy, dz, dux, duy, duz) +end + +""" + trace_normalized!(dy, y, p::TPNormalizedTuple, t) + +Normalized ODE equations for charged particle moving in static EM field with in-place form. +If the field is in 2D X-Y plane, periodic boundary should be applied for the field in z via +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) + + 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 + + return +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. +""" +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] + Ex, Ey, Ez = E(y, t) + Bx, By, Bz = B(y, t) + + 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) + + return +end \ No newline at end of file diff --git a/src/prepare.jl b/src/prepare.jl new file mode 100644 index 00000000..754fa4c1 --- /dev/null +++ b/src/prepare.jl @@ -0,0 +1,212 @@ +# Construction of tracing parameters. + +"Judge whether the field function is time dependent." +function is_time_dependent(f::Function) + applicable(f, zeros(6), 0.0) ? true : false +end + +abstract type AbstractField{itd} end + +""" + Field{itd, F} <: AbstractField{itd} + +A representation of a field function `f`, defined by: + +time-independent field +```math +\\mathbf{F} = F(\\mathbf{x}), +``` + +time-dependent field +```math +\\mathbf{F} = F(\\mathbf{x}, t). +``` + +# Arguments +- `field_function::Function`: the function of field. +- `itd::Bool`: whether the field function is time dependent. +- `F`: the type of `field_function`. +""" +struct Field{itd, F} <: AbstractField{itd} + field_function::F + Field{itd, F}(field_function::F) where {itd, F} = + isa(itd, Bool) ? new(field_function) : throw(ArgumentError("itd must be a boolean.")) +end + +Field(f::Function) = Field{is_time_dependent(f), typeof(f)}(f) + +(f::AbstractField{true})(xu, t) = f.field_function(xu, t) +(f::AbstractField{true})(xu) = throw(ArgumentError("Time-dependent field function must have a time argument.")) +(f::AbstractField{false})(xu, t) = f.field_function(xu) +(f::AbstractField{false})(xu) = f.field_function(xu) + +"The type of parameter tuple for full test particle problem." +FullTPTuple = Tuple{Float64, Float64, AbstractField, AbstractField, AbstractField} + +"The type of parameter tuple for normal test particle problem." +TPTuple = Tuple{Float64, AbstractField, AbstractField} + +"The type of parameter tuple for normalized test particle problem." +TPNormalizedTuple = Tuple{AbstractFloat, AbstractField, AbstractField} + + +""" + prepare(grid::CartesianGrid, E, B; species=Proton) -> (q2m, E, B) + +Return a tuple consists of particle charge-mass ratio for a prescribed `species` and +interpolated EM field functions. + +# keywords +- `order::Int=1`: order of interpolation in [1,2,3]. +- `bc::Int=1`: type of boundary conditions, 1 -> NaN, 2 -> periodic. + + prepare(grid::CartesianGrid, E, B, F; species=Proton, q=1.0, m=1.0) -> (q, m, E, B, F) + +Return a tuple consists of particle charge, mass for a prescribed `species` of charge `q` +and mass `m`, interpolated EM field functions, and external force `F`. + + prepare(x::AbstractRange, y::AbstractRange, z::AbstractRange, E, B) -> (q2m, E, B) + prepare(x, y, E, B) -> (q2m, E, B) + +Direct range input for uniform grid in 2/3D is also accepted. + + prepare(E, B; species=Proton, q=1.0, m=1.0) -> (q2m, E, B) + +Return a tuple consists of particle charge-mass ratio for a prescribed `species` of charge +`q` and mass `m` and analytic EM field functions. Prescribed `species` are `Electron` and +`Proton`; other species can be manually specified with `species=Ion/User`, `q` and `m`. + + prepare(E, B, F; species=Proton, q=1.0, m=1.0) -> (q, m, E, B, F) + +Return a tuple consists of particle charge, mass for a prescribed `species` of charge `q` +and mass `m`, analytic EM field functions, and external force `F`. +""" +function prepare(grid::CartesianGrid, E::TE, B::TB; species::Species=Proton, + q::AbstractFloat=1.0, m::AbstractFloat=1.0, order::Int=1, bc::Int=1) where {TE, TB} + + q, m = getchargemass(species, q, m) + + if embeddim(grid) == 3 + gridx, gridy, gridz = makegrid(grid) + E = TE <: AbstractArray ? getinterp(E, gridx, gridy, gridz, order, bc) : E + B = TB <: AbstractArray ? getinterp(B, gridx, gridy, gridz, order, bc) : B + elseif embeddim(grid) == 2 + gridx, gridy = makegrid(grid) + E = TE <: AbstractArray ? getinterp(E, gridx, gridy, order, bc) : E + B = TB <: AbstractArray ? getinterp(B, gridx, gridy, order, bc) : B + end + + q/m, Field(E), Field(B) +end + +function prepare(grid::CartesianGrid, E::TE, B::TB, F::TF; species::Species=Proton, + q::AbstractFloat=1.0, m::AbstractFloat=1.0, order::Int=1, bc::Int=1) where {TE, TB, TF} + + q, m = getchargemass(species, q, m) + + gridx, gridy, gridz = makegrid(grid) + + E = TE <: AbstractArray ? getinterp(E, gridx, gridy, gridz, order, bc) : E + B = TB <: AbstractArray ? getinterp(B, gridx, gridy, gridz, order, bc) : B + F = TF <: AbstractArray ? getinterp(F, gridx, gridy, gridz, order, bc) : F + + q, m, Field(E), Field(B), Field(F) +end + +function prepare(x::T, y::T, E::TE, B::TB; species::Species=Proton, q::AbstractFloat=1.0, + m::AbstractFloat=1.0, order::Int=1, bc::Int=1) where {T<:AbstractRange, TE, TB} + + q, m = getchargemass(species, q, m) + + E = TE <: AbstractArray ? getinterp(E, x, y, order, bc) : E + B = TB <: AbstractArray ? getinterp(B, x, y, order, bc) : B + + q/m, Field(E), Field(B) +end + +function prepare(x::T, y::T, z::T, E::TE, B::TB; + species::Species=Proton, q::AbstractFloat=1.0, m::AbstractFloat=1.0, order::Int=1, + bc::Int=1) where {T<:AbstractRange, TE, TB} + + q, m = getchargemass(species, q, m) + + E = TE <: AbstractArray ? getinterp(E, x, y, z, order, bc) : E + B = TB <: AbstractArray ? getinterp(B, x, y, z, order, bc) : B + + q/m, Field(E), Field(B) +end + +function prepare(E, B; species::Species=Proton, q::AbstractFloat=1.0, m::AbstractFloat=1.0) + q, m = getchargemass(species, q, m) + + q/m, Field(E), Field(B) +end + +function prepare(E, B, F; species::Species=Proton, q::AbstractFloat=1.0, + m::AbstractFloat=1.0) + q, m = getchargemass(species, q, m) + + q, m, Field(E), Field(B), Field(F) +end + +""" + guiding_center(xu, param::Union{TPTuple, FullTPTuple}) + +Calculate the coordinates of the guiding center according to the phase space coordinates of a particle. +Reference: https://en.wikipedia.org/wiki/Guiding_center + +A simple definition: +```math +\\mathbf{X}=\\mathbf{x}-m\\frac{\\mathbf{v}\\times\\mathbf{B}}{qB} +``` +""" +function guiding_center(xu, param::TPTuple) + q2m, _, B_field = param + t = xu[end] + v = @view xu[4:6] + Bv = B_field(xu, t) + B = hypot(Bv...) + # unit vector along B + b = Bv./B + # the vector of Larmor radius + ρ = (b × v) ./ (q2m*B) + + X = xu[1:3] - ρ +end + +function guiding_center(xu, param::FullTPTuple) + q, m, _, B_field = param + t = xu[end] + v = @view xu[4:6] + Bv = B_field(xu, t) + B = hypot(Bv...) + # unit vector along B + b = Bv./B + # the vector of Larmor radius + ρ = (b × v) ./ (q*B/m) + + X = xu[1:3] - ρ +end + +""" + get_gc(param::Union{TPTuple, FullTPTuple}) + +Get three functions for plotting the orbit of guiding center. + +For example: +```julia +param = prepare(E, B; species=Proton) +gc = get_gc(params) +# The definitions of stateinit, tspan, E and B are ignored. +prob = ODEProblem(trace!, stateinit, tspan, param) +sol = solve(prob, Vern7(); dt=2e-11) + +f = Figure(fontsize=18) +ax = Axis3(f[1, 1], aspect = :data) +gc_plot(x,y,z,vx,vy,vz) = (gc([x,y,z,vx,vy,vz])...,) +lines!(ax, sol, idxs=(gc_plot, 1, 2, 3, 4, 5, 6)) +``` +""" +function get_gc(param::Union{TPTuple, FullTPTuple}) + gc(xu) = guiding_center(xu, param) +end \ No newline at end of file diff --git a/src/sampler.jl b/src/sampler.jl new file mode 100644 index 00000000..25f1eabd --- /dev/null +++ b/src/sampler.jl @@ -0,0 +1,91 @@ +# Particle sampling + +""" +Abstract type for velocity distribution functions. +""" +abstract type VDF end + +""" +Type for Maxwellian velocity distributions. +""" +struct Maxwellian{T<:AbstractFloat} <: VDF + "Bulk velocity" + u0::Vector{T} + "Thermal speed" + uth::T + + function Maxwellian(u0::Vector{T}, uth::T) where T + @assert length(u0) == 3 "Bulk velocity must have length 3!" + new{T}(u0, uth) + end +end + +""" +Type for BiMaxwellian velocity distributions. +""" +struct BiMaxwellian{T<:AbstractFloat, U} <: VDF + "Bulk velocity" + u0::Vector{T} + "Parallel thermal speed" + uthpar::T + "Perpendicular thermal speed" + uthperp::T + "Unit magnetic field" + b0::Vector{U} + + function BiMaxwellian(u0::Vector{T}, upar::T, uperp::T, B::Vector{U}) where + {T <: AbstractFloat, U <: AbstractFloat} + @assert length(u0) == 3 && length(B) == 3 "The field vector must have length 3!" + b0 = B ./ hypot(B...) + new{T, U}(u0, upar, uperp, b0) + end +end + + +""" + sample(vdf::Maxwellian, nparticles::Int) + +Sample velocities from a [`Maxwellian`](@ref) distribution `vdf` with `npoints`. + + sample(vdf::BiMaxwellian, nparticles::Int) + +Sample velocities from a [`BiMaxwellian`](@ref) distribution `vdf` with `npoints`. +""" +function sample(vdf::Maxwellian, nparticles::Int) + sqr2 = typeof(vdf.uth)(√2) + # Convert from thermal speed to std + σ = vdf.uth / sqr2 + v = σ .* randn(typeof(vdf.uth), 3, nparticles) .+ vdf.u0 +end + +function sample(vdf::BiMaxwellian{T, U}, nparticles::Int) where {T, U} + sqr2 = T(√2) + # Convert from thermal speed to std + σpar = vdf.uthpar / sqr2 + σperp = vdf.uthperp / sqr2 + # Transform to Cartesian grid + v = fill(vdf.u0, (3, nparticles)) + vrand = σpar .* randn(T, nparticles) + vrand = reshape(vrand, 1, nparticles) + vpar = repeat(vrand, outer=3) + @inbounds for i in 1:3, ip in 1:nparticles + vpar[i,ip] = vpar[i,ip]*vdf.b0[i] + vdf.u0[i] + end + # Sample vectors on a 2D plane + μ = zeros(SVector{2,T}) + σ = SA[σperp 0; 0 σperp] + d = MvNormal(μ, σ) + vrand = rand(d, nparticles) + + vperp = zeros(T, (3, nparticles)) + # Rotate vectors to be perpendicular to b̂ + k = SVector{3, T}(0, 0, 1) + axis = vdf.b0 × k::SVector{3, T} + θ = acos(vdf.b0 ⋅ k) + R = get_rotation_matrix(axis, θ) + @inbounds for ip in 1:nparticles + vperp[:,ip] = R * SA[vrand[1,ip], vrand[2,ip], 0] + end + + v = vpar .+ vperp +end \ No newline at end of file diff --git a/src/utility/interpolation.jl b/src/utility/interpolation.jl new file mode 100644 index 00000000..913b5093 --- /dev/null +++ b/src/utility/interpolation.jl @@ -0,0 +1,95 @@ +# Field interpolations. + +""" + getinterp(A, gridx, gridy, gridz, order::Int=1, bc::Int=1) + getinterp(A, gridx, gridy, order::Int=1, bc::Int=1) + +Return a function for interpolating array `A` on the grid given by `gridx`, `gridy`, and +`gridz`. + +# Arguments +- `order::Int=1`: order of interpolation in [1,2,3]. +- `bc::Int=1`: type of boundary conditions, 1 -> NaN, 2 -> periodic. +""" +function getinterp(A, gridx, gridy, gridz, order::Int=1, bc::Int=1) + @assert size(A,1) == 3 && ndims(A) == 4 "Inconsistent 3D force field and grid!" + + Ax = @view A[1,:,:,:] + Ay = @view A[2,:,:,:] + Az = @view A[3,:,:,:] + + itpx, itpy, itpz = _getinterp(Ax, Ay, Az, order, bc) + + interpx = scale(itpx, gridx, gridy, gridz) + interpy = scale(itpy, gridx, gridy, gridz) + interpz = scale(itpz, gridx, gridy, gridz) + + # Return field value at a given location. + function get_field(xu) + r = @view xu[1:3] + + return SA[interpx(r...), interpy(r...), interpz(r...)] + end + + return get_field +end + +function getinterp(A, gridx, gridy, order::Int=1, bc::Int=2) + @assert size(A,1) == 3 && ndims(A) == 3 "Inconsistent 2D force field and grid!" + + Ax = @view A[1,:,:] + Ay = @view A[2,:,:] + Az = @view A[3,:,:] + + itpx, itpy, itpz = _getinterp(Ax, Ay, Az, order, bc) + + interpx = scale(itpx, gridx, gridy) + interpy = scale(itpy, gridx, gridy) + interpz = scale(itpz, gridx, gridy) + + # Return field value at a given location. + function get_field(xu) + r = @view xu[1:2] + + return SA[interpx(r...), interpy(r...), interpz(r...)] + end + + return get_field +end + +function _getinterp(Ax, Ay, Az, order::Int, bc::Int) + gtype = OnCell() + + if bc == 1 + if order == 1 + itpx = extrapolate(interpolate(Ax, BSpline(Linear())), NaN) + itpy = extrapolate(interpolate(Ay, BSpline(Linear())), NaN) + itpz = extrapolate(interpolate(Az, BSpline(Linear())), NaN) + elseif order == 2 + itpx = extrapolate(interpolate(Ax, BSpline(Quadratic())), NaN) + itpy = extrapolate(interpolate(Ay, BSpline(Quadratic())), NaN) + itpz = extrapolate(interpolate(Az, BSpline(Quadratic())), NaN) + elseif order == 3 + itpx = extrapolate(interpolate(Ax, BSpline(Cubic(Line(gtype)))), NaN) + itpy = extrapolate(interpolate(Ay, BSpline(Cubic(Line(gtype)))), NaN) + itpz = extrapolate(interpolate(Az, BSpline(Cubic(Line(gtype)))), NaN) + end + else + bctype = Periodic() + if order == 1 + itpx = extrapolate(interpolate(Ax, BSpline(Linear(bctype))), bctype) + itpy = extrapolate(interpolate(Ay, BSpline(Linear(bctype))), bctype) + itpz = extrapolate(interpolate(Az, BSpline(Linear(bctype))), bctype) + elseif order == 2 + itpx = extrapolate(interpolate(Ax, BSpline(Quadratic(Periodic(gtype)))), bctype) + itpy = extrapolate(interpolate(Ay, BSpline(Quadratic(Periodic(gtype)))), bctype) + itpz = extrapolate(interpolate(Az, BSpline(Quadratic(Periodic(gtype)))), bctype) + elseif order == 3 + itpx = extrapolate(interpolate(Ax, BSpline(Cubic(Periodic(gtype)))), bctype) + itpy = extrapolate(interpolate(Ay, BSpline(Cubic(Periodic(gtype)))), bctype) + itpz = extrapolate(interpolate(Az, BSpline(Cubic(Periodic(gtype)))), bctype) + end + end + + itpx, itpy, itpz +end \ No newline at end of file diff --git a/src/utility/utility.jl b/src/utility/utility.jl index 26cca24c..030a6a3a 100644 --- a/src/utility/utility.jl +++ b/src/utility/utility.jl @@ -12,6 +12,46 @@ include("current_sheet.jl") include("dipole.jl") include("confinement.jl") +""" + getchargemass(species::Species, q::AbstractFloat, m::AbstractFloat) + +Return charge and mass for `species`. if `species = User`, input `q` and `m` are returned. +""" +function getchargemass(species::Species, q::AbstractFloat, m::AbstractFloat) + if species == Proton + q = qᵢ + m = mᵢ + elseif species == Electron + q = qₑ + m = mₑ + end + q, m +end + +"Return uniform range from 2D/3D CartesianGrid." +function makegrid(grid::CartesianGrid{3, T}) where T + gridmin = coordinates(minimum(grid)) + gridmax = coordinates(maximum(grid)) + Δx = spacing(grid) + + gridx = range(gridmin[1], gridmax[1], step=Δx[1]) + gridy = range(gridmin[2], gridmax[2], step=Δx[2]) + gridz = range(gridmin[3], gridmax[3], step=Δx[3]) + + gridx, gridy, gridz +end + +function makegrid(grid::CartesianGrid{2, T}) where T + gridmin = coordinates(minimum(grid)) + gridmax = coordinates(maximum(grid)) + Δx = spacing(grid) + + gridx = range(gridmin[1], gridmax[1], step=Δx[1]) + gridy = range(gridmin[2], gridmax[2], step=Δx[2]) + + gridx, gridy +end + """ set_axes_equal(ax)