Hi all
I’m working on my first Turing model and am also very new to Julia. I have finally gotten my model to the point of running at all – that is, it failed with an error telling me to increase the number of maxiters after about 1.5 hours of sampling. I am aware of existing posts on this issue and that I can experiment with different solvers. My bigger concern at this point is that my model is much too slow for me to even try these different options. After ~1.5 hours, the sampling process was about 18% complete:
So as a first step, I need to speed up my model substantially. In fact, this is actually only a simplified version of the model I want to build – eventually I want to estimate multiple parameters, but here I am only trying to estimate a single parameter (a11):
using Parameters, LinearAlgebra, DifferentialEquations, Turing, MCMCChains, Distributions, Random, Plots, StatsPlots
"""Specify options"""
@with_kw mutable struct options
TR::Float64
TE::Float64
duration::Float64
dt::Float64
layers::Int
SNR::Float64
end
"""Define constants"""
@consts begin
λ = 0.5
ϑ₀ = 28.265*3
r₀ = 110
V₀ = 100*0.04
κ = 0.65
γ = 0.41
τ = 0.98
α = 0.32
E₀ = 0.34
ε = 0.4584
τᵈ = 0.5
pₕ = [κ γ τ α E₀ ε τᵈ]
end
"""Functions for the forward model"""
function f!(x, p, t)
a11, a12, a21, a22, c11, c12, c21, c22, u, dt = p
A = [a11 a21; a12 a22]
C = [c11 c21; c12 c22]
# Get input at time t
idx = max(Int(ceil(t/dt)), 1)
ut = u[idx,:]
J = A
J*x + C*ut
end
function h!(dz, z, p, t)
H, HC, xn, dt = p
idx = max(Int(ceil(t/dt)), 1)
H = pₕ .* exp.(H)
HC = [0 1; 0 0] .* exp.(HC)
fv = (abs.(z[:,3])).^(1 ./ H[:,4])
ff = ( 1 .- (1 .- H[:,5]).^(1 ./ z[:,2]) ) ./ H[:,5]
# Compute state equations
dz[:,1] = xn[idx] .- H[:,1].*z[:,1] - H[:,2].*(z[:,2] .- 1)
dz[:,2] = z[:,1]
dz[:,3] = (z[:,2] - fv + HC*z[:,5]) ./ H[:,3]
dz[:,4] = (ff.*z[:,2] - fv.*z[:,4]./z[:,3] + HC*z[:,6]) ./ H[:,3]
dz[:,5] = (-z[:,5] + z[:,3] .- 1)./(H[:,7])
dz[:,6] = (-z[:,6] + z[:,4] .- 1)./(H[:,7])
end
function g(z, H, TE)
E0 = E₀ .* exp.(H[:,5])
εₕ = ε .* exp.(H[2,6])
k₁ = 4.3.*ϑ₀.*E₀.*TE
k₂ = εₕ.*r₀.*E₀.*TE
k₃ = 1 - εₕ
v = map(i -> (z[i][:,3]), 1:12000)
v = transpose(reduce(hcat, v))
q = map(i -> (z[i][:,4]), 1:12000)
q = transpose(reduce(hcat, q))
y = V₀ * (k₁'.*(1 .- q) + k₂'.*(1 .- q./v) + k₃'.*(1 .- v))
end
# Forward model
function forward_model(params, states, opt)
a11, a12, a21, a22, c11, c12, c21, c22, u, H, HC = params
x0, z0 = states
ts = range(opt.dt, step=opt.dt, stop=opt.duration) # sample times
tspan = (0.0, opt.duration)
p = [a11, a12, a21, a22, c11, c12, c21, c22, u, opt.dt] # parameters
ode = ODEProblem(f!, typeof.(p[1]).(x0), tspan, p)
sol = solve(ode, AutoTsit5(Rosenbrock23()), saveat=ts, tstops=ts, dt=1e-2)
if sol.retcode != :Success
println(sol.retcode)
end
ph = [H, HC, sol.u, opt.dt] # parameters
odeh = ODEProblem(h!, typeof.(p[1]).(z0), tspan, ph)
solh = solve(odeh, AutoTsit5(Rosenbrock23()), saveat=ts, tstops=ts, dt = 1e-2)
if solh.retcode != :Success
println(solh.retcode)
end
y = g(solh.u, H, opt.TE)
end
"""Functions for model inversion"""
# Model inversion
Turing.setadbackend(:forwarddiff)
@model function invert_model(data, params, opt)
# Unpack parameters
a12, a21, a22, c11, c12, c21, c22, u, H, HC = params
# Specify priors
a11 ~ Normal(-0.5, 0.0625)
# Set initial conditions
x0 = zeros(opt.layers, 1) # initial condition
z0 = zeros(opt.layers, 6) # initial condition
z0[:,2:4] = exp.(z0[:,2:4]) # exponentiate haemodynamic state variables
states = [x0, z0]
P = [a11, a12, a21, a22, c11, c12, c21, c22, u, H, HC]
predicted = forward_model(P, states, opt)
σ ~ LogNormal(6, 1/128) # Check what this should actually be
for i = 1:size(predicted, 1)
data[i,:] ~ MvNormal(predicted[i,:], σ)
end
end
## 1. Simulate data
# Initialise options:
opt = options(TR = 2, TE = 0.03, duration = 600.0, dt = 0.05, layers = 2, SNR = 4.0) # keywords optional
# Generate a single impulse as input:
u = zeros(Int(opt.duration / opt.dt), 2)
u[6000:6004] .= 1
# Set parameters:
H = zeros(2,7)
HC = log(λ) * [0 1; 0 0]
parameters = [-1.1, 0, -1.1, 0, 0.3, 0, 0.3, 0, u, H, HC]
# Initial states:
x0 = zeros(opt.layers, 1)
z0 = zeros(opt.layers, 6)
z0[:,2:4] = exp.(z0[:,2:4])
states = [x0, z0]
# Run model:
Y = forward_model(parameters, states, opt)
Yn = Y + 0.01*randn(size(Y))
plot(Yn)
## 2. Estimate parameters
# Set parameters:
p_set = [0, -1.1, 0, 0.3, 0, 0.3, 0, u, H, HC] # only estimate a11 for now
model = invert_model(Yn, p_set, opt)
chain = mapreduce(c -> sample(model, NUTS(.65), 1000), chainscat, 1:3)
# Plot
plot(chain)
I would appreciate any help and suggestions on how to adjust my model so it runs faster.