Multiple covariance structures

Questions v3
1

Fitting models to vectors of data with two distinct means and distinct variance structures is no problem using indexing, e.g.:

import numpy as np
import pymc3 as pm
import theano.tensor as tt

#univariate data (uv)

numSubs = 2
numData = 100
idx = np.repeat(np.arange(0,numSubs),numData)

uv_sds = []
uv_mus=[]
uv_data=[]
for i in np.arange(0,numSubs):
    uv_mu = np.random.randn()
    uv_sd = np.random.rand()
    uv_data.append(np.random.normal(uv_mu, uv_sd, numData))
    uv_sds.append(uv_sd)
    uv_mus.append(uv_mu)

uv_y = np.hstack(uv_data)

with pm.Model() as uv_model:
    uv_mu_prior = pm.Normal('uv_mu_prior',mu=0,sigma=1,shape=numSubs)
    uv_sd_prior = pm.Exponential('uv_sd_prior',lam=1,shape=numSubs)
    
    #likelihood function (lf)
    uv_lf = pm.Normal('uv_lf',mu=uv_mu_prior[idx],sigma=uv_sd_prior[idx],observed = uv_y)
    
    uv_trace = pm.sample(1000,tune=1000)

This fits fine, however I can’t seem to apply the same indexing logic to data with separate mutivariate structures, e.g.:

#multivariate data (mv)
numSubs = 2
numData = 100
idx = np.repeat(np.arange(0,numSubs),numData)
numVars = 2

mv_covs = []
mv_mus=[]
mv_data=[]
for i in np.arange(0,numSubs):
    mv_mu = np.random.randn(numVars)
    #make a positive semi-definite matrix for covariance
    A = np.random.randn(numVars, numVars)
    mv_cov = np.dot(A, A.transpose())
    mv_data.append(np.random.multivariate_normal(mv_mu, mv_cov, numData))
    mv_covs.append(mv_cov)
    mv_mus.append(mv_mu)
    
    mv_y = np.vstack(mv_data)

with pm.Model() as mv_model:
    
    mu_prior = pm.Normal('mu_prior',mu=0,sigma=1,shape=(numSubs,numVars))
    sd_dist = pm.Exponential.dist(1.0, shape=(numSubs,numVars))
    chol,_,_ = pm.LKJCholeskyCov('chol', n=(numSubs,numVars), eta=1, sd_dist= sd_dist, compute_corr=True)
    vals = pm.MvNormal('vals', mu=mu_prior[idx], chol=chol[idx], observed=mv_y)
        
    trace = pm.sample(1000,tune=1000)

Played around with a bunch of different settings and just can’t seem to get it to work, any suggestions?