StackOverflowError trying to fit a difference equation with Turing

I am trying to sample a difference equation using Turing but encounter errors.

Basically, it’s a SIR model that solves fine, with a quick solver function written below (I wrote my own very simple solver to understand the basics and see how the solver could be altered/extended for modeling purposes).

Then, I add random noise to the solution and use it as data for the Turing model, in order to find the parameters of the SIR model using Turing. I used this strategy successfully with the Lokta Voltera model or the tutorial (see my modification of the tutorial here), but I can’t make it work with my own simple solver.

Any help is appreciated.

Julia version: 1.5.0.

Error

chain = sample(model, NUTS(.65),1000)

ERROR: StackOverflowError:
Stacktrace:
 [1] ScalMat at /Users/aamzallag/.julia/packages/PDMats/8bndI/src/scalmat.jl:5 [inlined] (repeats 2 times)
 [2] convert at /Users/aamzallag/.julia/packages/PDMats/8bndI/src/scalmat.jl:13 [inlined]
 [3] convert at /Users/aamzallag/.julia/packages/PDMats/8bndI/src/scalmat.jl:14 [inlined]
 [4] MvNormal(::Array{Real,1}, ::PDMats.ScalMat{Float64}) at /Users/aamzallag/.julia/packages/Distributions/jFoHB/src/multivariate/mvnormal.jl:202 (repeats 79984 times)

Code

using Turing, Distributions

# Import MCMCChain, Plots, and StatsPlots for visualizations and diagnostics.
using MCMCChains, Plots, StatsPlots

# Set a seed for reproducibility.
using Random
Random.seed!(12);

# Defining the dynamics
function sir_model(du,u,p,t)
    S, I, R = u
    β, γ = p
    du[1] = dS = -β*S*I/N
    du[2] = dI = β*S*I/N - γ*I
    du[3] = dR = γ*I
end


### trying a simple discrete time solver 

# this solver takes a function func_ describing
# the dynamics (a la du/dt = f(u,p,t)) and simulate
# the dynamics with a simple loop, and store the values
# in variable u
function diffsolve(func_, u0, p, n=100)
    du = [0., 0., 0.]
    u_t = copy(u0)
    #u = [u0]
    u = Array{Array{Real,1},1}() # if Real not specified, 
    # I get the error: in typeassert, expected Float64, got ForwardDiff.Dual{Nothing,Float64,3}
    push!(u, u0)
    for t in 1:n
        func_(du, u_t, p, 1)
        u_t = u_t + du
        push!(u, u_t)
    end
    #hcat(u...)
    u
end

N=10e3 # total population
conf=10 # confirmed cases at the beginning
u0 = [N-conf,conf,0.0] #.* N


## Trying the solver and plotting
x=diffsolve(sir_model, u0, [0.8, 0.1], 50)
x=hcat(x...)
# add noise
#odedata = x + 0.1*maximum(x) .* randn(size(x))
odedata = x .* exp.(0.1 .* randn(size(x)))

plot(x')
scatter!(odedata')


Turing.setadbackend(:forwarddiff)

@model function fitsir(data)
    σ ~ InverseGamma(1, 1)
    #σ ~ truncated(Normal(0., 0.5),0,5)
    β ~ truncated(Normal(0., 0.5),0,5)
    γ ~ truncated(Normal(0., 0.5),0,5)
    
    p = [β,γ]
    predicted = diffsolve(sir_model, u0, p, 10)

    for i = 1:length(predicted)
        data[:,i] ~ MvNormal(predicted[i], σ)
    end
end

model = fitsir(odedata)
chain = sample(model, NUTS(.65),1000)

# ERROR: StackOverflowError:
# Stacktrace:
#  [1] ScalMat at /Users/aamzallag/.julia/packages/PDMats/8bndI/src/scalmat.jl:5 [inlined] (repeats 2 times)
#  [2] convert at /Users/aamzallag/.julia/packages/PDMats/8bndI/src/scalmat.jl:13 [inlined]
#  [3] convert at /Users/aamzallag/.julia/packages/PDMats/8bndI/src/scalmat.jl:14 [inlined]
#  [4] MvNormal(::Array{Real,1}, ::PDMats.ScalMat{Float64}) at /Users/aamzallag/.julia/packages/Distributions/jFoHB/src/multivariate/mvnormal.jl:202 (repeats 79984 times)

The way that you write the array made it not concretely typed and then it hits a bad dispatch.

Thanks Chris,
when I try to declare u as

u = []

or

u = Array{Array{Float64,1},1}()

I get
ERROR: TypeError: in typeassert, expected Float64, got ForwardDiff.Dual{Nothing,Float64,3}
I don’t know else how to collect the results from the loop

You want to make this be dual valued, and only your parameters are dual valued, so the easiest thing to do is:

predicted = diffsolve(sir_model, eltype(p).(u0), p, 10)

and

u = [u0]

on the inside. That should propagate the duals in a type-stable way. I’ll just quickly note that DifferentialEquations.jl’s AD integration takes care of these cases for you.

1 Like

Thanks Chris,
unfortunately this makes no difference.

Good to know that DifferentialEquations.jl handles it for the user. It would be nice if Turing could be handling these kinds of problems, or return an informative error, for problems that are not differential equation problems.

Turing handles these issues using a specific syntax. You can find details on turing.ml in the documentation on AD.

Thank you @trappmartin - I am not sure there is something relevant to my problem in the AD section - isn’t it rather in the performance tips, the section on type stability? Performance Tips

1 Like

This was answered by @wupeifan on Julia Slack. Replacing the definition of diffsolve by what follows resolves the problem:

function diffsolve(func_, u0, p, n=100)
    u_t = deepcopy(u0)
    T = promote_type(eltype(u0), eltype(p))
    u = Array{Array{T,1},1}() 
    du = Vector{T}(undef, 3) 

    push!(u, u0)
    for t in 1:n
        func_(du, u_t, p, 1)
        u_t = u_t + du
        push!(u, u_t)
    end
    u
end

This solution converts both du and u0 to the same type as p (not only u0 as in a previous answer) and hence, it works.

1 Like