Field Interpolation

A robust field interpolation is the prerequisite for pushing particles. This example demonstrates the construction of scalar/vector field interpolators for Cartesian/Spherical grids. If the field is analytic, you can directly pass the generated function to prepare.

using TestParticle
using Meshes
using StaticArrays
using Chairmarks
import TestParticle as TP

function setup_spherical_field(ns = 16)
   r = logrange(0.1, 10.0, length = ns)
   r_uniform = range(0.1, 10.0, length = ns)
   θ = range(0, π, length = ns)
   ϕ = range(0, 2π, length = ns)

   B₀ = 1e-8 # [nT]
   B = zeros(3, length(r), length(θ), length(ϕ)) # vector
   A = zeros(length(r), length(θ), length(ϕ)) # scalar

   for (iθ, θ_val) in enumerate(θ)
      sinθ, cosθ = sincos(θ_val)
      B[1, :, iθ, :] .= B₀ * cosθ
      B[2, :, iθ, :] .= -B₀ * sinθ
      A[:, iθ, :] .= B₀ * sinθ
   end

   B_field_nu = build_interpolator(StructuredGrid, B, r, θ, ϕ)
   A_field_nu = build_interpolator(StructuredGrid, A, r, θ, ϕ)
   B_field = build_interpolator(StructuredGrid, B, r_uniform, θ, ϕ)
   A_field = build_interpolator(StructuredGrid, A, r_uniform, θ, ϕ)

   return B_field_nu, A_field_nu, B_field, A_field
end

function setup_cartesian_field(ns = 16)
   x = range(-10, 10, length = ns)
   y = range(-10, 10, length = ns)
   z = range(-10, 10, length = ns)
   B = zeros(3, length(x), length(y), length(z)) # vector
   B[3, :, :, :] .= 10e-9
   A = zeros(length(x), length(y), length(z)) # scalar
   A[:, :, :] .= 10e-9

   B_field = build_interpolator(B, x, y, z)
   A_field = build_interpolator(A, x, y, z)

   return B_field, A_field
end

function setup_cartesian_nonuniform_field()
   x = logrange(0.1, 10.0, length = 16)
   y = range(-10, 10, length = 16)
   z = range(-10, 10, length = 16)
   B = zeros(3, length(x), length(y), length(z)) # vector
   B[3, :, :, :] .= 10e-9
   A = zeros(length(x), length(y), length(z)) # scalar
   A[:, :, :] .= 10e-9

   B_field = build_interpolator(RectilinearGrid, B, x, y, z)
   A_field = build_interpolator(RectilinearGrid, A, x, y, z)

   return B_field, A_field
end

function setup_time_dependent_field(ns = 16)
   x = range(-10, 10, length = ns)
   y = range(-10, 10, length = ns)
   z = range(-10, 10, length = ns)

   # Create two time snapshots
   B0 = zeros(3, length(x), length(y), length(z))
   B0[3, :, :, :] .= 1.0 # Bz = 1 at t=0

   B1 = zeros(3, length(x), length(y), length(z))
   B1[3, :, :, :] .= 2.0 # Bz = 2 at t=1

   times = [0.0, 1.0]

   function loader(i)
       if i == 1
           # For demonstration, we assume we load from disk here
           return build_interpolator(CartesianGrid, B0, x, y, z)
       elseif i == 2
           return build_interpolator(CartesianGrid, B1, x, y, z)
       else
           error("Index out of bounds")
       end
   end

   # B_field_t(x, t)
   B_field_t = LazyTimeInterpolator(times, loader)

   return B_field_t
end

function setup_mixed_precision_field(
      ns = 11, order = 1, bc = FillExtrap(NaN);
      coeffs = OnTheFly(), store = StorePolicy()
   )
   x = range(0.0f0, 10.0f0, length = ns)
   y = range(0.0f0, 10.0f0, length = ns)
   z = range(0.0f0, 10.0f0, length = ns)
   B = fill(0.0f0, 3, ns, ns, ns)
   B[3, :, :, :] .= 1.0f-8

   itp = build_interpolator(B, x, y, z, order, bc; coeffs, store)
   return itp
end

B_sph_nu, A_sph_nu, B_sph, A_sph = setup_spherical_field();
B_car, A_car = setup_cartesian_field();
B_car_nu, A_car_nu = setup_cartesian_nonuniform_field();
B_td = setup_time_dependent_field();
itp_f32 = setup_mixed_precision_field();

loc = SA[1.0, 1.0, 1.0];
loc_f32 = SA[1.0f0, 1.0f0, 1.0f0];
loc_f64 = SA[1.0, 1.0, 1.0];

3-element StaticArraysCore.SVector{3, Float64} with indices SOneTo(3):
 1.0
 1.0
 1.0

Gridded spherical interpolation

Input Location

For spherical data, the input location is still in Cartesian coordinates!

julia> @be B_sph_nu($loc)
Benchmark: 3148 samples with 306 evaluations
 min    94.454 ns (2 allocs: 64 bytes)
 median 96.614 ns (2 allocs: 64 bytes)
 mean   97.695 ns (2 allocs: 64 bytes)
 max    225.761 ns (2 allocs: 64 bytes)

julia> @be A_sph_nu($loc)
Benchmark: 3151 samples with 388 evaluations
 min    73.977 ns (2 allocs: 48 bytes)
 median 75.861 ns (2 allocs: 48 bytes)
 mean   76.688 ns (2 allocs: 48 bytes)
 max    135.126 ns (2 allocs: 48 bytes)

Uniform spherical interpolation

julia> @be B_sph($loc)
Benchmark: 3143 samples with 306 evaluations
 min    93.049 ns (2 allocs: 64 bytes)
 median 95.471 ns (2 allocs: 64 bytes)
 mean   97.075 ns (2 allocs: 64 bytes)
 max    238.330 ns (2 allocs: 64 bytes)

julia> @be A_sph($loc)
Benchmark: 3108 samples with 397 evaluations
 min    71.592 ns (2 allocs: 48 bytes)
 median 74.118 ns (2 allocs: 48 bytes)
 mean   75.855 ns (2 allocs: 48 bytes)
 max    187.209 ns (2 allocs: 48 bytes)

Uniform Cartesian interpolation

julia> @be B_car($loc)
Benchmark: 3131 samples with 753 evaluations
 min    37.413 ns (2 allocs: 64 bytes)
 median 38.596 ns (2 allocs: 64 bytes)
 mean   39.717 ns (2 allocs: 64 bytes)
 max    97.664 ns (2 allocs: 64 bytes)

julia> @be A_car($loc)
Benchmark: 1147 samples with 529 evaluations
 min    50.321 ns (2 allocs: 48 bytes)
 median 52.329 ns (2 allocs: 48 bytes)
 mean   261.719 ns (2 allocs: 48 bytes, 0.09% gc time)
 max    236.095 μs (2 allocs: 48 bytes, 99.90% gc time)

Non-uniform Cartesian interpolation

julia> @be B_car_nu($loc)
Benchmark: 3233 samples with 752 evaluations
 min    37.024 ns (2 allocs: 64 bytes)
 median 38.182 ns (2 allocs: 64 bytes)
 mean   38.602 ns (2 allocs: 64 bytes)
 max    70.851 ns (2 allocs: 64 bytes)

julia> @be A_car_nu($loc)
Benchmark: 3240 samples with 780 evaluations
 min    34.899 ns (2 allocs: 48 bytes)
 median 36.671 ns (2 allocs: 48 bytes)
 mean   37.365 ns (2 allocs: 48 bytes)
 max    197.990 ns (2 allocs: 48 bytes)

Based on the benchmarks, for the same grid size, gridded interpolation (StructuredGrid with non-uniform ranges, RectilinearGrid) is 2x slower than uniform mesh interpolation (StructuredGrid with uniform ranges, CartesianGrid).

Mixed precision interpolation

Numerical field data from files is often stored in Float32. TestParticle supports constructing interpolators from Float32 data and ranges, which can then be queried with both Float32 and Float64 location vectors.

julia> @be itp_f32($loc_f32)
Benchmark: 3001 samples with 723 evaluations
 min    38.398 ns (2 allocs: 64 bytes)
 median 39.852 ns (2 allocs: 64 bytes)
 mean   43.148 ns (2 allocs: 64 bytes)
 max    80.661 ns (2 allocs: 64 bytes)

julia> @be itp_f32($loc_f64)
Benchmark: 34 samples with 444 evaluations
 min    61.241 ns (2 allocs: 64 bytes)
 median 65.380 ns (2 allocs: 64 bytes)
 mean   7.133 μs (2 allocs: 64 bytes, 2.94% gc time)
 max    240.336 μs (2 allocs: 64 bytes, 99.91% gc time)

Memory usage analysis

Large numerical field arrays can consume significant amounts of memory during interpolator construction. FastInterpolations.jl uses StorePolicy() by default, which copies grids and field data so that the interpolator owns a stable snapshot. This protects it from later caller mutations, but constructing an interpolator for a 45 GB field requires another approximately 45 GB allocation.

TestParticle forwards the store keyword to FastInterpolations.jl. Use StorePolicy(copy = false) to alias compatible grids and field data instead of copying them. This removes the field-sized construction allocation for linear interpolation and for cubic interpolation backed by dense scalar or SVector arrays.

julia> x_memory = range(0.0f0, 1.0f0, length = 32);

julia> B_memory = rand(Float32, 3, 32, 32, 32);

julia> B_memory_svector = fill(SA[1.0f0, 1.0f0, 1.0f0], 32, 32, 32);

julia> copy_store = StorePolicy();

julia> zero_copy_store = StorePolicy(copy = false);

julia> # Default owned-copy construction
       @be build_interpolator($B_memory, $x_memory, $x_memory, $x_memory, 1; store = $copy_store)
Benchmark: 802 samples with 1 evaluation
 min    27.512 μs (5 allocs: 384.398 KiB)
 median 37.441 μs (5 allocs: 384.398 KiB)
 mean   122.085 μs (5 allocs: 384.398 KiB, 0.37% gc time)
 max    34.737 ms (5 allocs: 384.398 KiB, 99.70% gc time)

julia> @be build_interpolator($B_memory_svector, $x_memory, $x_memory, $x_memory, 3; store = $copy_store)
Benchmark: 1139 samples with 1 evaluation
 min    15.814 μs (5 allocs: 384.430 KiB)
 median 22.084 μs (5 allocs: 384.430 KiB)
 mean   82.728 μs (5 allocs: 384.430 KiB, 0.44% gc time)
 max    28.766 ms (5 allocs: 384.430 KiB, 99.71% gc time)

julia> # Zero-copy construction
       @be build_interpolator($B_memory, $x_memory, $x_memory, $x_memory, 1; store = $zero_copy_store)
Benchmark: 861 samples with 485 evaluations
 min    44.852 ns (2 allocs: 320 bytes)
 median 53.524 ns (2 allocs: 320 bytes)
 mean   264.777 ns (2 allocs: 320 bytes, 0.23% gc time)
 max    84.314 μs (2 allocs: 320 bytes, 99.73% gc time)

julia> @be build_interpolator($B_memory_svector, $x_memory, $x_memory, $x_memory, 3; store = $zero_copy_store)
Benchmark: 1076 samples with 390 evaluations
 min    59.990 ns (2 allocs: 352 bytes)
 median 73.508 ns (2 allocs: 352 bytes)
 mean   222.522 ns (2 allocs: 352 bytes, 0.09% gc time)
 max    107.342 μs (2 allocs: 352 bytes, 99.79% gc time)

With zero-copy storage, the interpolator may alias both the field and grid arrays. They must remain alive and must not be mutated or resized while the interpolator is in use. If only the large field should be aliased, use StorePolicy(copy_values = false) to retain owned grid storage. Inputs that require an element-type conversion may still be copied.

Component-first vector fields with shape (3, nx, ny, nz) are exposed internally as a ReinterpretArray. FastInterpolations v0.4 currently cannot evaluate ND cardinal interpolation from that aliased wrapper, so TestParticle warns and falls back to owned storage for order = 3. Store the field directly as a dense 3D array of SVectors to use zero-copy cubic interpolation.

Base.summarysize reports all objects reachable through an interpolator, including aliased inputs, so construction allocation measurements are a better way to verify that the copy was eliminated.

On-the-fly vs Precomputed coefficients

Cubic interpolation (order = 3) requires high-order coefficients. TestParticle uses OnTheFly() coefficients, which are calculated during each query. Zero-copy storage reduces construction time and memory but does not change interpolation throughput. PreCompute() could trade additional memory for faster queries, but it is not yet supported for TestParticle's ND cardinal interpolation.

Status in FastInterpolations v0.4.19

PreCompute() is available for the global natural cubic spline (CubicInterp) in ND. However, it is not yet supported for the local Hermite cubic spline (CardinalInterp, i.e. order = 3 in TestParticle) in ND. Constructing such an interpolator with coeffs = PreCompute() raises an ArgumentError, so OnTheFly() remains the only option for cubic interpolation in TestParticle.

julia> # Benchmark evaluation time
       itp_fly = setup_mixed_precision_field(11, 3; coeffs = OnTheFly());

julia> # itp_pre = setup_mixed_precision_field(11, 3; coeffs = PreCompute()); # unsupported for ND cardinal cubic
       
       @be itp_fly($loc_f64)
Benchmark: 3127 samples with 142 evaluations
 min    206.585 ns (2 allocs: 64 bytes)
 median 209.113 ns (2 allocs: 64 bytes)
 mean   211.116 ns (2 allocs: 64 bytes)
 max    376.972 ns (2 allocs: 64 bytes)

julia> # Compare total object size
       Base.summarysize(itp_fly)
16196

As shown, OnTheFly() preserves memory efficiency while providing higher-order accuracy. Once PreCompute() support for ND local Hermite cubic interpolation lands, it will offer a faster alternative for memory-abundant systems.

Time-dependent field interpolation

For time-dependent fields, we can use LazyTimeInterpolator. It takes a list of time points and a loader function that returns a spatial interpolator for a given time index. The interpolator will linearly interpolate between the two nearest time points.

julia> @be B_td($loc, 0.5)
Benchmark: 3184 samples with 231 evaluations
 min    122.519 ns (2 allocs: 64 bytes)
 median 127.030 ns (2 allocs: 64 bytes)
 mean   128.188 ns (2 allocs: 64 bytes)
 max    244.221 ns (2 allocs: 64 bytes)

GPU interpolation

The interpolators above are built on FastInterpolations.jl and live on the host. Pushing a particle on a device asks for something different: a plain struct with Int32 indices whose axes and data can be copied to device memory and evaluated inside a kernel. TestParticle provides GPUSphericalGrid for spherical grids, and GPUGrid3D, GPUGrid2D and GPUGrid1D for Cartesian ones.

prepare performs the conversion by itself once the parameters are handed to a device backend, see GPU Ensemble Tracing. The very same grids can be built on the CPU() backend, which is what the following does with TP.CPU(), so that the device interpolator can be inspected and benchmarked where no device is available:

B_gpu = TP._to_gpu_spherical_grid(B_sph.itp, TP.CPU());
A_gpu = TP._to_gpu_spherical_grid(A_sph.itp, TP.CPU());
B_nu_gpu = TP._to_gpu_spherical_grid(B_sph_nu.itp, TP.CPU());

(::GPUSphericalGrid{Float64, GPUNonUniformAxis{Float64, Vector{Float64}, FillExtrap{StaticArraysCore.SVector{3, Float64}}}, GPUUniformAxis{Float64, FillExtrap{StaticArraysCore.SVector{3, Float64}}}, GPUUniformAxis{Float64, WrapExtrap}, StaticArraysCore.SVector{3, Float64}, Array{StaticArraysCore.SVector{3, Float64}, 3}, Tuple{FillExtrap{StaticArraysCore.SVector{3, Float64}}, FillExtrap{StaticArraysCore.SVector{3, Float64}}, WrapExtrap}}) (generic function with 6 methods)

A uniform grid vector becomes a GPUUniformAxis, which locates a cell with a multiply, and a non-uniform one a GPUNonUniformAxis, which brackets it with a binary search:

julia> typeof(B_gpu)
GPUSphericalGrid{Float64, GPUUniformAxis{Float64, FillExtrap{StaticArraysCore.SVector{3, Float64}}}, GPUUniformAxis{Float64, FillExtrap{StaticArraysCore.SVector{3, Float64}}}, GPUUniformAxis{Float64, WrapExtrap}, StaticArraysCore.SVector{3, Float64}, Array{StaticArraysCore.SVector{3, Float64}, 3}, Tuple{FillExtrap{StaticArraysCore.SVector{3, Float64}}, FillExtrap{StaticArraysCore.SVector{3, Float64}}, WrapExtrap}}

julia> B_gpu.axis_r
GPUUniformAxis{Float64, FillExtrap{StaticArraysCore.SVector{3, Float64}}}(0.1, 0.66, 1.5151515151515151, 16, 10.0, FillExtrap{StaticArraysCore.SVector{3, Float64}}([NaN, NaN, NaN]))

julia> typeof(B_nu_gpu.axis_r)
GPUNonUniformAxis{Float64, Vector{Float64}, FillExtrap{StaticArraysCore.SVector{3, Float64}}}

Queries are unchanged, only the backend behind them differs. The location is still Cartesian, and a vector result is rotated back into the Cartesian basis:

julia> @be B_gpu($loc)
Benchmark: 3156 samples with 283 evaluations
 min    101.604 ns (2 allocs: 64 bytes)
 median 103.689 ns (2 allocs: 64 bytes)
 mean   105.130 ns (2 allocs: 64 bytes)
 max    210.141 ns (2 allocs: 64 bytes)

julia> @be B_sph($loc)
Benchmark: 3177 samples with 303 evaluations
 min    94.597 ns (2 allocs: 64 bytes)
 median 96.449 ns (2 allocs: 64 bytes)
 mean   97.393 ns (2 allocs: 64 bytes)
 max    218.713 ns (2 allocs: 64 bytes)

julia> B_gpu($loc)
ERROR: syntax: "$" expression outside quote around REPL[3]:1

julia> A_gpu($loc)
ERROR: syntax: "$" expression outside quote around REPL[4]:1

Outside the grid the boundary condition decides, filling with NaN in r and in θ and wrapping periodically in ϕ:

julia> B_gpu(0.01, 0.0, 0.0)
3-element StaticArraysCore.SVector{3, Float64} with indices SOneTo(3):
 NaN
 NaN
 NaN

julia> B_gpu(20.0, 0.0, 0.0)
3-element StaticArraysCore.SVector{3, Float64} with indices SOneTo(3):
 NaN
 NaN
 NaN

One caveat comes with the device grids: they always interpolate linearly, whatever order the host interpolator was built with. A field prepared with order = 3 is therefore not reproduced on the device, and order = 1, the default, is the one to use when the host and the device are meant to agree.

Related API

TestParticle.build_interpolator Function

build_interpolator(gridtype, A, grids..., order::Int=1, bc=FillExtrap(NaN);
    coeffs=OnTheFly(), store=StorePolicy())
build_interpolator(A, grids..., order::Int=1, bc=FillExtrap(NaN);
    coeffs=OnTheFly(), store=StorePolicy())

Return a function for interpolating field array A on the given grids.

Arguments

• gridtype: CartesianGrid, RectilinearGrid or StructuredGrid. Usually determined by the number of grids.

• A: field array. For vector fields, the first dimension should be 3 unless the array contains SVectors.

• order::Int=1: order of interpolation in [0,1,3].

• bc=FillExtrap(NaN): boundary condition type from FastInterpolations.jl.

  • FillExtrap(NaN): Fill with NaN (default).

  • ClampExtrap(): Clamp (flat extrapolation).

  • WrapExtrap(): Exclusive periodic wrapping ($L=N\mathrm{\Delta }x$).

• coeffs=OnTheFly(): coefficient strategy for cubic interpolation (order=3).

• store=StorePolicy(): storage policy forwarded to FastInterpolations.jl. The default copies the grid and field data. Use StorePolicy(copy=false) to alias compatible inputs and avoid the construction-time data copy.

Notes

With StorePolicy(copy=false), the interpolator may alias A and the grid arrays. Do not mutate or resize aliased inputs while the interpolator is in use. When input conversion is required, FastInterpolations.jl may still copy the data.

source

TestParticle.prepare Function

prepare(args...; kwargs...) -> (q2m, m, E, B, F)
prepare(E, B, F = ZeroField(); kwargs...)
prepare(grid::CartesianGrid, E, B, F = ZeroField(); kwargs...)
prepare(x, E, B, F = ZeroField(); dir = 1, kwargs...)
prepare(x, y, E, B, F = ZeroField(); kwargs...)
prepare(x, y, z, E, B, F = ZeroField(); kwargs...)
prepare(B; E = ZeroField(), F = ZeroField(), kwargs...)

Return a tuple consists of particle charge-mass ratio for a prescribed species of charge q and mass m, mass m for a prescribed species, analytic/interpolated EM field functions, and external force F.

Prescribed species are Electron and Proton; other species can be manually specified with m and q keywords or species = Ion(m̄, q̄), where m̄ and q̄ are the mass and charge numbers respectively.

Direct range input for uniform grid in 1/2/3D is supported. The grid vectors must be sorted. For 1D grid, an additional keyword dir is used for specifying the spatial direction, 1 -> x, 2 -> y, 3 -> z. For 3D grid, the default grid type is CartesianGrid. To use StructuredGrid (spherical) grid, an additional keyword gridtype is needed. For StructuredGrid (spherical) grid, dimensions of field arrays should be (Br, Bθ, Bϕ).

Keywords

• order::Int=1: order of interpolation in [0,1,3].

• bc=FillExtrap(NaN): boundary condition type from FastInterpolations.jl.

• store=StorePolicy(): interpolation storage policy. Use StorePolicy(copy=false) for zero-copy construction when field and grid inputs will remain unchanged for the interpolator's lifetime.

• species=Proton: particle species.

• q=nothing: particle charge.

• m=nothing: particle mass.

• gridtype: CartesianGrid, RectilinearGrid, StructuredGrid.

source

prepare(pf::PotentialField; species = Proton, q = nothing, m = nothing, c = c, type = nothing)

Return parameters (q, m, c, pf) for TraceCanonicalProblem.

source