bayesian pca using PyMC

I'm trying to implement Bayesian PCA using PyMC library for python. But, I'm stuck where I define lower dimensional coordinates...

Model is

x = Wz + e

where x is observation vector, W is the transformation matrix, and z is lower dimensional coordinate vector.

First I define a distribution for the transformation matrix W. Each column is drawn from a normal distribution (zero mean, and identity covariance for simplicity)

def W_logp(value):
   logLikes = np.array([multivariate_normal.logpdf(value[:,i], mean=np.zeros(dimX), cov=1) for i in range(0, dimZ)])
   return logLikes.sum()

def W_random():
   W = np.zeros([dimX, dimZ])
   for i in range(0, dimZ):
      W[:,i] = multivariate_normal.rvs(mean=np.zeros(dimX), cov=1)
   return W

w0 = np.random.randn(dimX, dimZ)

W = pymc.Stochastic(
   logp = W_logp,
   doc = 'Transformation',
   name = 'W',
   parents = {},
   random = W_random,
   trace = True,
   value = w0,
   dtype = float,
   rseed = 116.,
   observed = False,
   cache_depth = 2,
   plot = False,
   verbose = 0)

Then, I want to define distribution for z that is again a multivariate normal (zero mean, and identity covariance). However, I need to draw a z for each observation separately while W is common for all of them. So, I tried

z = pymc.MvNormal('z', np.zeros(dimZ), np.eye(dimZ), size=N)

However, pymc.MvNormal does not have a size parameter. So it raises an error. Next step would be

m = Data.mean(axis=0) + np.dot(W, z)
obs = pymc.MvNormal('Obs', m, C, value=Data, observed=True)

I did not give the specification for C above since it is irrelevant for now. Any ideas how to implement?

Thanks

EDIT

After Chris Fonnesbeck's answer I changed my code as follows

numD, dimX = Data.shape
dimZ = 3
mm = Data.mean(axis=0)

tau = pymc.Gamma('tau', alpha=10, beta=2)
tauW = pymc.Gamma('tauW', alpha=20, beta=2, size=dimZ)

@pymc.deterministic(dtype=float)
def C(tau=tau):
    return (tau)*np.eye(dimX)

@pymc.deterministic(dtype=float)
def CW(tau=tauW):
    return np.diag(tau)

W = [pymc.MvNormal('W%i'%i, np.zeros(dimZ), CW) for i in range(dimX)]
z = [pymc.MvNormal('z%i'%i, np.zeros(dimZ), np.eye(dimZ)) for i in range(numD)]
mu = [pymc.Lambda('mu%i'%i, lambda W=W, z=z: mm + np.dot(np.array(W), np.array(z[i]))) for i in range(numD)]

obs = [pymc.MvNormal('Obs%i'%i, mu[i], C, value=Data[i,:], observed=True) for i in range(numD)]

model = pymc.Model([tau, tauW] + obs + W + z)
mcmc = pymc.MCMC(model)

But this time, it tries to allocate a huge amount of memory (more than 8GB) when running pymc.MCMC(model), with numD=45 and dimX=504. Even when I try it with only numD=1 (so creating only 1 z, mu, and obs), it does the same. Any idea why?

Solution

Unfortunately, PyMC does not easily let you define vectors of multivariate stochastics. Hopefully we can make this happen in PyMC 3. For now, you would have to specify this using a container. For example:

z = [pymc.MvNormal('z_%i' % i, np.zeros(dimZ), np.eye(dimZ)) for i in range(N)]