Writing Turing.jl model to avoid loops?

Specific Domains Statistics
,
1

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.

2

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
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:

# 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)
1 Like
3

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.