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

问题

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()

回答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.

来源：https://stackoverflow.com/questions/57024194/how-to-create-an-ax-legend-method-for-contourf-plots-that-doesnt-require-pass