genhpf 1.0.11__py3-none-any.whl

genhpf/scripts/preprocess/genhpf/main.py

@@ -0,0 +1,175 @@
+import argparse
+import logging
+import os
+import sys
+
+from ehrs import EHR_REGISTRY
+from pyspark.sql import SparkSession
+
+logging.basicConfig(
+    format="%(asctime)s | %(levelname)s | %(message)s",
+    datefmt="%Y-%m-%d %H:%M:%S",
+    level=os.environ.get("LOGLEVEL", "INFO").upper(),
+    stream=sys.stdout,
+)
+logger = logging.getLogger(__name__)
+
+
+def get_parser():
+    parser = argparse.ArgumentParser()
+    # task
+    parser.add_argument(
+        "--task",
+        type=str,
+        default=None,
+        help="specific task from the all_features pipeline." "if not set, run the all_features steps.",
+    )
+    parser.add_argument("--dest", default="outputs", type=str, metavar="DIR", help="output directory")
+
+    # data
+    parser.add_argument(
+        "--ehr",
+        type=str,
+        required=True,
+        choices=["mimiciii", "mimiciv", "eicu"],
+        help="name of the ehr system to be processed.",
+    )
+    parser.add_argument(
+        "--data",
+        metavar="DIR",
+        default=None,
+        help="directory containing data files of the given ehr (--ehr)."
+        "if not given, try to download from the internet.",
+    )
+    parser.add_argument(
+        "--ext",
+        type=str,
+        default=None,
+        help="extension for ehr data to look for. " "if not given, try to infer from --data",
+    )
+    parser.add_argument(
+        "--ccs",
+        type=str,
+        default=None,
+        help="path to `ccs_multi_dx_tool_2015.csv`" "if not given, try to download from the internet.",
+    )
+    parser.add_argument(
+        "--gem",
+        type=str,
+        default=None,
+        help="path to `icd10cmtoicd9gem.csv`" "if not given, try to download from the internet.",
+    )
+
+    parser.add_argument(
+        "-c", "--cache", action="store_true", help="whether to load data from cache if exists"
+    )
+
+    # misc
+    parser.add_argument("--max_event_size", type=int, default=256, help="max event size to crop to")
+    parser.add_argument(
+        "--min_event_size",
+        type=int,
+        default=5,
+        help="min event size to skip small samples",
+    )
+    parser.add_argument("--min_age", type=int, default=18, help="min age to skip too young patients")
+    parser.add_argument("--max_age", type=int, default=None, help="max age to skip too old patients")
+    parser.add_argument("--obs_size", type=int, default=12, help="observation window size by the hour")
+    parser.add_argument("--gap_size", type=int, default=12, help="time gap window size by the hour")
+    parser.add_argument("--pred_size", type=int, default=24, help="prediction window size by the hour")
+    parser.add_argument(
+        "--long_term_pred_size",
+        type=int,
+        default=336,
+        help="prediction window size by the hour (for long term mortality task)",
+    )
+    parser.add_argument(
+        "--first_icu",
+        action="store_true",
+        help="whether to use only the first icu or not",
+    )
+
+    # tasks
+    parser.add_argument("--mortality", action="store_true", help="whether to include mortality task or not")
+    parser.add_argument(
+        "--long_term_mortality",
+        action="store_true",
+        help="whether to include long term mortality task or not",
+    )
+    parser.add_argument("--los_3day", action="store_true", help="whether to include 3-day los task or not")
+    parser.add_argument("--los_7day", action="store_true", help="whether to include 7-day los task or not")
+    parser.add_argument(
+        "--readmission", action="store_true", help="whether to include readmission task or not"
+    )
+    parser.add_argument(
+        "--final_acuity", action="store_true", help="whether to include final acuity task or not"
+    )
+    parser.add_argument(
+        "--imminent_discharge", action="store_true", help="whether to include imminent discharge task or not"
+    )
+    parser.add_argument("--diagnosis", action="store_true", help="whether to include diagnosis task or not")
+    parser.add_argument("--creatinine", action="store_true", help="whether to include creatinine task or not")
+    parser.add_argument("--bilirubin", action="store_true", help="whether to include bilirubin task or not")
+    parser.add_argument("--platelets", action="store_true", help="whether to include platelets task or not")
+    parser.add_argument(
+        "--wbc", action="store_true", help="whether to include blood white blood cell count task or not"
+    )
+
+    parser.add_argument(
+        "--chunk_size",
+        type=int,
+        default=1024,
+        help="chunk size to read large csv files",
+    )
+    parser.add_argument("--bins", type=int, default=20, help="num buckets to bin time intervals by")
+
+    parser.add_argument(
+        "--max_event_token_len", type=int, default=192, help="max token length for each event (Hierarchical)"
+    )
+
+    parser.add_argument(
+        "--max_patient_token_len", type=int, default=8192, help="max token length for each patient (Flatten)"
+    )
+
+    parser.add_argument("--num_threads", type=int, default=8, help="number of threads to use")
+
+    # CodeEmb / Feature select
+    parser.add_argument(
+        "--emb_type",
+        type=str,
+        choices=["codebase", "textbase"],
+        default="textbase",
+        help="feature embedding type, codebase model = [SAND, Rajikomar], textbase model = [GenHPF, DescEmb]",
+    )
+    parser.add_argument(
+        "--feature",
+        choices=["select", "all_features"],
+        default="all_features",
+        help="pre-define feature select or all_features table use",
+    )
+    parser.add_argument("--bucket_num", type=int, default=10, help="feature bucket num")
+    return parser
+
+
+def main():
+    parser = get_parser()
+    args = parser.parse_args()
+
+    if not os.path.exists(args.dest):
+        os.makedirs(args.dest)
+
+    ehr = EHR_REGISTRY[args.ehr](args)
+
+    spark = (
+        SparkSession.builder.master(f"local[{args.num_threads}]")
+        .config("spark.driver.memory", "100g")
+        .config("spark.driver.maxResultSize", "10g")
+        .config("spark.network.timeout", "100s")
+        .appName("Main_Preprocess")
+        .getOrCreate()
+    )
+    ehr.run_pipeline(spark)
+
+
+if __name__ == "__main__":
+    main()

genhpf/scripts/preprocess/genhpf/manifest.py

@@ -0,0 +1,79 @@
+import argparse
+import os
+import random
+
+import h5pickle
+from tqdm import tqdm
+
+
+def get_parser():
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "root", metavar="DIR", type=str, help="root directory containing the data.h5 and label.csv"
+    )
+    parser.add_argument("--dest", default=".", type=str, metavar="DIR", help="output directory")
+    parser.add_argument("--prefix", default="", type=str, help="prefix for the output files")
+    parser.add_argument(
+        "--valid-percent",
+        default=0.1,
+        type=float,
+        metavar="D",
+        help="percentage of data to use as validation and test set (between 0 and 0.5)",
+    )
+    parser.add_argument("--seed", default=42, type=int, metavar="N", help="random seed")
+    return parser
+
+
+def main(args):
+    assert 0 <= args.valid_percent <= 0.5, (
+        f"Invalid valid-percent: {args.valid_percent}. " "Must be between 0 and 0.5."
+    )
+
+    if len(args.prefix) > 0 and not args.prefix.endswith("_"):
+        args.prefix += "_"
+
+    root_path = os.path.realpath(args.root)
+    data_path = os.path.join(root_path, "data.h5")
+    label_path = os.path.join(root_path, "label.csv")
+    rand = random.Random(args.seed)
+
+    if not os.path.exists(args.dest):
+        os.makedirs(args.dest)
+
+    with (
+        open(os.path.join(args.dest, f"{args.prefix}train.tsv"), "w") as train_f,
+        open(os.path.join(args.dest, f"{args.prefix}valid.tsv"), "w") as valid_f,
+        open(os.path.join(args.dest, f"{args.prefix}test.tsv"), "w") as test_f,
+    ):
+        print(data_path, file=train_f)
+        print(label_path, file=train_f)
+        print(data_path, file=valid_f)
+        print(label_path, file=valid_f)
+        print(data_path, file=test_f)
+        print(label_path, file=test_f)
+
+        def write(subjects, dest, split=None):
+            for subject in tqdm(subjects, total=len(subjects), desc=split):
+                print(subject, file=dest)
+
+        data = h5pickle.File(data_path, "r")["ehr"]
+        subjects = list(data.keys())
+        rand.shuffle(subjects)
+
+        valid_len = int(len(subjects) * args.valid_percent)
+        test_len = int(len(subjects) * args.valid_percent)
+        train_len = len(subjects) - valid_len - test_len
+
+        train = subjects[:train_len]
+        valid = subjects[train_len : train_len + valid_len]
+        test = subjects[train_len + valid_len :]
+
+        write(train, train_f, split=f"{args.prefix}train")
+        write(valid, valid_f, split=f"{args.prefix}valid")
+        write(test, test_f, split=f"{args.prefix}test")
+
+
+if __name__ == "__main__":
+    parser = get_parser()
+    args = parser.parse_args()
+    main(args)

genhpf/scripts/preprocess/genhpf/sample_dataset.py

@@ -0,0 +1,177 @@
+import argparse
+import os
+
+import h5py
+import torch
+from torch.nn.functional import one_hot
+from torch.nn.utils.rnn import pad_sequence
+from torch.utils.data import Dataset
+
+"""
+Sample Dataset for Integrated-EHR-Pipeline
+NOTE: The preprocessed token indices are stored as np.int16 type for efficienty,
+        so overflow can be caused when vocab size > 32767.
+ - ehr.cohort.labeled.index: dataframe pickle
+    - index: should be resetted by `df.reset_index()`
+    - [hi_start:hi_end] is the range of indices for the patient's history
+    - split: one of ['train', 'val', 'test']
+ - ehr.data: .np file with np.int16 type, (num_total_events, 3, max_word_len)
+    - second dimension
+        - 0: input_ids
+        - 1: type_ids
+        - 2: dpe_ids
+"""
+
+
+def get_parser():
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "--ehr",
+        type=str,
+        required=True,
+        choices=["mimiciii", "mimiciv", "eicu"],
+        help="Name of the EHR dataset",
+    )
+    parser.add_argument("--data", type=str, required=True, help="Path to the preprocessed data")
+    parser.add_argument(
+        "--pred_target",
+        type=str,
+        required=True,
+        choices=[
+            "mortality",
+            "readmission",
+            "los_3day",
+            "los_7day",
+            "final_acuity",
+            "imminent_discharge",
+            "diagnosis",
+        ],
+        help="Prediction target",
+    )
+    parser.add_argument("--max_event_size", type=int, default=256, help="max event size to crop to")
+    parser.add_argument(
+        "--max_event_token_len", type=int, default=128, help="max token length for each event (Hierarchical)"
+    )
+
+    parser.add_argument(
+        "--max_patient_token_len", type=int, default=8192, help="max token length for each patient (Flatten)"
+    )
+    return parser
+
+
+class BaseEHRDataset(Dataset):
+    def __init__(self, args, split):
+        super().__init__()
+
+        self.args = args
+        self.data_path = os.path.join(args.data, f"{args.ehr}.h5")
+
+        self.data = h5py.File(self.data_path, "r")["ehr"]
+        self.keys = []
+        for key in self.data.keys():
+            if self.data[key].attrs["split"] == split:
+                self.keys.append(key)
+        self.pred_target = args.pred_target
+
+        self.num_classes = {
+            "mimiciii": {
+                "mortality": 2,
+                "readmission": 2,
+                "los_3day": 2,
+                "los_7day": 2,
+                "final_acuity": 18,
+                "imminent_discharge": 18,
+                "diagnosis": 18,
+            },
+            "mimiciv": {
+                "mortality": 2,
+                "readmission": 2,
+                "los_3day": 2,
+                "los_7day": 2,
+                "final_acuity": 14,
+                "imminent_discharge": 14,
+                "diagnosis": 18,
+            },
+            "eicu": {
+                "mortality": 2,
+                "readmission": 2,
+                "los_3day": 2,
+                "los_7day": 2,
+                "final_acuity": 9,
+                "imminent_discharge": 9,
+                "diagnosis": 18,
+            },
+        }[self.args.ehr][self.args.pred_target]
+
+    def __len__(self):
+        return len(self.keys)
+
+    def __getitem__(self, idx):
+        raise NotImplementedError()
+
+    def collate_fn(self, out):
+        ret = dict()
+        if len(out) == 1:
+            for k, v in out[0].items():
+                if k == "label":
+                    ret[k] = (
+                        one_hot(torch.LongTensor([v]), self.num_classes)
+                        .float()
+                        .reshape((1, self.num_classes))
+                    )
+                else:
+                    ret[k] = v.unsqueeze(0)
+        else:
+            for k, v in out[0].items():
+                if k == "label":
+                    ret[k] = one_hot(torch.LongTensor([i[k] for i in out]), self.num_classes).float()
+                else:
+                    ret[k] = pad_sequence([i[k] for i in out], batch_first=True)
+        return ret
+
+
+class HierarchicalEHRDataset(BaseEHRDataset):
+    def __init__(self, args, split):
+        super().__init__(args, split)
+
+    def __getitem__(self, idx):
+        # NOTE: Warning occurs when converting np.int16 read-only array into tensor, but ignoreable
+
+        data = self.data[self.keys[idx]]["hi"]
+        label = self.data[self.keys[idx]].attrs[self.pred_target]
+        return {
+            "input_ids": torch.IntTensor(data[:, 0, :]),
+            "type_ids": torch.IntTensor(data[:, 1, :]),
+            "dpe_ids": torch.IntTensor(data[:, 2, :]),
+            "label": label,
+        }
+
+
+class FlattenEHRDataset(BaseEHRDataset):
+    def __init__(self, args, split):
+        super().__init__(args, split)
+        self.data_path = os.path.join(args.data, f"{args.ehr}.flat.npy")
+
+    def __getitem__(self, idx):
+        # NOTE: Warning occurs when converting np.int16 read-only array into tensor, but ignoreable
+        data = self.data[self.keys[idx]]["fl"]
+        label = self.data[self.keys[idx]].attrs[self.pred_target]
+        return {
+            "input_ids": torch.IntTensor(data[0, :]),
+            "type_ids": torch.IntTensor(data[1, :]),
+            "dpe_ids": torch.IntTensor(data[2, :]),
+            "label": label,
+        }
+
+
+def main():
+    args = get_parser().parse_args()
+    dataset = HierarchicalEHRDataset(args, "train")
+    print(dataset.__getitem__(1))
+    dataset = FlattenEHRDataset(args, "train")
+    print(dataset.__getitem__(1))
+    pass
+
+
+if __name__ == "__main__":
+    main()

genhpf/scripts/preprocess/genhpf/utils/__init__.py

@@ -0,0 +1,3 @@
+from .utils import col_name_add, q_cut
+
+__all__ = ["q_cut", "col_name_add"]

genhpf/scripts/preprocess/genhpf/utils/utils.py

@@ -0,0 +1,16 @@
+import numpy as np
+import pandas as pd
+
+
+def col_name_add(x, cate_col):
+    if not (x == "nan" or x == pd.isnull(x)):
+        return cate_col + "_" + str(x)
+    else:
+        return x  # return nan
+
+
+def q_cut(x, cuts):
+    unique_var = len(np.unique([i for i in x]))
+    nunique = len(pd.qcut(x, min(unique_var, cuts), duplicates="drop").cat.categories)
+    output = pd.qcut(x, min(unique_var, cuts), labels=range(1, min(nunique, cuts) + 1), duplicates="drop")
+    return output

genhpf/scripts/preprocess/manifest.py

@@ -0,0 +1,83 @@
+import argparse
+import random
+import os
+import logging
+import h5pickle
+from tqdm import tqdm
+
+logger = logging.getLogger(__name__)
+
+def get_parser():
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "data", type=str, help="path to the .h5 file containing the data"
+    )
+    parser.add_argument(
+        "label", type=str, help="path to the .csv file containing the labels corresponded to the data"
+    )
+    parser.add_argument(
+        "--dest", default=".", type=str, metavar="DIR", help="output directory"
+    )
+    parser.add_argument(
+        "--prefix", default="", type=str, help="prefix for the output files"
+    )
+    parser.add_argument(
+        "--valid-percent",
+        default=0.1,
+        type=float,
+        metavar="D",
+        help="percentage of the data to use as validation and test set (between 0 and 0.5)"
+    )
+    parser.add_argument("--seed", default=42, type=int, metavar="N", help="random seed")
+    return parser
+
+def main(args):
+    assert 0 <= args.valid_percent <= 0.5
+
+    data_path = os.path.realpath(args.data)
+    label_path = os.path.realpath(args.label)
+    prefix = args.prefix
+    if len(prefix) > 0 and not prefix.endswith("_"):
+        prefix += "_"
+    rand = random.Random(args.seed)
+
+    if not os.path.exists(args.dest):
+        os.makedirs(args.dest)
+    
+    with (
+        open(os.path.join(args.dest, f"{prefix}train.tsv"), "w") as train_f,
+        open(os.path.join(args.dest, f"{prefix}valid.tsv"), "w") as valid_f,
+        open(os.path.join(args.dest, f"{prefix}test.tsv"), "w") as test_f
+    ):
+        print(data_path, file=train_f)
+        print(label_path, file=train_f)
+        print(data_path, file=valid_f)
+        print(label_path, file=valid_f)
+        print(data_path, file=test_f)
+        print(label_path, file=test_f)
+
+        def write(subject_ids, dest, split):
+            for subject_id in tqdm(subject_ids, total=len(subject_ids), desc=split):
+                print(subject_id, file=dest)
+
+        data = h5pickle.File(data_path, "r")["ehr"]
+        subject_ids = list(data.keys())
+
+        rand.shuffle(subject_ids)
+
+        valid_len = int(len(subject_ids) * args.valid_percent)
+        test_len = int(len(subject_ids) * args.valid_percent)
+        train_len = len(subject_ids) - valid_len - test_len
+
+        train = subject_ids[:train_len]
+        valid = subject_ids[train_len: train_len + valid_len]
+        test = subject_ids[train_len + valid_len:]
+
+        write(train, train_f, split=f"{prefix}train")
+        write(valid, valid_f, split=f"{prefix}valid")
+        write(test, test_f, split=f"{prefix}test")
+
+if __name__ == "__main__":
+    parser = get_parser()
+    args = parser.parse_args()
+    main(args)