How to create an ax.legend() method for contourf plots that doesn't require passing of legend handles from a user?

Question

Desired feature

I would like to be able to call

ax.legend()

on an axis containing a contourf plot and automatically get the legend (see plot below for an example).

More Detail

I know how to create legend entries for contourf plots using proxies, see code below and which is already discussed in this Q&A. However, I would be interested in a solution where the final call to axes[0][-1].legend() does not require any handles being passed.

The plot generation (more complex plots than in this example) is happening in a package and the user will have access to fig and axes and depending on the plots might prefer some axis over the others to plot the legend in. It would be nice if the call to ax.legend() could be simple and would not require the use of proxies and explicit passing of handles. This works automatically for normal plots, scatter plots, hists, etc., but contourf does not accept label as a kwarg and does not come with its own handle so I need to create a proxy (Rectangle patch in this case).

But how could I attach/attribute/... the proxy alongside a label to the contourf plot or to the axes such that ax.legend() can automatically access them the way it does for other types of plots?

Example Image

Example Code

import numpy as np
from scipy import stats
import matplotlib.pyplot as plt
from matplotlib.patches import Rectangle
from matplotlib.colors import LinearSegmentedColormap


########################
# not accessed by User #
########################

def basic_cmap(color):
    return LinearSegmentedColormap.from_list(color, ['#ffffff', color])
cmap1 = basic_cmap('C0')
cmap2 = basic_cmap('C1')

x = np.linspace(0, 10, 50)
mvn1 = stats.multivariate_normal(mean=[4, 4])
mvn2 = stats.multivariate_normal(mean=[6, 7])
X, Y = np.meshgrid(x, x)
Z1 = [[mvn1.pdf([x1, x2]) for x1 in x] for x2 in x]
Z2 = [[mvn2.pdf([x1, x2]) for x1 in x] for x2 in x]
Z1 = Z1 / np.max(Z1)
Z2 = Z2 / np.max(Z2)

fig, axes = plt.subplots(2, 2, sharex='col', sharey='row')
for i, row in enumerate(axes):
    for j, ax in enumerate(row):
        cont1 = ax.contourf(X, Y, Z1, [0.05, 0.33, 1], cmap=cmap1, alpha=0.7)
        cont2 = ax.contourf(X, Y, Z2, [0.05, 0.33, 1], cmap=cmap2, alpha=0.7)


###################################
# User has access to fig and axes #
###################################

proxy1 = plt.Rectangle((0, 0), 1, 1, fc=cmap1(0.999), ec=cmap1(0.33), alpha=0.7, linewidth=3)
proxy2 = plt.Rectangle((0, 0), 1, 1, fc=cmap2(0.999), ec=cmap2(0.33), alpha=0.7, linewidth=3)

# would like this without passing of handles and labels
axes[0][-1].legend(handles=[proxy1, proxy2], labels=['foo', 'bar'])  


plt.savefig("contour_legend.png")
plt.show()

Panter

It pretty rare to see contours with a legend, so no there is not automatic way to call legend. How would legend know which contours to put in the legend. Many people call contour with hundreds of levels. – Turnbull

I know this could not be done for general types of contour plots. However, in astronomy and cosmology it is very common to see the 68% and 95% contours for multiple data sets or multiple theoretical models to see how well they agree/disagree. So this kind of forms a subgroup of contour plots used to plot multiple posterior "blobs" in one plot. Could this maybe be turned into a new plot command inheriting most of plt.contourf's features but with a different treatment of legends where a new standard legend handle is provided? – Panter

Answer 1

Well, I dappled a bit more and found a solution after all that's surprisingly simple, but I had to dig much deeper into matplotlib.legend to get the right idea. In _get_legend_handles it shows how it collects the handles:

    for ax in axs:
        handles_original += (ax.lines + ax.patches +
                             ax.collections + ax.containers)

So all I was lacking was to pass the labels to the proxies and the proxies to ax.patches

Example Code with Solution

changes

        # pass labels to proxies and place proxies in loop
        proxy1 = plt.Rectangle((0, 0), 1, 1, fc=cmap1(0.999), ec=cmap1(0.33), 
                               alpha=0.7, linewidth=3, label='foo')
        proxy2 = plt.Rectangle((0, 0), 1, 1, fc=cmap2(0.999), ec=cmap2(0.33), 
                               alpha=0.7, linewidth=3, label='bar')

        # pass proxies to ax.patches
        ax.patches += [proxy1, proxy2]


###################################
# User has access to fig and axes #
###################################

# no passing of handles and labels anymore
axes[0][-1].legend()

full code

import numpy as np
from scipy import stats
import matplotlib.pyplot as plt
from matplotlib.patches import Rectangle
from matplotlib.colors import LinearSegmentedColormap


########################
# not accessed by User #
########################

def basic_cmap(color):
    return LinearSegmentedColormap.from_list(color, ['#ffffff', color])
cmap1 = basic_cmap('C0')
cmap2 = basic_cmap('C1')

x = np.linspace(0, 10, 50)
mvn1 = stats.multivariate_normal(mean=[4, 4])
mvn2 = stats.multivariate_normal(mean=[6, 7])
X, Y = np.meshgrid(x, x)
Z1 = [[mvn1.pdf([x1, x2]) for x1 in x] for x2 in x]
Z2 = [[mvn2.pdf([x1, x2]) for x1 in x] for x2 in x]
Z1 = Z1 / np.max(Z1)
Z2 = Z2 / np.max(Z2)

fig, axes = plt.subplots(2, 2, sharex='col', sharey='row')
for i, row in enumerate(axes):
    for j, ax in enumerate(row):
        cont1 = ax.contourf(X, Y, Z1, [0.05, 0.33, 1], cmap=cmap1, alpha=0.7)
        cont2 = ax.contourf(X, Y, Z2, [0.05, 0.33, 1], cmap=cmap2, alpha=0.7)

        # pass labels to proxies and place proxies in loop
        proxy1 = plt.Rectangle((0, 0), 1, 1, fc=cmap1(0.999), ec=cmap1(0.33), 
                               alpha=0.7, linewidth=3, label='foo')
        proxy2 = plt.Rectangle((0, 0), 1, 1, fc=cmap2(0.999), ec=cmap2(0.33), 
                               alpha=0.7, linewidth=3, label='bar')

        # pass proxies to ax.patches
        ax.patches += [proxy1, proxy2]


###################################
# User has access to fig and axes #
###################################

# no passing of handles and labels anymore
axes[0][-1].legend()  


plt.savefig("contour_legend_solved.png")
plt.show()

This produces the same image as shown in the question.

Sorry, was able to come up with a solution on my own after all, but maybe this will be helpful for someone else in the future.

Panter