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.

Construct each numerical field explicitly, then pass its callable to prepare:

using TestParticle
using TestParticle: PeriodicBC, FillExtrap
B = build_interpolator(Bdata, (x, y, z);
    bc = (PeriodicBC(; endpoint = :exclusive),
        PeriodicBC(; endpoint = :exclusive), FillExtrap(NaN)))
p = prepare(; B)

The tuple constructor requires bc; see build_interpolator for endpoint, array-layout, and device contracts. E and B can use independent grids and policies. Array inputs belong in build_interpolator, not prepare. Foreign interpolators need only implement f(x) or f(x, t). For callables with both signatures, select explicitly with TestParticle.Field(f; time_dependent=false) or true; automatic detection prefers the two-argument signature.

Periodic fields do not wrap particle coordinates. Trajectories remain unwrapped for displacement analysis; wrap coordinates separately for plotting. Open-field NaNs do not terminate a trajectory cleanly: configure the solver's boundary termination. For spherical fields use gridtype=StructuredGrid and provide a tuple of radial, colatitude, and longitude policies. Longitude must be periodic with period 2π; choose inclusive or exclusive endpoints explicitly. Radial and colatitude policies support fill or clamp. Vector input components are spherical; queries and outputs are Cartesian. Nonzero fill vectors are also rotated from the spherical basis.

using TestParticle
using TestParticle: FillExtrap, LazyTimeInterpolator
using TestParticle: OnTheFly, PreCompute, StorePolicy
using Meshes
using StaticArrays
using Chairmarks
using KernelAbstractions: CPU
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: 3126 samples with 378 evaluations
 min    75.085 ns (2 allocs: 64 bytes)
 median 77.603 ns (2 allocs: 64 bytes)
 mean   79.772 ns (2 allocs: 64 bytes)
 max    146.571 ns (2 allocs: 64 bytes)

julia> @be A_sph_nu($loc)
Benchmark: 3166 samples with 515 evaluations
 min    55.637 ns (2 allocs: 48 bytes)
 median 57.194 ns (2 allocs: 48 bytes)
 mean   58.471 ns (2 allocs: 48 bytes)
 max    138.559 ns (2 allocs: 48 bytes)

Uniform spherical interpolation

julia> @be B_sph($loc)
Benchmark: 3179 samples with 379 evaluations
 min    74.702 ns (2 allocs: 64 bytes)
 median 76.844 ns (2 allocs: 64 bytes)
 mean   78.666 ns (2 allocs: 64 bytes)
 max    228.631 ns (2 allocs: 64 bytes)

julia> @be A_sph($loc)
Benchmark: 2566 samples with 522 evaluations
 min    54.680 ns (2 allocs: 48 bytes)
 median 55.910 ns (2 allocs: 48 bytes)
 mean   71.022 ns (2 allocs: 48 bytes, 0.04% gc time)
 max    37.336 μs (2 allocs: 48 bytes, 99.41% gc time)

Uniform Cartesian interpolation

julia> @be B_car($loc)
Benchmark: 2620 samples with 1028 evaluations
 min    25.524 ns (2 allocs: 64 bytes)
 median 27.186 ns (2 allocs: 64 bytes)
 mean   35.660 ns (2 allocs: 64 bytes, 0.04% gc time)
 max    17.130 μs (2 allocs: 64 bytes, 99.28% gc time)

julia> @be A_car($loc)
Benchmark: 2577 samples with 1248 evaluations
 min    21.026 ns (2 allocs: 48 bytes)
 median 22.454 ns (2 allocs: 48 bytes)
 mean   30.083 ns (2 allocs: 48 bytes, 0.04% gc time)
 max    12.874 μs (2 allocs: 48 bytes, 99.36% gc time)

Non-uniform Cartesian interpolation

julia> @be B_car_nu($loc)
Benchmark: 2579 samples with 1071 evaluations
 min    25.163 ns (2 allocs: 64 bytes)
 median 27.203 ns (2 allocs: 64 bytes)
 mean   34.802 ns (2 allocs: 64 bytes, 0.04% gc time)
 max    16.343 μs (2 allocs: 64 bytes, 99.26% gc time)

julia> @be A_car_nu($loc)
Benchmark: 2690 samples with 1066 evaluations
 min    23.224 ns (2 allocs: 48 bytes)
 median 24.916 ns (2 allocs: 48 bytes)
 mean   33.379 ns (2 allocs: 48 bytes, 0.04% gc time)
 max    20.319 μs (2 allocs: 48 bytes, 99.55% gc time)

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: 2654 samples with 1052 evaluations
 min    25.152 ns (2 allocs: 64 bytes)
 median 26.933 ns (2 allocs: 64 bytes)
 mean   34.261 ns (2 allocs: 64 bytes, 0.04% gc time)
 max    16.253 μs (2 allocs: 64 bytes, 99.28% gc time)

julia> @be itp_f32($loc_f64)
Benchmark: 2502 samples with 1012 evaluations
 min    26.314 ns (2 allocs: 64 bytes)
 median 28.808 ns (2 allocs: 64 bytes)
 mean   38.141 ns (2 allocs: 64 bytes, 0.04% gc time)
 max    18.620 μs (2 allocs: 64 bytes, 99.45% 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: 1842 samples with 1 evaluation
 min    11.918 μs (12 allocs: 385.117 KiB)
 median 17.056 μs (12 allocs: 385.117 KiB)
 mean   42.167 μs (12 allocs: 385.117 KiB, 1.26% gc time)
 max    17.041 ms (12 allocs: 385.117 KiB, 99.01% gc time)

julia> @be build_interpolator($B_memory_svector, $x_memory, $x_memory, $x_memory, 3; store = $copy_store)
Benchmark: 1241 samples with 1 evaluation
 min    9.885 μs (12 allocs: 385.164 KiB)
 median 13.620 μs (12 allocs: 385.164 KiB)
 mean   42.566 μs (12 allocs: 385.164 KiB, 0.86% gc time)
 max    3.683 ms (12 allocs: 385.164 KiB, 99.24% gc time)

julia> # Zero-copy construction
       @be build_interpolator($B_memory, $x_memory, $x_memory, $x_memory, 1; store = $zero_copy_store)
Benchmark: 1486 samples with 123 evaluations
 min    174.984 ns (9 allocs: 1.047 KiB)
 median 439.728 ns (9 allocs: 1.047 KiB)
 mean   662.974 ns (9 allocs: 1.047 KiB, 0.13% gc time)
 max    241.758 μs (9 allocs: 1.047 KiB, 99.73% gc time)

julia> @be build_interpolator($B_memory_svector, $x_memory, $x_memory, $x_memory, 3; store = $zero_copy_store)
Benchmark: 1628 samples with 120 evaluations
 min    179.933 ns (9 allocs: 1.078 KiB)
 median 450.050 ns (9 allocs: 1.078 KiB)
 mean   486.300 ns (9 allocs: 1.078 KiB, 0.06% gc time)
 max    225.182 μs (9 allocs: 1.078 KiB, 99.63% 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: 2873 samples with 144 evaluations
 min    148.069 ns (2 allocs: 64 bytes)
 median 153.569 ns (2 allocs: 64 bytes)
 mean   164.418 ns (2 allocs: 64 bytes)
 max    1.469 μs (2 allocs: 64 bytes)

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

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: 3079 samples with 350 evaluations
 min    81.526 ns (2 allocs: 64 bytes)
 median 84.386 ns (2 allocs: 64 bytes)
 mean   87.274 ns (2 allocs: 64 bytes)
 max    194.580 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 GPUGrid, a multilinear interpolator over Cartesian or spherical grids.

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, so that the device interpolator can be inspected and benchmarked where no device is available:

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

(::TestParticle.GPUGrid{3, TestParticle.SphericalCoords, Tuple{TestParticle.GPUNonUniformAxis{Float64, Vector{Float64}, FastInterpolations.FillExtrap{StaticArraysCore.SVector{3, Float64}}}, TestParticle.GPUUniformAxis{Float64, FastInterpolations.FillExtrap{StaticArraysCore.SVector{3, Float64}}}, TestParticle.GPUUniformAxis{Float64, FastInterpolations.WrapExtrap}}, StaticArraysCore.SVector{3, Float64}, Array{StaticArraysCore.SVector{3, Float64}, 3}}) (generic function with 4 methods)

An equally spaced grid vector becomes a GPUUniformAxis, which locates a cell with a multiply, and any other one a GPUNonUniformAxis, which brackets it with a binary search:

julia> typeof(B_gpu)
TestParticle.GPUGrid{3, TestParticle.SphericalCoords, Tuple{TestParticle.GPUUniformAxis{Float64, FastInterpolations.FillExtrap{StaticArraysCore.SVector{3, Float64}}}, TestParticle.GPUUniformAxis{Float64, FastInterpolations.FillExtrap{StaticArraysCore.SVector{3, Float64}}}, TestParticle.GPUUniformAxis{Float64, FastInterpolations.WrapExtrap}}, StaticArraysCore.SVector{3, Float64}, Array{StaticArraysCore.SVector{3, Float64}, 3}}

julia> B_gpu.axes[1]
TestParticle.GPUUniformAxis{Float64, FastInterpolations.FillExtrap{StaticArraysCore.SVector{3, Float64}}}(0.1, 1.5151515151515151, 16, 10.0, FastInterpolations.FillExtrap{StaticArraysCore.SVector{3, Float64}}([NaN, NaN, NaN]))

julia> typeof(B_nu_gpu.axes[1])
TestParticle.GPUNonUniformAxis{Float64, Vector{Float64}, FastInterpolations.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: 2347 samples with 396 evaluations
 min    73.040 ns (2 allocs: 64 bytes)
 median 91.528 ns (2 allocs: 64 bytes)
 mean   101.072 ns (2 allocs: 64 bytes)
 max    743.621 ns (2 allocs: 64 bytes)

julia> @be B_sph($loc)
Benchmark: 1868 samples with 330 evaluations
 min    72.445 ns (2 allocs: 64 bytes)
 median 76.661 ns (2 allocs: 64 bytes)
 mean   153.396 ns (2 allocs: 64 bytes, 0.05% gc time)
 max    133.780 μs (2 allocs: 64 bytes, 99.63% gc time)

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

The device grids implement a subset of the host interpolators: linear interpolation (order = 1, the default) and the boundary conditions FillExtrap, ClampExtrap, WrapExtrap and an inclusive PeriodicBC. Moving any other interpolator to a device throws an ArgumentError.

Related API

TestParticle.build_interpolator Function

build_interpolator(A, grids::Tuple; bc, order=1, gridtype=CartesianGrid, kwargs...)
build_interpolator([gridtype,] A, grids..., order=1, bc=FillExtrap(NaN); kwargs...)

Build a callable field for prepare. Vector data has shape (3, length.(grids)...) or contains SVector{3} values. Orders are 0, 1, and 3. The tuple API requires bc: one policy for all axes or a tuple of per-axis policies. FillExtrap values must agree across axes.

PeriodicBC(endpoint=:inclusive) repeats the endpoint; :exclusive omits it (uniform-grid period NΔx). WrapExtrap() wraps over the grid span without imposing spline periodicity. FillExtrap(NaN) fills outside; ClampExtrap() extends constantly. Field lookup neither wraps particle state nor terminates trajectories.

gridtype=StructuredGrid uses spherical (r, θ, ϕ) grids and (Br, Bθ, Bϕ) components; queries and outputs are Cartesian. Supply fill/clamp for radius and colatitude and PeriodicBC with period 2π for longitude. Its positional default is (FillExtrap(NaN), FillExtrap(NaN), PeriodicBC(check = false)).

coeffs selects cubic coefficient storage. store=StorePolicy(copy=false) may alias inputs, which must remain unchanged. Device conversion requires order 1 and fill, clamp, wrap, or inclusive periodic policies; exclusive periodic fields are CPU-only.

source

TestParticle.prepare Function

prepare(; E=nothing, B=nothing, F=nothing, species=Proton, q, m, c, type, erased=false)

Build particle parameters from callable electric, magnetic, and external force fields. Omitted fields are zero. Each callable accepts x or (x, t) and returns three components. Construct numerical fields with build_interpolator before calling prepare. See Field to select the calling convention. q and m override species; c sets the relativistic speed of light. type converts particle constants.

source

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

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

source