# Writing Turing.jl model to avoid loops?

I’m trying to estimate an 2-dimension Multivariate Ornstein-Uhlenbeck process in Turing.jl using the NUTS algorithm.

The solution to this model involves the computation of a matrix exponential. As far as I can deduce, the only autodiff backend that can handle this is Zygote.jl. Thus, I need to write my model avoiding loops. but I am having trouble implementing it.

A snapshot of my model with loops is as follows:

``````
@model function fit_toy_model_observed(data; t_start = t_start, t_end = t_end, dt = dt)
#Set Prior Distributions
σ_y ~ truncated(InverseGamma(2.0,0.2); lower = 0.001)
σ_z ~ truncated(InverseGamma(2.0,0.2); lower = 0.001)
μ_y ~ Normal(mean(data), std(data))
κ_y ~ truncated(Normal(0.0, 0.3); lower =  0.01, upper = 0.99)
κ_z ~ truncated(Normal(0.0,0.3), lower = 0.01,  upper = 0.99)

#Set parameters and compute covariance
p = [κ_y * μ_y, κ_y, κ_z, σ_y, σ_z]
Σ = [σ_y^2 0.0; 0.0 σ_z^2]
β = [κ_y -1.0; 0.0 κ_z]
vecΣ_r = kron(β, β) \ (I - exp(-kron(β, β) * dt)) * ((Σ)[:])
Σ_r = [vecΣ_r[1] vecΣ_r[2]; vecΣ_r[2] vecΣ_r[4]]

#Simulate SDE
tspan = (t_start, t_end)
prob = SDEProblem(drift_toy_latent_model, diffusion_toy_latent_model, data[1], tspan, p)
sol = solve(prob, SOSRI(); saveat = dt, verbose = false)

#Calculate Likelihood
for t =1:length(data)-1
data[t+1] ~ MvNormal(sol.u[t+1], Σ_r)
end

end
``````

I cannot figure out how to use `arraydist()` or `filldist()` with type `Vector{Vector{Float64}}` which is the type of `sol.u`. I have had success by first defining `data_xform = data - sol.u`, and then writing the loop as `data_xform[2:end] ~ filldist(MvNormal([0.0,0.0], Σ_r), length(data)-1)`, but this treats each element of `data_xform` as a parameter, which I do not want.

There are several approaches you could take here.

First, you could tell ForwardDiff or ReverseDiff to use the ChainRules AD rules for the matrix exponential:

ForwardDiff:

``````using LinearAlgebra, ForwardDiff, ForwardDiffChainRules, ChainRules
@ForwardDiff_frule LinearAlgebra.exp!(x::AbstractMatrix{<:ForwardDiff.Dual})
``````

ReverseDiff:

``````using ReverseDiff, ChainRules
``````

Without an MWE, it’s not clear to me why your model needs to avoid loops. But here’s how you would use `arraydist`:

``````# outside of model
using PDMats
data_mat = stack(data[2:end])
Σ_r = PDMat([vecΣ_r[1] vecΣ_r[2]; vecΣ_r[2] vecΣ_r[4]])

# inside of model
F_Σ_r = cholesky(Σ_r)
data_mat ~ @views arraydist(MvNormal.(sol.u[2:end], Ref(Σ_r)))
``````

I haven’t tested it, but a `MatrixNormal` might also work:

``````# inside of model
sol_mat = @views stack(sol.u[2:end])
n = size(data_mat, 1)
data_mat' ~ MatrixNormal(sol_mat', I(n), Σ_r)
``````

FWIW, even though the Bayesian Estimation of Differential Equations demos sampling a model containing an SDE with NUTS, I don’t recommend this. All HMC variants require that the transition be reversible, and this assumption is violated by the stochasticity of the SDE. I wouldn’t expect that your draws target the correct posterior.

However, I don’t know what inference methods are applicable in this scenario.