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

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.

CodePudding user response：

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