[WIP] Add contours of highest posterior density for 2D KDE plot

arviz/plots/kdeplot.py

@@ -4,7 +4,7 @@
 
 from ..data import InferenceData
 from ..rcparams import rcParams
-from ..stats.density_utils import _fast_kde_2d, kde
+from ..stats.density_utils import _fast_kde_2d, kde, _find_hdi_contours
 from .plot_utils import get_plotting_function
 
 
@@ -20,6 +20,7 @@ def plot_kde(
     quantiles=None,
     rotated=False,
     contour=True,
+    levels=None,
     fill_last=False,
     figsize=None,
     textsize=None,
@@ -74,6 +75,8 @@ def plot_kde(
     contour : bool
         If True plot the 2D KDE using contours, otherwise plot a smooth 2D KDE.
         Defaults to True.
+    levels : list
+        Confidence levels for highest density interval contours of a 2D KDE.
     fill_last : bool
         If True fill the last contour of the 2D KDE plot. Defaults to False.
     figsize : tuple
@@ -261,6 +264,29 @@ def plot_kde(
         gridsize = (128, 128) if contour else (256, 256)
         density, xmin, xmax, ymin, ymax = _fast_kde_2d(values, values2, gridsize=gridsize)
 
+        if levels is not None:
+            # Check hdi levels are within bounds (0, 1)
+            if min(levels) <= 0 or max(levels) >= 1:
+                raise ValueError(
+                    "Highest density interval confidence levels must be between 0 and 1"
+                )
+
+            # Calculate contour levels and sort for matplotlib
+            contour_levels = _find_hdi_contours(density, levels)
+            contour_levels.sort()
+
+            contour_level_list = [0] + list(contour_levels) + [density.max()]
+
+            # Add keyword arguments to contour, contourf
+            if contour_kwargs is None:
+                contour_kwargs = {"levels": contour_level_list}
+            else:
+                contour_kwargs.setdefault("levels", contour_level_list)
+            if contourf_kwargs is None:
+                contourf_kwargs = {"levels": contour_level_list}
+            else:
+                contourf_kwargs.setdefault("levels", contour_level_list)
+
         lower, upper, density_q = [None] * 3
 
     kde_plot_args = dict(

arviz/stats/density_utils.py

@@ -1036,3 +1036,34 @@ def histogram(data, bins, range_hist=None):
     hist, bin_edges = np.histogram(data, bins=bins, range=range_hist)
     hist_dens = hist / (hist.sum() * np.diff(bin_edges))
     return hist, hist_dens, bin_edges
+
+
+def _find_hdi_contours(density, hdi_levels):
+    """
+    Find contours enclosing regions of highest posterior density.
+
+    Parameters
+    ----------
+    density : array-like
+        A gridded 2D KDE.
+    hdi_levels : array-like
+        An array of highest density interval confidence levels.
+
+    Returns
+    -------
+    contour_levels : array
+        The contour levels corresponding to the given HDI levels.
+    """
+    # Using the algorithm from corner.py
+    sorted_density = np.sort(density, axis=None)[::-1]
+    sm = sorted_density.cumsum()
+    sm /= sm[-1]
+
+    contours = np.empty_like(hdi_levels)
+    for idx, hdi_level in enumerate(hdi_levels):
+        try:
+            contours[idx] = sorted_density[sm <= hdi_level][-1]
+        except IndexError:
+            contours[idx] = sorted_density[0]
+
+    return contours

arviz/tests/base_tests/test_plots_matplotlib.py

@@ -395,6 +395,9 @@ def test_plot_joint_bad(models):
         {"is_circular": "radians"},
         {"is_circular": "degrees"},
         {"adaptive": True},
+        {"levels": [0.3, 0.9, 0.6]},
+        {"levels": [0.3, 0.6, 0.9], "contourf_kwargs": {"cmap": "Blues"}},
+        {"levels": [0.9, 0.6, 0.3], "contour_kwargs": {"alpha": 0}},
     ],
 )
 def test_plot_kde(continuous_model, kwargs):
@@ -404,6 +407,20 @@ def test_plot_kde(continuous_model, kwargs):
     assert axes is axes1
 
 
+@pytest.mark.parametrize(
+    "kwargs",
+    [
+        {"levels": [1, 2, 3]},
+        {"levels": [-0.3, 0.6, 0.9]},
+        {"levels": [0, 0.3, 0.6]},
+        {"levels": [0.3, 0.6, 1]},
+    ],
+)
+def test_plot_kde_levels_bad(continuous_model, kwargs):
+    with pytest.raises(ValueError):
+        plot_kde(continuous_model["x"], continuous_model["y"], **kwargs)
+
+
 @pytest.mark.parametrize("shape", [(8,), (8, 8), (8, 8, 8)])
 def test_cov(shape):
     x = np.random.randn(*shape)

arviz/tests/base_tests/test_utils.py

@@ -7,7 +7,7 @@
 import scipy.stats as st
 
 from ...data import dict_to_dataset, from_dict, load_arviz_data
-from ...stats.density_utils import _circular_mean, _normalize_angle
+from ...stats.density_utils import _circular_mean, _normalize_angle, _find_hdi_contours
 from ...utils import (
     _stack,
     _subset_list,
@@ -291,3 +291,48 @@ def test_normalize_angle(mean):
 
     values = _normalize_angle(rvs, zero_centered=False)
     assert ((values >= 0) & (values <= 2 * np.pi)).all()
+
+
+@pytest.mark.parametrize("mean", [[0, 0], [1, 1]])
+@pytest.mark.parametrize(
+    "cov",
+    [
+        np.diag([1, 1]),
+        np.diag([0.5, 0.5]),
+        np.diag([0.25, 1]),
+        np.array([[0.4, 0.2], [0.2, 0.8]]),
+    ],
+)
+@pytest.mark.parametrize("contour_sigma", [np.array([1, 2, 3])])
+def test_find_hdi_contours(mean, cov, contour_sigma):
+    """Test `_find_hdi_contours()` against SciPy's multivariate normal distribution."""
+    # Set up scipy distribution
+    rv = st.multivariate_normal(mean, cov)
+
+    # Find standard deviations and eigenvectors
+    eigenvals, eigenvecs = np.linalg.eig(cov)
+    eigenvecs = eigenvecs.T
+    stdevs = np.sqrt(eigenvals)
+
+    # Find min and max for grid at 7-sigma contour
+    extremes = np.empty((4, 2))
+    for i in range(4):
+        extremes[i] = mean + (-1) ** i * 7 * stdevs[i // 2] * eigenvecs[i // 2]
+    x_min, y_min = np.amin(extremes, axis=0)
+    x_max, y_max = np.amax(extremes, axis=0)
+
+    # Create 256x256 grid
+    x = np.linspace(x_min, x_max, 256)
+    y = np.linspace(y_min, y_max, 256)
+    grid = np.dstack(np.meshgrid(x, y))
+
+    density = rv.pdf(grid)
+
+    contour_sp = np.empty(contour_sigma.shape)
+    for idx, sigma in enumerate(contour_sigma):
+        contour_sp[idx] = rv.pdf(mean + sigma * stdevs[0] * eigenvecs[0])
+
+    hdi_probs = 1 - np.exp(-0.5 * contour_sigma ** 2)
+    contour_az = _find_hdi_contours(density, hdi_probs)
+
+    np.testing.assert_allclose(contour_sp, contour_az, rtol=1e-2, atol=1e-4)