# Writing Turing.jl model to avoid loops?

**URL:** <https://discourse.julialang.org/t/writing-turing-jl-model-to-avoid-loops/104593>\
**Category:** Statistics\
**Tags:** turing, zygote\
**Created:** [October 4, 2023, 9:54pm UTC](https://discourse.julialang.org/t/writing-turing-jl-model-to-avoid-loops/104593 "2023-10-04T21:54:04Z")\
**Posts on this page:** 3\
**Page:** 1

<div class="post-metadata">

**Author:** ![dleather](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/dleather/32/43762_2.png) [@dleather](https://discourse.julialang.org/u/dleather)\
**Post date:** [October 4, 2023, 9:54pm UTC](https://discourse.julialang.org/t/writing-turing-jl-model-to-avoid-loops/104593/1 "2023-10-04T21:54:04Z")

</div>

I’m trying to estimate an 2-dimension [Multivariate Ornstein-Uhlenbeck](https://www2.stat.duke.edu/~scs/Projects/StructuralPhylogeny/multivariateOU.pdf) 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:

```julia

@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.

---

<div class="post-metadata">

**Author:** ![sethaxen](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/sethaxen/32/35604_2.png) [@sethaxen](https://discourse.julialang.org/u/sethaxen)\
**Post date:** [October 5, 2023, 9:24pm UTC](https://discourse.julialang.org/t/writing-turing-jl-model-to-avoid-loops/104593/2 "2023-10-05T21:24:33Z")

</div>

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:

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

```

ReverseDiff:

```julia
using ReverseDiff, ChainRules
ReverseDiff.@grad_from_chainrules Base.exp(x::ReverseDiff.TrackedMatrix)

```

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

```julia
# 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:

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

```

---

<div class="post-metadata">

**Author:** ![sethaxen](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/sethaxen/32/35604_2.png) [@sethaxen](https://discourse.julialang.org/u/sethaxen)\
**Post date:** [October 6, 2023, 7:39am UTC](https://discourse.julialang.org/t/writing-turing-jl-model-to-avoid-loops/104593/3 "2023-10-06T07:39:37Z")

</div>

FWIW, even though the [Bayesian Estimation of Differential Equations](https://turinglang.org/dev/tutorials/10-bayesian-differential-equations/#inference-of-a-stochastic-differential-equation) 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.