CODEWITHSUNDEEP
matplotlibseabornheatmapmaskcolorbar
Senna
Senna

Reputation: 169

How to mask seaborn heatmap while keeping the colorbar range

I am trying to create a visualization of square matrix with seaborn heatmap, where all elements range between 0 and 1. However, I want to only show those greater than some threshold (ex. 0.5) and set other values to 0. Moreover, I want to set the range of colorbar to be shown between 0.5 and 1, but I do not want to adjust the full colormap to range between 0.5 and 1, but keep the original colormap range.

For example, I attach two examples that I tried:

1st example

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

# Make a random matrix composed of elements between 0 and 1
matrix=np.random.rand(20, 20)
mask_threshold=0.5
matrix_masked=np.zeros_like(matrix)
for i in range(len(matrix_masked)):
    for j in range(len(matrix_masked)):
        if matrix[i][j]<mask_threshold:
            matrix[i][j]=0
        else:
            matrix_masked[i][j]=matrix[i][j]

# Plot random matrix with seaborn heatmap, however I want to mask elements whose values are less than 0.5
# while keeping the colorbar range consistent
fig, ax=plt.subplots(figsize=(24, 20))
sns.heatmap(matrix, cmap="seismic", ax=ax, vmin=0, vmax=1)

cbar=ax.collections[0].colorbar
cbar.ax.tick_params(labelsize=30)
cbar.ax.yaxis.label.set_size(30)

cbar.outline.set_linewidth(1.0)
plt.tight_layout()
plt.show();

In this first example, the visualized matrix exactly corresponds to what I originally wanted, but not the colorbar because it ranges between 0 and 1. Therefore, I tried changing the range of colorbar like in the second example.

2nd example

# Plot random matrix with seaborn heatmap, however I want to mask elements whose values are less than 0.5
# while keeping the colorbar range consistent
fig, ax=plt.subplots(figsize=(24, 20))
sns.heatmap(matrix_masked, cmap="seismic", ax=ax, vmin=0, vmax=1)

cbar=ax.collections[0].colorbar
cbar.ax.set_ylim(0.5, 1.0)
cbar.ax.tick_params(labelsize=30)
cbar.ax.set_yticklabels(labels=[0.5, 0.6, 0.7, 0.8, 0.9, 1.0], weight="bold")
cbar.ax.yaxis.label.set_size(30)

cbar.outline.set_linewidth(1.0)
plt.tight_layout()
plt.show();

In this 2nd example, the range of colorbar indeed changes like what I want, but the height of figure and colorbar does not match, which is serious problem.

How can I solve this problem?

Upvotes: 3

Views: 1657

Answers (1)

tdy
tdy

Reputation: 41327

Your second example seems fine on my end, but here is a simpler alternative:

  1. Use the Reds cmap with vmin=0.5 (no need to alter the cbar's ylim and yticklabels)
  2. Use the mask param to automatically mask thresholded values (no need to zero them out)
  3. Set the "bad" and "under" colors for displaying masked values on the heatmap and colorbar
matrix = np.random.default_rng(0).random(size=(20, 20))

cmap = plt.cm.get_cmap('Reds').copy()
cmap.set_bad('midnightblue')      # color of mask on heatmap
cmap.set_under('midnightblue')    # color of mask on cbar

sns.heatmap(matrix,
    cmap=cmap, vmin=0.5, vmax=1,  # set cbar range from 0.5 to 1
    mask=matrix < mask_threshold, # use "bad" color for thresholded values
    cbar_kws={'extend': 'min'})   # extend cbar to show "under" color

Upvotes: 3

Related Questions