import pandas as pd import numpy as np import pyreadstat from sklearn.experimental import enable_iterative_imputer from sklearn.impute import IterativeImputer if __name__ == "__main__": CHARLS_data = pd.read_csv("CHARLS_data_pollutants_p_n_m_nd_h.csv") cavariates = ["last_year_NO2", "before_last_NO2", "last_year_O3", "before_last_O3", "last_year_pm1", "before_last_pm1", "last_year_pm2.5", "before_last_pm2.5", "last_year_pm10", 'before_last_pm10', 'last_year_SO4', 'last_year_NO3', 'last_year_NH4', 'last_year_OM', 'last_year_BC', 'before_last_SO4', 'before_last_NO3', 'before_last_NH4', 'before_last_OM', 'before_last_BC', 'last_year_nl', 'before_last_nl'] #挑出需要的字段 data = CHARLS_data[["ID", "rgender", "age", "marital_status", "education", "Physical_activity", "Psychiatric_score", "BMI", "ADL", "Smoke", "Drink", 'Hypertension','Dyslipidemia','Disabetes_or_High_Blood_Sugar','Cancer_or_Malignant_Tumor','Chronic_Lung_Diseases', 'Liver_Disease', 'Heart_Problems', 'Stroke', 'Kidney_Diease','Stomach_or_Other_Digestive_Disease', 'Emotional_Nervous_or_Psychiatric_Problems', 'Memory_Related_Disease','Arthritis_or_Rheumatism','Asthma', 'wave' ]+cavariates] #处理共病状态 data["state"] = data['Hypertension'] + data['Dyslipidemia']+ data['Disabetes_or_High_Blood_Sugar'] + data['Cancer_or_Malignant_Tumor'] + data['Chronic_Lung_Diseases'] + data['Liver_Disease'] \ + data['Heart_Problems'] + data['Stroke'] + data['Kidney_Diease'] + data['Stomach_or_Other_Digestive_Disease'] + data['Emotional_Nervous_or_Psychiatric_Problems'] + data['Memory_Related_Disease'] \ + data['Arthritis_or_Rheumatism'] + data['Asthma'] data["state"] = data['state'].apply(lambda x : 1 if x == 0 else 2 if x == 1 else 3 if x >= 2 else np.nan) data = data.drop(columns=['Hypertension','Dyslipidemia','Disabetes_or_High_Blood_Sugar','Cancer_or_Malignant_Tumor','Chronic_Lung_Diseases', 'Liver_Disease', 'Heart_Problems', 'Stroke', 'Kidney_Diease','Stomach_or_Other_Digestive_Disease', 'Emotional_Nervous_or_Psychiatric_Problems', 'Memory_Related_Disease','Arthritis_or_Rheumatism','Asthma']) #增加一列死亡状态 #0：未死亡 #1：死亡 #读取2013年的死亡数据 exit, meta = pyreadstat.read_dta("/root/r_base/CHARLS/CHARLS2013/Exit_Interview.dta") exit['ID'] = pd.to_numeric(exit['ID'], errors='coerce').astype('Int64') exit["exit_year"] = exit["exb001_1"] data = pd.merge(data, exit[['ID', "exit_year"]], on = "ID", how="left") #读取2020年的死亡数据 exit, meta = pyreadstat.read_dta("/root/r_base/CHARLS/CHARLS2020/Exit_Module.dta") exit['ID'] = pd.to_numeric(exit['ID'], errors='coerce').astype('Int64') exit["exit_year"] = exit["exb001_1"] data = pd.merge(data, exit[['ID', "exit_year"]], on = "ID", how="left") #合并两次死亡数据 data["exit_year"] = data["exit_year_x"].fillna(data["exit_year_y"]) # 定义随访年的列表 followup_years = [2011, 2013, 2015, 2018, 2020] # 将 'id' 和 'wave' 转换为整型，确保排序不会出问题 data['ID'] = data['ID'].astype(int) data['wave'] = data['wave'].astype(int) # 找到每个死亡个体的最后一次随访数据 last_followup = data.dropna(subset=['exit_year']).groupby('ID').apply(lambda x: x[x['wave'] == x['wave'].max()]) # 创建一个布尔掩码，用于标记需要将 'state' 修改为 4 的行 mask = last_followup['wave'] > last_followup['exit_year'] # 将 MultiIndex 转换为单一的 ID 索引 last_followup_ids = last_followup[mask].index.get_level_values('ID') # 使用布尔索引直接在 data 中修改对应行的 'state' data.loc[data['ID'].isin(last_followup_ids) & (data['wave'] == data.groupby('ID')['wave'].transform('max')), 'state'] = 4 # 创建新的记录并为每个死亡个体设置下一次随访年 new_rows = last_followup[last_followup['wave'] <= last_followup['exit_year']].copy() new_rows['wave'] = new_rows['wave'].apply(lambda x: next((year for year in followup_years if year > x), None)) new_rows['state'] = 4 # 将新行添加到原始 DataFrame 中 data = pd.concat([data, new_rows], ignore_index=True).sort_values(by=['ID', 'wave']).reset_index(drop=True) #删除多余列 data = data.drop(columns=["exit_year_x", "exit_year_y", "exit_year"]) # 统计唯一用户的个数 unique_user_count = data['ID'].nunique() print(f"删除空状态前的用户数：{unique_user_count}") # 将状态为空的数据删除 # 查找所有有空值 state 的用户 ID users_with_na = data[data['state'].isna()]['ID'].unique() # 删除这些用户的所有数据 data = data[~data['ID'].isin(users_with_na)] unique_user_count = data['ID'].nunique() print(f"删除空状态后的用户数：{unique_user_count}") #获取参加全部批次的纵向数据 # 1. 统计每个用户的批次数 user_counts = data.groupby('ID')['wave'].count().reset_index() user_counts.columns = ['ID', 'wave_count'] # 2. 找到每个用户的最大批次的状态 max_wave_state = data.loc[data.groupby('ID')['wave'].idxmax()][['ID', 'state']] max_wave_state.columns = ['ID', 'max_wave_state'] # 3. 将用户的批次数和最大批次状态合并回原始数据 data = data.merge(user_counts, on='ID').merge(max_wave_state, on='ID') # 4. 筛选满足条件的用户 condition_1 = (data['wave_count'] == 5) condition_2 = (data['max_wave_state'] == 4) & (data['wave_count'] > 1) data = data[condition_1 | condition_2] # 5. 清除多余的列 data = data.drop(columns=['wave_count', 'max_wave_state']).reset_index(drop=True) unique_user_count = data['ID'].nunique() print(f"参加全部批次的用户数：{unique_user_count}") #删除45岁以下的用户ID users_with_45_age = data[data['age']<45]['ID'].unique() # 删除这些用户的所有数据 data = data[~data['ID'].isin(users_with_45_age)] unique_user_count = data['ID'].nunique() print(f"删除45岁后的用户数：{unique_user_count}") # 查找所有有空值age 的用户 ID users_with_na_age = data[data['age'].isna()]['ID'].unique() # 删除这些用户的所有数据 data = data[~data['ID'].isin(users_with_na_age)] unique_user_count = data['ID'].nunique() print(f"删除空年龄后的用户数：{unique_user_count}") # 查找所有有空值education 的用户 ID users_with_na_education = data[data['education'].isna()]['ID'].unique() # 删除这些用户的所有数据 data = data[~data['ID'].isin(users_with_na_education)] unique_user_count = data['ID'].nunique() print(f"删除空教育后的用户数：{unique_user_count}") #删除异常的BMI users_with_BMI_wr = data[data["BMI"]>=200]['ID'].unique() # 删除这些用户的所有数据 data = data[~data['ID'].isin(users_with_BMI_wr)] unique_user_count = data['ID'].nunique() print(f"删除异常BMI的用户数：{unique_user_count}") #多重插补 # 初始化多重插补模型 imputer = IterativeImputer(max_iter=10, random_state=0) # 进行多重插补 imputed_data = imputer.fit_transform(data) # 将插补后的数据转换为 DataFrame data = pd.DataFrame(imputed_data, columns=data.columns) # 将分类变量列取整 categorical_columns = ['Physical_activity', 'Psychiatric_score', 'ADL', 'Smoke', 'Drink'] # 分类变量列名 data[categorical_columns] = data[categorical_columns].round().astype(int) # 修正负值，将所有小于 0 的值替换为 0 或指定的最小值 data[data < 0] = 0 #排序将ID相同的放到一起 data.to_csv("paper_data.csv", index=False)