Turing, Zygote, and Interpolations

Specific Domains Machine Learning
1

This is a follow up to a previous post I did, in which I thought everything was resolved, but it is not. I want to use Turing with both FFTW and some type of interpolation method. I am trying something of the form:

using Turing
using Distributions
using FFTW
using Random
using Zygote
using Interpolations

Turing.setadbackend(:zygote)

N = 8;
γ = 0.01;
x = LinRange(0,1,N+1)[2:end-1];

# recover on particular points
x_data = [0.3, 0.4, 0.7];

n_data = length(x_data);

# true value that we wish to recover
uᵗ(x) = x*(1-x);

# generate noisy data
Random.seed!(500); # set a seed for reproducibility
y_data = @. uᵗ(x_data) + γ * randn()

function build_field(ξ; α=1.)
    k = 1:N;
    λ = @. 1/(π*k)^(2*α);

    # fill in the nonzero entries
    # NOTE we need to multiply by 2 *N for FFT scaling
    # @. uhat[2:N+1] = 2 * N * sqrt(λ) * sqrt(2) * ξ;
    uhat = [0; @. 2 * N * sqrt(λ) * sqrt(2) * ξ; zeros(N-1)];

    # invert and get the relevant imaginary part
    u = @views imag.(ifft(uhat)[N+2:end]);
    return u
end

@model function mean_recovery(x_data, y_data)
    ξ ~ MvNormal(zeros(N), 1.)
    u = build_field(ξ);
    u_spl = LinearInterpolation(x, u);
    u_pred = u_spl(x_data);
    n_data = length(x_data)
    
    for i in 1:n_data
       y_data[i]~Normal(u_pred[i], γ)
    end
    
end
model=mean_recovery(x_data, y_data);
chain = sample(model, HMC(0.01, 10), 10^4)

which generates the error:

gradient of 7-element extrapolate(interpolate((::Vector{Float64},), ::Vector{Float64}, Gridded(Linear())), Throw()) with element type Float64:
  0.6506507222175089
  0.32945144279414945
  0.872387463894707
  0.6718748183867302
 -0.25268069044424946
 -0.5033626843984641
 -0.4305674717804948 not supported for position ([0.3, 0.4, 0.7],)

Stacktrace:
  [1] error(s::String)
    @ Base ./error.jl:33
#...continues

If, instead, I don’t use any interpolations (by, for instance, x_data = x and u_pred = u), the code works as expected. I am not particularly wedded to using Interpolations, so if this would work with a competing interpolation package, that’d be fine by me. Thanks for any suggestions.

2

I did get something working with DataInterpolations.jl but it would be nice if some of the other packages worked too.