TypeError when sampling for a forced differential equation problem

I am new to Julia and Turing and am trying to fit a forced 0-D box ODE to data, but I get type error when doing sampling.

The tutorial (https://turing.ml/dev/tutorials/10-bayesian-differential-equations/) gives an example for unforced LV model. Following this page (https://stackoverflow.com/questions/49428939/solve-system-of-odes-with-read-in-external-forcing), I added an interpolation handle of the forcing as a parameter of the ODE, and the ODE solver runs without problems.

When I combine the forced ODE with Turing, however, I get a type error during HMC's gradient calculation, quoting ”TypeError: in typeassert, expected Float64, got a value of type ForwardDiff.Dual{Nothing, Float64, 3}”.

Below are my codes that should reproduce the error I got.

I appreciate any help, Thanks.

Below are my codes:

using Interpolations, DifferentialEquations, Plots, Turing, LinearAlgebra

# Define ODE
function f(du,u,p,t)
      α, β, F = p
      du[1] = α * u[1] + β * F(t) # Interpolations.jl interpolates via ()
end

# Define initial value problem
begin
    # Define forcing 
    time_forcing = -1.:9.
    data_forcing = [10,0,0,0,0,0,0,0,0, 0, 0]
    F = Interpolations.scale(interpolate(data_forcing, BSpline(Linear())), time_forcing)
    α = -0.5
    β = 1
    p_lin = (α, β, F) 

    # Define u0 and tspan
    u0 = [0.]
    tspan = (-1.,9.) # Note, that we would need to extrapolate beyond 
    ode_lin = ODEProblem(f,u0,tspan,p_lin)
end

# Generate data
sol  = solve(ode_lin, Tsit5(); saveat=1)
data = Array(sol) + 0.2 * randn(size(Array(sol)))

@model function fit_simple_box(data, F, ode_lin)
    # Prior distributions.
    σ ~ InverseGamma(2, 3)
    α ~ Normal(0, 3)
    β ~ Normal(0, 3)

    # Simulate Lotka-Volterra model. 
    p = (α, β, F)
    predicted = solve(ode_lin, Tsit5(); p=p, saveat=1)

    # Observations.
    for i in 1:length(predicted)
        data ~ Normal(predicted[i][1], σ^2)
    end

    return nothing
end

model = fit_simple_box(data, F, ode_lin)
chain = sample(model, NUTS(0.65), MCMCSerial(), 1000, 2)

TuringLang / Turing.jl

TypeError when sampling for a forced differential equation problem #1837