Caroline Gebara
Caroline Gebara

Reputation: 55

How to set fixed color ranges with Seaborn heatmap?

I am trying to make a heatmap with the Seaborn package, where I define my own color ranges. However, I cannot get the legend to show the values in a non-continous way. I would like the following color indication:

0-0.001: green

0.001-0.25: yellow

0.25-0.50: orange

0.50-0.75: red

0.75-1.00: dark red

But I get this:

enter image description here

I suspect that the first range disturbs the picture, because it is smaller. However, I would like the legens "boxes" or area to be the same sizes. Is it therefore possible to have something like this or similar:

enter image description here

The code that I currently use is the one below. Any hint or suggestions would be highly appreciated. Thanks!

my_colors=['#02ab2e','gold','orange','red', 'darkred']


grid_kws = {"height_ratios": (.9, .025), "hspace": .1}
f, (ax, cbar_ax) = plt.subplots(2, gridspec_kw=grid_kws)
ax = sns.heatmap(STEdata.iloc[:,3:13].reindex(ste_order_reg.sort_values().index, axis=0), 
                 yticklabels=2, ax=ax,
                 cmap = my_colors,
                 cbar_ax=cbar_ax, 
                 cbar_kws={"orientation": "horizontal"})
# sns.set(rc = {'figure.figsize':(8, 18)})

colorbar = ax.collections[0].colorbar
colorbar.set_ticks([0, 0, 0.25, 0.5, .75])
colorbar.set_ticklabels(['0',']0-0.25]', ']0.25-0.50]',']0.50-0.75]', ']0.75-1.00]'])

Upvotes: 2

Views: 8506

Answers (1)

JohanC
JohanC

Reputation: 80339

To set uneven color ranges, a BoundaryNorm can be used. The colorbar ticks can be positioned at the center of each range. A ListedColormap creates a colormap from a list of colors.

import matplotlib.pyplot as plt
from matplotlib.colors import BoundaryNorm, ListedColormap
import seaborn as sns
import numpy as np

my_colors = ['#02ab2e', 'gold', 'orange', 'red', 'darkred']
my_cmap = ListedColormap(my_colors)
bounds = [0, 0.0001, 0.25, 0.50, 0.75, 1]
my_norm = BoundaryNorm(bounds, ncolors=len(my_colors))

grid_kws = {"height_ratios": (.9, .025), "hspace": .1}
fig, (ax, cbar_ax) = plt.subplots(nrows=2, figsize=(8,18), gridspec_kw=grid_kws)
sns.heatmap(np.clip(np.random.rand(21, 12) - 0.1, 0, 1),
            yticklabels=2, 
            ax=ax,
            cmap=my_cmap,
            norm=my_norm,
            cbar_ax=cbar_ax,
            cbar_kws={"orientation": "horizontal"})

colorbar = ax.collections[0].colorbar
colorbar.set_ticks([(b0+b1)/2 for b0, b1 in zip(bounds[:-1], bounds[1:])])
colorbar.set_ticklabels(['0', ']0-0.25]', ']0.25-0.50]', ']0.50-0.75]', ']0.75-1.00]'])

plt.show()

sns.heatmap with BoundaryNorm

Instead of a colormap, a custom legend could be created. The same BoundaryNorm will be needed to assign the correct colors in the heatmap. ax.text() and ax.hlines() can be used to place text and lines for the grouping. The y-axis transform uses the y-coordinates of the data and x coordinates as a fraction of the axes. clip_on=False allows drawing outside the main plot area.

import matplotlib.pyplot as plt
from matplotlib.colors import BoundaryNorm, ListedColormap
from matplotlib.lines import Line2D
import seaborn as sns
import numpy as np

unavail_color = 'lightgray'
my_colors = ['#02ab2e', 'gold', 'orange', 'red', 'darkred']
my_cmap = ListedColormap(my_colors)
bounds = [0, 0.0001, 0.25, 0.50, 0.75, 1]
my_norm = BoundaryNorm(bounds, ncolors=len(my_colors))

data = np.exp(np.random.rand(21, 12)) - 1
data[data > 1] = np.nan

fig, ax = plt.subplots(figsize=(8, 18))
ax.set_facecolor(unavail_color)
sns.heatmap(data,
            yticklabels=2, ax=ax,
            cmap=my_cmap,
            norm=my_norm,
            cbar=False)
ax.tick_params(labelsize=16)

group_edges = np.array([0, 5, 6, 13, 14, 21])
group_labels = ['group A', '', 'group B', '', 'group C']
ax.hlines(group_edges, np.zeros(len(group_edges)), np.zeros(len(group_edges)) - 0.12,
          color='navy', lw=2, clip_on=False, transform=ax.get_yaxis_transform())
for label, b0, b1 in zip(group_labels, group_edges[:-1], group_edges[1:]):
    ax.text(-0.12, (b0 + b1) / 2, label, color='navy', fontsize=20, ha='left', va='center',
            rotation=90, transform=ax.get_yaxis_transform())

handles = [Line2D([], [], lw=10, color=color, label=label)
           for color, label in zip(my_colors + [unavail_color],
                                   ['0', ']0-0.25]', ']0.25-0.50]', ']0.50-0.75]', ']0.75-1.00]', 'unavailable'])]
ax.legend(handles=handles, handlelength=0.5, ncol=len(my_colors) + 1,
          bbox_to_anchor=(0.5, -0.02), loc='upper center', frameon=False)

plt.tight_layout()
plt.show()

sns.heatmap with custom legend

Upvotes: 4

Related Questions