Gradient error in `trace_gc_drifts!`

[henry2004y]

When the magnetic field is given in a numerical form, ForwardDiff cannot figure out its gradient:

using TestParticle
using OrdinaryDiffEq
using StaticArrays
using CircularArrays
using LinearAlgebra

"""
	 trilinear_interp(x, y, z, Q)

Trilinear interpolation for x1,y1,z1=(0,0,0) and x2,y2,z2=(1,1,1).
Q's are surrounding points such that Q000 = F[0,0,0], Q100 = F[1,0,0], etc.
"""
function trilinear_interp(
	x::T,
	y::T,
	z::T,
	Q000,
	Q100,
	Q010,
	Q110,
	Q001,
	Q101,
	Q011,
	Q111,
) where {T <: Real}
	oneT = one(T)
	mx = oneT - x
	my = oneT - y
	mz = oneT - z
	fout::T =
		Q000 * mx * my * mz +
		Q100 * x * my * mz +
		Q010 * y * mx * mz +
		Q110 * x * y * mz +
		Q001 * mx * my * z +
		Q101 * x * my * z +
		Q011 * y * mx * z +
		Q111 * x * y * z
end

"""
	 grid_interp(x, y, z, field, ix, iy, iz, xsize, ysize)

Interpolate a value at (x,y,z) in a field. `ix`,`iy` and `iz` are indexes for x, y and z
locations (0-based). `xsize` and `ysize` are the sizes of field in X and Y.
"""
grid_interp(
	x::T,
	y::T,
	z::T,
	ix::Int,
	iy::Int,
	iz::Int,
	field::AbstractArray{U, 3},
) where
	{T <: Real, U <: Number} =
	trilinear_interp(x - ix, y - iy, z - iz,
		field[ix+1, iy+1, iz+1],
		field[ix+2, iy+1, iz+1],
		field[ix+1, iy+2, iz+1],
		field[ix+2, iy+2, iz+1],
		field[ix+1, iy+1, iz+2],
		field[ix+2, iy+1, iz+2],
		field[ix+1, iy+2, iz+2],
		field[ix+2, iy+2, iz+2],
	)

function get_field_interp(gridx, gridy, gridz, Bx, By, Bz; B0 = 1)
	dx = gridx[2] - gridx[1]
	dy = gridy[2] - gridy[1]
	dz = gridz[2] - gridz[1]
	# Return field value at a given location.
	function get_field(xu)
		r = @view xu[1:3]

		x = (r[1] - gridx[1]) / dx
		y = (r[2] - gridy[1]) / dy
		z = (r[3] - gridz[1]) / dz

		# Find surrounding points (0-based indices)
		ix = floor(Int, x)
		iy = floor(Int, y)
		iz = floor(Int, z)

		SVector{3, Float64}(
			grid_interp(x, y, z, ix, iy, iz, Bx) / B0,
			grid_interp(x, y, z, ix, iy, iz, By) / B0,
			grid_interp(x, y, z, ix, iy, iz, Bz) / B0,
		)
	end

	return get_field
end

using Batsrus

function load_B()
	filedir = "/nobackupp18/cdong1/run_Turbulence_and_Reconnection_git_2/test-2000s_1000_S_8e4_Roe_Bg4/GM/IO2"

	fnamebx = joinpath(filedir, "3d__var_3_n00006288.h5")
	fnameby = joinpath(filedir, "3d__var_4_n00006288.h5")
	fnamebz = joinpath(filedir, "3d__var_5_n00006288.h5")

	fbx = BatsrusHDF5Uniform(fnamebx)
	bx, (xl_new, yl_new, zl_new), (xu_new, yu_new, zu_new) = extract_var(fbx, "bx")
	bx = CircularArray(bx)

	fby = BatsrusHDF5Uniform(fnameby)
	by, _, _ = extract_var(fby, "by")
	by = CircularArray(by)

	fbz = BatsrusHDF5Uniform(fnamebz)
	bz, _, _ = extract_var(fbz, "bz")
	bz = CircularArray(bz)

	bx, by, bz
end

bx, by, bz = load_B()
nx, ny, nz = size(bx)

const Lx = nx
const Ly = ny
const Lz = nz

gridx = range(0.5, nx - 0.5, step = 1.0)
gridy = range(0.5, ny - 0.5, step = 1.0)
gridz = range(0.5, nz - 0.5, step = 1.0)
# Background mean field
B0 = 4.102587

#q/m, Field(E), Field(B)
B = get_field_interp(gridx, gridy, gridz, bx, by, bz; B0) |> TestParticle.Field

Bmag(x) = √(B(x) ⋅ B(x))

∇B = ForwardDiff.gradient(Bmag, [0.0,0.0,0.0])

julia> ∇B = ForwardDiff.gradient(Bmag, [0.0,0.0,0.0])
ERROR: MethodError: no method matching Float64(::ForwardDiff.Dual{ForwardDiff.Tag{typeof(Bmag), Float64}, Float64, 3})
The type `Float64` exists, but no method is defined for this combination of argument types when trying to construct it.

Closest candidates are:
  (::Type{T})(::Real, ::RoundingMode) where T<:AbstractFloat
   @ Base rounding.jl:265
  (::Type{T})(::T) where T<:Number
   @ Core boot.jl:900
  Float64(::IrrationalConstants.Twoinvπ)
   @ IrrationalConstants ~/.julia/packages/IrrationalConstants/lWTip/src/macro.jl:131
  ...

[henry2004y]

With the direct wrapper of the interpolating function using Interpolations.jl, this AD issue may be resolved.