Mask two groups of values in matlotlib colormap with gradual colorscheme

Question

Say I want to distinguish the NaNs in a matplotlib colormap. Then:

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

# create a (4,5) matrix with values ranging from 0 to 19
np_data = np.arange(20).reshape((4,5)).astype(float)
# add a row with NaNs in the middle
np_data = np.insert(np_data,2,[np.nan for x in range(0,5)],axis=0)
# mask invalid data (NaNs)
np_data = np.ma.masked_invalid(np_data)

 # get figure and ax objects from plot
fig, ax = plt.subplots() 
 # Draw an "X" on transparent values (masked values)
ax.patch.set(hatch='x', edgecolor='blue')

# get a predefined color scheme
reds_cm = plt.get_cmap("Reds")
# Plot heatmap, add a colorbar and show it 
heatmap = ax.pcolor(np_data, cmap=reds_cm)
cbar = fig.colorbar(heatmap)
plt.show()

Plots:

Now NaNs are easily identifiable in the plot.

Now, say I want to be able to easily tell apart between NaNs, 0s and the rest of the values.

If I now mask the 0s, I won't be able to tell the NaNs and the 0s apart.

How can I differentiate 2 groups of values in a colormap? In this case NaNs on one hand and 0s in the other.

Answer 1

In case you want to tell appart the first or last value of your colormap the following solution is a good way to go. You can modify the colormap such that those values become a different color quite easily

reds_cm = plt.get_cmap("name of colormap")
# init colormap such that its members are available
reds_cm._init()
# set the first value to black
reds_cm._lut[0,: ] = (0,0,0,1) #this is an RGBA tuple 
# set the last value to lightgreen
reds_cm._lut[-4:,: ] = np.array([149,238,58,255])/255.

Here is a full solution.

import numpy as np
import matplotlib.pyplot as plt

# create a (4,5) matrix with values ranging from 0 to 19
np_data = np.arange(20).reshape((4,5)).astype(float)
# add a row with NaNs in the middle
np_data = np.insert(np_data,2,[np.nan for x in range(0,5)],axis=0)
# mask invalid data (NaNs)
np_data = np.ma.masked_invalid(np_data)

 # get figure and ax objects from plot
fig, ax = plt.subplots() 
 # Draw an "X" on transparent values (masked values)
ax.patch.set(hatch='x', edgecolor='blue')

# get a predefined color scheme
reds_cm = plt.get_cmap("Reds")
# init colormap such that its members are available
reds_cm._init()
# set the first value to black
reds_cm._lut[0,: ] = (0,0,0,1)
# set the last value to lightgreen
reds_cm._lut[-4:,: ] = np.array([149,238,58,255])/255.

# Plot heatmap, add a colorbar and show it 
heatmap = ax.pcolor(np_data, cmap=reds_cm)
cbar = fig.colorbar(heatmap)
plt.show()

producing

Answer 2

I found this answer from @unutbu in an unrelated question. I adapted his answer to my problem and also fixed the issue that the new hatches are also included in the NaN cells. To avoid this, just get the cells with value 0 before masking the numpy array (I would comment on his answer to point this out in context but I don't have the required reputation). I only include code changed from my question.

# (previous imports)
# Import to add patches to "non transparent" cells
import matplotlib.patches as mpatches


# (generate np_data)

# Get mask positions of 0 values before masking NaNs so NaN cells aren't included
cells_with_0 = np_data == 0
# mask invalid data (NaNs)
np_data = np.ma.masked_invalid(np_data)

# (get color scheme, plot heatmap, plot colorbar)

#set the background color as gray so the transparent values (NaNs here) use that color
ax.patch.set_facecolor((0.6, 0.6, 0.6, 1.0))
# Draw an "X" on transparent values (masked values)
ax.patch.set(hatch='x', edgecolor='black')
# Put an x over cells which have value 0
for j, i in np.column_stack(np.where(cells_with_0)):
      ax.add_patch(
          mpatches.Rectangle(
              (i, j),     # (x,y)
              1,          # width
              1,          # height
              fill=False, 
              edgecolor='blue',
              snap=False,
              hatch='x' # the more slashes, the denser the hash lines 
          ))

plt.show()

New heatmap: