Parallel Tempering in Julia

using DifferentialEquations
using Random
using Distributions
using Turing
using DataFrames
using StatsPlots
using MCMCTempering

function sir_ode!(du,u,p,t)
    (S,I,R,C) = u
    (β,c,γ, δ) = p
    N = S+I+R
    infection = β*c*I*S/N
    recovery = γ*I
    wane = δ * R
    @inbounds begin
        du[1] = -infection + wane
        du[2] = infection - recovery
        du[3] = recovery - wane
        du[4] = infection 
    end
    nothing
end;


tmax = 365.0
tspan = (0.0,tmax)
obstimes = 1.0:1.0:tmax
u0 = [5990.0,10.0,0.0,0.0] # S,I.R,C
p = [0.05,10.0,0.20, 1.0/180]; # β, c, γ, δ


prob_ode = ODEProblem(sir_ode!,u0,tspan,p);


sol_ode = solve(prob_ode,
            Tsit5(),
            saveat = 1.0);

@model bayes_sir(y) = begin
  # Calculate number of timepoints
  l = length(y)
  i₀  ~ Uniform(0.0,.2)
  β ~ Uniform(0, .1)
  immunity ~ Uniform(90,365)

  I = i₀*6000.0
  u0=[6000.0-I,I,0.0,0.0]
  p=[β,10.0,0.2, 1.0/immunity]
  tspan = (0.0,float(l))
  prob = ODEProblem(sir_ode!,
          u0,
          tspan,
          p)
  sol = solve(prob,
              Tsit5(),
              saveat = 1.0)
  sol_C = Array(sol)[4,:] # Cumulative cases
  sol_X = sol_C[2:end] - sol_C[1:(end-1)]
  l = length(y)
  for i in 1:l
    y[i] ~ Poisson(ifelse(sol_X[i] <0 , 1e-6, sol_X[i] ))
  end
end;

mod = bayes_sir(Y[1:200]);

ode_nuts = sample(mod, NUTS(0.65),10000);

const temperature_steps = 4
using MCMCTempering
sampler2 = NUTS()
tempered_sampler = tempered(sampler2, temperature_steps)

chain = sample(mod, tempered_sampler, n_samples; discard_initial = n_adapts)

@online{dewitt2022,
  author = {Michael E. DeWitt},
  title = {Parallel Tempering in Julia},
  date = {2022-08-28},
  url = {https://michaeldewittjr.com/blog/2022-08-28-parallel-tempering-in-julia/},
  langid = {en}
}
Copy