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

CHANGELOG.md

@@ -12,6 +12,7 @@
 * Add more helpful error message for HDF5 problems reading `InferenceData` from NetCDF ([1637](https://github.com/arviz-devs/arviz/pull/1637))
 * Added `data.log_likelihood`, `stats.ic_compare_method` and `plot.density_kind` to `rcParams` ([1611](https://github.com/arviz-devs/arviz/pull/1611))
 * Improve error messages in `stats.compare()`, and `var_name` parameter. ([1616](https://github.com/arviz-devs/arviz/pull/1616))
+* Added ability to plot HDI contours to `plot_kde` with the new `hdi_probs` parameter. ([1665](https://github.com/arviz-devs/arviz/pull/1665))
 
 ### Maintenance and fixes
 * Fixed conversion of Pyro output fit using GPUs ([1659](https://github.com/arviz-devs/arviz/pull/1659))

arviz/plots/backends/bokeh/pairplot.py

@@ -1,5 +1,6 @@
 """Bokeh pairplot."""
 import warnings
+from copy import deepcopy
 from uuid import uuid4
 
 import bokeh.plotting as bkp
@@ -291,7 +292,7 @@ def get_width_and_height(jointplot, rotate):
                         backend="bokeh",
                         backend_kwargs={},
                         show=False,
-                        **kde_kwargs,
+                        **deepcopy(kde_kwargs),
                     )
 
                 if "hexbin" in kind:

arviz/plots/backends/matplotlib/pairplot.py

@@ -1,5 +1,6 @@
 """Matplotlib pairplot."""
 import warnings
+from copy import deepcopy
 
 import matplotlib.pyplot as plt
 import numpy as np
@@ -284,7 +285,7 @@ def plot_pair(
                             var1,
                             var2,
                             ax=ax[j, i],
-                            **kde_kwargs,
+                            **deepcopy(kde_kwargs),
                         )
 
                     if "hexbin" in kind:

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,
+    hdi_probs=None,
     fill_last=False,
     figsize=None,
     textsize=None,
@@ -74,6 +75,9 @@ def plot_kde(
     contour : bool
         If True plot the 2D KDE using contours, otherwise plot a smooth 2D KDE.
         Defaults to True.
+    hdi_probs : list
+        Plots highest density credibility regions for the provided probabilities for a 2D KDE.
+        Defaults to matplotlib chosen levels with no fixed probability associated.
     fill_last : bool
         If True fill the last contour of the 2D KDE plot. Defaults to False.
     figsize : tuple
@@ -216,6 +220,13 @@ def plot_kde(
         ...     contour_kwargs={"levels":3}, contourf_kwargs={"levels":3}
         ... );
 
+    Plot 2d contour KDE with 30%, 60% and 90% HDI contours.
+
+    .. plot::
+        :context: close-figs
+
+        >>> az.plot_kde(mu_posterior, values2=tau_posterior, hdi_probs=[0.3, 0.6, 0.9])
+
     Plot 2d smooth KDE
 
     .. plot::
@@ -261,6 +272,27 @@ def plot_kde(
         gridsize = (128, 128) if contour else (256, 256)
         density, xmin, xmax, ymin, ymax = _fast_kde_2d(values, values2, gridsize=gridsize)
 
+        if hdi_probs is not None:
+            # Check hdi probs are within bounds (0, 1)
+            if min(hdi_probs) <= 0 or max(hdi_probs) >= 1:
+                raise ValueError("Highest density interval probabilities must be between 0 and 1")
+
+            # Calculate contour levels and sort for matplotlib
+            contour_levels = _find_hdi_contours(density, hdi_probs)
+            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

@@ -1049,3 +1049,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_probs):
+    """
+    Find contours enclosing regions of highest posterior density.
+
+    Parameters
+    ----------
+    density : array-like
+        A 2D KDE on a grid with cells of equal area.
+    hdi_probs : array-like
+        An array of highest density interval confidence probabilities.
+
+    Returns
+    -------
+    contour_levels : array
+        The contour levels corresponding to the given HDI probabilities.
+    """
+    # 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_probs)
+    for idx, hdi_prob in enumerate(hdi_probs):
+        try:
+            contours[idx] = sorted_density[sm <= hdi_prob][-1]
+        except IndexError:
+            contours[idx] = sorted_density[0]
+
+    return contours

arviz/tests/base_tests/test_plots_bokeh.py

@@ -266,6 +266,11 @@ def test_plot_kde_1d(continuous_model):
         {"contour": True, "contourf_kwargs": {"cmap": "plasma"}},
         {"contour": False},
         {"contour": False, "pcolormesh_kwargs": {"cmap": "plasma"}},
+        {"contour": True, "contourf_kwargs": {"levels": 3}},
+        {"contour": True, "contourf_kwargs": {"levels": [0.1, 0.2, 0.3]}},
+        {"hdi_probs": [0.3, 0.9, 0.6]},
+        {"hdi_probs": [0.3, 0.6, 0.9], "contourf_kwargs": {"cmap": "Blues"}},
+        {"hdi_probs": [0.9, 0.6, 0.3], "contour_kwargs": {"alpha": 0}},
     ],
 )
 def test_plot_kde_2d(continuous_model, kwargs):

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},
+        {"hdi_probs": [0.3, 0.9, 0.6]},
+        {"hdi_probs": [0.3, 0.6, 0.9], "contourf_kwargs": {"cmap": "Blues"}},
+        {"hdi_probs": [0.9, 0.6, 0.3], "contour_kwargs": {"alpha": 0}},
     ],
 )
 def test_plot_kde(continuous_model, kwargs):
@@ -404,6 +407,21 @@ def test_plot_kde(continuous_model, kwargs):
     assert axes is axes1
 
 
+@pytest.mark.parametrize(
+    "kwargs",
+    [
+        {"hdi_probs": [1, 2, 3]},
+        {"hdi_probs": [-0.3, 0.6, 0.9]},
+        {"hdi_probs": [0, 0.3, 0.6]},
+        {"hdi_probs": [0.3, 0.6, 1]},
+    ],
+)
+def test_plot_kde_hdi_probs_bad(continuous_model, kwargs):
+    """Ensure invalid hdi probabilities are rejected."""
+    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
+    prob_dist = 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 = prob_dist.pdf(grid)
+
+    contour_sp = np.empty(contour_sigma.shape)
+    for idx, sigma in enumerate(contour_sigma):
+        contour_sp[idx] = prob_dist.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)

examples/bokeh/bokeh_plot_kde_2d.py

@@ -1,6 +1,6 @@
 """
-2d KDE Bokeh
-============
+2d KDE
+======
 
 _thumb: .1, .8
 """

examples/bokeh/bokeh_plot_kde_2d_bis.py

@@ -0,0 +1,19 @@
+"""
+2d KDE
+======
+
+_thumb: .1, .8
+"""
+import numpy as np
+
+import arviz as az
+
+az.style.use("arviz-darkgrid")
+
+az.plot_kde(
+    np.random.beta(2, 5, size=100),
+    np.random.beta(2, 5, size=100),
+    contour_kwargs={"levels": 30},
+    contourf_kwargs={"alpha": 0.5, "levels": 30, "cmap": "viridis"},
+    backend="bokeh",
+)

examples/bokeh/bokeh_plot_kde_2d_hdi.py

@@ -0,0 +1,23 @@
+"""
+2d KDE
+======
+
+_thumb: .5, .8
+_example_title: 2d KDE with HDI Contours
+"""
+import numpy as np
+
+import arviz as az
+
+az.style.use("arviz-darkgrid")
+
+rng = np.random.default_rng()
+data = rng.multivariate_normal([2, 2], [[1, 0.4], [0.4, 0.8]], 1000000)
+
+az.plot_kde(
+    data[:, 0],
+    data[:, 1],
+    hdi_probs=[0.393, 0.865, 0.989],  # 1, 2 and 3 sigma contours
+    contourf_kwargs={"cmap": "Blues"},
+    backend="bokeh",
+)

examples/bokeh/bokeh_plot_pair_kde_hdi.py

@@ -0,0 +1,27 @@
+"""
+KDE Pair Plot
+=============
+
+_thumb: .2, .5
+_example_title: KDE Pair Plot with HDI Contours
+"""
+import arviz as az
+
+az.style.use("arviz-darkgrid")
+
+centered = az.load_arviz_data("centered_eight")
+
+coords = {"school": ["Choate", "Deerfield"]}
+az.plot_pair(
+    centered,
+    var_names=["theta", "mu", "tau"],
+    kind="kde",
+    coords=coords,
+    divergences=True,
+    textsize=22,
+    kde_kwargs={
+        "hdi_probs": [0.3, 0.6, 0.9],  # Plot 30%, 60% and 90% HDI contours
+        "contourf_kwargs": {"cmap": "Blues"},
+    },
+    backend="bokeh",
+)

examples/matplotlib/mpl_plot_kde_2d_hdi.py

@@ -0,0 +1,26 @@
+"""
+2d KDE
+======
+
+_thumb: .5, .8
+_example_title: 2d KDE with HDI Contours
+"""
+import matplotlib.pyplot as plt
+import numpy as np
+
+import arviz as az
+
+az.style.use("arviz-darkgrid")
+
+rng = np.random.default_rng()
+data = rng.multivariate_normal([2, 2], [[1, 0.4], [0.4, 0.8]], 1000000)
+
+ax = az.plot_kde(
+    data[:, 0],
+    data[:, 1],
+    hdi_probs=[0.393, 0.865, 0.989],  # 1, 2 and 3 sigma contours
+    contourf_kwargs={"cmap": "Blues"},
+)
+
+ax.set_aspect("equal")
+plt.show()

examples/matplotlib/mpl_plot_pair_kde_hdi.py

@@ -0,0 +1,29 @@
+"""
+KDE Pair Plot
+=============
+
+_thumb: .2, .5
+_example_title: KDE Pair Plot with HDI Contours
+"""
+import matplotlib.pyplot as plt
+
+import arviz as az
+
+az.style.use("arviz-darkgrid")
+
+centered = az.load_arviz_data("centered_eight")
+
+coords = {"school": ["Choate", "Deerfield"]}
+az.plot_pair(
+    centered,
+    var_names=["theta", "mu", "tau"],
+    kind="kde",
+    coords=coords,
+    divergences=True,
+    textsize=22,
+    kde_kwargs={
+        "hdi_probs": [0.3, 0.6, 0.9],  # Plot 30%, 60% and 90% HDI contours
+        "contourf_kwargs": {"cmap": "Blues"},
+    },
+)
+plt.show()