ODE model in Turing

Hi everyone,

I’m new to Julia and would like to use the Turing package for Bayesian inference on high dimensional ordinary differential equations. With this in mind, efficiency is key.
To get started I implemented a linear toy model with two variables based on lotka volterra tutorial provided here: Bayesian Estimation of Differential Equations

I also tried to work with with the DataFrames packages in a long dataformat as this is how I will import my own data at a later time point.

Unfortunately, the sampler does not identify the correct parameters and neither does it. Furthermore, the sampling is very slow.

What am I doing wrong / what could be improved? Again, an efficient implementation is key as I’d like to potentially extent this code to a system of up to 30 variables (with say a 100 parameters).

I first define the ODEs and generate fake data:

using Turing, Distributions, DifferentialEquations, MCMCChains, Plots, StatsPlots, Random, DataFrames, CSV
using Statistics: mean, std
using Random:seed!
seed!(123)

### Define the ODEs
function linear_sdm(du,u,p,t)
  x, y = u
  α, β, γ, δ  = p
  du[1] = α * x + β * y 
  du[2] = γ * x + δ * y 
end

### Define ``real'' parameter values
p_mean = [0.025, -0.0075, -0.05, 0.03]
randomness_size = 0.2

### Create initial conditions for simulated individuals
u0_1 = [-0.8, -0.6, -0.4, -0.2, 0, 0.2, 0.4, 0.6, 0.8, 1.0];
u0_2 = [1, 0.8, 0.6, 0.4, 0.2, 0.2, 0.4, 0.6, 0.8, 1.0];

t_span = (0.0, 48.0)
N = 10
dt = 3
add_missings = false
df = []  # Allocate memory for this dataframe
prob = ODEProblem(linear_sdm, Array([0,0]), t_span, p_mean) 

for i = 1:N
    #println(i)
    u0_i = [u0_1[i], u0_2[i]]
    prob1 = remake(prob, u0=u0_i)
    sol = solve(prob1, Tsit5())
    #display(plot(sol))  # Plot the set of ODEs

    ### Generate data from ODEs
    sol1 = solve(prob1, Tsit5(), saveat=dt)
    odedata = Array(sol1) + randomness_size * randn(size(Array(sol1)))
    
    if add_missings   # Generate missings in the simulated data
        odedata[1,8] = NaN
        odedata[1,3] = NaN
        odedata[1,9] = NaN
        odedata[1,13] = NaN

        odedata[2,1] = NaN
        odedata[2,6] = NaN
        odedata[2,12] = NaN
        odedata[2,16] = NaN
    end

    # Store in dataframe odedata 
    df_i = DataFrame(ADAS13 = odedata[1,:], 
                     FDG = odedata[2,:], 
                     Time = sol1.t,
                     RID = repeat([i], 
                     length(sol1.t)))
    
    if i == 1
        df = deepcopy(df_i)
    else
        df = reduce(vcat, [df, df_i])  # Join dataframe with original dataframe
    end
    
end

Then I run the model, which gives very poor results

### Run Bayesian model
Turing.setadbackend(:forwarddiff)
using Suppressor

@suppress_err begin  # Silence warnings
    @model function sdm_model(data, prob, N, t_span, dt)
        σ ~ InverseGamma(0.2, 0.05) 
        α ~ Normal(0.1, 0.05)
        β ~ Normal(-0.1, 0.05) 
        γ ~ Normal(-0.1, 0.05)
        δ ~ Normal(0.1, 0.05) 

        p = [α,β,γ,δ]
        
        for k in 1:N
            u0_i = Array(df[in([k]).(df.RID), ["ADAS13", "FDG"]][1,:])  # Set baseline
            prob_i = remake(prob, p=p, u0=u0_i)
            predicted = solve(prob_i, Tsit5(), saveat=dt)

            for i = 2:length(predicted) # where {T <: Real}
                data = Array(df[in([k]).(df.RID), ["ADAS13", "FDG"]])[i,:]
                data ~ MvNormal(predicted[i], σ) #
            end
        end

    end

    model = sdm_model(df, prob, N, t_span, dt) 
    chain = sample(model, NUTS(.65), MCMCThreads(), 1000, 3, progress=true)
end

Any suggestions? Thank you

Try decreasing the solver tolerance a bit? That can usually help stabilize it.

Thanks for your reply @ChrisRackauckas .

I lowered the solver tolerance to “reltol=1e-10, abstol=1e-10” but this did not improve the sampler. Is that what you meant?

Yeah then everything seems fine on the solver side so I’ll let someone on the Turing side investigate.