Generalised Wishart Processes - stuck with accessing variables in a loop

Questions
1

Hello,

I’m new to PyMC3 and have bitten off a bit more than I can chew. I hope you can help me out. I’m trying to implement Generalized Wishart Processes. This model uses multiple latent draws from a Gaussian process to construct a covariance matrix for every observed time point (essentially covariance regression).

Most of the model is straightforward to implement in PyMC3, but I’m stuck at the likelihood. Here, I need to construct this covariance matrix Sigma(t) at every time point, but I do not know how to access the latent GP variables in a loop. My minimal not-really-working example is:

import matplotlib.pyplot as plt
import numpy as np
import pymc3 as pm
import theano
import theano.tensor as tt

D   = 2
nu  = D+1
N   = 50
x   = np.linspace(0, 1, N)    
L   = np.identity(D)

X   = x[:, None]

with pm.Model() as model:

    def gwp_mvn_likelihood(y):  
        # y ~ GWP = \prod_{t=1}^N MvN(0, Sigma(t))
        # with Sigma(t) = \sum_{i=1}^nu L u_i(t) u^T_i(t) L^T
        # u_i(t) is of size D x 1 
        # L is of size D x D
        logp = 0.0
        for t in range(N):
            Sigma_t = tt.dmatrix('Sigma_t')
            for i in range(nu):
#                Sigma_t += ? # I don't know how to acces u_{i,d} here!
            logp += pm.MvNormal.dist(mu=0, cov=Sigma_t).logp(y).eval()
        return logp
                    
    lengthscale     = pm.HalfNormal('lengthscale', sd=1.0)
    cov             = pm.gp.cov.Matern32(1, lengthscale)    
    gp_prior        = pm.gp.Latent(cov)
    
    # u_id ~ GP(0, k(x,x'))    
    for i in range(nu):
        for d in range(D):
            gp_prior.prior('u_{:d}{:d}'.format(i,d), X=X)
            
    pm.DensityDist('y', gwp_mvn_likelihood, observed=y)
    trace = pm.sample(cores=1, njobs=1)

# plot some sample draws from u
plt.figure()
for i in range(1,10):    
    L = trace.get_values(varname='u_00')[i, :, :]
    z = trace.get_values(varname='u_00_rotated_')[i,:]
    f = np.matmul(L, z)
    plt.plot(np.squeeze(X), f)

I assume that I’m not constructing the likelihood correctly anyway, since it’s written as a numpy-function rather than in symbolic Theano. Any points for that would be greatly appreciated as well, but foremost: how do I access the variables u_{i,d} in my likelihood definition? If it was a single variable u I had direct access to it, but now I have no idea how to do this.

Thanks a bunch!