preprocess.py

import os

import argparse

import numpy as np
import pandas as pd


class Reader:
    def __init__(self, data_path):
        self.data_path = data_path

    def read_admissions_table(self):
        df = pd.read_csv(
            os.path.join(self.data_path, "ADMISSIONS.csv.gz"),
            compression="gzip",
            header=0,
            index_col=0,
        )
        df.columns = map(str.upper, df.columns)
        df.ADMITTIME = pd.to_datetime(df.ADMITTIME)
        df.DISCHTIME = pd.to_datetime(df.DISCHTIME)
        df = df[["SUBJECT_ID", "HADM_ID", "ADMITTIME", "DISCHTIME", "ETHNICITY"]]
        return df

    def read_icustay_table(self):
        df = pd.read_csv(
            os.path.join(self.data_path, "ICUSTAYS.csv.gz"),
            compression="gzip",
            header=0,
            index_col=0,
        )
        df.columns = map(str.upper, df.columns)
        df.INTIME = pd.to_datetime(df.INTIME)
        df.OUTTIME = pd.to_datetime(df.OUTTIME)
        df = df[["SUBJECT_ID", "HADM_ID", "ICUSTAY_ID", "INTIME", "OUTTIME", "LOS"]]
        return df

    def read_d_icd_diagnoses_table(self):
        d_icd_diagnoses = pd.read_csv(
            os.path.join(self.data_path, "D_ICD_DIAGNOSES.csv.gz"),
            compression="gzip",
            header=0,
            index_col=0,
        )
        d_icd_diagnoses.columns = map(str.upper, d_icd_diagnoses.columns)
        return d_icd_diagnoses

    def read_d_items_table(self):
        d_items = pd.read_csv(
            os.path.join(self.data_path, "D_ITEMS.csv.gz"),
            compression="gzip",
            header=0,
            index_col=0,
        )
        d_items.columns = map(str.upper, d_items.columns)
        d_items = d_items[["ITEMID", "LABEL", "DBSOURCE", "PARAM_TYPE"]]
        return d_items

    def read_d_labitems_table(self):
        d_labitems = pd.read_csv(
            os.path.join(self.data_path, "D_LABITEMS.csv.gz"),
            compression="gzip",
            header=0,
            index_col=0,
        )
        d_labitems.columns = map(str.upper, d_labitems.columns)
        d_labitems = d_labitems[["ITEMID", "LABEL", "FLUID", "CATEGORY"]]
        return d_labitems

    def read_patients_table(self):
        patients = pd.read_csv(
            os.path.join(self.data_path, "PATIENTS.csv.gz"),
            compression="gzip",
            header=0,
            index_col=0,
        )
        patients.columns = map(str.upper, patients.columns)
        patients = patients[["SUBJECT_ID", "GENDER", "DOB"]]
        return patients

    def read_diagnoses_icd_table(self):
        diagnoses_icd = pd.read_csv(
            os.path.join(self.data_path, "DIAGNOSES_ICD.csv.gz"),
            compression="gzip",
            header=0,
            index_col=0,
        )
        diagnoses_icd.columns = map(str.upper, diagnoses_icd.columns)
        diagnoses_icd = diagnoses_icd[["SUBJECT_ID", "HADM_ID", "ICD9_CODE", "SEQ_NUM"]]
        return diagnoses_icd


cache = {}


def get_info_admissions(reader: Reader, formatted_path: str):

    df = reader.read_admissions_table()
    df["STAYTIME"] = (
        df["DISCHTIME"] - df["ADMITTIME"]
    )  # stay time : discharge time - admission time
    df["STAYTIME"] = df["STAYTIME"] / np.timedelta64(1, "h")

    # formula to calcultate the age of patiens in MIMIC3

    patients = reader.read_patients_table()
    df = pd.merge(df, patients, how="left", on="SUBJECT_ID")
    df["DOB"] = pd.to_datetime(df["DOB"])
    df["ADMITTIME"] = pd.to_datetime(df["ADMITTIME"])
    df["AGE"] = df["ADMITTIME"].dt.year - df["DOB"].dt.year

    # Patients who are older than 89 years old at any time in the database
    # have had their date of birth shifted to obscure their age and comply with HIPAA.
    # The date of birth was then set to exactly 300 years before their first admission.
    df.loc[((df.AGE > 89) | (df.AGE < 0)), "AGE"] = 90

    icustays = reader.read_icustay_table()

    # merge on the HADM_ID, unique, represents a single patient's admission to the hospital
    # while subject_id can be redundant meaning that a patient had many stays at the hospital
    df = pd.merge(df, icustays, how="right", on="HADM_ID")

    # the elapsed time between the admission to the hospital and the tranfer to the ICU
    df["Time go ICU"] = (df["INTIME"] - df["ADMITTIME"]) / np.timedelta64(1, "h")

    # the elapsed time in the ICU
    df["Time in ICU"] = (df["OUTTIME"] - df["INTIME"]) / np.timedelta64(1, "h")

    # the elapsed time between the admission to the ICU and the final discharge from the hospital
    df["Time after go ICU"] = (df["DISCHTIME"] - df["INTIME"]) / np.timedelta64(1, "h")

    # number of times the patient has been transferred to the ICU during one admission
    df["Count times go ICU"] = df.groupby("HADM_ID")["ICUSTAY_ID"].transform("count")

    with open(os.path.join(formatted_path, "ADMISSIONS.csv"), "w") as f:
        df.to_csv(f, encoding="utf-8", header=True)


def check_AKI_before(hadm_id, dataset_path: str):
    key = "check_AKI_before"
    global cache

    if key not in cache:
        diagnoses = pd.read_csv(
            os.path.join(dataset_path, "DIAGNOSES_ICD.csv.gz"), compression="gzip"
        )
        diagnoses.columns = map(str.upper, diagnoses.columns)
        diagnoses = diagnoses.loc[
            diagnoses["ICD9_CODE"].isin(["5845", "5846", "5847", "5848"])
        ]
        cache[key] = diagnoses

    diagnoses = cache[key]

    if not diagnoses[diagnoses["HADM_ID"].isin(hadm_id)].empty:
        return True

    return False


def check_CKD(hadm_id, dataset_path: str):
    key = "check_CKD"
    global cache

    if key not in cache:
        diagnoses = pd.read_csv(
            os.path.join(dataset_path, "DIAGNOSES_ICD.csv.gz"), compression="gzip"
        )
        diagnoses.columns = map(str.upper, diagnoses.columns)
        diagnoses = diagnoses.loc[
            diagnoses["ICD9_CODE"].isin(["5851", "5852", "5853", "5854", "5855"])
        ]
        cache[key] = diagnoses

    diagnoses = cache[key]

    if not diagnoses[diagnoses["HADM_ID"].isin(hadm_id)].empty:
        return True

    return False


def check_renal_failure(hadm_id, formatted_path: str):
    key = "check_renal_failure"
    global cache

    if key not in cache:
        diagnoses = pd.read_csv(
            os.path.join(formatted_path, "comorbidities_DBSOURCE.csv")
        )
        diagnoses.columns = map(str.upper, diagnoses.columns)
        diagnoses = diagnoses.loc[diagnoses["RENAL_FAILURE"] == 1]
        cache[key] = diagnoses

    diagnoses = cache[key]

    if not diagnoses[diagnoses["HADM_ID"].isin(hadm_id)].empty:
        return True

    return False


def caculate_eGFR_MDRD_equation(cr, gender, eth, age):
    temp = 186 * (cr ** (-1.154)) * (age ** (-0.203))
    if gender == "F":
        temp = temp * 0.742
    if eth == "BLACK/AFRICAN AMERICAN":
        temp = temp * 1.21
    return temp


def get_aki_patients_7days_creatinine(reader: Reader, formatted_path: str):
    dataset_path = reader.data_path

    df = pd.read_csv(os.path.join(formatted_path, "ADMISSIONS.csv"))
    df = df.sort_values(by=["SUBJECT_ID_x", "HADM_ID", "ICUSTAY_ID"])

    print("admissions info", df.shape)
    print("number of unique subjects in admission: ", df["SUBJECT_ID_x"].nunique())
    print("number of icustays info in admissions: ", df["ICUSTAY_ID"].nunique())

    info_save = df.drop_duplicates(subset=["ICUSTAY_ID"])
    info_save["AKI"] = -1
    info_save["EGFR"] = -1

    print(
        "the biggest number of ICU stays for a patient: ",
        info_save["Count times go ICU"].max(),
    )

    c_aki_7d = pd.read_csv(
        os.path.join(formatted_path, "AKI_KIDIGO_7D_SQL_CREATININE_DBSOURCE.csv")
    )
    c_aki_7d.columns = map(str.upper, c_aki_7d.columns)
    c_aki_7d = c_aki_7d.drop(columns=["UNNAMED: 0"])
    print("c_aki_7d infos")
    print("Total icustays: ", c_aki_7d["ICUSTAY_ID"].nunique())
    print(
        "NORMAL Patients in 7DAY: {}".format(
            c_aki_7d.loc[c_aki_7d["AKI_STAGE_7DAY"] == 0]["ICUSTAY_ID"].count()
        )
    )
    print(
        "AKI patients STAGE 1 within 7DAY: {}".format(
            c_aki_7d.loc[c_aki_7d["AKI_STAGE_7DAY"] == 1]["ICUSTAY_ID"].count()
        )
    )
    print(
        "AKI Patients STAGE 2 in 7DAY: {}".format(
            c_aki_7d.loc[c_aki_7d["AKI_STAGE_7DAY"] == 2]["ICUSTAY_ID"].count()
        )
    )
    print(
        "AKI Patients STAGE 3 7DAY: {}".format(
            c_aki_7d.loc[c_aki_7d["AKI_STAGE_7DAY"] == 3]["ICUSTAY_ID"].count()
        )
    )
    print(
        "NAN patients within 7DAY: {}".format(c_aki_7d["AKI_STAGE_7DAY"].isna().sum())
    )
    c_aki_7d = c_aki_7d.dropna(subset=["AKI_STAGE_7DAY"])

    print("Total icustays: ", c_aki_7d["ICUSTAY_ID"].nunique())

    df_save = pd.merge(info_save, c_aki_7d, how="inner", on="ICUSTAY_ID")
    df_save.columns = map(str.upper, df_save.columns)
    icustays_data = [frame for season, frame in df_save.groupby(["ICUSTAY_ID"])]

    count_ckd_normal = 0
    count_ckd_aki = 0
    count_akibefore_normal = 0
    count_akibefore_aki = 0
    count_normal = 0
    count_aki = 0
    count_renalfailure_normal = 0
    count_renalfailure_aki = 0

    for temp in icustays_data:

        temp = temp.sort_values(by=["ICUSTAY_ID"])

        first_row = temp.iloc[0]
        gender = first_row["GENDER"]
        age = first_row["AGE"]
        eth = first_row["ETHNICITY"]
        cr = first_row["CREAT"]
        icustay_id = first_row["ICUSTAY_ID"]

        eGFR = caculate_eGFR_MDRD_equation(cr=cr, gender=gender, age=age, eth=eth)

        df_save.loc[df_save["ICUSTAY_ID"] == icustay_id, "EGFR"] = eGFR
        df_save.loc[df_save["ICUSTAY_ID"] == icustay_id, "AKI"] = c_aki_7d.loc[
            c_aki_7d["ICUSTAY_ID"] == icustay_id
        ]["AKI_7DAY"].values[0]

        if df_save.loc[df_save["ICUSTAY_ID"] == icustay_id, "AKI"].values[0] == 1:
            count_aki = count_aki + 1
        else:
            count_normal = count_normal + 1

        has_aki = (
            info_save.loc[info_save["ICUSTAY_ID"] == icustay_id, "AKI"].values[0] == 1
        )

        if check_CKD(temp["HADM_ID"], dataset_path) == True:
            df_save.loc[df_save["ICUSTAY_ID"] == icustay_id, "AKI"] = 2
            if has_aki:
                count_ckd_aki = count_ckd_aki + 1
            else:
                count_ckd_normal = count_ckd_normal + 1

        if check_AKI_before(temp["HADM_ID"], dataset_path) == True:
            df_save.loc[df_save["ICUSTAY_ID"] == icustay_id, "AKI"] = 3
            if has_aki:
                count_akibefore_aki = count_akibefore_aki + 1
            else:
                count_akibefore_normal = count_akibefore_normal + 1

        if check_renal_failure(temp["HADM_ID"], formatted_path) == True:
            df_save.loc[df_save["ICUSTAY_ID"] == icustay_id, "AKI"] = 4
            if has_aki:
                count_renalfailure_aki = count_renalfailure_aki + 1
            else:
                count_renalfailure_normal = count_renalfailure_normal + 1

    lab = pd.read_csv(os.path.join(formatted_path, "labstay_DBSOURCE.csv"))
    lab.columns = map(str.upper, lab.columns)
    info_save = pd.merge(df_save, lab, how="left", on="ICUSTAY_ID")
    cols_to_drop = set(info_save.columns).intersection(
        set(["UNNAMED: 0_x", "UNNAMED: 0_y", "SUBJECT_ID"])
    )
    if len(cols_to_drop) > 0:
        info_save = info_save.drop(columns=list(cols_to_drop))
    info_save = info_save.rename(
        columns={"SUBJECT_ID_X": "SUBJECT_ID", "HADM_ID_x": "HADM_ID"}
    )

    chart = pd.read_csv(os.path.join(formatted_path, "chart_vitals_stay_DBSOURCE.csv"))
    chart.columns = map(str.upper, chart.columns)
    df_save = pd.merge(info_save, chart, how="left", on="ICUSTAY_ID")
    df_save = df_save.drop(
        columns=["UNNAMED: 0", "HADM_ID_y", "HADM_ID_y", "SUBJECT_ID_Y", "SUBJECT_ID_y"]
    )
    df_save = df_save.rename(
        columns={"SUBJECT_ID_X": "SUBJECT_ID", "HADM_ID_x": "HADM_ID"}
    )

    comorbidities = pd.read_csv(
        os.path.join(formatted_path, "comorbidities_DBSOURCE.csv")
    )
    comorbidities.columns = map(str.upper, comorbidities.columns)
    info_save = pd.merge(df_save, comorbidities, how="left", on="HADM_ID")
    info_save = info_save.drop(columns=["UNNAMED: 0"])

    print(
        "NORMAL Patients in 7DAY: {}".format(
            c_aki_7d.loc[c_aki_7d["AKI_STAGE_7DAY"] == 0]["ICUSTAY_ID"].count()
        )
    )
    print(
        "AKI patients STAGE 1 within 7DAY: {}".format(
            c_aki_7d.loc[c_aki_7d["AKI_STAGE_7DAY"] == 1]["ICUSTAY_ID"].count()
        )
    )
    print("CKD counted as normal: {}".format(count_ckd_normal))
    print("CKD counted as aki: {}".format(count_ckd_aki))
    print("AKI on admission counted as normal: {}".format(count_akibefore_normal))
    print("AKI on admission counted as aki: {}".format(count_akibefore_aki))
    print("RENAL FAILURE counted as normal: {}".format(count_renalfailure_normal))
    print("RENAL FAILURE counted as aki: {}".format(count_renalfailure_aki))
    print("normal: {}".format(count_normal))
    print("aki: {}".format(count_aki))

    with open(
        os.path.join(formatted_path, "INFO_DATASET_7days_creatinine.csv"), "w"
    ) as f:
        info_save.to_csv(f, encoding="utf-8", header=True)


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "--formatted_path",
        type=str,
        default="./formatted",
        help='Path to formatted MIMIC III data from "parse_db.py"',
    )
    parser.add_argument(
        "--data_path",
        type=str,
        default="./data",
        help="Path to gzipped MIMIC III data",
    )

    args = parser.parse_args()

    formatted_path = args.formatted_path
    data_path = args.data_path

    reader = Reader(data_path=data_path)

    os.makedirs(formatted_path, exist_ok=True)

    get_info_admissions(reader, formatted_path)
    get_aki_patients_7days_creatinine(reader, formatted_path)


if __name__ == "__main__":
    main()