How to mask seaborn heatmap while keeping the colorbar range-CodePudding

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?

CodePudding user response：

Your