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: 3155 samples with 233 evaluations
 min    123.146 ns (2 allocs: 64 bytes)
 median 125.124 ns (2 allocs: 64 bytes)
 mean   126.802 ns (2 allocs: 64 bytes)
 max    243.330 ns (2 allocs: 64 bytes)

julia> @be A_sph_nu($loc)
Benchmark: 3071 samples with 282 evaluations
 min    100.610 ns (2 allocs: 48 bytes)
 median 104.057 ns (2 allocs: 48 bytes)
 mean   107.223 ns (2 allocs: 48 bytes)
 max    1.189 μs (2 allocs: 48 bytes)

Uniform spherical interpolation

julia> @be B_sph($loc)
Benchmark: 3178 samples with 230 evaluations
 min    122.709 ns (2 allocs: 64 bytes)
 median 124.970 ns (2 allocs: 64 bytes)
 mean   126.897 ns (2 allocs: 64 bytes)
 max    242.452 ns (2 allocs: 64 bytes)

julia> @be A_sph($loc)
Benchmark: 3123 samples with 307 evaluations
 min    93.495 ns (2 allocs: 48 bytes)
 median 95.550 ns (2 allocs: 48 bytes)
 mean   97.491 ns (2 allocs: 48 bytes)
 max    208.622 ns (2 allocs: 48 bytes)

Uniform Cartesian interpolation

julia> @be B_car($loc)
Benchmark: 3161 samples with 602 evaluations
 min    46.515 ns (2 allocs: 64 bytes)
 median 48.262 ns (2 allocs: 64 bytes)
 mean   48.993 ns (2 allocs: 64 bytes)
 max    100.834 ns (2 allocs: 64 bytes)

julia> @be A_car($loc)
Benchmark: 3160 samples with 662 evaluations
 min    42.344 ns (2 allocs: 48 bytes)
 median 44.160 ns (2 allocs: 48 bytes)
 mean   44.777 ns (2 allocs: 48 bytes)
 max    97.140 ns (2 allocs: 48 bytes)

Non-uniform Cartesian interpolation

julia> @be B_car_nu($loc)
Benchmark: 2203 samples with 576 evaluations
 min    49.241 ns (2 allocs: 64 bytes)
 median 76.714 ns (2 allocs: 64 bytes)
 mean   74.562 ns (2 allocs: 64 bytes)
 max    411.748 ns (2 allocs: 64 bytes)

julia> @be A_car_nu($loc)
Benchmark: 3163 samples with 613 evaluations
 min    45.697 ns (2 allocs: 48 bytes)
 median 47.788 ns (2 allocs: 48 bytes)
 mean   48.441 ns (2 allocs: 48 bytes)
 max    101.964 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: 3338 samples with 570 evaluations
 min    47.263 ns (2 allocs: 64 bytes)
 median 48.723 ns (2 allocs: 64 bytes)
 mean   49.563 ns (2 allocs: 64 bytes)
 max    129.458 ns (2 allocs: 64 bytes)

julia> @be itp_f32($loc_f64)
Benchmark: 3176 samples with 575 evaluations
 min    48.840 ns (2 allocs: 64 bytes)
 median 50.685 ns (2 allocs: 64 bytes)
 mean   51.559 ns (2 allocs: 64 bytes)
 max    121.557 ns (2 allocs: 64 bytes)

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: 148 samples with 1 evaluation
 min    27.522 μs (5 allocs: 384.398 KiB)
 median 39.595 μs (5 allocs: 384.398 KiB)
 mean   675.525 μs (5 allocs: 384.398 KiB, 0.67% gc time)
 max    94.286 ms (5 allocs: 384.398 KiB, 99.85% gc time)

julia> @be build_interpolator($B_memory_svector, $x_memory, $x_memory, $x_memory, 3; store = $copy_store)
Benchmark: 740 samples with 1 evaluation
 min    19.018 μs (5 allocs: 384.430 KiB)
 median 26.415 μs (5 allocs: 384.430 KiB)
 mean   132.366 μs (5 allocs: 384.430 KiB, 0.40% gc time)
 max    42.174 ms (5 allocs: 384.430 KiB, 99.73% gc time)

julia> # Zero-copy construction
       @be build_interpolator($B_memory, $x_memory, $x_memory, $x_memory, 1; store = $zero_copy_store)
Benchmark: 1136 samples with 401 evaluations
 min    53.571 ns (2 allocs: 320 bytes)
 median 61.864 ns (2 allocs: 320 bytes)
 mean   206.661 ns (2 allocs: 320 bytes, 0.09% gc time)
 max    122.464 μ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: 1551 samples with 179 evaluations
 min    73.631 ns (2 allocs: 352 bytes)
 median 116.715 ns (2 allocs: 352 bytes)
 mean   335.617 ns (2 allocs: 352 bytes, 0.06% gc time)
 max    283.047 μ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: 3121 samples with 102 evaluations
 min    284.451 ns (2 allocs: 64 bytes)
 median 289.941 ns (2 allocs: 64 bytes)
 mean   293.190 ns (2 allocs: 64 bytes)
 max    637.824 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: 3162 samples with 171 evaluations
 min    164.871 ns (2 allocs: 64 bytes)
 median 171.251 ns (2 allocs: 64 bytes)
 mean   173.220 ns (2 allocs: 64 bytes)
 max    339.170 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: 3233 samples with 222 evaluations
 min    130.509 ns (2 allocs: 64 bytes)
 median 133.581 ns (2 allocs: 64 bytes)
 mean   135.099 ns (2 allocs: 64 bytes)
 max    269.279 ns (2 allocs: 64 bytes)

julia> @be B_sph($loc)
Benchmark: 3122 samples with 240 evaluations
 min    120.971 ns (2 allocs: 64 bytes)
 median 122.892 ns (2 allocs: 64 bytes)
 mean   124.312 ns (2 allocs: 64 bytes)
 max    255.846 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 ( ).

• 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