CODEWITHSUNDEEP
pythonseabornheatmap
kutu
kutu

Reputation: 53

Assign specific color to seaborn heatmap

I'm trying to make heatmap using seaborn, but got stuck to change color on specific values. Suppose, the value 0 should be white, and value 1 should be grey, then over that uses the palette as provided by cmap.

Was trying to use mask, but got confused.

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

df = pd.read_csv('/home/test.csv', index_col=0)

fig, ax = plt.subplots()
sns.heatmap(df, cmap="Reds", vmin=0, vmax=15)
plt.show()

this for the sample data

TAG     A   B   C   D   E   F   G   H   I   J
TAG_1   1   0   0   5   0   7   1   1   0   10
TAG_2   0   1   0   6   0   6   0   0   0   7
TAG_3   0   1   0   2   0   4   0   0   1   4
TAG_4   0   0   0   3   1   3   0   0   0   10
TAG_5   1   0   1   5   0   2   1   1   0   11
TAG_6   0   0   0   0   0   0   0   0   0   12
TAG_7   0   1   0   0   1   0   0   0   0   0
TAG_8   0   0   0   1   0   0   1   0   1   0
TAG_9   0   0   1   0   0   0   0   0   0   0
TAG_10  0   0   0   0   0   0   0   0   0   0

Upvotes: 5

Views: 5839

Answers (1)

JohanC
JohanC

Reputation: 80319

df.set_index('TAG', inplace=True) tells seaborn that the tags should be used as tags, not as data.

The 'binary' colormap goes smoothly from white for the lower values to dark black for the highest. Playing with vmin and vmax, setting vmin=0 and vmax to a value between 1.5 and about 5, value 0 will be white and 1 will be any desired type of gray.

To set a mask, the dataframe should be converted to a 2D numpy array and be of type float.

import matplotlib.pyplot as plt
import seaborn as sns
import pandas as pd
from io import StringIO

data_str = StringIO('''TAG     A   B   C   D   E   F   G   H   I   J
TAG_1   1   0   0   5   0   7   1   1   0   10
TAG_2   0   1   0   6   0   6   0   0   0   7
TAG_3   0   1   0   2   0   4   0   0   1   4
TAG_4   0   0   0   3   1   3   0   0   0   10
TAG_5   1   0   1   5   0   2   1   1   0   11
TAG_6   0   0   0   0   0   0   0   0   0   12
TAG_7   0   1   0   0   1   0   0   0   0   0
TAG_8   0   0   0   1   0   0   1   0   1   0
TAG_9   0   0   1   0   0   0   0   0   0   0
TAG_10  0   0   0   0   0   0   0   0   0   0''')

df = pd.read_csv(data_str, delim_whitespace=True)
df.set_index('TAG', inplace=True)
values = df.to_numpy(dtype=float)
ax = sns.heatmap(values, cmap='Reds', vmin=0, vmax=15, square=True)
sns.heatmap(values, xticklabels=df.columns, yticklabels=df.index,
            cmap=plt.get_cmap('binary'), vmin=0, vmax=2, mask=values > 1, cbar=False, ax=ax)
plt.show()

example plot

Alternatively, a custom colormap could be created. That way the colorbar will also show the adapted colors.

from matplotlib.colors import LinearSegmentedColormap

cmap_reds = plt.get_cmap('Reds')
num_colors = 15
colors = ['white', 'grey'] + [cmap_reds(i / num_colors) for i in range(2, num_colors)]
cmap = LinearSegmentedColormap.from_list('', colors, num_colors)
ax = sns.heatmap(df, cmap=cmap, vmin=0, vmax=num_colors, square=True, cbar=False)
cbar = plt.colorbar(ax.collections[0], ticks=range(num_colors + 1))
plt.show()

custom colormap

Upvotes: 6

Related Questions