matrix-normal-failing-because-cov-matrix

matrix-normal-failing-because-cov-matrix.py

import numpy as np
import aesara.tensor as at
import pymc3 as pm

np.random.seed(20090425)

n = 1
p = 10
k = 5
t = 200
W    = np.random.normal(.1,.5,size=(p,k))
lamb = np.diag(np.random.normal(10,.2,size=(k)))
C    = np.matmul(np.matmul(W,lamb),W.T)
y    =  np.random.normal(.1,.5,size=(n,t,p))

def beta_spike_slab(shape,spike):
    inclusion_prop = 0.05
    beta_spike = pm.Normal('beta_spike', 0, spike, shape=shape)
    beta_slab = pm.Normal('beta_slab', 10, shape=shape)
    gamma = pm.Bernoulli('gamma', inclusion_prop, shape=shape)
    beta_spike_slab = pm.Deterministic('beta_spike_slab',(beta_spike * (1-gamma)) + ((beta_slab * gamma)))
    return beta_spike_slab

with pm.Model() as model:
    beta  = pm.HalfCauchy("beta", 4)
    alpha = pm.HalfCauchy("alpha", 4)
   
    W_t = beta_spike_slab((k,p),0.1)
    A = pm.MvNormal("A", mu=np.zeros(k) , cov=alpha*np.eye(k), shape =(n,k), testval=np.random.normal(1,2,size=(n,k)))

    for i in range(n):
       
        z = pm.MvNormal("Z"+str(i), mu=np.zeros(k), cov=at.diag(at.squeeze(A[i,:])), shape=(t,k), testval=np.random.normal(.01,.02,size=(t,k)))
        obs = np.squeeze(y[i,:,:])
        
        y_obs = pm.MatrixNormal("y_obs"+str(i), mu=(pm.math.dot(z,W_t)), colcov=0.1*(1/beta)*np.eye(p), rowcov=np.eye(t), observed=obs.reshape(t, p),shape=(t,p))

    tr = pm.sample(1000, tune=1000)