Unable to correctly fit parameters for SIR (ODE simulation) example

[TorkelE]

I am trying to get Turing to work for bayesian parameter optimisation in combination with DiffEq. Unfortunately, I cannot really get it to mesh properly. In this case, I can go through the full process, however, the outputs do not look like expected (i.e. the posterior distributions are very far from the true values).

I am not super familiar with what is going on (and got little experience with Bayesian stuff), so I might make a stupid mistake. Either way, some hint on whether seeing something like this is as bad as I presume would be useful!

I start by delcaring the model and generating some synthetic data where I know the parameters (two for the models and one for the noise distribution).

# Fetch packages
using Distributions
using Turing
using OrdinaryDiffEq
using StatsPlots
using Random
Random.seed!(14);

# Define SIR model.
function sir(du, u, p, t)
    S, I, R = u
    γ, ν = p
    infection = γ * S * I
    recovery = ν * I
    du[1] = -infection
    du[2] = infection - recovery
    du[3] = recovery
    return nothing
end

# Create ODEProblem for true model parameters.
u0_known = [999, 1, 0]
tend_known = 100.0
γ_true = 0.0003; ν_true = 0.1; σI_true = 0.2;
p_true = [γ_true, ν_true]
oprob_true = ODEProblem(sir, u0_known, tend_known, p_true)

# Generate synthetic data (and plot it).
sol_true = solve(oprob_true)
measured_t = 5:5:100
measured_I = [rand(Normal(I, σI_true*I)) for I in sol_true(measured_t, idxs = 2).u]
plot(sol_true; label = "Vals (true)")
plot!(measured_t, measured_I; label = "Vals (measured)", seriestype = :scatter)

Next I I create and run a Turing model.

# Creates Turing model.
@model function fit_sir(data, prob)
    # Sets prior distributions.
    γ ~ LogUniform(0.00001, 0.001)
    ν ~ LogUniform(0.01, 0.9)
    σI ~ LogUniform(0.1, 1)

    # Simulate the model.
    prob = remake(prob; p = [γ, ν])
    predicted = solve(prob; saveat = measured_t, verbose = false, maxiters = 10000)

    # If simulation was unsuccessful, the likelihood is -Inf.
    if !SciMLBase.successful_retcode(predicted)
        Turing.@addlogprob! -Inf
        return nothing
    end

    # Computes the likelihood of the observations.
    for i in eachindex(predicted)
        I = max(predicted[i][2], 0.0)
        data[i] ~ Normal(I, σI *I)
    end

    return nothing
end

# Fit model to data.
model = fit_sir(sol_true[2,:], oprob_true)
chain = sample(model, NUTS(), MCMCSerial(), 100000, 3; progress = false)

Finally, I plot the posterior distributions, including the true values:

# Plots the result of the fits (doesn't look correct).
plt1 = begin
    plt = plot(plot(chain)[1,2])
    plot!([γ_true, γ_true], collect(ylims(plt)); color = :black, lw = 5, la = 0.5)
end
plt2 = begin
    plt = plot(plot(chain)[2,2])
    plot!([ν_true, ν_true], collect(ylims(plt)); color = :black, lw = 5, la = 0.5)
end
plt3 = begin
    plt = plot(plot(chain)[3,2])
    plot!([σI_true, σI_true], collect(ylims(plt)); color = :black, lw = 5, la = 0.5)
end
plot(plt1, plt2, plt3; layout = (1,3), size = (1200, 400), title = ["γ" "ν" "σI"], yguide = "")

I also tried seeding the sample with initial values corresponding to the correct ones, but this did not help

model = fit_sir(sol_true[2,:], oprob_true)
initial_params = fill([0.0005, 0.05, 0.35], 3)
chain = sample(model, NUTS(), MCMCSerial(), 100000, 3; progress = false, initial_params)
plot(chain)

[rzhli]

Might be helpful

[TorkelE]

Thanks 🙂

initially had a look at it when I started doing my example. However, I tried to rework it to fit the case that I have (above), and I am not really sure why this won't work 🙁

[devmotion]

Where does

        I = max(predicted[i][2], 0.0)
        data[i] ~ Normal(I, σI *I)

come from? Intuitively that seems problematic - the ODE solver can always end up with small negative values (which does not necessarily point to any major problem as long as they are within the prescribed error tolerances) but in this case you'd evaluate the likelihood with Normal(0, 0) which yields an infinite log-likelihood for any non-zero observation at this point.

Generally, since you're working with numbers of individuals, susceptibles etc. at the initial time point, IMO you might want to work with discrete observations and model them with a discrete distribution (such as the Poisson distribution in the linked example above). If you're working with continuous observations, you could disallow negative values in your model by using a distribution that has only support on the non-negative real numbers.

[TorkelE]

So I did

I = max(predicted[i][2], 0.0)

because when the ODE becomes (negligibly) negative,

data[i] ~ Normal(I, σI *I)

got a negative std, and started throwing an error, so I needed to force non-negativity somehow. Since I figured the std would be really really low in these cases, I figured setting it to 0 should be ok, but maybe this is a problem as you say?

In reality, the disease dynamic has integer-valued individuals, but when I approximate it as an ODE I kind of make the approximation that it can be considered as continuous concentrations (although the initial condition is individual-inspired, I guess). Had I made a true individual-based model I would have to use Gillespie-style jump simulations.

[rzhli]

Even though I is non-negative, the sample from distribution of Normal(I, σI *I) can also be negative, this violates the true value of data[i], you should try non-negative distributions like truncated(Normal()) or Poisson, etc.

[TorkelE]

I tried making an updated example, where I simplified things so that I don't have to worry about negative stuff in the distribution. The two main parameters are not found quite well, however, the noise parameter is just pushed to basically smallest possible value.

# Fetch packages
using Distributions
using Turing
using OrdinaryDiffEq
using StatsPlots
using Random
Random.seed!(14);

# Define SIR model.
function sir(du, u, p, t)
    S, I, R = u
    γ, ν = p
    infection = γ * S * I
    recovery = ν * I
    du[1] = -infection
    du[2] = infection - recovery
    du[3] = recovery
    return nothing
end

# Create ODEProblem for true model parameters.
u0_known = [999, 1, 0]
tend_known = 100.0
γ_true = 0.0003; ν_true = 0.1; σS_true = 20.0;
p_true = [γ_true, ν_true]
oprob_true = ODEProblem(sir, u0_known, tend_known, p_true)

# Generate synthetic data (and plot it).
measured_t = 1:1:100
sol_true = solve(oprob_true; saveat = measured_t)
measured_S = [rand(Normal(S, σS_true)) for S in sol_true(measured_t, idxs = 1).u]
plot(sol_true; label = ["S" "I" "R"])
plot!(measured_t, measured_S; label = "S (measured)", seriestype = :scatter, color = 1)

# Creates Turing model.
@model function fit_sir(data, prob)
    # Sets prior distributions.
    γ ~ LogUniform(0.00001, 0.001)
    ν ~ LogUniform(0.01, 1.0)
    σS ~ LogUniform(10.0, 30.0)

    # Simulate the model.
    prob = remake(prob; p = [γ, ν])
    predicted = solve(prob; saveat = measured_t, verbose = false, maxiters = 10000)

    # If simulation was unsuccessful, the likelihood is -Inf.
    if !SciMLBase.successful_retcode(predicted)
        Turing.@addlogprob! -Inf
        return nothing
    end

    # Computes the likelihood of the observations.
    for i in eachindex(predicted)
        data[i] ~ Normal(predicted[i][1], σS)
    end

    return nothing
end

# Fit model to data.
model = fit_sir(sol_true[1,:], oprob_true)
chain = sample(model, NUTS(), MCMCSerial(), 1000, 4; progress = false)
plot(chain)

# Plots the result of the fits (doesn't look correct).
plt1 = begin
    plt = plot(plot(chain)[1,2])
    plot!([γ_true, γ_true], collect(ylims(plt)); color = :black, lw = 5, la = 0.5)
end
plt2 = begin
    plt = plot(plot(chain)[2,2])
    plot!([ν_true, ν_true], collect(ylims(plt)); color = :black, lw = 5, la = 0.5)
end
plt3 = begin
    plt = plot(plot(chain)[3,2])
    plot!([σS_true, σS_true], collect(ylims(plt)); color = :black, lw = 5, la = 0.5)
end
plot(plt1, plt2, plt3; layout = (1,3), size = (1200, 400), title = ["γ" "ν" "σS"], yguide = "")

Hhere, the vertical grey lines are the true values. 10 is the minimum value of the prior distribution of sigmaS. If the threshold becomes lower, the posterior distribution follows.

[rzhli]

fit_sir(sol_true[1,:], oprob_true)

why sol_true[1,:] instead of measured_S？

[TorkelE]

You are totally right, I entirely mixed up how to use that one. That works perfectly, thanks a lot! 😃