# note: to get stacktrace, run with julia -p 1 using Distributed @everywhere using Turing, ForwardDiff, OrdinaryDiffEq, StatsPlots, AdvancedVI, LinearAlgebra, Random @everywhere Random.seed!(1) using Plots; plotly() # toggle autodiff on/off in ODE solver @everywhere autodiff = true # toy ODE system @everywhere function lotka_volterra(du, u, p, t) α, β, γ, δ = p x, y = u try du[1] = (α - β*y) * x du[2] = (δ*x - γ) * y catch e @error e @error "Error" exception = (e, catch_backtrace()) println("") @show α β δ γ x y sleep(3600) # pause run end return nothing end # simulation output @everywhere function gen_predict(p, jobid) @info "Started job $jobid" u0 = [1.0, 1.0] tspan = (0.0, 10.0) prob = ODEProblem{true, SciMLBase.FullSpecialize}(lotka_volterra, u0, tspan, p) sol = solve(prob, TRBDF2(; autodiff); dt = 0.1, saveat = 0.1) return vcat(sol.u...) end # note: ADVI requires @everywhere to be placed here @everywhere @model function fit_lv(data) # sampling σ² ~ filldist(truncated(Normal(0.05, 0.0125); lower = 0, upper = 1.0), 1) α ~ filldist(truncated(Normal(1.5, 0.5); lower = 0.5, upper = 2.5), 1) β ~ filldist(truncated(Normal(1.2, 0.5); lower = 0, upper = 2), 1) γ ~ filldist(truncated(Normal(3.0, 0.5); lower = 1, upper = 4), 1) δ ~ filldist(truncated(Normal(1.0, 0.5); lower = 0, upper = 2), 1) p = [α[1], β[1], γ[1], δ[1]] # distribute simulation runs over two tests joblist = 1:2 predict = pmap(jobid -> gen_predict(p, jobid), joblist) y_sim = vcat(predict...) # simulation vs experiment data ~ MvNormal(y_sim, σ²[1] * I) return nothing end # target parameter values p = [1.5, 1.0, 3.0, 1.0] # generate mock data points covering two tests u0 = [1.0, 1.0] tspan = (0.0, 10.0) prob = ODEProblem{true, SciMLBase.FullSpecialize}(lotka_volterra, u0, tspan, p) sol = solve(prob, TRBDF2(; autodiff); saveat = 0.1) u = vcat(sol.u...) data = vcat(0.9.*u..., 1.1.*u...) # configure inference model and fit parameters model = fit_lv(data) advi = ADVI(1, 1000) optimizer = Turing.Variational.TruncatedADAGrad(0.01, 1.0, 10) # @time res = vi(model, advi; optimizer) # ADVI @time res = sample(model, NUTS(0.65), MCMCSerial(), 100, 1) # NUTS println("\ndone")