Python：转换患者事件数据以生成桑基图(Sankey plot)

需求说明 #

将包含患者事件时间区间、索引事件的原始DataFrame，按**索引事件前一周（code1）、当周（code2）、后一周（code3）**分组，转换为宽表格式：

• code1：事件完全落在索引日期前7天至前1天的Value值
• code2：固定为索引事件的Value_Index_event值
• code3：事件完全落在索引日期后1天至后7天的Value值
• 无对应时间段事件的位置填充NaN

原始数据 #

import pandas as pd

# 原始DataFrame
df = pd.DataFrame({
    'patient_id': [1, 2, 3],
    'start_date': ['28-12-1999', '28-12-2000', '28-12-2001'],
    'end_date': ['02-01-2000', '02-12-2001', '02-01-2002'],
    'Value': ['A', 'B', 'A'],
    'Index_event_date': ['01-01-2000', '01-01-2001', '01-01-2002'],
    'Value_Index_event': ['X', 'X', 'X']
})

目标格式 #

# 转换后的目标DataFrame
target_df = pd.DataFrame({
    'patient_id': [1, 2, 3],
    'code1': ['A', 'B', 'A'],
    'code2': ['X', 'X', 'X'],
    'code3': ['A', pd.NA, 'A']
})

实现步骤与代码 #

1. 转换日期格式：将所有日期列从字符串转为datetime类型，方便时间计算
2. 计算时间区间：为每个患者计算索引事件的前一周、后一周时间范围
3. 匹配对应code的Value：按患者分组，判断事件所属时间段，提取对应值
4. 合并为宽表：将分组结果整理成目标格式

# 1. 转换日期格式，指定日-月-年的格式
df['start_date'] = pd.to_datetime(df['start_date'], format='%d-%m-%Y')
df['end_date'] = pd.to_datetime(df['end_date'], format='%d-%m-%Y')
df['Index_event_date'] = pd.to_datetime(df['Index_event_date'], format='%d-%m-%Y')

# 2. 定义时间段判断函数，按患者分组处理
def process_patient(group):
    idx_date = group['Index_event_date'].iloc[0]
    # 前一周：索引日期前7天到前1天
    pre_week_start = idx_date - pd.Timedelta(days=7)
    pre_week_end = idx_date - pd.Timedelta(days=1)
    # 后一周：索引日期后1天到后7天
    post_week_start = idx_date + pd.Timedelta(days=1)
    post_week_end = idx_date + pd.Timedelta(days=7)
    
    # 提取code1：事件完全落在前一周
    code1 = group[(group['start_date'] >= pre_week_start) & (group['end_date'] <= pre_week_end)]['Value'].iloc[0] if len(group[(group['start_date'] >= pre_week_start) & (group['end_date'] <= pre_week_end)]) > 0 else pd.NA
    # code2固定为索引事件值
    code2 = group['Value_Index_event'].iloc[0]
    # 提取code3：事件完全落在后一周
    code3 = group[(group['start_date'] >= post_week_start) & (group['end_date'] <= post_week_end)]['Value'].iloc[0] if len(group[(group['start_date'] >= post_week_start) & (group['end_date'] <= post_week_end)]) > 0 else pd.NA
    
    return pd.Series({'code1': code1, 'code2': code2, 'code3': code3})

# 3. 分组处理并合并
result_df = df.groupby('patient_id').apply(process_patient).reset_index()

# 查看结果
print(result_df)

输出结果 #

运行代码后得到的结果与目标格式一致：

patient_id code1 code2 code3
0           1     A     X     A
1           2     B     X  <NA>
2           3     A     X     A

内容的提问来源于stack exchange，提问作者Economist_Ayahuasca