在数据框的其他列上使用多个日期列和条件注释热图

Posted 2023-03-12

【中文标题】在数据框的其他列上使用多个日期列和条件注释热图

【英文标题】：Annotate heatmap with multiple date columns and conditions on other columns of a dataframe

【发布时间】：2022-01-11 07:55:45

【问题描述】：

我有一个示例数据框，如下所示。

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)

我们的目标是创建一个热图，其中 2 列日期组合在一起，按顺序形成单个 x 轴热图。热图的 y 轴具有 ID 数字。热图应该有 3 个类别，如下所示：

0 - 如果当天的两个 C/I 列中都没有干预。 1 - 如果当天Aim2 C/I 列中有“干预”。 2 - 如果当天Aim3 C/I 列中有“干预”。

下面给出了输出图像的外观。

我真的很难检查条件、创建类别并在热图中对其进行注释。

【参考方案1】：

首先重塑数据框；那么热图本身就很简单了：

# 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'])

---

详细分类

使用np.where 映射Intervention 代码：

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

concat 将Date 和C/I 列转换为长格式：

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

创建一个pivot_table 的日期作为列：

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  

用discretized colorbar 绘制带注释的heatmap：

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'])

【讨论】：

感谢您的详细解释。我错过了旋转数据框的技巧，因此一直在苦苦挣扎。非常感谢。