Bayesian Inference of ODE

For this tutorial, we will show how to do Bayesian inference to infer the parameters of the Lotka-Volterra equations using each of the three backends:

  • Turing.jl
  • Stan.jl
  • DynamicHMC.jl

Setup

First, let's set up our ODE and the data. For the data, we will simply solve the ODE and take that solution at some known parameters as the dataset. This looks like the following:

using DiffEqBayes, ParameterizedFunctions, OrdinaryDiffEq, RecursiveArrayTools,
      Distributions
f1 = @ode_def LotkaVolterra begin
    dx = a * x - x * y
    dy = -3 * y + x * y
end a

p = [1.5]
u0 = [1.0, 1.0]
tspan = (0.0, 10.0)
prob1 = ODEProblem(f1, u0, tspan, p)

σ = 0.01                         # noise, fixed for now
t = collect(1.0:10.0)   # observation times
sol = solve(prob1, Tsit5())
priors = [Normal(1.5, 1)]
randomized = VectorOfArray([(sol(t[i]) + σ * randn(2)) for i in 1:length(t)])
data = convert(Array, randomized)
2×10 Matrix{Float64}:
 2.7752    6.76937  0.962997  1.8979   …  4.35316  3.25413  1.04712
 0.250245  2.00973  1.90896   0.31747     0.32149  4.55979  0.905685

Inference Methods

Stan

using StanSample #required for using the Stan backend
bayesian_result_stan = stan_inference(prob1, :rk45, t, data, priors)
1000×3 DataFrame
  Row │ sigma1.1  sigma1.2  theta_1 
      │ Float64   Float64   Float64 
──────┼─────────────────────────────
    1 │ 0.325445  0.227392  1.50776
    2 │ 0.272685  0.295252  1.50289
    3 │ 0.171984  0.201156  1.5045
    4 │ 0.171984  0.201156  1.5045
    5 │ 0.305768  0.28925   1.49564
    6 │ 0.226879  0.2106    1.50313
    7 │ 0.267993  0.273367  1.4935
    8 │ 0.212334  0.367141  1.50032
  ⋮   │    ⋮         ⋮         ⋮
  994 │ 0.283118  0.232546  1.49621
  995 │ 0.231955  0.257723  1.49795
  996 │ 0.293731  0.238898  1.505
  997 │ 0.217553  0.165812  1.50117
  998 │ 0.343947  0.435806  1.50536
  999 │ 0.299487  0.406141  1.49717
 1000 │ 0.205624  0.168844  1.4985
                    985 rows omitted

Turing

bayesian_result_turing = turing_inference(prob1, Tsit5(), t, data, priors)
╭─FlexiChain (1000 iterations, 1 chain) ───────────────────────────────────────╮
 ↓ iter  = 501:1500                                                           
 → chain = 1:1                                                                
                                                                              
 Parameters (2) ── AbstractPPL.VarName                                        
  Float64          theta[1]                                                   
  Vector{Float64}  σ                                                          
                                                                              
 Extras (14)                                                                  
  Int64    n_steps, tree_depth                                                
  Bool     is_accept, numerical_error                                         
  Float64  acceptance_rate, log_density, hamiltonian_energy,                  
           hamiltonian_energy_error, max_hamiltonian_energy_error, step_size, 
           nom_step_size, logprior, loglikelihood, logjoint                   
╰──────────────────────────────────────────────────────────────────────────────╯

DynamicHMC

We can use DynamicHMC.jl as the backend for sampling with the dynamic_inference function. It is similarly used as follows:

bayesian_result_hmc = dynamichmc_inference(prob1, Tsit5(), t, data, priors)
(posterior = [(parameters = [1.5007402491735486], σ = [0.007343914127771412, 0.005307632754874232]), (parameters = [1.5001520155654253], σ = [0.007300576593088145, 0.005789981279715627]), (parameters = [1.5005784187301527], σ = [0.012306594496606283, 0.012662542488325584]), (parameters = [1.500757491374837], σ = [0.010177011157450495, 0.008612670549235214]), (parameters = [1.5001848061359315], σ = [0.011103298228719902, 0.013000718634647659]), (parameters = [1.5002553844575701], σ = [0.007372775571214522, 0.008925270794091694]), (parameters = [1.5002116340598948], σ = [0.011955807931512972, 0.008625727886865752]), (parameters = [1.500356512422847], σ = [0.007113858374474057, 0.008099289368169631]), (parameters = [1.5003564521686352], σ = [0.009206893500280899, 0.009286603400275798]), (parameters = [1.5003187427272173], σ = [0.009754771057901908, 0.00698660410417593])  …  (parameters = [1.5005176584027842], σ = [0.006520061599068119, 0.007258616937905013]), (parameters = [1.5005578589863915], σ = [0.00859552935791887, 0.008047626205475071]), (parameters = [1.5003061189574871], σ = [0.00673895015880558, 0.007771409981429232]), (parameters = [1.5004260083891405], σ = [0.012672625536715701, 0.010035347835984965]), (parameters = [1.500512653471777], σ = [0.009254651458311029, 0.010562692848966803]), (parameters = [1.5004294084350658], σ = [0.00860580541004711, 0.011235632010464484]), (parameters = [1.5001997588727245], σ = [0.0091187826621989, 0.008038169723290186]), (parameters = [1.5003525532401254], σ = [0.009145588745623543, 0.009243185896665826]), (parameters = [1.500340897612104], σ = [0.00868662701186652, 0.007134581960195547]), (parameters = [1.5001774921631394], σ = [0.008686559520240977, 0.006906879143303816])], posterior_matrix = [0.4059584858263915 0.40556644668352104 … 0.4056923473621053 0.4055834292167056; -4.913883318573893 -4.91980194859243 … -4.745970560983171 -4.745978330614311; -5.238609352005519 -5.151626220612196 … -4.942801619791447 -4.975237386740301], tree_statistics = DynamicHMC.TreeStatisticsNUTS[DynamicHMC.TreeStatisticsNUTS(42.045436074842826, 3, turning at positions -6:1, 0.9535964197119698, 7, DynamicHMC.Directions(0xa6a68291)), DynamicHMC.TreeStatisticsNUTS(39.62081746771642, 2, turning at positions -3:0, 0.2889627796647019, 3, DynamicHMC.Directions(0x3775a658)), DynamicHMC.TreeStatisticsNUTS(49.580249425113024, 3, turning at positions -2:5, 0.9999999999999999, 7, DynamicHMC.Directions(0x608dd405)), DynamicHMC.TreeStatisticsNUTS(49.313122722402625, 2, turning at positions -3:-6, 0.9919483584949365, 7, DynamicHMC.Directions(0x2c54f579)), DynamicHMC.TreeStatisticsNUTS(49.415142997115595, 3, turning at positions -2:5, 0.9453661167484345, 7, DynamicHMC.Directions(0x8f66f0fd)), DynamicHMC.TreeStatisticsNUTS(47.31998578190226, 3, turning at positions -6:1, 0.9201054647812636, 7, DynamicHMC.Directions(0x273fd981)), DynamicHMC.TreeStatisticsNUTS(51.61484481827187, 3, turning at positions -6:1, 0.9965268744914948, 7, DynamicHMC.Directions(0x8d121011)), DynamicHMC.TreeStatisticsNUTS(51.70444462129945, 3, turning at positions -2:5, 0.9951295027471412, 7, DynamicHMC.Directions(0xbbb690f5)), DynamicHMC.TreeStatisticsNUTS(52.07462081601487, 3, turning at positions -2:5, 0.9778102136589206, 7, DynamicHMC.Directions(0x0ea6bc5d)), DynamicHMC.TreeStatisticsNUTS(52.29947552879216, 2, turning at positions 3:6, 0.9886004189533185, 7, DynamicHMC.Directions(0x51d755de))  …  DynamicHMC.TreeStatisticsNUTS(51.01949659538698, 2, turning at positions -3:0, 0.9229198534303866, 3, DynamicHMC.Directions(0x59d52434)), DynamicHMC.TreeStatisticsNUTS(49.63666296155436, 2, turning at positions 0:3, 0.8014114884640016, 3, DynamicHMC.Directions(0x9f33f5df)), DynamicHMC.TreeStatisticsNUTS(50.11373537964046, 2, turning at positions 2:5, 0.8824945760618922, 7, DynamicHMC.Directions(0x62bc0d85)), DynamicHMC.TreeStatisticsNUTS(51.06641318273166, 3, turning at positions -4:3, 0.9463371572279738, 7, DynamicHMC.Directions(0x2a751863)), DynamicHMC.TreeStatisticsNUTS(50.07054688552282, 3, turning at positions -3:4, 0.9758655078796811, 7, DynamicHMC.Directions(0x403cf774)), DynamicHMC.TreeStatisticsNUTS(51.848648745562286, 2, turning at positions -3:0, 0.9965711971943484, 3, DynamicHMC.Directions(0x6b36a46c)), DynamicHMC.TreeStatisticsNUTS(49.56265587865348, 3, turning at positions -6:1, 0.9397356031291143, 7, DynamicHMC.Directions(0x2ce7a0b9)), DynamicHMC.TreeStatisticsNUTS(51.75466783311041, 2, turning at positions 3:4, 0.9502444695465748, 5, DynamicHMC.Directions(0xc3752256)), DynamicHMC.TreeStatisticsNUTS(52.70376325307364, 3, turning at positions -1:6, 0.9990487956030628, 7, DynamicHMC.Directions(0x3d18e05e)), DynamicHMC.TreeStatisticsNUTS(51.80398599982209, 2, turning at positions -2:1, 0.9297837174740766, 3, DynamicHMC.Directions(0x2e5e85ad))], logdensities = [46.506522038909154, 50.3717112219624, 50.44808729832697, 50.842372472651356, 50.7717734306387, 52.325211416763004, 52.057229740703505, 52.47841138126405, 52.77931417315911, 52.585260029442644  …  51.237039520761826, 52.30224515509972, 52.0793020577323, 51.5851221797741, 52.14443306479037, 52.146307419388144, 52.59526221013158, 52.793313402528106, 52.74579161178275, 52.19239504433013], κ = Gaussian kinetic energy (Diagonal), √diag(M⁻¹): [0.0001401130713064973, 0.2575032327456374, 0.22560542282892063], ϵ = 0.6296960317982011)

More Information

For a better idea of the summary statistics and plotting, you can take a look at the benchmarks.