Python Matplotlib: how to add "bad" color to the legend?

Question

The color map in matplotlib allows to mark "bad" values, i.e. NaNs, with a specific color. When we plot the color bar afterwards, this color is not included. Is there a preferred approach to have both the contiuous color bar and a discrete legend for the specific color for bad values?

Edit: Certainly, it's possible to make use of the "extend" functionality. However, this solution is not satisfactory. The function of the legend/colorbar is to clarify the meaning of colors to the user. In my opinion, this solution does not communicate that the value is a NaN.

Code example:

import numpy as np
import matplotlib.pyplot as plt
import matplotlib.cm as cm

data = np.random.rand(10, 10)
data[0:3, 0:3] = np.nan  # some bad values for set_bad

colMap = cm.RdBu
colMap.set_bad(color='black')

plt.figure(figsize=(10, 9))
confusion_matrix = plt.imshow(data, cmap=colMap, vmin=0, vmax=1)
plt.colorbar(confusion_matrix)
plt.show()

Which produces:

Answer 1

A legend element could be created and used as follows:

from matplotlib.patches import Patch

legend_elements = [Patch(facecolor=colMap(np.nan), label='Bad values')]
plt.legend(handles=legend_elements)

Answer 2

While adding it to the legend is an option (see JohanC's answer), the color semantically belongs to the colorbar and I usually don't want a separate legend on a color-coded plot with a colorbar.

After experimenting for way too long, the best way I found is to add a separate single-color colorbar that associates with the proper colorbar:

import numpy as np
import matplotlib as mpl
import matplotlib.pyplot as plt
import matplotlib.cm as cm

data = np.random.rand(10, 10)
data[0:3, 0:3] = np.nan  # some bad values for set_bad

colMap = cm.RdBu
colMap.set_bad(color='purple')

plt.figure(figsize=(10, 9))
plt.grid(False)
confusion_matrix = plt.imshow(data, cmap=colMap, vmin=0, vmax=1)
cbar= plt.colorbar(confusion_matrix)

sm = cm.ScalarMappable(cmap= mpl.colors.ListedColormap([colMap.get_bad()]))
divider = make_axes_locatable(cbar.ax)  # for tight_layout compatibility
nan_ax = divider.append_axes("bottom", size= "5%", pad= "3%", aspect= 1, anchor= cbar.ax.get_anchor())
nan_ax.grid(visible=False, which='both', axis='both')  # required for Colorbar constructor below
nan_cbar = mpl.colorbar.Colorbar(ax=nan_ax, mappable=sm, orientation='vertical')
nan_cbar.set_ticks([0.5], labels=['NaN'])
nan_cbar.ax.tick_params(length= 0)  # optional to drop the unnecessary tick line

Size and distance of the NaN colorbar can be adjusted via size and pad above.

Elaboration on proposed alternative

As you wrote in your edit, another option is to abuse extend, but the color will be added directly below/above the colorbar with no visible gap and there is no clear way of associating a label with it in the API:

import numpy as np
import matplotlib as mpl
import matplotlib.pyplot as plt
import matplotlib.cm as cm

data = np.random.rand(10, 10)
data[0:3, 0:3] = np.nan  # some bad values for set_bad

colMap = cm.RdBu
colMap.set_under(color='purple')

plt.figure(figsize=(10, 9))
plt.grid(False)
confusion_matrix = plt.imshow(data, cmap=colMap, vmin=0, vmax=1)
cbar= plt.colorbar(confusion_matrix, extend='min', extendrect=True)

The third option would be to create a new colormap and reserve a fraction for the NaN value, but that is also cumbersome and allows for no clear gap to the real colors.