Support for gradients w.r.t. data

[lxvm]

Hi,

Thank you for developing this very capable package. It would be a very useful feature in applications such as topology optimization to be able to take derivatives w.r.t. the data for an interpolant. I made this example trying to use Zygote.jl to take such a derivative:

using FastInterpolations
import Zygote
using Random

Random.seed!(0)

Nx = Ny = 5
data = rand(Nx, Ny)
x = range(0, 1, length=Nx)
y = range(0, 1, length=Ny)

x0 = (0.2, 0.3)

function f0(data, x0)
    itp = cubic_interp((x, y), data)
    return itp(x0)
end

@show f0(data, x0) # 0.5931453845794765
Zygote.withgradient(f0, data, x0)

However, the autodiff fails with the following error:

ERROR: LoadError: Compiling Tuple{typeof(FastInterpolations._get_bank), FastInterpolations.GlobalRegistry, Type{FastInterpolations.CacheBank{FastInterpolations.CacheEntry{Float64, Deriv1{Float64}, Deriv1{Float64}, StepRangeLen{Float64, Base.TwicePrecision{Float64}, Base.TwicePrecision{Float64}, Int64}, FastInterpolations.ScalarSpacing{Float64}}}}}: ArgumentError: array must be non-empty

In addition, some topology optimization algorithms make use of gradients and hessians in which case it would be helpful to get gradients of the data with respect to those, like in this code:

function f1(data, x0)
    itp = cubic_interp((x, y), data)
    return sum(abs2, gradient(itp, x0))
end

f1(data, x0)
Zygote.withgradient(f1, data, x0)

function f2(data, x0)
    itp = cubic_interp((x, y), data)
    return sum(abs2, hessian(itp, x0))
end

f2(data, x0)
Zygote.withgradient(f2, data, x0)

Would supporting these kinds of gradients be possible? I would be interested in helping to write rrules.

[mgyoo86]

Hi @lxvm,
thank you for the detailed issue and your interest in contributing!
I'm not an AD expert, but would like share the current status, and some workarounds as far as I understand.

Supporting ∂f/∂data (adjoint of the interpolation operator) is on my roadmap and I'll be prioritizing it as a next development item — as also discussed in issue (#56). The goal of a dedicated adjoint API would compute this directly without any AD backend and also enable full Zygote support via rrule.

In the meantime, here are some workarounds that might help. You can use ForwardDiff to compute ∂f/∂data — it should work for all methods (constant, linear, quadratic, cubic) in both 1D and ND. If you need Zygote specifically, it currently supports ∂f/∂data only for 1D linear/constant; ND and higher-order methods hit Zygote's array mutation limitation.

Also, as you'll see in the examples below, for query-coordinate derivatives (∂f/∂xq) we'd recommend using our built-in analytical gradient/hessian API where possible — it's much faster than AD backends and zero-allocation.
See the Vector-Calculus-Functions docs for details.

∂f/∂data via ForwardDiff (Available Now)

using FastInterpolations, ForwardDiff

Nx = Ny = 5
data = rand(Nx, Ny)
x = range(0, 1, length=Nx)
y = range(0, 1, length=Ny)
x0 = (0.2, 0.3)

# ∂f/∂data — works for linear, quadratic, cubic
ForwardDiff.gradient(d -> cubic_interp((x, y), d, x0), data)

Your three use cases from the issue all work:

# (A) gradient of interpolant w.r.t. data
ForwardDiff.gradient(d -> cubic_interp((x, y), d, x0), data)

# (B) gradient of ‖∇f‖² w.r.t. data
ForwardDiff.gradient(data) do d
    itp = cubic_interp((x, y), d)
    sum(abs2, FastInterpolations.gradient(itp, x0))
end

# (C) gradient of ‖Hf‖² w.r.t. data
ForwardDiff.gradient(data) do d
    itp = cubic_interp((x, y), d)
    sum(abs2, FastInterpolations.hessian(itp, x0))
end

For more details, see the Adjoint via AD (∂f/∂y) documentation page.

Zygote Status

For ∂f/∂data, Zygote currently only works with 1D linear/constant. ND and quadratic/cubic construction paths involve in-place array mutations.

For ∂f/∂xq (query-coordinate sensitivity on a pre-built interpolant), Zygote works for all methods (except SeriesInterpolant) and dimensions via our ChainRulesCore extension.

Built-in Analytical ∂f/∂xq (Faster Alternative)

If your workflow also involves derivatives w.r.t. query coordinates, we provide native analytical gradient and hessian functions that are much faster than AD backends and zero-allocation:

itp = cubic_interp((x, y), data)

FastInterpolations.gradient(itp, (0.2, 0.3))   # (∂f/∂x, ∂f/∂y)  ~21 ns
FastInterpolations.hessian(itp, (0.2, 0.3))     # 2×2 Hessian       ~34 ns

# vs Zygote.gradient: ~40 ns (via rrule), ForwardDiff.gradient: ~335 ns
# vs ForwardDiff.hessian: ~659 ns

These work for all methods (linear, quadratic, cubic) in any dimension. See the AD Support (ND) Docs.

Roadmap

A dedicated adjoint operator that computes ∂f/∂data directly (without AD) is planned as a next development item. I'll update this issue when progress is made.

If you'd like to contribute rrules for the adjoint operator, that would be very welcome!

[lxvm]

Thank you very much for the status update and it is great to have this already supported in 1D! I will look at the implementation of cubic_adjoint later this week to see if I can extend it to ND.

[mgyoo86]

FYI, the N-dimensional cubic_adjoint implementation and the corresponding rrules for AD packages are almost ready.
I'll clean things up a bit before opening a PR, and then release a new version afterward.

Just to note that the 1D cubic_adjoint has not been released yet either. I plan to include it in the same release after the ND version is merged. In the meantime, if you'd like to try it, you can use the master branch.

For now, the adjoint is only implemented for the cubic-spline first (hopefully, this could suffice your needs).
I’m planning to add the same functionality for quadratic and linear/constant interpolation in a subsequent release.

[stevengj]

Note that ForwardDiff is not a great option here because forward-mode AD is extremely inefficient for computing the derivative with respect to a large number of data points. We really have to have an rrule here … especially for 2+ dimensions where the amount of data is huge. I agree that for performance reasons you probably want an analytical adjoint that you can then plug in as an rrule into AD systems rather than relying on AD (even if AD worked).

See also #56.

[mgyoo86]

Quick Update:
I've implemented ND cubic adjoint, and baked rrules for Zygote and Enzyme in PR #71, which I believe resolves this issue for cubic spline.

As @stevengj commented, I confirm that ForwardDiff is extremly slower than other methods. The larger grid the slower performance, like ~1000x times slower than others for (100x100) grid.

Although Zygote and Enzyme uses the native API through rrules, but they are ~2x slower than using native API due to some allocations and overhead via rrule path.

After completing some final tasks like documenting and adding simple features, I'll release a new version and update here.

[stevengj]

This looks great, thanks!

[mgyoo86]

v0.4.2 has just been released.
It adds native Adjoint API and AD support for cubic interpolation, which covers the main use case, so I’m closing this issue for now.
Please give it a try, and let us know if you run into anything unexpected.

Support for the other methods (constant, linear, and quadratic) is planned for upcoming releases.