Selecting arbitrary subsets of 2d random variables for plotting

Questions version agnostic
1

Suppose I estimate a model with a matrix-valued parameter:

import pymc as pm
import arviz as az
from scipy import stats
import numpy as np

test_data = stats.multivariate_normal(mean=np.zeros(3), cov=np.eye(3) * np.random.exponential(1)).rvs(100)
with pm.Model() as mod:
    mu = pm.Normal('mu', size=3)
    chol, *_ = pm.LKJCholeskyCov('chol_cov', n=3, eta=1, sd_dist=pm.Exponential.dist(1))
    cov = pm.Determinstic('cov', chol @ chol.T)
    
    obs = pm.MvNormal('obs', mu=mu, cov=cov, observed=test_data)
    idata = pm.sampling_jax.sample_blackjax_nuts()

(Imagine for example that I was doing a non-centered MVN model or something, so it wouldn’t be so silly to do the intermediate covariance computation)

Is there a way to then plot only an arbitrary subset of the elements of the variable cov? I can select entire rows or columns using the coords keyword, for example:

az.plot_posterior(idata, var_list=['cov'], coords={'cov_dim_0':[0]})

But is it possible to somehow plot only the diagonal, only the lower triangle, or some arbitrary set of index pairs? I could extract them in the PyMC model and save them as deterministics for plotting later, but this seems like overkill, and I’m hoping I’m missing something!

2

Yes, you can use 1d DataArrays with a different and common dimension (e.g. pointwise_sel) to specify any arbitrary subset. There is an example of this in Label guide — ArviZ 0.15.1 documentation