coords for pairs of variables.

Short Description

HI! I’m trying to understand how I can make the coords keyword of many arviz plotting routines to behave in a certain way. I hope explain the problem well enough, I’ll gladly try to clarify if not.

To outline my problem, I’m inferring the contents of a correlation matrix with PyMC3. Now, because the matrix is symmetrical, I only want to use the upper (or lower, if it makes a difference) off-diagonal elements, that is to use elements [0,1], [1,2], [0,2] if it’s a 3-dimensional problem.

However, plot_pair seems to interpret coords as a “slice” along a dimension, such that if I set coords to [0,1] on dimension one of the correlation matrix and [1,2] on dimension two, I will also get the plots for [1,1]. Is there any way that I can leave this one out automatically, or do I have to manually iterate over a set of axes to get this?

Code Example or link

Here’s a short example, you can find a more complete version of the code for context in this repo.

import numpy as np
import pymc3 as pm
import arviz as az


mu = np.array([10., 30.])
sigma = np.array([20., 40.])
rho = -0.7
cov = np.zeros((len(mu), len(mu)))
cov[0,:] = [sigma[0]**2, rho*sigma[0]*sigma[1]]
cov[1,:] = [rho*sigma[1]*sigma[0], sigma[1]**2]

data = np.random.multivariate_normal(mu, cov, size=100)
ndim = data.shape[1]

with pm.Model() as model:
            #we put weakly informative priors on the means and standard deviations of the multivariate normal distribution
            mu = pm.Normal("mu", mu=mu_prior[0], sigma=mu_prior[1], shape=ndim)
            sigma = pm.HalfCauchy.dist(sigma_prior)
            #and a prior on the covariance matrix which weakly penalises strong correlations
            chol, corr, stds = pm.LKJCholeskyCov("chol", n=ndim, eta=2.0, sd_dist=sigma, compute_corr=True)
            #the prior gives us the Cholesky Decomposition of the covariance matrix, so for completeness we can calculate that determinisitically
            cov = pm.Deterministic("cov", chol.dot(chol.T))

            #and now we can put our observed values into a multivariate normal to complete the model
            vals = pm.MvNormal('vals', mu=mu, chol=chol, observed=data)

trace = pm.sample(
                steps, tune=tune, target_accept=0.9, compute_convergence_checks=False,return_inferencedata=True
            )

coords = {"chol_corr_dim_0":[0], "chol_corr_dim_1":[1]} #For 2-D data I can do this
                        #But it's not clear how I can just get the upper triangle in 3 or more dimensions

plot_vars = ['mu', 'chol_corr']
az.plot_pair(trace,
                     var_names = plot_vars,
                     coords = coords,
                     kind="kde",
                     marginals=True,
                     point_estimate="mean",
                     show=True,
            )

Thanks in advance for any input you can give me!

[pscicluna]

Ah, okay, thanks for explaining that. If I get the chance, I will attempt the PR, but seeing as this is quite a new package I might as well also force an update to the latest versions of pymc and arviz once the custom labellers are available and I have figured out how to use them 😃

Thanks for all the help!

[OriolAbril]

I’m trying to do some clean up, so I’ll close this as it seems resolved, feel free to reopen if it is not the case.