Automatic Differentiation

Both operators and interpolators can be differentiated with reverse-mode AD through the Enzyme.jl package extension, which provides native EnzymeRules (augmented_primal/reverse) for every supported component.

All examples use DifferentiationInterface.jl, which provides a unified API over AD backends.

Implementation Status

Enzyme is the supported AD backend. The evaluation kernels are multithreaded, so they cannot be traced generically by other reverse-mode backends — differentiation goes through the provided rules. Differentiate the out-of-place forms (op(x), weights(op) * x); the in-place forms (op(y, x), mul!(y, op, x)) have no AD rules and are not supported under differentiation.

Differentiating Through Operators

The most common use case is differentiating a loss function with respect to field values while keeping the operator fixed. Create the operator once outside the loss function, then differentiate through its application.

using RadialBasisFunctions
using StaticArrays
import DifferentiationInterface as DI
import Enzyme

# Create points and operator (outside loss function)
points = [SVector{2}(0.1 + 0.8 * i / 7, 0.1 + 0.8 * j / 7) for i in 1:7 for j in 1:7]
values = sin.(getindex.(points, 1)) .+ cos.(getindex.(points, 2))

lap = laplacian(points)

# Loss function: minimize squared Laplacian
function loss(v)
    result = lap(v)
    return sum(result .^ 2)
end

# Compute gradient using DifferentiationInterface
backend = DI.AutoEnzyme(; function_annotation=Enzyme.Const)
grad = DI.gradient(loss, backend, values)
grad[1:5]  # Show first 5 gradient values

5-element Vector{Float64}:
 -895.3558338641012
  320.6587608218761
 -586.7957170069684
 -441.1121212423749
 -459.14643578997976

The function_annotation=Enzyme.Const tells Enzyme that data captured by the loss closure (here the operator lap) is constant — we differentiate w.r.t. the input values, never the captures. This is required for the operator-capture pattern above.

This works with any operator type:

# Gradient operator (vector-valued)
∇f = gradient(points)

function loss_grad(v)
    result = ∇f(v)
    return sum(result .^ 2)
end

grad = DI.gradient(loss_grad, backend, values)
grad[1:5]

5-element Vector{Float64}:
 -27.12103246485196
 -24.840445463329793
 -32.08909142649798
 -27.70040059366693
 -34.349507250253154

# Partial derivative operator
∂x = partial(points, 1, 1)

function loss_partial(v)
    result = ∂x(v)
    return sum(result .^ 2)
end

grad = DI.gradient(loss_partial, backend, values)
grad[1:5]

5-element Vector{Float64}:
 -34.550520724617165
 -24.171745925500087
 -34.79386557444219
 -29.245023400774606
 -33.51450706876448

Differentiating Through Interpolators

When differentiating through interpolation, the Interpolator must be constructed inside the loss function since changing the input values changes the interpolation weights.

N_interp = 30
points_interp = [SVector{2}(0.5 + 0.4 * cos(2π * i / N_interp), 0.5 + 0.4 * sin(2π * i / N_interp)) for i in 1:N_interp]
values_interp = sin.(getindex.(points_interp, 1))
eval_points = [SVector{2}(0.5, 0.5), SVector{2}(0.6, 0.6)]

# Loss function - must rebuild interpolator inside
function loss_interp(v)
    interp = Interpolator(points_interp, v)
    result = interp(eval_points)
    return sum(result .^ 2)
end

grad = DI.gradient(loss_interp, backend, values_interp)
grad[1:5]

5-element Vector{Float64}:
 -180.6580738786631
  205.53075789759478
 -150.9770244184113
   34.060282950687025
   93.49176118063092

Differentiating Basis Functions Directly

For low-level control, you can differentiate basis function evaluations directly. This is useful for custom applications or understanding the underlying derivatives.

x = [0.5, 0.5]
xi = [0.3, 0.4]

# PHS basis
phs = PHS(3)
function loss_phs(xv)
    return phs(xv, xi)^2
end

grad = DI.gradient(loss_phs, backend, x)

2-element Vector{Float64}:
 0.003
 0.0014999999999999998

All basis types are supported:

# IMQ basis
imq = IMQ(1.0)
function loss_imq(xv)
    return imq(xv, xi)^2
end

grad = DI.gradient(loss_imq, backend, x)

2-element Vector{Float64}:
 -0.36281179138321995
 -0.18140589569160992

# Gaussian basis
gauss = Gaussian(1.0)
function loss_gauss(xv)
    return gauss(xv, xi)^2
end

grad = DI.gradient(loss_gauss, backend, x)

2-element Vector{Float64}:
 -0.7238699344287678
 -0.3619349672143838

Differentiating Weight Construction

For advanced use cases like mesh optimization or shape parameter tuning, you can differentiate through the weight construction process using the internal _build_weights function. It returns a StencilWeights, whose dense matrix of stencil weight values is accessed with parent(W). The AD path is host-resident and stencil-major by design — weight construction (and therefore its pullback) always runs on CPU over the k × N_eval layout, regardless of any device orientation the applied operator's weights may use.

points_weights = [SVector{2}(0.1 + 0.8 * i / 5, 0.1 + 0.8 * j / 5) for i in 1:5 for j in 1:5]
N_weights = length(points_weights)
adjl = RadialBasisFunctions.find_neighbors(points_weights, 10)
basis = PHS(3; poly_deg=2)
ℒ = Partial(1, 1)  # First derivative in x

# Loss function w.r.t. point positions
function loss_weights(pts)
    pts_vec = [SVector{2}(pts[2*i-1], pts[2*i]) for i in 1:N_weights]
    W = RadialBasisFunctions._build_weights(ℒ, pts_vec, pts_vec, adjl, basis)
    return sum(parent(W) .^ 2)
end

pts_flat = reduce(vcat, points_weights)
grad = DI.gradient(loss_weights, backend, pts_flat)
grad[1:6]  # Gradients for first 3 points (x,y pairs)

6-element Vector{Float64}:
 1359.8433242375106
 -160.1093957698763
 1271.584760206346
 -142.32197568606665
 1001.5986906332257
  -42.813602864985924

This also works with the Laplacian operator and different basis types:

ℒ_lap = Laplacian()
basis_imq = IMQ(1.0; poly_deg=2)

function loss_weights_lap(pts)
    pts_vec = [SVector{2}(pts[2*i-1], pts[2*i]) for i in 1:N_weights]
    W = RadialBasisFunctions._build_weights(ℒ_lap, pts_vec, pts_vec, adjl, basis_imq)
    return sum(parent(W) .^ 2)
end

grad = DI.gradient(loss_weights_lap, backend, pts_flat)
grad[1:6]

6-element Vector{Float64}:
      -5.435978381515452e6
 -780158.3431338153
       1.1478421407828113e7
 -385860.5538132653
      -5.479421531195566e6
  590390.441464169

Supported Components

Component	Enzyme
Operator evaluation (op(values))	✓
Weight matvec (weights(op) * x, W * x)	✓
Interpolator construction	✓
Interpolator evaluation	✓
Basis functions (PHS, IMQ, Gaussian)	✓
Weight construction (_build_weights)	✓
Shape parameter (ε) differentiation	✓

In-place evaluation (op(y, x), mul!(y, op, x)) is not differentiable — use the out-of-place forms inside losses.