How to label a seaborn contour plot

Question

So I'm using seaborn to make a kdeplot with sns.kdeplot(x, y, ax=plt.gca(), cmap="coolwarm").

I can change the levels with the levels kwarg but I want to be able to label the contours as well. In matplotlib you would simply do plt.clabel(CS, CS.levels, inline=True) but seaborn doesn't return the contour collection CS.

How would I do this? Or do I just have to do it all from scratch myself?

Edit: Is there maybe a way to make a wrapper which will also return CS? I can't see how though...

Answer 1

Unfortunately, seaborn does everything to keep the countourset secret from the user. Apart from drawing a plt.contour plot from the data, which isn't actually too hard, you have the obtion to monkey patch the seaborn _bivariate_kdeplot and let it return the countourset for further use.

This could look as follows:

import matplotlib.pyplot as plt
import numpy as np; np.random.seed(10)
import seaborn as sns
import seaborn.distributions as sd
from seaborn.palettes import color_palette, blend_palette
from six import string_types


def _bivariate_kdeplot(x, y, filled, fill_lowest,
                       kernel, bw, gridsize, cut, clip,
                       axlabel, cbar, cbar_ax, cbar_kws, ax, **kwargs):
    """Plot a joint KDE estimate as a bivariate contour plot."""
    # Determine the clipping
    if clip is None:
        clip = [(-np.inf, np.inf), (-np.inf, np.inf)]
    elif np.ndim(clip) == 1:
        clip = [clip, clip]

    # Calculate the KDE
    if sd._has_statsmodels:
        xx, yy, z = sd._statsmodels_bivariate_kde(x, y, bw, gridsize, cut, clip)
    else:
        xx, yy, z = sd._scipy_bivariate_kde(x, y, bw, gridsize, cut, clip)

    # Plot the contours
    n_levels = kwargs.pop("n_levels", 10)
    cmap = kwargs.get("cmap", "BuGn" if filled else "BuGn_d")
    if isinstance(cmap, string_types):
        if cmap.endswith("_d"):
            pal = ["#333333"]
            pal.extend(color_palette(cmap.replace("_d", "_r"), 2))
            cmap = blend_palette(pal, as_cmap=True)
        else:
            cmap = plt.cm.get_cmap(cmap)

    kwargs["cmap"] = cmap
    contour_func = ax.contourf if filled else ax.contour
    cset = contour_func(xx, yy, z, n_levels, **kwargs)
    if filled and not fill_lowest:
        cset.collections[0].set_alpha(0)
    kwargs["n_levels"] = n_levels

    if cbar:
        cbar_kws = {} if cbar_kws is None else cbar_kws
        ax.figure.colorbar(cset, cbar_ax, ax, **cbar_kws)

    # Label the axes
    if hasattr(x, "name") and axlabel:
        ax.set_xlabel(x.name)
    if hasattr(y, "name") and axlabel:
        ax.set_ylabel(y.name)

    return ax, cset

# monkey patching
sd._bivariate_kdeplot = _bivariate_kdeplot

# some data
mean, cov = [0, 2], [(1, .5), (.5, 1)]
x, y = np.random.multivariate_normal(mean, cov, size=50).T

# plot
fig, ax = plt.subplots()
_, cs = sns.kdeplot(x, y, ax=ax, cmap="coolwarm")
# label the contours
plt.clabel(cs, cs.levels, inline=True)
# add a colorbar
fig.colorbar(cs)

plt.show()