How to use the earthpy.plot.draw_legend function in earthpy

# You can plot the classified NDVI with a categorical legend using the
# ``draw_legend`` function from the ``earthpy.plot`` module.

# Define color map
nbr_colors = ["gray", "y", "yellowgreen", "g", "darkgreen"]
nbr_cmap = ListedColormap(nbr_colors)

fig, ax = plt.subplots(figsize=(12, 12))

im = ax.imshow(ndvi_landsat_class, cmap=nbr_cmap)

# Get list of classes
classes = np.unique(ndvi_landsat_class)
classes = classes.tolist()

ep.draw_legend(im_ax=im, classes=classes, titles=ndvi_cat_names)

ax.set_title(
    "Landsat 8 - Normalized Difference Vegetation Index (NDVI) Classes",
    fontsize=14,
)
ax.set_axis_off()

# Auto adjust subplot to fit figure size
plt.tight_layout()

# You can plot the classified NDVI with a categorical legend using the
# ``draw_legend`` function from the ``earthpy.plot`` module.

# Define color map
nbr_colors = ["gray", "y", "yellowgreen", "g", "darkgreen"]
nbr_cmap = ListedColormap(nbr_colors)

fig, ax = plt.subplots(figsize=(12, 12))

im = ax.imshow(ndvi_landsat_class, cmap=nbr_cmap)

# Get list of classes
classes = np.unique(ndvi_landsat_class)
classes = classes.tolist()

ep.draw_legend(im_ax=im, classes=classes, titles=ndvi_cat_names)

ax.set_title(
    "Landsat 8 - Normalized Difference Vegetation Index (NDVI) Classes",
    fontsize=16,
)
ax.set_axis_off()
plt.show()

# Create a numpy array. Let's pretend this is what you want to plot.
arr = np.random.randint(4, size=(5, 5))

# When plot_bands is updated a cbar will be here as well
ep.plot_bands(arr)
plt.show()

###############################################################################
# Create Custom Discrete Legends with Earthpy
# -------------------------------------------
# If you want to create a custom categorical legend, you can use the ``ep.draw_legend()`` function.

f, ax = plt.subplots(figsize=(8, 5))
im = ax.imshow(arr)
ep.draw_legend(im)
plt.tight_layout()

###############################################################################
# Customize Discrete Legends
# ---------------------------
# By default the draw_legend function creates a legend with default categories.
# You can customize the legend by adding titles.

f, ax = plt.subplots(figsize=(8, 5))
im = ax.imshow(arr)
ep.draw_legend(im, titles=["Small", "Bigger", "Even Bigger", "Ginormous"])
plt.tight_layout()

###############################################################################
# Discrete Legends With Missing Values in the Array
# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~

)
plt.tight_layout()

###############################################################################
# Custom Colormaps and Ensuring Cmaps Apply to All Valid Classes
# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
#
# You can customize the color map used in your plot too. Notice that in this example,
# 4 category colors are rendered. Yet, the image only contains three values and thus will be
# rendered using three colors. The three colors used to render the image are incorrect by default.
# The colors begin at white and end at black even though the value of 0 which should be black
# is missing from the data.

f, ax = plt.subplots(figsize=(8, 5))
im = ax.imshow(new_arr, cmap="Greys_r")
ep.draw_legend(
    im,
    titles=["Small", "Bigger", "Even Bigger", "Ginormous"],
    classes=[0, 1, 2, 3],
)
plt.tight_layout()

###############################################################################
# Specify vmin and vmax to set the colormap range
# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
#
# In this case, you can use the ``vmin`` and ``vmax`` arguments to set the range of values to use for the colormap.

f, ax = plt.subplots(figsize=(8, 5))
im = ax.imshow(new_arr, cmap="Greys_r", vmin=0, vmax=3)
ep.draw_legend(
    im,

ep.draw_legend(
    im,
    titles=["Small", "Bigger", "Even Bigger", "Ginormous"],
    classes=[0, 1, 2, 3],
)
plt.tight_layout()

###############################################################################
# Specify vmin and vmax to set the colormap range
# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
#
# In this case, you can use the ``vmin`` and ``vmax`` arguments to set the range of values to use for the colormap.

f, ax = plt.subplots(figsize=(8, 5))
im = ax.imshow(new_arr, cmap="Greys_r", vmin=0, vmax=3)
ep.draw_legend(
    im,
    titles=["Small", "Bigger", "Even Bigger", "Ginormous"],
    classes=[0, 1, 2, 3],
)
plt.tight_layout()

# If you want to create a custom categorical legend, you can use the ``ep.draw_legend()`` function.

f, ax = plt.subplots(figsize=(8, 5))
im = ax.imshow(arr)
ep.draw_legend(im)
plt.tight_layout()

###############################################################################
# Customize Discrete Legends
# ---------------------------
# By default the draw_legend function creates a legend with default categories.
# You can customize the legend by adding titles.

f, ax = plt.subplots(figsize=(8, 5))
im = ax.imshow(arr)
ep.draw_legend(im, titles=["Small", "Bigger", "Even Bigger", "Ginormous"])
plt.tight_layout()

###############################################################################
# Discrete Legends With Missing Values in the Array
# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
#
# Now let's pretend that you are creating a series of classified plots. You may have a range of
# values that are expected between 0-4. However not all of your data has all values
# In this case, your legend won't be able to by default create 4 categories because one
# doesn't exist in your data. In this instance, you can specify the values explicitly:

new_arr = arr.copy()
new_arr[new_arr == 0] = 1

f, ax = plt.subplots(figsize=(8, 5))
im = ax.imshow(new_arr)