Automatic Differentiation

Both operators and interpolators can be differentiated with reverse-mode AD. Two backends are supported through package extensions:

• Enzyme.jl - Native EnzymeRules for high-performance reverse-mode AD; the recommended default backend

• Mooncake.jl - Reverse-mode AD with support for mutation; a fully supported alternative

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

Implementation Status

Both backends have native AD rule implementations. Enzyme.jl uses EnzymeRules (augmented_primal/reverse) and Mooncake.jl uses native rrule!! with @is_primitive.

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}:
 -955.3201507351854
  395.9998781032441
 -714.7851809604294
 -379.07044853290387
 -462.56171290353996

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.121032464851965
 -24.84044546332979
 -32.089091426498015
 -27.70040059366695
 -34.34950725025317

# 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.55052072461721
 -24.171745925500097
 -34.7938655744422
 -29.2450234007746
 -33.51450706876449

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.

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(W.nzval .^ 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(W.nzval .^ 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	Mooncake
Operator evaluation (op(values))	✓	✓
Interpolator construction	✓	✓
Interpolator evaluation	✓	✓
Basis functions (PHS, IMQ, Gaussian)	✓	✓
Weight construction (_build_weights)	✓	✓
Shape parameter (ε) differentiation	✓	✓

Using Mooncake Backend

Switch to Mooncake by changing the backend:

import DifferentiationInterface as DI
import Mooncake

backend = DI.AutoMooncake(; config=nothing)
grad = DI.gradient(loss, backend, values)