Parameterizing matrix in Turing model without mutating arrays

New to Julia
1

Hi All.

So I’ve got this compartment model which uses a matrix exponential,
\hat y = e^{A\,t} x_0,
where the elements of the N by N matrix A are flow rates between N compartments, and x_0 are the initial compartment concentrations. I hope to estimate these flow rates in A.

The matrix A should be only positive numbers on the off-diagonals, and the diagonals should be their negative column sums. How do I achieve this without mutating arrays?

I can produce A using the Zygote backend by producing a multivariate normal and then reshaping it, as I’ve done below (for the case of 4 compartments).

Is there a way to parameterize matrices without mutating arrays?

Or should I perhaps implement the ODE behind this model instead?

Thanks a bunch for checking my question.

Here’s what I can do:

@model function makemodel(data, init, time_range)
    N = size(data)[1]
    σ1 ~ InverseGamma(2, 2)
    σ2 ~ InverseGamma(10, 9)
    σ3 ~ InverseGamma(3, 0.5)
    Sigma = zeros(16, 16) + Diagonal([σ2, σ1, σ1, σ1,
                                      σ1, σ2, σ1, σ1,
                                      σ1, σ1, σ2, σ1,
                                      σ1, σ1, σ1, σ2])
    K ~ MvNormal(Sigma)
    A = reshape(K, N, N)
    
    predicted = reduce(hcat, [real(exp(A*t))*init for t in time_range])
    
    data = transpose(data)[:]
    data ~ arraydist([Normal(predicted[i], σ3) for i=1:length(predicted)])
end

Here’s what I wish I could do, for the case of 3 compartments:

@model function makemodel(data, init, time_range)
    N = 3
    σ1 ~ InverseGamma(2, 2)
    σ2 ~ InverseGamma(3, 0.5)
    K ~ MvNormal(3, σ2)
    K_diags ~ TruncatedNormal([0.0, 0.0, 0.0], 1.0, -100.0, 0.0)
    A = zeros(3, 3)
    A[2, 1] = K[1]
    A[3, 1] = K[2]
    A[3, 2] = K[2]
    for i=1:3
        A[i, i] = -sum(A[:, i]) + K_diags[i]
    
    predicted = reduce(hcat, [real(exp(A*t))*init for t in time_range])
    
    data = transpose(data)[:]
    data ~ arraydist([Normal(predicted[i], σ3) for i=1:length(predicted)])
end