VDF extension

Project.toml

@@ -22,13 +22,14 @@ SpecialFunctions = "276daf66-3868-5448-9aa4-cd146d93841b"
 StaticArrays = "90137ffa-7385-5640-81b9-e52037218182"
 Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2"
 Tensors = "48a634ad-e948-5137-8d70-aa71f2a747f4"
-VelocityDistributionFunctions = "63a0b508-f74a-48f7-b6e7-0e65849b0078"
 
 [weakdeps]
 RecursiveArrayTools = "731186ca-8d62-57ce-b412-fbd966d074cd"
+VelocityDistributionFunctions = "63a0b508-f74a-48f7-b6e7-0e65849b0078"
 
 [extensions]
 TestParticleRecursiveArrayToolsExt = "RecursiveArrayTools"
+TestParticleVDFExt = "VelocityDistributionFunctions"
 
 [compat]
 Adapt = "4.4"

docs/examples/applications/demo_fermi_foreshock.jl

@@ -24,6 +24,7 @@
 import DisplayAs #hide
 using TestParticle, OrdinaryDiffEqVerner, StaticArrays
 import TestParticle as TP
+using VelocityDistributionFunctions
 using TestParticle: mₑ, Rₑ, qₑ
 using TestParticle: get_EField, get_BField
 using Random
@@ -99,7 +100,7 @@ function prob_func(prob, i, repeat)
     u0 = SA[0.0, 0.0, 0.0]
     T₀ = 10 # [eV]
     vth = √(2T₀ * abs(qₑ) / mₑ) # [m/s]
-    vdf = Maxwellian(u0, vth)
+    vdf = TP.Maxwellian(u0, vth)
     v0 = rand(vdf)
 
     return prob = remake(prob, u0 = [x0..., v0...])

docs/examples/applications/demo_shock.jl

@@ -5,6 +5,8 @@
 import DisplayAs #hide
 using TestParticle, OrdinaryDiffEqVerner
 using TestParticle: mᵢ, kB
+import TestParticle as TP
+using VelocityDistributionFunctions
 using LinearAlgebra
 using Statistics: mean, std
 using Printf
@@ -60,8 +62,8 @@ mid_ = length(x) ÷ 2
 B[:, 1:mid_] .= B₂
 E[:, 1:mid_] .= E₂
 
-const vdf₁ = Maxwellian(V₁, Pth₁, n₁; m = mᵢ)
-vdf₂ = Maxwellian(V₂, Pth₂, n₂; m = mᵢ)
+const vdf₁ = TP.Maxwellian(V₁, Pth₁, n₁; m = mᵢ)
+vdf₂ = TP.Maxwellian(V₂, Pth₂, n₂; m = mᵢ)
 
 println(vdf₁)
 println(vdf₂)
@@ -348,8 +350,8 @@ mid_ = length(x) ÷ 2
 B[:, 1:mid_] .= B₂
 E[:, 1:mid_] .= E₂
 
-const vdf₁_para = Maxwellian(V₁, Pth₁, n₁; m = mᵢ)
-vdf₂_para = Maxwellian(V₂, Pth₂, n₂; m = mᵢ)
+const vdf₁_para = TP.Maxwellian(V₁, Pth₁, n₁; m = mᵢ)
+vdf₂_para = TP.Maxwellian(V₂, Pth₂, n₂; m = mᵢ)
 
 trajectories = 2
 weight₁ = n₁ / trajectories

docs/examples/features/demo_distributions.jl

@@ -5,6 +5,8 @@
 # We compare the sampled distributions with their theoretical probability density functions (PDFs).
 
 using TestParticle
+import TestParticle as TP
+using VelocityDistributionFunctions
 using SpecialFunctions: gamma
 using Random
 using StaticArrays
@@ -35,7 +37,7 @@ n = 1.0 # Number density
 m = 1.0 # Mass
 
 ## Construct distribution
-vdf = Maxwellian(u0, p, n; m)
+vdf = TP.Maxwellian(u0, p, n; m)
 println(vdf)
 
 ## Sample
@@ -73,7 +75,7 @@ ppar = 2.0
 pperp = 0.5
 
 ## Construct distribution
-vdf_bi = BiMaxwellian(B, u0, ppar, pperp, n; m)
+vdf_bi = TP.BiMaxwellian(B, u0, ppar, pperp, n; m)
 println(vdf_bi)
 
 ## Sample
@@ -120,7 +122,7 @@ f = DisplayAs.PNG(f) #hide
 kappa = 3.0
 
 ## Construct distribution
-vdf_kappa = Kappa(u0, p, n, kappa; m)
+vdf_kappa = TP.Kappa(u0, p, n, kappa; m)
 println(vdf_kappa)
 
 ## Sample
@@ -159,7 +161,7 @@ pperp = 1.0
 kappa = 4.0
 
 ## Construct distribution
-vdf_bikappa = BiKappa(B, u0, ppar, pperp, n, kappa; m)
+vdf_bikappa = TP.BiKappa(B, u0, ppar, pperp, n, kappa; m)
 println(vdf_bikappa)
 
 ## Sample

docs/examples/features/demo_ensemble.jl

@@ -9,7 +9,8 @@
 
 import DisplayAs #hide
 using TestParticle
-using TestParticle: get_BField
+import TestParticle as TP
+using VelocityDistributionFunctions
 using OrdinaryDiffEqVerner
 using StaticArrays
 using Statistics
@@ -66,7 +67,7 @@ Random.seed!(1234)
 ## Define a new prob_func that samples from a Maxwellian
 function prob_func_maxwellian(prob, i, repeat)
     ## Sample from a Maxwellian with bulk speed 0 and thermal speed 1.0
-    vdf = Maxwellian([0.0, 0.0, 0.0], 1.0)
+    vdf = TP.Maxwellian([0.0, 0.0, 0.0], 1.0)
     v = rand(vdf)
     return remake(prob; u0 = [prob.u0[1:3]..., v...])
 end
@@ -134,7 +135,7 @@ prob_custom = ODEProblem(trace_normalized!, stateinit_custom, tspan_custom, para
 
 ## Define prob_func to initialize particles with different pitch angles
 function prob_func_custom(prob, i, repeat)
-    B0 = get_BField(prob)(prob.u0)
+    B0 = TP.get_BField(prob)(prob.u0)
     B0 = normalize(B0)
 
     Bperp1 = normalize(SA[0.0, -B0[3], B0[2]])
@@ -150,7 +151,7 @@ end
 
 ## Define output_func to save specific data
 function output_func_custom(sol, i)
-    getB = get_BField(sol)
+    getB = TP.get_BField(sol)
     b = getB.(sol.u)
 
     ## Calculate cosine of pitch angle

docs/examples/features/demo_number_density.jl

@@ -5,6 +5,8 @@
 
 import DisplayAs #hide
 using TestParticle
+import TestParticle as TP
+using VelocityDistributionFunctions
 using Meshes
 using OrdinaryDiffEq
 using StaticArrays
@@ -33,7 +35,7 @@ x0 = [SVector(0.0, 0.0, 0.0) for _ in 1:N];
 
 # Maxwellian velocity distribution
 # ``u_0 = 0``, ``p = n k_B T = 1`` (assuming ``n=1`` effectively for distribution shape)
-vdf = TestParticle.Maxwellian([0.0, 0.0, 0.0], T, 1.0; m = m)
+vdf = TP.Maxwellian([0.0, 0.0, 0.0], T, 1.0; m = m)
 # Use the vectorized rand to get a Vector of SVectors
 v0 = rand(vdf, N);
 
@@ -45,7 +47,7 @@ prob_func(prob, i, repeat) = remake(prob, u0 = vcat(x0[i], v0[i]))
 # Define a single problem template
 u0_dummy = vcat(x0[1], v0[1])
 ## Zero fields
-param = prepare(TestParticle.ZeroField(), TestParticle.ZeroField(); q, m)
+param = prepare(TP.ZeroField(), TP.ZeroField(); q, m)
 tspan = (0.0, t_end)
 prob = ODEProblem(trace, u0_dummy, tspan, param);
 
@@ -66,7 +68,7 @@ grid = CartesianGrid((-L, -L, -L), (L, L, L); dims)
 
 # Calculate density at ``t_{end}``
 # `get_number_density` returns count / volume
-density = TestParticle.get_number_density(sols, grid, t_end);
+density = TP.get_number_density(sols, grid, t_end);
 
 # Extract a 1D slice along x-axis (approx. ``y=0, z=0``)
 mid_y = dims[2] ÷ 2
@@ -75,7 +77,7 @@ density_x = density[:, mid_y, mid_z];
 
 # Grid coordinates for plotting
 # `get_cell_centers` returns ranges for each dimension
-grid_x, grid_y, grid_z = TestParticle.get_cell_centers(grid)
+grid_x, grid_y, grid_z = TP.get_cell_centers(grid)
 xs_plot = collect(grid_x);
 
 # ## Analytical Solution

ext/TestParticleRecursiveArrayToolsExt.jl

@@ -1,4 +1,5 @@
 module TestParticleRecursiveArrayToolsExt
+
 import TestParticle: trace, trace!
 using TestParticle: get_dv
 using RecursiveArrayTools: ArrayPartition

ext/TestParticleVDFExt.jl

@@ -0,0 +1,65 @@
+module TestParticleVDFExt
+
+using TestParticle
+import TestParticle as TP
+import VelocityDistributionFunctions as VDF
+import TestParticle: Maxwellian, BiMaxwellian, Kappa, BiKappa
+using PrecompileTools: @setup_workload, @compile_workload
+
+function Maxwellian(args...; kwargs...)
+    return VDF.Maxwellian(args...; kwargs...)
+end
+
+function Maxwellian(u0, p, n; m = TP.mᵢ)
+    vth = TP.get_thermal_speed(p, n, m)
+
+    return VDF.Maxwellian(vth; u0)
+end
+
+function BiMaxwellian(args...; kwargs...)
+    return VDF.BiMaxwellian(args...; kwargs...)
+end
+
+function BiMaxwellian(B, u0, ppar, pperp, n; m = TP.mᵢ)
+    vpar = TP.get_thermal_speed(ppar, n, m)
+    vperp = TP.get_thermal_speed(pperp, n, m)
+
+    return VDF.BiMaxwellian(vperp, vpar, B; u0)
+end
+
+Kappa(args...; kwargs...) = VDF.Kappa(args...; kwargs...)
+
+function Kappa(u0, p, n, kappa; m = TP.mᵢ)
+    vth = TP.get_thermal_speed(p, n, m)
+
+    return VDF.Kappa(vth, kappa; u0)
+end
+
+BiKappa(args...; kwargs...) = VDF.BiKappa(args...; kwargs...)
+
+function BiKappa(B, u0, ppar, pperp, n, kappa; m = TP.mᵢ)
+    vpar = TP.get_thermal_speed(ppar, n, m)
+    vperp = TP.get_thermal_speed(pperp, n, m)
+
+    return VDF.BiKappa(vperp, vpar, kappa, B; u0)
+end
+
+@setup_workload begin
+    @compile_workload begin
+        u0 = [0.0, 0.0, 0.0]
+        p = 1.0e-9
+        n = 1.0e6
+        vdf = Maxwellian(u0, p, n)
+        v = rand(vdf)
+        B = [1.0, 0.0, 0.0]
+        vdf = BiMaxwellian(B, u0, p, p, n)
+        v = rand(vdf)
+        kappa = 4.0
+        vdf = Kappa(u0, p, n, kappa)
+        v = rand(vdf)
+        vdf = BiKappa(B, u0, p, p, n, kappa)
+        v = rand(vdf)
+    end
+end
+
+end

src/TestParticle.jl

@@ -31,8 +31,7 @@ export trace!, trace_relativistic!, trace_normalized!, trace_relativistic_normal
     trace_gc_drifts!, trace_gc_flr!, trace_gc_exb!, trace_fieldline!, trace_fieldline,
     get_gc_velocity, full_to_gc, gc_to_full
 export Proton, Electron, Ion
-export Maxwellian, BiMaxwellian
-export Kappa, BiKappa
+export Maxwellian, BiMaxwellian, Kappa, BiKappa
 export AdaptiveBoris, AdaptiveHybrid
 export CurrentLoop, getB_loop
 export get_gyrofrequency,

src/precompile.jl

@@ -17,10 +17,6 @@
             dims = (length(x) - 1, length(y) - 1, length(z) - 1)
         )
 
-        vdf = Maxwellian([0.0, 0.0, 0.0], 1.0e-9, 1.0e6)
-        v = rand(vdf)
-        vdf = BiMaxwellian([1.0, 0.0, 0.0], [0.0, 0.0, 0.0], 1.0e-9, 1.0e-9, 1.0e6)
-        v = rand(vdf)
         # numerical field
         param = prepare(x, y, z, E, B)
         param = prepare(mesh, E, B)

src/sampler.jl

@@ -1,6 +1,5 @@
 # Particle sampling
 
-import VelocityDistributionFunctions
 
 # Wrapper functions to avoid type piracy when adding convenience constructors
 
@@ -14,15 +13,7 @@ Construct a `Maxwellian` distribution. Forwards to `VelocityDistributionFunction
 Construct a `Maxwellian` distribution with bulk velocity `u0`, thermal pressure `p`, and
 number density `n` in SI units. The default particle is proton.
 """
-function Maxwellian(args...; kwargs...)
-    return VelocityDistributionFunctions.Maxwellian(args...; kwargs...)
-end
-
-function Maxwellian(u0, p, n; m = mᵢ)
-    vth = get_thermal_speed(p, n, m)
-
-    return VelocityDistributionFunctions.Maxwellian(vth; u0)
-end
+function Maxwellian end
 
 """
      BiMaxwellian(args...; kw...)
@@ -35,16 +26,7 @@ Construct a `BiMaxwellian` distribution with magnetic field `B`, bulk velocity `
 thermal pressure `ppar`, perpendicular thermal pressure `pperp`, and number density `n` in
 SI units. The default particle is proton.
 """
-function BiMaxwellian(args...; kwargs...)
-    return VelocityDistributionFunctions.BiMaxwellian(args...; kwargs...)
-end
-
-function BiMaxwellian(B, u0, ppar, pperp, n; m = mᵢ)
-    vpar = get_thermal_speed(ppar, n, m)
-    vperp = get_thermal_speed(pperp, n, m)
-
-    return VelocityDistributionFunctions.BiMaxwellian(vperp, vpar, B; u0)
-end
+function BiMaxwellian end
 
 """
      Kappa(args...; kw...)
@@ -56,13 +38,7 @@ Construct a `Kappa` distribution. Forwards to `VelocityDistributionFunctions.Kap
 Construct a `Kappa` distribution with bulk velocity `u0`, thermal pressure `p`, number density
 `n`, and spectral index `kappa` in SI units. The default particle is proton.
 """
-Kappa(args...; kwargs...) = VelocityDistributionFunctions.Kappa(args...; kwargs...)
-
-function Kappa(u0, p, n, kappa; m = mᵢ)
-    vth = get_thermal_speed(p, n, m)
-
-    return VelocityDistributionFunctions.Kappa(vth, kappa; u0)
-end
+function Kappa end
 
 """
      BiKappa(args...; kw...)
@@ -75,13 +51,6 @@ Construct a `BiKappa` distribution with magnetic field `B`, bulk velocity `u0`,
 thermal pressure `ppar`, perpendicular thermal pressure `pperp`, number density `n`, and
 spectral index `kappa` in SI units. The default particle is proton.
 """
-BiKappa(args...; kwargs...) = VelocityDistributionFunctions.BiKappa(args...; kwargs...)
-
-function BiKappa(B, u0, ppar, pperp, n, kappa; m = mᵢ)
-    vpar = get_thermal_speed(ppar, n, m)
-    vperp = get_thermal_speed(pperp, n, m)
-
-    return VelocityDistributionFunctions.BiKappa(vperp, vpar, kappa, B; u0)
-end
+function BiKappa end
 
 get_thermal_speed(p, n, m) = √(2 * p / (n * m))

test/Project.toml

@@ -1,6 +1,5 @@
 [deps]
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
-Meshes = "eacbb407-ea5a-433e-ab97-5258b1ca43fa"
 OrdinaryDiffEq = "1dea7af3-3e70-54e6-95c3-0bf5283fa5ed"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 SciMLBase = "0bca4576-84f4-4d90-8ffe-ffa030f20462"

test/runtests.jl

@@ -1,10 +1,12 @@
 using TestParticle, OrdinaryDiffEq, StaticArrays
 using TestParticle: Field, qᵢ, mᵢ, qₑ, mₑ, c
 import TestParticle as TP
-using Meshes: CartesianGrid, StructuredGrid
 using Random, StableRNGs
+using VelocityDistributionFunctions
 using Test
 
+import TestParticle: Maxwellian, BiMaxwellian, Kappa, BiKappa
+
 """
 Initial state perturbation for EnsembleProblem.
 """

test/test_density.jl

@@ -1,5 +1,4 @@
 using TestParticle
-using Meshes
 using StaticArrays
 using Test
 using OrdinaryDiffEq