def test_colors(arr_plot_blues):
"""Test that the correct colors appear in the patches of the legend"""
im_ax, _ = arr_plot_blues
the_legend = ep.draw_legend(im_ax=im_ax)
legend_cols = [i.get_facecolor() for i in the_legend.get_patches()]
# Get the array and cmap from axis object
cmap_name = im_ax.axes.get_images()[0].get_cmap().name
unique_vals = np.unique(im_ax.get_array().data)
image_colors = ep.make_col_list(unique_vals, cmap=cmap_name)
assert image_colors == legend_cols
plt.close()
def test_colors(binned_array_3bins):
"""Test that the correct colors appear in the patches of the legend"""
bins, im_arr_bin = binned_array_3bins
f, ax = plt.subplots()
im = ax.imshow(im_arr_bin, cmap="Blues")
the_legend = ep.draw_legend(im_ax=im)
# NOTE: Do I know for sure things are rendering in the right order?
legend_cols = [i.get_facecolor() for i in the_legend.get_patches()]
# Get the array and cmap from axis object
cmap_name = im.axes.get_images()[0].get_cmap().name
unique_vals = np.unique(im.get_array().data)
image_colors = ep.make_col_list(unique_vals, cmap=cmap_name)
assert image_colors == legend_cols
def draw_legend(im_ax, bbox=(1.05, 1), titles=None, cmap=None, classes=None):
"""Create a custom legend with a box for each class in a raster.
This is an exact copy of the `earthpy.plot.draw_legend` function except
that it uses a relative font size in the legend to allow the text to
scale with image size, plus a couple tweaks to allow it to run outside of
its origianl context.
Parameters
----------
im_ax : matplotlib image object
This is the image returned from a call to imshow().
bbox : tuple (default = (1.05, 1))
This is the bbox_to_anchor argument that will place the legend
anywhere on or around your plot.
titles : list (optional)
A list of a title or category for each unique value in your raster.
This is the label that will go next to each box in your legend. If
nothing is provided, a generic "Category x" will be populated.
cmap : str (optional)
Colormap name to be used for legend items.
classes : list (optional)
A list of unique values found in the numpy array that you wish to plot.
Returns
----------
matplotlib.pyplot.legend
A matplotlib legend object to be placed on the plot.
"""
# Lazy load imports to keep this function self-contained so that it can
# be removed cleanly if a better solution presents itself
from matplotlib import patches as mpatches
from matplotlib.colors import ListedColormap
try:
im_ax.axes
except AttributeError:
raise AttributeError(
"The legend function requires a matplotlib axis object to "
"run properly. You have provided a {}.".format(type(im_ax)))
# If classes not provided, get them from the im array in the ax object
# Else use provided vals
if classes is not None:
# Get the colormap from the mpl object
cmap = im_ax.cmap.name
# If the colormap is manually generated from a list
if cmap == "from_list":
cmap = ListedColormap(im_ax.cmap.colors)
colors = ep.make_col_list(nclasses=len(classes),
unique_vals=classes,
cmap=cmap)
# If there are more colors than classes, raise value error
if len(set(colors)) < len(classes):
raise ValueError(
"There are more classes than colors in your cmap. "
"Please provide a ListedColormap with the same number "
"of colors as classes.")
else:
classes = list(np.unique(im_ax.axes.get_images()[0].get_array()))
# Remove masked values, could next this list comp but keeping it simple
classes = [
aclass for aclass in classes if aclass is not np.ma.core.masked
]
colors = [im_ax.cmap(im_ax.norm(aclass)) for aclass in classes]
# If titles are not provided, create filler titles
if not titles:
titles = ["Category {}".format(i + 1) for i in range(len(classes))]
if not len(classes) == len(titles):
raise ValueError(
"The number of classes should equal the number of "
"titles. You have provided {0} classes and {1} titles.".format(
len(classes), len(titles)))
patches = [
mpatches.Patch(color=colors[i], label="{lab}".format(lab=titles[i]))
for i in range(len(titles))
]
# Get the axis for the legend
ax = im_ax.axes
return ax.legend(
handles=patches,
bbox_to_anchor=bbox,
loc=2,
borderaxespad=0.0,
prop={"size": "small"},
)