cehrbert-data 0.0.4__tar.gz → 0.0.6__tar.gz

{cehrbert_data-0.0.4 → cehrbert_data-0.0.6}/.github/workflows/python-build.yml

@@ -68,7 +68,7 @@ jobs:
         name: python-package-distributions
         path: dist/
     - name: Sign the dists with Sigstore
-      uses: sigstore/gh-action-sigstore-python@v2.1.1
+      uses: sigstore/gh-action-sigstore-python@v3.0.0
       with:
         inputs: >-
           ./dist/*.tar.gz
@@ -90,4 +90,4 @@ jobs:
       run: >-
         gh release upload
         '${{ github.ref_name }}' dist/**
-        --repo '${{ github.repository }}'
+        --repo '${{ github.repository }}'

{cehrbert_data-0.0.4 → cehrbert_data-0.0.6}/.github/workflows/tests.yml

@@ -23,6 +23,11 @@ jobs:
         uses: actions/setup-python@v3
         with:
           python-version: "3.10"
+      - name: Set up Java 11
+        uses: actions/setup-java@v3
+        with:
+          java-version: "11" # specify the Java version here
+          distribution: "temurin" # or use 'adopt' or 'zulu', depending on your preference
       - name: Install dependencies
         run: |
           python -m pip install --upgrade pip

{cehrbert_data-0.0.4 → cehrbert_data-0.0.6}/PKG-INFO

@@ -1,6 +1,6 @@
-Metadata-Version: 2.1
+Metadata-Version: 2.2
 Name: cehrbert_data
-Version: 0.0.4
+Version: 0.0.6
 Summary: The Spark ETL tools for generating the CEHR-BERT and CEHR-GPT pre-training and finetuning data
 Author-email: Chao Pang <chaopang229@gmail.com>, Xinzhuo Jiang <xj2193@cumc.columbia.edu>, Krishna Kalluri <kk3326@cumc.columbia.edu>, Nishanth Parameshwar Pavinkurve <np2689@cumc.columbia.edu>, Karthik Natarajan <kn2174@cumc.columbia.edu>
 License: MIT License

{cehrbert_data-0.0.4 → cehrbert_data-0.0.6}/src/cehrbert_data/apps/generate_training_data.py

@@ -41,13 +41,15 @@ def main(
         apply_age_filter: bool,
         include_death: bool,
         att_type: AttType,
+        inpatient_att_type: AttType,
         include_sequence_information_content: bool = False,
         exclude_demographic: bool = False,
         use_age_group: bool = False,
         with_drug_rollup: bool = True,
         include_inpatient_hour_token: bool = False,
         continue_from_events: bool = False,
-        refresh_measurement: bool = False
+        refresh_measurement: bool = False,
+        aggregate_by_hour: bool = True,
 ):
     spark = SparkSession.builder.appName("Generate CEHR-BERT Training Data").getOrCreate()
 
@@ -71,6 +73,7 @@ def main(
         f"use_age_group: {use_age_group}\n"
         f"with_drug_rollup: {with_drug_rollup}\n"
         f"refresh_measurement: {refresh_measurement}\n"
+        f"aggregate_by_hour: {aggregate_by_hour}\n"
     )
 
     domain_tables = []
@@ -128,7 +131,8 @@ def main(
         processed_measurement = get_measurement_table(
             spark,
             input_folder,
-            refresh=refresh_measurement
+            refresh=refresh_measurement,
+            aggregate_by_hour=aggregate_by_hour,
         )
         if patient_events:
             # Union all measurement records together with other domain records
@@ -165,9 +169,12 @@ def main(
             patient_demographic=person if gpt_patient_sequence else None,
             death=death,
             att_type=att_type,
+            inpatient_att_type=inpatient_att_type,
             exclude_demographic=exclude_demographic,
             use_age_group=use_age_group,
             include_inpatient_hour_token=include_inpatient_hour_token,
+            spark=spark,
+            persistence_folder=output_folder,
         )
     else:
         sequence_data = create_sequence_data(
@@ -325,12 +332,23 @@ if __name__ == "__main__":
         dest="refresh_measurement",
         action="store_true"
     )
+    parser.add_argument(
+        "--aggregate_by_hour",
+        dest="aggregate_by_hour",
+        action="store_true"
+    )
     parser.add_argument(
         "--att_type",
         dest="att_type",
         action="store",
         choices=[e.value for e in AttType],
     )
+    parser.add_argument(
+        "--inpatient_att_type",
+        dest="inpatient_att_type",
+        action="store",
+        choices=[e.value for e in AttType],
+    )
 
     ARGS = parser.parse_args()
 
@@ -352,10 +370,12 @@ if __name__ == "__main__":
         ARGS.apply_age_filter,
         ARGS.include_death,
         AttType(ARGS.att_type),
+        AttType(ARGS.inpatient_att_type),
         exclude_demographic=ARGS.exclude_demographic,
         use_age_group=ARGS.use_age_group,
         with_drug_rollup=ARGS.with_drug_rollup,
         include_inpatient_hour_token=ARGS.include_inpatient_hour_token,
         continue_from_events=ARGS.continue_from_events,
         refresh_measurement=ARGS.refresh_measurement,
+        aggregate_by_hour=ARGS.aggregate_by_hour,
     )

{cehrbert_data-0.0.4 → cehrbert_data-0.0.6}/src/cehrbert_data/cohorts/spark_app_base.py

@@ -5,6 +5,7 @@ import shutil
 from abc import ABC
 from typing import List
 
+from numpy.random import permutation
 from pandas import to_datetime
 from pyspark.sql import DataFrame, SparkSession
 from pyspark.sql import functions as F
@@ -246,6 +247,7 @@ class BaseCohortBuilder(ABC):
     def _add_demographics(self, cohort: DataFrame):
         return (
             cohort.join(self._dependency_dict[PERSON], "person_id")
+            .withColumn("year_of_birth", F.coalesce(F.year("birth_datetime"), F.col("year_of_birth")))
             .withColumn("age", F.year("index_date") - F.col("year_of_birth"))
             .select(
                 F.col("person_id"),
@@ -295,6 +297,7 @@ class NestedCohortBuilder:
             is_drug_roll_up_concept: bool = True,
             include_concept_list: bool = True,
             refresh_measurement: bool = False,
+            aggregate_by_hour: bool = True,
             is_new_patient_representation: bool = False,
             gpt_patient_sequence: bool = False,
             is_hierarchical_bert: bool = False,
@@ -306,9 +309,14 @@ class NestedCohortBuilder:
             is_observation_window_unbounded: bool = False,
             is_population_estimation: bool = False,
             att_type: AttType = AttType.CEHR_BERT,
+            inpatient_att_type: AttType = AttType.MIX,
+            include_inpatient_hour_token: bool = False,
             exclude_demographic: bool = True,
             use_age_group: bool = False,
             single_contribution: bool = False,
+            exclude_features: bool = True,
+            meds_format: bool = False,
+            cache_events: bool = False,
     ):
         self._cohort_name = cohort_name
         self._input_folder = input_folder
@@ -340,15 +348,21 @@ class NestedCohortBuilder:
         self._is_prediction_window_unbounded = is_prediction_window_unbounded
         self._include_concept_list = include_concept_list
         self._refresh_measurement = refresh_measurement
+        self._aggregate_by_hour = aggregate_by_hour
         self._allow_measurement_only = allow_measurement_only
         self._output_data_folder = os.path.join(
             self._output_folder, re.sub("[^a-z0-9]+", "_", self._cohort_name.lower())
         )
         self._is_population_estimation = is_population_estimation
         self._att_type = att_type
+        self._inpatient_att_type = inpatient_att_type
+        self._include_inpatient_hour_token = include_inpatient_hour_token
         self._exclude_demographic = exclude_demographic
         self._use_age_group = use_age_group
         self._single_contribution = single_contribution
+        self._exclude_features = exclude_features
+        self._meds_format = meds_format
+        self._cache_events = cache_events
 
         self.get_logger().info(
             f"cohort_name: {cohort_name}\n"
@@ -377,12 +391,18 @@ class NestedCohortBuilder:
             f"is_prediction_window_unbounded: {is_prediction_window_unbounded}\n"
             f"include_concept_list: {include_concept_list}\n"
             f"refresh_measurement: {refresh_measurement}\n"
+            f"aggregate_by_hour: {aggregate_by_hour}\n"
             f"is_observation_window_unbounded: {is_observation_window_unbounded}\n"
             f"is_population_estimation: {is_population_estimation}\n"
             f"att_type: {att_type}\n"
+            f"inpatient_att_type: {inpatient_att_type}\n"
+            f"include_inpatient_hour_token: {include_inpatient_hour_token}\n"
             f"exclude_demographic: {exclude_demographic}\n"
             f"use_age_group: {use_age_group}\n"
             f"single_contribution: {single_contribution}\n"
+            f"extract_features: {exclude_features}\n"
+            f"meds_format: {meds_format}\n"
+            f"cache_events: {cache_events}\n"
         )
 
         self.spark = SparkSession.builder.appName(f"Generate {self._cohort_name}").getOrCreate()
@@ -508,13 +528,30 @@ class NestedCohortBuilder:
             )
             cohort = cohort.withColumn("record_rank", record_rank).where("record_rank == 1").drop("record_rank")
 
-        ehr_records_for_cohorts = self.extract_ehr_records_for_cohort(cohort)
-        # ehr_records_for_cohorts.show()
-        cohort = (
-            cohort.join(ehr_records_for_cohorts, ["person_id", "cohort_member_id"])
-            .where(F.col("num_of_visits") >= self._num_of_visits)
-            .where(F.col("num_of_concepts") >= self._num_of_concepts)
-        )
+        if not self._exclude_features:
+            ehr_records_for_cohorts = self.extract_ehr_records_for_cohort(cohort)
+            cohort = (
+                cohort.join(ehr_records_for_cohorts, ["person_id", "cohort_member_id"])
+                .where(F.col("num_of_visits") >= self._num_of_visits)
+                .where(F.col("num_of_concepts") >= self._num_of_concepts)
+            )
+
+        person_id_column = "person_id"
+        index_date_column = "index_date"
+        if self._meds_format:
+            cohort = cohort.withColumnRenamed(
+                "person_id", "subject_id"
+            ).withColumnRenamed(
+                "index_date", "prediction_time"
+            ).withColumnRenamed(
+                "label", "boolean_value"
+            ).withColumn(
+                "prediction_time", F.to_timestamp("prediction_time")
+            ).withColumn(
+                "boolean_value", F.col("boolean_value").cast("boolean")
+            )
+            person_id_column = "subject_id"
+            index_date_column = "prediction_time"
 
         if self._is_prediction_window_unbounded:
             observation_period = self._dependency_dict[OBSERVATION_PERIOD]
@@ -529,16 +566,28 @@ class NestedCohortBuilder:
             # Add time_to_event
             cohort = cohort.withColumn(
                 "study_end_date",
-                F.coalesce(F.col("outcome_date"), F.date_add(cohort.index_date, self._prediction_window))
+                F.coalesce(
+                    F.col("outcome_date"),
+                    F.date_add(
+                        cohort[index_date_column], self._prediction_window
+                    )
+                )
             )
-        cohort = cohort.withColumn("time_to_event", F.datediff("study_end_date", "index_date"))
+        cohort = cohort.withColumn("time_to_event", F.datediff("study_end_date", index_date_column))
 
         # if patient_splits is provided, we will
         if self._patient_splits_folder:
             patient_splits = self.spark.read.parquet(self._patient_splits_folder)
-            cohort.join(patient_splits, "person_id").orderBy("person_id", "cohort_member_id").write.mode(
+            cohort.join(
+                patient_splits,
+                cohort[person_id_column] == patient_splits.person_id
+            ).select(
+                [cohort[c] for c in cohort.columns] + [patient_splits.split]
+            ).orderBy(person_id_column, index_date_column).write.mode(
                 "overwrite"
-            ).parquet(os.path.join(self._output_data_folder, "temp"))
+            ).parquet(
+                os.path.join(self._output_data_folder, "temp")
+            )
             # Reload the data from the disk
             cohort = self.spark.read.parquet(os.path.join(self._output_data_folder, "temp"))
             cohort.where('split="train"').write.mode("overwrite").parquet(
@@ -547,7 +596,7 @@ class NestedCohortBuilder:
             cohort.where('split="test"').write.mode("overwrite").parquet(os.path.join(self._output_data_folder, "test"))
             shutil.rmtree(os.path.join(self._output_data_folder, "temp"))
         else:
-            cohort.orderBy("person_id", "cohort_member_id").write.mode("overwrite").parquet(self._output_data_folder)
+            cohort.orderBy(person_id_column, index_date_column).write.mode("overwrite").parquet(self._output_data_folder)
 
     def extract_ehr_records_for_cohort(self, cohort: DataFrame):
         """
@@ -566,47 +615,54 @@ class NestedCohortBuilder:
             with_drug_rollup=self._is_drug_roll_up_concept,
             include_concept_list=self._include_concept_list,
             refresh_measurement=self._refresh_measurement,
+            aggregate_by_hour=self._aggregate_by_hour,
         )
 
         # Duplicate the records for cohorts that allow multiple entries
-        ehr_records = ehr_records.join(cohort, "person_id").select(
-            [ehr_records[field_name] for field_name in ehr_records.schema.fieldNames()] + ["cohort_member_id"]
-        )
+        ehr_records = ehr_records.alias("ehr").join(
+            cohort.alias("cohort"), F.col("ehr.person_id") == F.col("cohort.person_id")
+        ).select(
+            [F.col("ehr." + col) for col in ehr_records.columns] + [F.col("cohort.cohort_member_id")]
+        ).selectExpr("*")
 
         # Only allow the data records that occurred between the index date and the prediction window
         if self._is_population_estimation:
             if self._is_prediction_window_unbounded:
-                record_window_filter = ehr_records["date"] <= F.current_date()
+                record_window_filter = F.col("ehr.date") <= F.current_date()
             else:
-                record_window_filter = ehr_records["date"] <= F.date_add(cohort["index_date"], self._prediction_window)
+                record_window_filter = F.col("ehr.date") <= F.date_add(
+                    F.col("cohort.index_date"), self._prediction_window
+                )
         else:
             # For patient level prediction, we remove all records post index date
             if self._is_observation_post_index:
-                record_window_filter = ehr_records["date"].between(
-                    cohort["index_date"],
-                    F.date_add(cohort["index_date"], self._observation_window),
+                record_window_filter = F.col("ehr.date").between(
+                    F.col("cohort.index_date"),
+                    F.date_add(F.col("cohort.index_date"), self._observation_window),
                 )
             else:
                 if self._is_observation_window_unbounded:
-                    record_window_filter = ehr_records["date"] <= F.date_sub(
-                        cohort["index_date"], self._hold_off_window
+                    record_window_filter = F.col("ehr.date") <= F.date_sub(
+                        F.col("cohort.index_date"), self._hold_off_window
                     )
                 else:
-                    record_window_filter = ehr_records["date"].between(
+                    record_window_filter = F.col("ehr.date").between(
                         F.date_sub(
-                            cohort["index_date"],
+                            F.col("cohort.index_date"),
                             self._observation_window + self._hold_off_window,
                         ),
-                        F.date_sub(cohort["index_date"], self._hold_off_window),
+                        F.date_sub(F.col("cohort.index_date"), self._hold_off_window),
                     )
 
-        cohort_ehr_records = (
-            ehr_records.join(
-                cohort,
-                (ehr_records.person_id == cohort.person_id) & (ehr_records.cohort_member_id == cohort.cohort_member_id),
-            )
-            .where(record_window_filter)
-            .select([ehr_records[field_name] for field_name in ehr_records.schema.fieldNames()])
+        # Somehow the dataframe join does not work without using the alias
+        cohort_ehr_records = ehr_records.alias("ehr").join(
+            cohort.alias("cohort"),
+            (F.col("ehr.person_id") == F.col("cohort.person_id")) &
+            (F.col("ehr.cohort_member_id") == F.col("cohort.cohort_member_id")),
+        ).where(
+            record_window_filter
+        ).select(
+            [F.col("ehr." + field_name) for field_name in ehr_records.schema.fieldNames()]
         )
 
         if self._is_hierarchical_bert:
@@ -647,8 +703,12 @@ class NestedCohortBuilder:
                 exclude_visit_tokens=self._exclude_visit_tokens,
                 patient_demographic=(patient_demographic if self._gpt_patient_sequence else None),
                 att_type=self._att_type,
+                inpatient_att_type=self._inpatient_att_type,
                 exclude_demographic=self._exclude_demographic,
                 use_age_group=self._use_age_group,
+                include_inpatient_hour_token=self._include_inpatient_hour_token,
+                spark=self.spark if self._cache_events else None,
+                persistence_folder=self._output_data_folder if self._cache_events else None,
             )
 
         return create_sequence_data(
@@ -720,6 +780,7 @@ def create_prediction_cohort(
             age_upper_bound=spark_args.age_upper_bound,
             prior_observation_period=0,
             post_observation_period=0,
+            continue_job=spark_args.continue_job,
         )
         .build()
         .load_cohort()
@@ -748,6 +809,7 @@ def create_prediction_cohort(
         is_drug_roll_up_concept=spark_args.is_drug_roll_up_concept,
         include_concept_list=spark_args.include_concept_list,
         refresh_measurement=spark_args.refresh_measurement,
+        aggregate_by_hour=spark_args.aggregate_by_hour,
         is_new_patient_representation=spark_args.is_new_patient_representation,
         gpt_patient_sequence=spark_args.gpt_patient_sequence,
         is_hierarchical_bert=spark_args.is_hierarchical_bert,
@@ -763,7 +825,11 @@ def create_prediction_cohort(
         is_observation_window_unbounded=spark_args.is_observation_window_unbounded,
         is_population_estimation=spark_args.is_population_estimation,
         att_type=AttType(spark_args.att_type),
+        inpatient_att_type=AttType(spark_args.inpatient_att_type),
+        include_inpatient_hour_token=spark_args.include_inpatient_hour_token,
         exclude_demographic=spark_args.exclude_demographic,
         use_age_group=spark_args.use_age_group,
         single_contribution=spark_args.single_contribution,
+        exclude_features=spark_args.exclude_features,
+        meds_format=spark_args.meds_format,
     ).build()