Annotate heatmap with multiple date columns and conditions on other columns of a dataframe

Question

I have a sample dataframe as given below.

import pandas as pd

data = {'ID':['001', '002', '003', '004', '005'],
    'Aim2 Date':['9/3/19', '9/3/19', '9/13/19', '9/20/19', '9/23/19'],
    'Aim2 C/I' :['Intervention', 'Intervention', 'Intervention', 'Control', 'Intervention'],
    'Aim3 Date':['9/17/19','9/19/19','9/27/19','9/30/19','10/8/19'],
    'Aim3 C/I' :['Control', 'Intervention', 'Control', 'Withdraw', 'Withdraw']}

df = pd.DataFrame(data)

The goal is to create a heatmap where the 2 columns of date are combined to form a single x-axis of heatmap in order. The y-axis of the heatmap has the ID number. The heatmap should have 3 categories as given below:

• 0 - if there are no Interventions in both the C/I columns for that day.
• 1 - if there is 'Intervention' in Aim2 C/I column for that day.
• 2 - if there is 'Intervention' in Aim3 C/I column for that day.

The image of the output as in how it should look is given below.

I am really struggling with checking the conditions, creating the categories and annotating them within the heatmap.

Answer 1

First reshape the dataframe; then the heatmap itself will be straightforward:

# map the intervention codes
df['Aim2 C/I'] = np.where(df['Aim2 C/I'] == 'Intervention', 1, 0)
df['Aim3 C/I'] = np.where(df['Aim3 C/I'] == 'Intervention', 2, 0)

# concat the date and C/I columns into long form
df = pd.concat([
    df[['ID', 'Aim2 Date', 'Aim2 C/I']].rename(columns={'Aim2 Date': 'Date', 'Aim2 C/I': 'C/I'}),
    df[['ID', 'Aim3 Date', 'Aim3 C/I']].rename(columns={'Aim3 Date': 'Date', 'Aim3 C/I': 'C/I'}),
])

# pivot the dates into columns
df['Date'] = pd.to_datetime(df['Date']).dt.strftime('%Y-%m-%d')
df = df.pivot_table(index='ID', columns='Date', fill_value=0).droplevel(0, axis=1)

# plot the annotated heatmap with discretized colorbar
ax = sns.heatmap(df, annot=True, cmap=sns.color_palette('cividis', 3))
colorbar = ax.collections[0].colorbar
colorbar.set_ticks([0.33, 1, 1.67])
colorbar.set_ticklabels(['0: No interventions', '1: Aim2 intervention', '2: Aim3 intervention'])

---

Detailed breakdown

1. Use np.where to map the Intervention codes:

   df['Aim2 C/I'] = np.where(df['Aim2 C/I'] == 'Intervention', 1, 0)
   df['Aim3 C/I'] = np.where(df['Aim3 C/I'] == 'Intervention', 2, 0)
   
   #     ID Aim2 Date  Aim2 C/I Aim3 Date  Aim3 C/I
   # 0  001    9/3/19         1   9/17/19         0
   # 1  002    9/3/19         1   9/19/19         2
   # 2  003   9/13/19         1   9/27/19         0
   # 3  004   9/20/19         0   9/30/19         0
   # 4  005   9/23/19         1   10/8/19         0
   

2. concat the Date and C/I columns into long form:

   df = pd.concat([
       df[['ID', 'Aim2 Date', 'Aim2 C/I']].rename(columns={'Aim2 Date': 'Date', 'Aim2 C/I': 'C/I'}),
       df[['ID', 'Aim3 Date', 'Aim3 C/I']].rename(columns={'Aim3 Date': 'Date', 'Aim3 C/I': 'C/I'}),
   ])
   
   #     ID     Date  C/I
   # 0  001   9/3/19    1
   # 1  002   9/3/19    1
   # 2  003  9/13/19    1
   # 3  004  9/20/19    0
   # 4  005  9/23/19    1
   # 0  001  9/17/19    0
   # 1  002  9/19/19    2
   # 2  003  9/27/19    0
   # 3  004  9/30/19    0
   # 4  005  10/8/19    0
   

3. Create a pivot_table of the dates as columns:

   df['Date'] = pd.to_datetime(df['Date']).dt.strftime('%Y-%m-%d')
   df = df.pivot_table(index='ID', columns='Date', fill_value=0).droplevel(0, axis=1)
   
   # Date  2019-09-03  2019-09-13  2019-09-17  2019-09-19  2019-09-20  2019-09-23  2019-09-27  2019-09-30  2019-10-08  
   # ID                                                                                                                
   # 001            1           0           0           0           0           0           0           0           0  
   # 002            1           0           0           2           0           0           0           0           0  
   # 003            0           1           0           0           0           0           0           0           0  
   # 004            0           0           0           0           0           0           0           0           0  
   # 005            0           0           0           0           0           1           0           0           0  
   

4. Plot the annotated heatmap with a discretized colorbar:

   ax = sns.heatmap(df, annot=True, cmap=sns.color_palette('cividis', 3))
   colorbar = ax.collections[0].colorbar
   colorbar.set_ticks([0.33, 1, 1.67])
   colorbar.set_ticklabels(['0: No interventions', '1: Aim2 intervention', '2: Aim3 intervention'])