How to plot trace of either chol_corr[0,1] or chol_corr[1,0] from model with LKJCholeskyCov prior?

I’m adapting the tutorial from here LKJ Cholesky Covariance Priors for Multivariate Normal Models — PyMC example gallery

For my own use case, I want to plot the trace of “chol_corr”. Now, chol_corr is a 2x2 correlation matrix and the diagonals are always 1’s in every sample of the trace – the off-diagonal terms are identical and will vary from sample to sample. When I do

az.plot_trace(trace,var_names="chol_corr")

I get an error:

File /opt/anaconda3/envs/pymc/lib/python3.11/site-packages/arviz/plots/kdeplot.py:267, in plot_kde(values, values2, cumulative, rug, label, bw, adaptive, quantiles, rotated, contour, hdi_probs, fill_last, figsize, textsize, plot_kwargs, fill_kwargs, rug_kwargs, contour_kwargs, contourf_kwargs, pcolormesh_kwargs, is_circular, ax, legend, backend, backend_kwargs, show, return_glyph, **kwargs)
    264 if bw == "default":
    265     bw = "taylor" if is_circular else "experimental"
--> 267 grid, density = kde(values, is_circular, bw=bw, adaptive=adaptive, cumulative=cumulative)
    268 lower, upper = grid[0], grid[-1]
    270 density_q = density if cumulative else density.cumsum() / density.sum()

File /opt/anaconda3/envs/pymc/lib/python3.11/site-packages/arviz/stats/density_utils.py:498, in kde(x, circular, **kwargs)
    495 else:
    496     kde_fun = _kde_linear
--> 498 return kde_fun(x, **kwargs)

File /opt/anaconda3/envs/pymc/lib/python3.11/site-packages/arviz/stats/density_utils.py:590, in _kde_linear(x, bw, adaptive, extend, bound_correction, extend_fct, bw_fct, bw_return, custom_lims, cumulative, grid_len, **kwargs)
    588     grid, pdf = _kde_adaptive(x, bw, grid_edges, grid_counts, grid_len, bound_correction)
    589 else:
--> 590     grid, pdf = _kde_convolution(x, bw, grid_edges, grid_counts, grid_len, bound_correction)
    592 if cumulative:
    593     pdf = pdf.cumsum() / pdf.sum()

File /opt/anaconda3/envs/pymc/lib/python3.11/site-packages/arviz/stats/density_utils.py:708, in _kde_convolution(x, bw, grid_edges, grid_counts, grid_len, bound_correction, **kwargs)
    704 # See: https://stackoverflow.com/questions/2773606/gaussian-filter-in-matlab
    706 grid = (grid_edges[1:] + grid_edges[:-1]) / 2
--> 708 kernel_n = int(bw * 2 * np.pi)
    709 if kernel_n == 0:
    710     kernel_n = 1

OverflowError: cannot convert float infinity to integer

I think it’s because plot_trace is trying to do the kdeplot of the chol_corr[0,0] element trace which is always 1. How do I tell plot_trace to only plot the trace of the chol_corr[0,1] or chol_corr[1,0] element (just one of them, not both, since they are identical)?

FYI, I also raised this issue previously here: Cholesky Factor Invariance and Traceplot Errors

Seems like others have encountered this issue too.

Thanks – is there an obvious easy solution that you use?

@jeti wrote:

As a workaround I plot the correlations separately and only the dimensions of interest:

az.plot_trace(trace, var_names=['chol_cov_corr'], filter_vars='like',
                       coords={"chol_cov_corr_dim_0": 0, "chol_cov_corr_dim_1": 1, })

It seems the problem is the constant ones on the diagonal going into the KDE what cause the problem.