Incorporating forcing functions in the ODE model

Specific Domains Probabilistic Programming
,
3

Thank you for the quick reply @tomaklutfu! Please check the complete code below:

#=
Section 1: Import required packages
=#

using Turing, Distributions, DifferentialEquations, Interpolations
using MCMCChains, Plots, StatsPlots
using CSV, XLSX, DataFrames
using Random
Random.seed!(18431)

#=
Section 2: Read the data file containing observation data and get the NPI data into arrays
=#

my_data = DataFrame(XLSX.readtable("observation_data.xlsx","Sheet1")...);

total_weeks = 36;   # Total number of time points
N = 67081000;       # Population

y_time = 1:1:total_weeks;               # Timepoints (weeks)

y_S = Float64.(my_data.Susceptible);    # Susceptible
y_S = y_S[1:total_weeks];

y_D = Float64.(my_data.Deceased);       # Deceased
y_D = y_D[1:total_weeks];

y_HC = Float64.(my_data.Hosp_critical); # Critical hospitalizations
y_HC = y_HC[1:total_weeks];

y_T = Float64.(my_data.Hosp_total);     # Total hospitalizations
y_T = y_T[1:total_weeks];
y_HNC = y_T - y_HC;                     # Non-critical hospitalizations

observation_data = [y_S y_D y_HC y_HNC];

wet_data = DataFrame(XLSX.readtable("Wetdata.xlsx","Wetdata")...);
# IPTCC is a forcing function
IPTCC = wet_data.Normalized_IPTCC;
IPTCC = IPTCC[1:total_weeks];

mobil_data = DataFrame(XLSX.readtable("Mobdata.xlsx","Mobdata")...);
# mobil is another forcing function
mobil = mobil_data.Mean;
mobil = mobil[1:total_weeks];

wet_forcing = interpolate(IPTCC, BSpline(Linear()));
mobil_forcing = interpolate(mobil, BSpline(Linear()));

forcing_params = (wet_forcing, mobil_forcing);

#=
Section 3: Define the model and the respective parameters
=#

function epidemic_wildtype(dy, y, p, t)
    S, E, I, Hᵪ, Hₙ, R, D = y;
    β, λ, α, γ, θᵪ, θₙ, γᵪ, γₙ, δᵪ, w, m = p;
    N = 67081000;

    dy[1] = -β*w(t)*m(t)*I*S/N + λ*R;  # S
    dy[2] = β*w(t)*m(t)*I*S/N - α*E;   # E
    dy[3] = α*E - (γ + θᵪ + θₙ)*I;     # I
    dy[4] = θₙ*I - γₙ*Hᵪ;               # HNC
    dy[5] = θᵪ*I - (γᵪ + δᵪ)*Hₙ;       # HC
    dy[6] = γ*I + γₙ*Hₙ + γᵪ*Hᵪ - λ*R;  # R
    dy[7] = δᵪ*Hᵪ;                     # D
end

#=
Section 4: Define the priors and the Bayesian model
=#

Turing.setadbackend(:forwarddiff)
@model function fitting_epidemic_wildtype(observ_data, w_forcing, m_forcing)
    # Priors of model parameters
    β ~ truncated(Normal(0.65, 0.1), 0, 2)
    λ ~ truncated(Normal(0.5, 0.1), 0, 5)
    α ~ truncated(Normal(0.25, 0.1), 0.1, 0.5)
    γ ~ truncated(Normal(0.05, 0.1), 0, 5)
    γₙ ~ Uniform(0.05, 0.1)
    γᵪ ~ Uniform(0.05, 0.1)
    θₙ ~ Uniform(0.09, 0.75)
    θᵪ ~ Uniform(0.09, 0.75)
    δᵪ ~ Uniform(0.1, 0.8)

    p = [β, λ, α, γ, θᵪ, θₙ, γᵪ, γₙ, δᵪ, w_forcing, m_forcing];

    # Priors of standard deviations
    σ₁ ~ InverseGamma(1, 1) # Susceptible
    σ₂ ~ InverseGamma(1, 1) # Deceased
    σ₃ ~ InverseGamma(2, 3) # Critically hospitalized
    σ₄ ~ InverseGamma(1, 1) # Non-critically hospitalized

    # Initial conditions
    N = 67081000;
    S0 = N;
    I0 = 100;
    y0 = [S0, 0, I0, 0, 0, 0, 0];
    y0 = eltype(p).(y0);

    # Solve the model and compare with observed data
    problem = ODEProblem(epidemic_wildtype, y0, (1, 36), p)
    predicted = solve(problem, Tsit5(), saveat=1.0)

    for i = 1:length(predicted)
        observ_data[i,1] ~ Normal(predicted[1,i], σ₁)
        observ_data[i,2] ~ Normal(predicted[7,i], σ₂)
        observ_data[i,3] ~ Normal(predicted[5,i], σ₃)
        observ_data[i,4] ~ Normal(predicted[4,i], σ₄)
    end
end

#=
Section 5: Run the model-inference system and save the chains
=#

model = fitting_epidemic_wildtype(observation_data, wet_forcing, mobil_forcing);
number_of_chains = 1;
chain = sample(model, NUTS(0.65), MCMCThreads(), 10000, number_of_chains);

Hope it helps to debug! Thanks in advance.