Multivariate normal with diagonal covariance

I would like to use a multivariate normal distribution with a diagonal covariance matrix in my model. What is the correct way to do it? At the moment I am using the following code:

sigma = pm.InverseGamma('sigma', alpha, beta, shape=(num_dimensions))
mu = pm.MvNormal('mu', mu=means, cov=tt.nlinalg.diag(sigma), shape=(num_dimensions,))

which works but is very slow and feels not like an elegant way to do things.

1 Like

Why not use the pm.Normal directly?

mu = pm.Normal('mu', mu=means, sd = sigma, shape=(num_dimensions,))
1 Like

That would indeed be great. But is it possible? Documentation of Normal says that mu and sd are of type float.

Can it then be also vector-valued and each of the variables will pick up its parameters from the vector?

1 Like

Yes they can be n-d tensor. We should update the docstrings if this point is not clear.

1 Like

I’d like to follow up on this even though it’s been months since the original post. I’ve been working with data with N observations of d dimensional multivariate Normal variables and want to ensure I’m using shape correctly; I also want to see if I can use univariate Normal instead for speed if the covariance is diagonal.

I was doing some simple simulations to test this and found some odd behavior. I just updated to PyMC 3.6 by the way so it seems like the refactoring of shape hasn’t happened yet? Code is below, but the basic idea is that when I use MvNormal with cov instead of tau or chol, then try to sample from the distribution, it doesn’t work (wrong shape and get nan). If I build PyMC models around these cases and use NUTS to get the posterior for the parameters, it seems to work correctly in all cases, and is faster to use univariate Normal if the covariance is diagonal.

import pymc3
import numpy as np

d = 5
N = 10

noise_sd_true = 0.5
mu_true = np.tile(np.linspace(1,50,N),(d,1)).T
noise_cov_true = noise_sd_true**2 * np.eye(d)
noise_tau_true = 1./noise_sd_true**2 * np.eye(d)

# Generate (N,d) data using univariate Normal
samp = pymc3.Normal.dist(mu=mu_true, sd=noise_sd_true, shape=(N,d)).random()
print(samp.shape) # (N, d) vector with correct values

# Generate (N,d) data using MvNormal with tau
samp = pymc3.MvNormal.dist(mu=mu_true, tau=noise_tau_true, shape=(N,d)).random()
print(samp.shape) # (N, d) vector with correct values

# Generate (N,d) data using MvNormal with chol
samp = pymc3.MvNormal.dist(mu=mu_true, chol = noise_sd_true * np.eye(d), shape=(N,d)).random()
print(samp.shape) # (N, d) vector with correct values

# Generate (N,d) data using MvNormal with cov
samp = pymc3.MvNormal.dist(mu=mu_true, cov=noise_cov_true, shape=(N,d)).random()
print(samp.shape) # (d, ) vector of nan?
1 Like

Had the exact same problem and luckily found your solution in this thread. Agreed an update to the docstrings would be helpful.

1 Like