Discrete legend in seaborn heatmap plot

Question

I am using the data present here to construct this heat map using seaborn and pandas.

Code:

    import pandas     import seaborn.apionly as sns      # Read in csv file     df_trans = pandas.read_csv('LUH2_trans_matrix.csv')      sns.set(font_scale=0.8)     cmap = sns.cubehelix_palette(start=2.8, rot=.1, light=0.9, as_cmap=True)     cmap.set_under('gray')  # 0 values in activity matrix are shown in gray (inactive transitions)     df_trans = df_trans.set_index(['Unnamed: 0'])     ax = sns.heatmap(df_trans, cmap=cmap, linewidths=.5, linecolor='lightgray')      # X - Y axis labels     ax.set_ylabel('FROM')     ax.set_xlabel('TO')      # Rotate tick labels     locs, labels = plt.xticks()     plt.setp(labels, rotation=0)     locs, labels = plt.yticks()     plt.setp(labels, rotation=0)      # revert matplotlib params     sns.reset_orig() 

As you can see from csv file, it contains 3 discrete values: 0, -1 and 1. I want a discrete legend instead of the colorbar. Labeling 0 as A, -1 as B and 1 as C. How can I do that?

Answer 1

Well, there's definitely more than one way to accomplish this. In this case, with only three colors needed, I would pick the colors myself by creating a LinearSegmentedColormap instead of generating them with cubehelix_palette. If there were enough colors to warrant using cubehelix_palette, I would define the segments on colormap using the boundaries option of the cbar_kws parameter. Either way, the ticks can be manually specified using set_ticks and set_ticklabels.

The following code sample demonstrates the manual creation of LinearSegmentedColormap, and includes comments on how to specify boundaries if using a cubehelix_palette instead.

import matplotlib.pyplot as plt import pandas import seaborn.apionly as sns from matplotlib.colors import LinearSegmentedColormap  sns.set(font_scale=0.8) dataFrame = pandas.read_csv('LUH2_trans_matrix.csv').set_index(['Unnamed: 0'])  # For only three colors, it's easier to choose them yourself. # If you still really want to generate a colormap with cubehelix_palette instead, # add a cbar_kws={"boundaries": linspace(-1, 1, 4)} to the heatmap invocation # to have it generate a discrete colorbar instead of a continous one. myColors = ((0.8, 0.0, 0.0, 1.0), (0.0, 0.8, 0.0, 1.0), (0.0, 0.0, 0.8, 1.0)) cmap = LinearSegmentedColormap.from_list('Custom', myColors, len(myColors))  ax = sns.heatmap(dataFrame, cmap=cmap, linewidths=.5, linecolor='lightgray')  # Manually specify colorbar labelling after it's been generated colorbar = ax.collections[0].colorbar colorbar.set_ticks([-0.667, 0, 0.667]) colorbar.set_ticklabels(['B', 'A', 'C'])  # X - Y axis labels ax.set_ylabel('FROM') ax.set_xlabel('TO')  # Only y-axis labels need their rotation set, x-axis labels already have a rotation of 0 _, labels = plt.yticks() plt.setp(labels, rotation=0)  plt.show() 

Answer 2

Here's a simple solution based on the other answers that generalizes beyond 3 categories and uses a dict (vmap) to define the labels.

import seaborn as sns import numpy as np  # This just makes some sample 2D data and a corresponding vmap dict with labels for the values in the data data = [[1, 2, 3], [4, 5, 6], [7, 8, 9]] vmap = {i: chr(65 + i) for i in range(len(np.ravel(data)))} n = len(vmap)  print(vmap)  cmap = sns.color_palette("deep", n) ax = sns.heatmap(data, cmap=cmap)  # Get the colorbar object from the Seaborn heatmap colorbar = ax.collections[0].colorbar # The list comprehension calculates the positions to place the labels to be evenly distributed across the colorbar r = colorbar.vmax - colorbar.vmin colorbar.set_ticks([colorbar.vmin + 0.5 * r / (n) + r * i / (n) for i in range(n)]) colorbar.set_ticklabels(list(vmap.values()))