Using dims with LKJCholeskyCov

1

If I want to sample a multivariate normal, and I’d like the results of the idata object to have appropriately named dimensions and coordinates, how do I achieve that?

This code runs, but I’m not sure what to do with the coords to get them attached to the various parameters (mu, chol, Rho_, sigma_cafe):

import numpy as np
import pymc as pm

k = 3
n = 100

mean = np.random.randn(k)
covariance = np.random.rand(k, k)
covariance = np.dot(covariance, covariance.T)

data = np.random.multivariate_normal(mean, covariance, n)

coords = {"vars": np.arange(k)}

with pm.Model(coords=coords) as model:
    mu = pm.MvNormal('mu', mu=np.zeros([k,]), cov=np.eye(k))
    chol, Rho_, sigma_cafe = pm.LKJCholeskyCov('L_chol', n=k, eta=2, sd_dist=pm.Exponential.dist(1.0), compute_corr=True)

    y_obs = pm.MvNormal('y_obs', mu=mu, cov=chol, observed=data)

    trace = pm.sample(1000, tune=1000, random_seed=1234)
    pm.summary(trace, var_names=['mu', 'L_chol'])
2

There was some draft work on this by @jessegrabowski

3

This is helpful. It’s nice to know that work is being done.

I think it’s also possible to use xarray.rename_dims and xarray.rename_vars and xarray.assign_coords to do a lot of this after sampling.

The one thing that I still can’t figure out is how to get xarray to combine dimensions that are supposed to be the same, like chol_cov_corr_dim_0 and chol_cov_stds_dim_0.

4

If two dimensions have the same name, they will be combined automatically.

I would say the current recommended way to get coords onto LKJCholeskyCov is to use the idata_kwargs kwarg in pm.sample. This tutorial shows how to do it:

    # Lots of other stuff snipped

    sd_dist = pm.Exponential.dist(0.5, shape=(2,))
    chol, corr, stds = pm.LKJCholeskyCov("chol", n=2, eta=2.0, sd_dist=sd_dist)
   
    # More snipping

    covariation_intercept_slope_trace = pm.sample(
        1000,
        tune=3000,
        target_accept=0.95,
        idata_kwargs={"dims": {"chol_stds": ["param"], "chol_corr": ["param", "param_bis"]}},
    )
5

Thanks. That was actually a huge help.

Opher