pythonpandasmatplotlibseaborn

How to generate a hierarchical colourmap in matplotlib?


I have a hierarchical dataset that I wish to visualise in this manner. I've been able to construct a heatmap for it.

enter image description here

I want to generate a colormap in matplotlib such that Level 1 get categorical colours while Level 2 get different shades of the Level 1 colour. I was able to get Level 1 colours from a "tab20" palette but I can't figure out how to generate shades of the base Level 1 colour.

EDIT: Just to be clear, this needs to be a generic script. So I can't hard code the colormap.

MWE

At the moment this just creates a colormap based on the level 1 values. I am not sure how to generate the shades for the level 2 colours:

import pandas as pd
import seaborn as sns
import matplotlib.pyplot as plt
import matplotlib as mpl

df = pd.DataFrame({"Level 2": [4, 5, 6, 6, 7], "Level 1": [0, 0, 1, 1, 1]}).T

colours = mpl.colormaps["tab20"].resampled(len(df.loc["Level 1"].unique())).colors

colour_dict = {
    item: colour for item, colour in zip(df.loc["Level 1"].unique(), colours)
}

sns.heatmap(
    df,
    cmap=mpl.colors.ListedColormap([colour_dict[item] for item in colour_dict.keys()]),
)
colours

In this example, 4 and 5 should be shades of the colour for 0 and 6 and 7 should be shades of the colour for 1.

Edit 2

Applying @mozway's answer below, this is the heatmap I see:

enter image description here

This is with 423 values in level 2 and n=500.


Solution

  • What about combining several gradients to form a multi-colored cmap, then rescaling your data?

    import matplotlib as mpl
    from matplotlib.colors import LinearSegmentedColormap
    
    df = pd.DataFrame({"Level 2": [1, 2, 1, 2, 3, 2, 3, 4, 0, 1, 5],
                       "Level 1": [0, 0, 1, 1, 1, 2, 2, 2, 3, 3, 3]}).T
    
    n = 5 # max value per level
    level1 = pd.factorize(df.loc['Level 1'])[0]*n
    n_levels = df.loc['Level 1'].nunique()
    
    cmap = mpl.colormaps['tab10']
    
    i = np.linspace(0, 1, num=n_levels+1)
    colors = list(zip(np.sort(np.r_[i, i[1:-1]-0.001]),
                      [x for c in cmap.colors[:n_levels+1]
                       for x in (c, 'w')]))
    
    multi_cmap = LinearSegmentedColormap.from_list('hierachical', colors)
    
    tmp = pd.DataFrame({'Level 2': level1+df.loc['Level 2'],
                        'Level 1': level1
                        }).T
    
    sns.heatmap(tmp, cmap=multi_cmap, vmin=0, vmax=n*n_levels,
                square=True, cbar_kws={'orientation': 'horizontal'})
    

    Output:

    matplotlib hierarchical colormap

    If you want to annotate with the real values:

    sns.heatmap(tmp, annot=df, cmap=multi_cmap, vmin=0, vmax=n*n_levels,
                square=True, cbar_kws={'orientation': 'horizontal'})
    

    Output:

    matplotlib hierarchical colormap

    further customization

    If you want to reverse the order of the gradients:

    import matplotlib as mpl
    from matplotlib.colors import LinearSegmentedColormap
    
    df = pd.DataFrame({"Level 2": [1, 2, 1, 2, 3, 2, 3, 4, 0, 1, 5],
                       "Level 1": [0, 0, 1, 1, 1, 2, 2, 2, 3, 3, 3]}).T
    
    n = 5 # max value per level
    level1 = pd.factorize(df.loc['Level 1'])[0]*n
    n_levels = df.loc['Level 1'].nunique()
    
    cmap = mpl.colormaps['tab10']
    
    i = np.linspace(0, 1, num=n_levels+1)
    colors = list(zip(np.sort(np.r_[i, i[1:-1]-0.001]),
                      [x for c in cmap.colors[:n_levels+1]
                       for x in ('w', c)]))
    
    multi_cmap = LinearSegmentedColormap.from_list('hierachical', colors)
    
    tmp = pd.DataFrame({'Level 2': level1+df.loc['Level 2'],
                        'Level 1': level1+0.999*n
                        }).T
    
    sns.heatmap(tmp, annot=df, cmap=multi_cmap, vmin=0, vmax=n*n_levels,
                square=True, cbar_kws={'orientation': 'horizontal'})
    

    Output:

    matplotlib hierarchical colormap

    Output with n = 10

    matplotlib hierarchical colormap