Search code examples
pythonmatplotlibseaborndata-visualizationheatmap

How to combine two heatmaps in Seaborn in Python so both are shown in the same heatmap?


This is link to the data I'm using: https://github.com/fivethirtyeight/data/tree/master/drug-use-by-age

I'm using Jupyter Lab, and here's the code:

from matplotlib import pyplot as plt
import pandas as pd
import numpy as np
import seaborn as sb

url = 'https://raw.githubusercontent.com/fivethirtyeight/data/master/drug-use-by-age/drug-use-by-age.csv'

df = pd.read_csv(url, index_col = 0)

df.dtypes
df.replace('-', np.nan, inplace=True)
df = df.iloc[:,:].astype(float)
df = df.loc[:, df.columns != 'n']
#df.columns = df.columns.str.rstrip('-use')
df

fig, axes = plt.subplots(1,2, figsize=(20, 8))

fig.subplots_adjust(wspace=0.1)
fig.colorbar(ax.collections[0], ax=ax,location="right", use_gridspec=False, pad=0.2)
#plt.figure(figsize=(16, 16))
df_percentage = df.iloc[:,range(0,26,2)]
plot_precentage = sb.heatmap(df_percentage, cmap='Reds', ax=axes[0], cbar_kws={'format': '%.0f%%', 'label': '% used in past 12 months'})
df_frequency = df.iloc[:,range(1,27,2)]
plot_frequency = sb.heatmap(df_frequency, cmap='Blues', ax=axes[1], cbar_kws= dict(label = 'median frequency a user used'))

I can just show two of them in a subplot in separate diagrams.

I want to make it look like this (this is made in paint):

enter image description here

Also show the data side by side. Is there a simple way to achieve that?


Solution

  • A pretty simple solution with mask option:

    mask = np.vstack([np.arange(df.shape[1])]* df.shape[0]) % 2
    
    fig, axes = plt.subplots()
    
    plot_precentage = sns.heatmap(df,mask=mask, cmap='Reds', ax=axes, 
                                  cbar_kws={'format': '%.0f%%',
                                            'label': '% used in past 12 months'}
                                 )
    
    plot_frequency = sns.heatmap(df, mask=1-mask, cmap='Blues', ax=axes,
                                 cbar_kws= dict(label = 'median frequency a user used')
                                )
    

    Output:

    enter image description here