Data extraction for specific tasks on mimiciv

{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"name":"python","version":"3.10.14","mimetype":"text/x-python","codemirror_mode":{"name":"ipython","version":3},"pygments_lexer":"ipython3","nbconvert_exporter":"python","file_extension":".py"},"kaggle":{"accelerator":"none","dataSources":[{"sourceId":194747480,"sourceType":"kernelVersion"}],"dockerImageVersionId":30761,"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":false}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"markdown","source":"# Data extraction based on tasks\n\nThese tasks for evaluating data-science pipeline, more see [paper](#TODO). We focus on two tasks, `binary classification` and `regression` tasks.","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19"}},{"cell_type":"code","source":"#!pip install --ignore-installed -U -q datasets==2.18.0","metadata":{"execution":{"iopub.status.busy":"2024-09-09T14:54:18.197800Z","iopub.execute_input":"2024-09-09T14:54:18.198424Z","iopub.status.idle":"2024-09-09T14:54:18.204825Z","shell.execute_reply.started":"2024-09-09T14:54:18.198360Z","shell.execute_reply":"2024-09-09T14:54:18.203572Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import os\nfrom huggingface_hub import login\nfrom kaggle_secrets import UserSecretsClient\n","metadata":{"execution":{"iopub.status.busy":"2024-09-09T14:54:59.321843Z","iopub.execute_input":"2024-09-09T14:54:59.322923Z","iopub.status.idle":"2024-09-09T14:54:59.696650Z","shell.execute_reply.started":"2024-09-09T14:54:59.322849Z","shell.execute_reply":"2024-09-09T14:54:59.695302Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Parameters for specific tasks\n\n","metadata":{}},{"cell_type":"code","source":"time=0\ndisease_label=\"No Disease Filter\"\nlabel='mortality'\n\ncohort_output=\"cohort_\" + 'icu' + \"_\" + label + \"_\" + str(time) +\"_\" + disease_label\nsummary_output=\"summary_\" + 'icu' + \"_\" + label + \"_\" + str(time) +\"_\" + disease_label\n\n\ngroup_col='subject_id'\nvisit_col='stay_id'\nadmit_col='intime'\ndisch_col='outtime'\ndeath_col='dod'\nadm_visit_col='hadm_id'\nuse_mort=True,\nuse_los=False\nlos=0\nuse_admn=False\ndisease_label=\"\"\nuse_ICU=True","metadata":{"execution":{"iopub.status.busy":"2024-09-09T14:55:02.092645Z","iopub.execute_input":"2024-09-09T14:55:02.093683Z","iopub.status.idle":"2024-09-09T14:55:02.101591Z","shell.execute_reply.started":"2024-09-09T14:55:02.093630Z","shell.execute_reply":"2024-09-09T14:55:02.100048Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"icu_data_path='/kaggle/input/pre-download-mimic-iv-v2-0-dataset/physionet.org/files/mimiciv/2.0/icu/icustays.csv.gz'\nicu_patents_data_path='/kaggle/input/pre-download-mimic-iv-v2-0-dataset/physionet.org/files/mimiciv/2.0/hosp/patients.csv.gz'\nhosp_admissions='/kaggle/input/pre-download-mimic-iv-v2-0-dataset/physionet.org/files/mimiciv/2.0/hosp/admissions.csv.gz'","metadata":{"execution":{"iopub.status.busy":"2024-09-09T14:55:05.937124Z","iopub.execute_input":"2024-09-09T14:55:05.937626Z","iopub.status.idle":"2024-09-09T14:55:05.943406Z","shell.execute_reply.started":"2024-09-09T14:55:05.937576Z","shell.execute_reply":"2024-09-09T14:55:05.941971Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"Combines the MIMIC-IV core/patients table information with either the icu/icustays or core/admissions data.","metadata":{}},{"cell_type":"code","source":"import polars as pl\n\ndef get_visit_pts(group_col:str, visit_col:str, admit_col:str, disch_col:str, adm_visit_col:str, use_mort:bool, use_los:bool, los:int, use_admn:bool, disease_label:str, use_ICU:bool):\n    # 读取数据\n    visit = pl.read_csv(icu_data_path, compression='gzip', dtypes={admit_col: pl.Datetime, disch_col: pl.Datetime})\n    \n    if use_admn:\n        pts = pl.read_csv(icu_patients_data_path, compression='gzip', usecols=['subject_id', 'dod'], dtypes={'dod': pl.Datetime})\n        visit = visit.join(pts, on='subject_id', how='inner')\n        visit = visit.filter(pl.col('dod').is_null() | (pl.col('dod') >= pl.col(disch_col)))\n        \n        if len(disease_label):\n            hids = disease_cohort.extract_diag_cohort(visit['hadm_id'].to_pandas(), disease_label, mimic4_path)\n            visit = visit.filter(pl.col('hadm_id').is_in(hids['hadm_id'].to_pandas()))\n            print(\"[ READMISSION DUE TO \"+disease_label+\" ]\")\n\n    pts = pl.read_csv(\n        icu_patients_data_path, \n        compression='gzip', \n        usecols=[group_col, 'anchor_year', 'anchor_age', 'anchor_year_group', 'dod', 'gender']\n    )\n    \n    pts = pts.with_columns([\n        (pl.col('anchor_year') - pl.col('anchor_age')).alias('yob'),\n        (pl.col('anchor_year') + (2019 - pl.col('anchor_year_group').str.slice(-4).cast(pl.Int32))).alias('min_valid_year')\n    ])\n    \n    visit_pts = visit.join(pts, on=group_col, how='inner')\n    visit_pts = visit_pts.with_columns([\n        pl.col('anchor_age').alias('Age')\n    ]).filter(pl.col('Age') >= 18)\n    \n    # 添加 Demo 数据\n    eth = pl.read_csv(hosp_admissions, compression='gzip', usecols=['hadm_id', 'insurance', 'race'])\n    visit_pts = visit_pts.join(eth, on='hadm_id', how='inner')\n    \n    return visit_pts.select([group_col, visit_col, adm_visit_col, admit_col, disch_col, 'los', 'min_valid_year', 'dod', 'Age', 'gender', 'race', 'insurance'])\n\n# 调用函数\npts = get_visit_pts(\n    group_col=group_col,\n    visit_col=visit_col,\n    admit_col=admit_col,\n    disch_col=disch_col,\n    adm_visit_col=adm_visit_col,\n    use_mort=use_mort,\n    use_los=use_los,\n    los=los,\n    use_admn=use_admn,\n    disease_label=disease_label,\n    use_ICU=use_ICU\n)\n","metadata":{"execution":{"iopub.status.busy":"2024-09-09T14:54:19.061210Z","iopub.execute_input":"2024-09-09T14:54:19.061651Z","iopub.status.idle":"2024-09-09T14:54:20.876944Z","shell.execute_reply.started":"2024-09-09T14:54:19.061607Z","shell.execute_reply":"2024-09-09T14:54:20.874528Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"pts.info()","metadata":{"execution":{"iopub.status.busy":"2024-09-09T14:54:20.877972Z","iopub.status.idle":"2024-09-09T14:54:20.878419Z","shell.execute_reply.started":"2024-09-09T14:54:20.878199Z","shell.execute_reply":"2024-09-09T14:54:20.878220Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"pts.head(5)","metadata":{"execution":{"iopub.status.busy":"2024-09-09T14:54:20.879873Z","iopub.status.idle":"2024-09-09T14:54:20.880351Z","shell.execute_reply.started":"2024-09-09T14:54:20.880131Z","shell.execute_reply":"2024-09-09T14:54:20.880155Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Extraction columns","metadata":{}},{"cell_type":"code","source":"import numpy as np\nimport datetime\n\ndef partition_by_los(df: pl.DataFrame, los: int, group_col: str, visit_col: str, admit_col: str, disch_col: str, valid_col: str):\n    invalid = df.filter(\n        (pl.col(admit_col).is_null()) |\n        (pl.col(disch_col).is_null()) |\n        (pl.col('los').is_null())\n    )\n    \n    cohort = df.filter(\n        (~pl.col(admit_col).is_null()) &\n        (~pl.col(disch_col).is_null()) &\n        (~pl.col('los').is_null())\n    )\n    \n    pos_cohort = cohort.filter(pl.col('los') > los)\n    neg_cohort = cohort.filter(pl.col('los') <= los)\n    \n    pos_cohort = pos_cohort.with_column(pl.lit(1).alias('label'))\n    neg_cohort = neg_cohort.with_column(pl.lit(0).alias('label'))\n    \n    cohort = pl.concat([pos_cohort, neg_cohort], how='vertical')\n    cohort = cohort.sort(by=[group_col, admit_col])\n    \n    print(\"[ LOS LABELS FINISHED ]\")\n    return cohort, invalid\n\ndef partition_by_readmit(df: pl.DataFrame, gap: datetime.timedelta, group_col: str, visit_col: str, admit_col: str, disch_col: str, valid_col: str):\n    \"\"\"Applies labels to individual visits according to whether or not a readmission has occurred within the specified `gap` days.\"\"\"\n    \n    case = pl.DataFrame()  # hadm_ids with readmission within the gap period\n    ctrl = pl.DataFrame()  # hadm_ids without readmission within the gap period\n    invalid = pl.DataFrame()  # hadm_ids that are not considered in the cohort\n\n    grouped = df.sort(by=[group_col, admit_col]).groupby(group_col)\n    for subject, group in grouped:\n        group = group.sort(by=admit_col)\n        max_year = group.select(pl.col(disch_col).dt.year.max()).to_numpy()[0]\n\n        if group.shape[0] <= 1:\n            ctrl = ctrl.vstack(group)\n        else:\n            for idx in range(group.shape[0] - 1):\n                visit_time = group[idx, disch_col]  # For each index (a unique hadm_id), get its timestamp\n                if group.filter(\n                    (pl.col(admit_col) > visit_time) &\n                    (pl.col(admit_col) - visit_time <= pl.duration(days=gap.days))\n                ).shape[0] >= 1:\n                    case = case.vstack(group[idx])\n                else:\n                    ctrl = ctrl.vstack(group[idx])\n            ctrl = ctrl.vstack(group[-1])\n    \n    print(\"[ READMISSION LABELS FINISHED ]\")\n    return case, ctrl, invalid\n\ndef partition_by_mort(df: pl.DataFrame, group_col: str, visit_col: str, admit_col: str, disch_col: str, death_col: str):\n    \"\"\"Applies labels to individual visits according to whether or not a death has occurred within the times of the specified admit_col and disch_col\"\"\"\n    \n    invalid = df.filter(\n        (pl.col(admit_col).is_null()) |\n        (pl.col(disch_col).is_null())\n    )\n    \n    cohort = df.filter(\n        (~pl.col(admit_col).is_null()) &\n        (~pl.col(disch_col).is_null())\n    )\n    \n    cohort = cohort.with_column(pl.lit(0).alias('label'))\n    \n    pos_cohort = cohort.filter(~pl.col(death_col).is_null())\n    pos_cohort = pos_cohort.with_column(\n        (pl.col(death_col).dt.strptime('%Y-%m-%d') >= pl.col(admit_col)) &\n        (pl.col(death_col).dt.strptime('%Y-%m-%d') <= pl.col(disch_col))\n    ).alias('label')\n    \n    pos_cohort = pos_cohort.with_column(pl.col('label').cast(pl.Int32))\n    \n    cohort = pl.concat([pos_cohort, cohort.filter(pl.col(death_col).is_null())], how='vertical')\n    cohort = cohort.sort(by=[group_col, admit_col])\n    \n    print(\"[ MORTALITY LABELS FINISHED ]\")\n    return cohort, invalid\n\ndef get_case_ctrls(df: pl.DataFrame, gap: int, group_col: str, visit_col: str, admit_col: str, disch_col: str, valid_col: str, death_col: str, use_mort=False, use_admn=False, use_los=False) -> pl.DataFrame:\n    \"\"\"Handles logic for creating the labelled cohort based on arguments passed to extract().\"\"\"\n    \n    if use_mort:\n        return partition_by_mort(df, group_col, visit_col, admit_col, disch_col, death_col)\n    elif use_admn:\n        gap = datetime.timedelta(days=gap)\n        case, ctrl, invalid = partition_by_readmit(df, gap, group_col, visit_col, admit_col, disch_col, valid_col)\n        case = case.with_column(pl.lit(1).alias('label'))\n        ctrl = ctrl.with_column(pl.lit(0).alias('label'))\n        return pl.concat([case, ctrl], how='vertical'), invalid\n    elif use_los:\n        return partition_by_los(df, gap, group_col, visit_col, admit_col, disch_col, death_col)","metadata":{"execution":{"iopub.status.busy":"2024-09-09T14:54:20.881999Z","iopub.status.idle":"2024-09-09T14:54:20.882498Z","shell.execute_reply.started":"2024-09-09T14:54:20.882245Z","shell.execute_reply":"2024-09-09T14:54:20.882269Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"cols = [group_col, visit_col, admit_col, disch_col, 'Age','gender','ethnicity','insurance','label']\n\nif use_mort:\n    cols.append(death_col)\n    cohort, invalid = get_case_ctrls(pts, None, group_col, visit_col, admit_col, disch_col,'min_valid_year', death_col, use_mort=True,use_admn=False,use_los=False)\nelif use_admn:\n    interval = time\n    cohort, invalid = get_case_ctrls(pts, interval, group_col, visit_col, admit_col, disch_col,'min_valid_year', death_col, use_mort=False,use_admn=True,use_los=False)\nelif use_los:\n    cohort, invalid = get_case_ctrls(pts, los, group_col, visit_col, admit_col, disch_col,'min_valid_year', death_col, use_mort=False,use_admn=False,use_los=True)","metadata":{"execution":{"iopub.status.busy":"2024-09-09T14:54:20.884052Z","iopub.status.idle":"2024-09-09T14:54:20.884538Z","shell.execute_reply.started":"2024-09-09T14:54:20.884278Z","shell.execute_reply":"2024-09-09T14:54:20.884314Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"if use_ICU:\n    cols.append(adm_visit_col)\nprint(cohort.info())","metadata":{"execution":{"iopub.status.busy":"2024-09-09T14:54:20.885811Z","iopub.status.idle":"2024-09-09T14:54:20.886260Z","shell.execute_reply.started":"2024-09-09T14:54:20.886030Z","shell.execute_reply":"2024-09-09T14:54:20.886061Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"cohort=cohort.rename(columns={\"race\":\"ethnicity\"})\nprint(cohort[cols].info())","metadata":{"execution":{"iopub.status.busy":"2024-09-09T14:54:20.888183Z","iopub.status.idle":"2024-09-09T14:54:20.888680Z","shell.execute_reply.started":"2024-09-09T14:54:20.888423Z","shell.execute_reply":"2024-09-09T14:54:20.888445Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"cohort.head(5)","metadata":{"execution":{"iopub.status.busy":"2024-09-09T14:54:20.890290Z","iopub.status.idle":"2024-09-09T14:54:20.890770Z","shell.execute_reply.started":"2024-09-09T14:54:20.890544Z","shell.execute_reply":"2024-09-09T14:54:20.890567Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"cohort[cols].head(5)","metadata":{"execution":{"iopub.status.busy":"2024-09-09T14:54:20.892722Z","iopub.status.idle":"2024-09-09T14:54:20.893318Z","shell.execute_reply.started":"2024-09-09T14:54:20.893013Z","shell.execute_reply":"2024-09-09T14:54:20.893043Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"summary = \"\\n\".join([\n    f\"{label} FOR ICU DATA\",\n    f\"# Admission Records: {cohort.shape[0]}\",\n    f\"# Patients: {cohort[group_col].nunique()}\",\n    f\"# Positive cases: {cohort[cohort['label']==1].shape[0]}\",\n    f\"# Negative cases: {cohort[cohort['label']==0].shape[0]}\"\n])\n\nprint(summary)","metadata":{"execution":{"iopub.status.busy":"2024-09-09T14:54:20.895464Z","iopub.status.idle":"2024-09-09T14:54:20.895931Z","shell.execute_reply.started":"2024-09-09T14:54:20.895722Z","shell.execute_reply":"2024-09-09T14:54:20.895745Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"csv_file_path = f'{\"mimic_iv_classification_regression_tasks\"}.csv'\ncohort[cols].to_csv(csv_file_path, index=False)","metadata":{"execution":{"iopub.status.busy":"2024-09-09T14:54:20.897643Z","iopub.status.idle":"2024-09-09T14:54:20.898061Z","shell.execute_reply.started":"2024-09-09T14:54:20.897855Z","shell.execute_reply":"2024-09-09T14:54:20.897877Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# from datasets import Dataset\n\n# ds = Dataset.from_pandas(pd.read_csv(csv_file_path))\n# ds.push_to_hub(os.getenv('DATASET'))","metadata":{"execution":{"iopub.status.busy":"2024-09-09T14:54:20.899677Z","iopub.status.idle":"2024-09-09T14:54:20.900271Z","shell.execute_reply.started":"2024-09-09T14:54:20.899965Z","shell.execute_reply":"2024-09-09T14:54:20.899997Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Acknowledgements\n\n* https://github.com/healthylaife/MIMIC-IV-Data-Pipeline","metadata":{}}]}