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

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

@@ -1,6 +1,6 @@
-Metadata-Version: 2.2
+Metadata-Version: 2.4
 Name: cehrbert_data
-Version: 0.0.6
+Version: 0.0.7
 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
@@ -25,6 +25,7 @@ Requires-Dist: pytest-subtests; extra == "dev"
 Requires-Dist: rootutils; extra == "dev"
 Requires-Dist: hypothesis; extra == "dev"
 Requires-Dist: black; extra == "dev"
+Dynamic: license-file
 
 # cehrbert_data
 

{cehrbert_data-0.0.6 → cehrbert_data-0.0.7}/src/cehrbert_data/apps/generate_included_concept_list.py

@@ -18,6 +18,7 @@ Command-line Arguments:
 """
 
 import os
+from typing import List
 
 from pyspark.sql import SparkSession
 from pyspark.sql import functions as F
@@ -29,7 +30,13 @@ from cehrbert_data.utils.spark_utils import join_domain_tables, preprocess_domai
 DOMAIN_TABLE_LIST = ["condition_occurrence", "procedure_occurrence", "drug_exposure"]
 
 
-def main(input_folder, output_folder, min_num_of_patients, with_drug_rollup: bool = True):
+def main(
+        input_folder: str,
+        output_folder: str,
+        min_num_of_patients: int,
+        with_drug_rollup: bool = True,
+        ehr_table_list: List[str] = None
+):
     """
     Main function to generate a qualified concept list based on patient event data from multiple.
 
@@ -39,8 +46,9 @@ def main(input_folder, output_folder, min_num_of_patients, with_drug_rollup: boo
         input_folder (str): The directory where the input data is stored.
         output_folder (str): The directory where the output (qualified concept list) will be saved.
         min_num_of_patients (int): Minimum number of patients that a concept must be linked to for
-        nclusion.
+        inclusion.
         with_drug_rollup (bool): If True, applies drug rollup logic to the drug_exposure domain.
+        ehr_table_list (List[str]): List of patient event tables to include in the concept list.
 
     The function processes patient event data across various domain tables, excludes low-frequency
     concepts, and saves the filtered concepts to a specified output folder.
@@ -50,7 +58,8 @@ def main(input_folder, output_folder, min_num_of_patients, with_drug_rollup: boo
     domain_tables = []
     # Exclude measurement from domain_table_list if exists because we need to process measurement
     # in a different way
-    for domain_table_name in DOMAIN_TABLE_LIST:
+    domain_table_list = ehr_table_list if ehr_table_list else DOMAIN_TABLE_LIST
+    for domain_table_name in domain_table_list:
         if domain_table_name != MEASUREMENT:
             domain_tables.append(
                 preprocess_domain_table(
@@ -105,6 +114,14 @@ if __name__ == "__main__":
         required=False,
     )
     parser.add_argument("--with_drug_rollup", dest="with_drug_rollup", action="store_true")
+    parser.add_argument(
+        "--ehr_table_list",
+        dest="ehr_table_list",
+        nargs="+",
+        action="store",
+        help="The list of domain tables you want to include for feature extraction",
+        required=False,
+    )
 
     ARGS = parser.parse_args()
 
@@ -113,4 +130,5 @@ if __name__ == "__main__":
         ARGS.output_folder,
         ARGS.min_num_of_patients,
         ARGS.with_drug_rollup,
+        ARGS.ehr_table_list,
     )

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

@@ -22,6 +22,7 @@ from cehrbert_data.utils.spark_utils import (
     preprocess_domain_table,
     get_measurement_table,
     validate_table_names,
+    construct_artificial_visits
 )
 from cehrbert_data.utils.logging_utils import add_console_logging
 
@@ -50,6 +51,7 @@ def main(
         continue_from_events: bool = False,
         refresh_measurement: bool = False,
         aggregate_by_hour: bool = True,
+        should_construct_artificial_visits: bool = False,
 ):
     spark = SparkSession.builder.appName("Generate CEHR-BERT Training Data").getOrCreate()
 
@@ -74,6 +76,7 @@ def main(
         f"with_drug_rollup: {with_drug_rollup}\n"
         f"refresh_measurement: {refresh_measurement}\n"
         f"aggregate_by_hour: {aggregate_by_hour}\n"
+        f"should_construct_artificial_visits: {should_construct_artificial_visits}\n"
     )
 
     domain_tables = []
@@ -94,6 +97,7 @@ def main(
         "visit_start_date",
         "visit_start_datetime",
         "visit_end_date",
+        "visit_end_datetime",
         "visit_concept_id",
         "person_id",
         "discharged_to_concept_id",
@@ -159,6 +163,15 @@ def main(
 
     patient_events = spark.read.parquet(os.path.join(output_folder, "all_patient_events"))
 
+    if should_construct_artificial_visits:
+        # Construct artificial visits or re-link the visits for the problem list events
+        patient_events, visit_occurrence_person = construct_artificial_visits(
+            patient_events,
+            visit_occurrence_person,
+            spark=spark,
+            persistence_folder=output_folder,
+        )
+
     if is_new_patient_representation:
         sequence_data = create_sequence_data_with_att(
             patient_events,
@@ -274,7 +287,7 @@ if __name__ == "__main__":
         type=lambda s: datetime.datetime.strptime(s, "%Y-%m-%d"),
         action="store",
         required=False,
-        default="2018-01-01",
+        default="1985-01-01",
     )
     parser.add_argument(
         "-iv",
@@ -349,6 +362,13 @@ if __name__ == "__main__":
         action="store",
         choices=[e.value for e in AttType],
     )
+    parser.add_argument(
+        "--should_construct_artificial_visits",
+        dest="should_construct_artificial_visits",
+        action="store_true",
+        help="Indicate whether we should construct artificial visits for "
+             "the problem list records that could occur years ahead",
+    )
 
     ARGS = parser.parse_args()
 
@@ -378,4 +398,5 @@ if __name__ == "__main__":
         continue_from_events=ARGS.continue_from_events,
         refresh_measurement=ARGS.refresh_measurement,
         aggregate_by_hour=ARGS.aggregate_by_hour,
+        should_construct_artificial_visits=ARGS.should_construct_artificial_visits,
     )

cehrbert_data-0.0.7/src/cehrbert_data/cohorts/atrial_fibrillation.py

@@ -0,0 +1,63 @@
+from cehrbert_data.cohorts.query_builder import AncestorTableSpec, QueryBuilder, QuerySpec
+from cehrbert_data.const.common import CONDITION_OCCURRENCE, PERSON, VISIT_OCCURRENCE
+
+COHORT_QUERY_TEMPLATE = """
+SELECT
+    DISTINCT
+    c.person_id,
+    c.index_date,
+    c.visit_occurrence_id
+FROM
+(
+    SELECT
+        co.person_id,
+        vo.visit_occurrence_id,
+        CASE
+            WHEN co.condition_start_datetime < vo.visit_start_datetime AND datediff(vo.visit_start_date, co.condition_start_date) <= 1 
+                THEN to_timestamp(concat(date_format(vo.visit_start_date, 'yyyy-MM-dd'), ' 23:59:00'), 'yyyy-MM-dd HH:mm:ss')
+            ELSE 
+                to_timestamp(concat(date_format(co.condition_start_date, 'yyyy-MM-dd'), ' 23:59:00'), 'yyyy-MM-dd HH:mm:ss')
+        END AS index_date,
+        ROW_NUMBER() OVER(
+            PARTITION BY co.person_id
+            ORDER BY co.condition_start_datetime,
+                co.condition_start_date,
+                co.visit_occurrence_id
+        ) as r_number
+    FROM global_temp.condition_occurrence AS co
+    JOIN global_temp.visit_occurrence AS vo
+        ON co.visit_occurrence_id = vo.visit_occurrence_id
+    JOIN global_temp.{atrial_fibrillation_concepts} AS c
+        ON co.condition_concept_id = c.concept_id
+) AS c
+WHERE c.r_number = 1
+"""
+
+ATRIAL_FIBRILLATION_CONCEPT_ID = [313217]
+
+DEPENDENCY_LIST = [PERSON, VISIT_OCCURRENCE, CONDITION_OCCURRENCE]
+
+DEFAULT_COHORT_NAME = "atrial_fibrillation"
+ATRIAL_FIBRILLATION_CONCEPTS = "atrial_fibrillation_concepts"
+
+
+def query_builder():
+    query = QuerySpec(
+        table_name=DEFAULT_COHORT_NAME,
+        query_template=COHORT_QUERY_TEMPLATE,
+        parameters={"atrial_fibrillation_concepts": ATRIAL_FIBRILLATION_CONCEPTS},
+    )
+
+    ancestor_table_specs = [
+        AncestorTableSpec(
+            table_name=ATRIAL_FIBRILLATION_CONCEPTS,
+            ancestor_concept_ids=ATRIAL_FIBRILLATION_CONCEPT_ID,
+            is_standard=True,
+        )
+    ]
+    return QueryBuilder(
+        cohort_name=DEFAULT_COHORT_NAME,
+        dependency_list=DEPENDENCY_LIST,
+        query=query,
+        ancestor_table_specs=ancestor_table_specs,
+    )

{cehrbert_data-0.0.6 → cehrbert_data-0.0.7}/src/cehrbert_data/cohorts/cabg.py

@@ -10,10 +10,18 @@ FROM
 (
     SELECT DISTINCT
         vo.person_id,
-        FIRST(DATE(vo.visit_start_date)) OVER (PARTITION BY po.person_id
-            ORDER BY DATE(vo.visit_start_date), vo.visit_occurrence_id) AS index_date,
-        FIRST(vo.visit_occurrence_id) OVER (PARTITION BY po.person_id
-            ORDER BY DATE(vo.visit_start_date), vo.visit_occurrence_id) AS visit_occurrence_id
+        vo.visit_occurrence_id,
+        CASE
+            WHEN po.procedure_datetime < vo.visit_start_datetime AND datediff(vo.visit_start_datetime, po.procedure_datetime) <= 1 
+                THEN to_timestamp(concat(date_format(vo.visit_start_date, 'yyyy-MM-dd'), ' 23:59:00'), 'yyyy-MM-dd HH:mm:ss')
+            ELSE to_timestamp(concat(date_format(po.procedure_date, 'yyyy-MM-dd'), ' 23:59:00'), 'yyyy-MM-dd HH:mm:ss')
+        END AS index_date,
+        ROW_NUMBER() OVER(
+            PARTITION BY po.person_id
+            ORDER BY po.procedure_datetime,
+                po.procedure_date,
+                po.visit_occurrence_id
+        ) as r_number
     FROM global_temp.procedure_occurrence AS po
     JOIN global_temp.visit_occurrence AS vo
         ON po.visit_occurrence_id = vo.visit_occurrence_id
@@ -23,7 +31,7 @@ FROM
         WHERE po.procedure_concept_id = ie.concept_id
     )
 ) c
-WHERE c.index_date >= '{date_lower_bound}'
+WHERE c.index_date >= '{date_lower_bound}' AND c.r_number = 1
 """
 
 DEFAULT_COHORT_NAME = "cabg"

{cehrbert_data-0.0.6 → cehrbert_data-0.0.7}/src/cehrbert_data/cohorts/coronary_artery_disease.py

@@ -21,10 +21,18 @@ FROM
 (
     SELECT DISTINCT
         vo.person_id,
-        FIRST(DATE(vo.visit_start_date)) OVER (PARTITION BY co.person_id
-            ORDER BY DATE(vo.visit_start_date), vo.visit_occurrence_id) AS index_date,
-        FIRST(vo.visit_occurrence_id) OVER (PARTITION BY co.person_id
-            ORDER BY DATE(vo.visit_start_date), vo.visit_occurrence_id) AS visit_occurrence_id
+        vo.visit_occurrence_id,
+        CASE
+            WHEN co.condition_start_datetime < vo.visit_start_datetime AND datediff(vo.visit_start_datetime, co.condition_start_datetime) <= 1 
+                THEN to_timestamp(concat(date_format(vo.visit_start_date, 'yyyy-MM-dd'), ' 23:59:00'), 'yyyy-MM-dd HH:mm:ss')
+            ELSE to_timestamp(concat(date_format(co.condition_start_date, 'yyyy-MM-dd'), ' 23:59:00'), 'yyyy-MM-dd HH:mm:ss')
+        END AS index_date,
+        ROW_NUMBER() OVER(
+            PARTITION BY co.person_id
+            ORDER BY co.condition_start_datetime,
+                co.condition_start_date,
+                co.visit_occurrence_id
+        ) as r_number
     FROM global_temp.condition_occurrence AS co
     JOIN global_temp.visit_occurrence AS vo
         ON co.visit_occurrence_id = vo.visit_occurrence_id
@@ -41,7 +49,9 @@ WHERE NOT EXISTS (
     FROM prior_graft_stent AS exclusion
     WHERE exclusion.person_id = c.person_id
         AND c.index_date > exclusion.procedure_date
-) AND c.index_date >= '{date_lower_bound}'
+)
+    AND c.r_number = 1
+    AND c.index_date >= '{date_lower_bound}'
 """
 
 DEFAULT_COHORT_NAME = "coronary_artery_disease"

cehrbert_data-0.0.7/src/cehrbert_data/cohorts/ischemic_stroke.py

@@ -0,0 +1,61 @@
+from cehrbert_data.cohorts.query_builder import AncestorTableSpec, QueryBuilder, QuerySpec
+from cehrbert_data.const.common import CONDITION_OCCURRENCE, PERSON, VISIT_OCCURRENCE
+
+COHORT_QUERY_TEMPLATE = """
+SELECT DISTINCT
+    c.person_id,
+    c.index_date,
+    c.visit_occurrence_id
+FROM
+(
+    SELECT
+        co.person_id,
+        vo.visit_occurrence_id,
+        CASE
+            WHEN co.condition_start_datetime < vo.visit_start_datetime AND datediff(vo.visit_start_datetime, co.condition_start_datetime) <= 1 
+                THEN to_timestamp(concat(date_format(vo.visit_start_date, 'yyyy-MM-dd'), ' 23:59:00'), 'yyyy-MM-dd HH:mm:ss')
+            ELSE to_timestamp(concat(date_format(co.condition_start_date, 'yyyy-MM-dd'), ' 23:59:00'), 'yyyy-MM-dd HH:mm:ss')
+        END AS index_date,
+        ROW_NUMBER() OVER(
+            PARTITION BY co.person_id
+            ORDER BY co.condition_start_datetime,
+                co.condition_start_date,
+                co.visit_occurrence_id
+        ) as r_number
+    FROM global_temp.condition_occurrence AS co
+    JOIN global_temp.visit_occurrence AS vo
+        ON co.visit_occurrence_id = vo.visit_occurrence_id
+    JOIN global_temp.{ischemic_stroke_concepts} AS c
+        ON co.condition_concept_id = c.concept_id
+) AS c
+WHERE c.r_number = 1
+"""
+
+ISCHEMIC_STROKE_CONCEPT_ID = [443454]
+
+DEPENDENCY_LIST = [PERSON, VISIT_OCCURRENCE, CONDITION_OCCURRENCE]
+
+DEFAULT_COHORT_NAME = "ischemic_stroke"
+ISCHEMIC_STROKE_CONCEPTS = "ischemic_stroke_concepts"
+
+
+def query_builder():
+    query = QuerySpec(
+        table_name=DEFAULT_COHORT_NAME,
+        query_template=COHORT_QUERY_TEMPLATE,
+        parameters={"ischemic_stroke_concepts": ISCHEMIC_STROKE_CONCEPTS},
+    )
+
+    ancestor_table_specs = [
+        AncestorTableSpec(
+            table_name=ISCHEMIC_STROKE_CONCEPTS,
+            ancestor_concept_ids=ISCHEMIC_STROKE_CONCEPT_ID,
+            is_standard=True,
+        )
+    ]
+    return QueryBuilder(
+        cohort_name=DEFAULT_COHORT_NAME,
+        dependency_list=DEPENDENCY_LIST,
+        query=query,
+        ancestor_table_specs=ancestor_table_specs,
+    )

{cehrbert_data-0.0.6 → cehrbert_data-0.0.7}/src/cehrbert_data/cohorts/last_visit_discharged_home.py

@@ -11,7 +11,7 @@ FROM
     SELECT
         v.person_id,
         v.visit_occurrence_id,
-        v.visit_end_date AS index_date,
+        coalesce(v.visit_end_datetime, v.visit_end_date) AS index_date,
         v.discharged_to_concept_id,
         ROW_NUMBER() OVER(PARTITION BY v.person_id ORDER BY DATE(v.visit_end_date) DESC) AS rn
     FROM global_temp.visit_occurrence AS v

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

@@ -5,7 +5,6 @@ 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
@@ -22,6 +21,7 @@ from cehrbert_data.utils.spark_utils import (
     extract_ehr_records,
     get_descendant_concept_ids,
     preprocess_domain_table,
+    construct_artificial_visits
 )
 from cehrbert_data.utils.logging_utils import add_console_logging
 
@@ -232,8 +232,8 @@ class BaseCohortBuilder(ABC):
         FROM global_temp.cohort AS c
         JOIN global_temp.observation_period AS p
             ON c.person_id = p.person_id
-                AND DATE_ADD(c.index_date, -{prior_observation_period}) >= p.observation_period_start_date
-                AND DATE_ADD(c.index_date, {post_observation_period}) <= p.observation_period_end_date
+                AND c.index_date - INTERVAL {prior_observation_period} DAY >= p.observation_period_start_date
+                AND c.index_date + INTERVAL {post_observation_period} DAY <= p.observation_period_end_date
         """.format(
                 prior_observation_period=self._prior_observation_period,
                 post_observation_period=self._post_observation_period,
@@ -317,6 +317,7 @@ class NestedCohortBuilder:
             exclude_features: bool = True,
             meds_format: bool = False,
             cache_events: bool = False,
+            should_construct_artificial_visits: bool = False,
     ):
         self._cohort_name = cohort_name
         self._input_folder = input_folder
@@ -363,6 +364,7 @@ class NestedCohortBuilder:
         self._exclude_features = exclude_features
         self._meds_format = meds_format
         self._cache_events = cache_events
+        self._should_construct_artificial_visits = should_construct_artificial_visits
 
         self.get_logger().info(
             f"cohort_name: {cohort_name}\n"
@@ -403,6 +405,7 @@ class NestedCohortBuilder:
             f"extract_features: {exclude_features}\n"
             f"meds_format: {meds_format}\n"
             f"cache_events: {cache_events}\n"
+            f"should_construct_artificial_visits: {should_construct_artificial_visits}\n"
         )
 
         self.spark = SparkSession.builder.appName(f"Generate {self._cohort_name}").getOrCreate()
@@ -434,7 +437,7 @@ class NestedCohortBuilder:
             FROM global_temp.target_cohort AS t
             LEFT JOIN global_temp.{entry_cohort} AS o
                 ON t.person_id = o.person_id
-                    AND DATE_ADD(t.index_date, {prediction_start_days}) > o.index_date
+                AND t.index_date + INTERVAL {prediction_start_days} DAY > o.index_date
             WHERE o.person_id IS NULL
             """.format(
                     entry_cohort=ENTRY_COHORT,
@@ -468,7 +471,7 @@ class NestedCohortBuilder:
             LEFT JOIN global_temp.outcome_cohort AS exclusion
                 ON t.person_id = exclusion.person_id
                     AND exclusion.index_date BETWEEN t.index_date
-                        AND DATE_ADD(t.index_date, {prediction_start_days})
+                        AND t.index_date + INTERVAL {prediction_start_days} DAY
             WHERE exclusion.person_id IS NULL
             """.format(
                     prediction_start_days=max(prediction_start_days - 1, 0)
@@ -485,7 +488,7 @@ class NestedCohortBuilder:
             FROM global_temp.target_cohort AS t
             LEFT JOIN global_temp.outcome_cohort AS o
                 ON t.person_id = o.person_id
-                    AND o.index_date >= DATE_ADD(t.index_date, {prediction_start_days})
+                    AND o.index_date >= t.index_date + INTERVAL {prediction_start_days} DAY
             """
         else:
             query_template = """
@@ -496,11 +499,11 @@ class NestedCohortBuilder:
             FROM global_temp.target_cohort AS t
             LEFT JOIN global_temp.observation_period AS op
                 ON t.person_id = op.person_id
-                    AND DATE_ADD(t.index_date, {prediction_window}) <= op.observation_period_end_date
+                    AND t.index_date + INTERVAL {prediction_window} DAY <= op.observation_period_end_date
             LEFT JOIN global_temp.outcome_cohort AS o
                 ON t.person_id = o.person_id
-                    AND o.index_date BETWEEN DATE_ADD(t.index_date, {prediction_start_days})
-                        AND DATE_ADD(t.index_date, {prediction_window})
+                    AND o.index_date BETWEEN t.index_date + INTERVAL {prediction_start_days} DAY
+                        AND t.index_date + INTERVAL {prediction_window} DAY
             WHERE op.person_id IS NOT NULL OR o.person_id IS NOT NULL
             """
 
@@ -557,20 +560,24 @@ class NestedCohortBuilder:
             observation_period = self._dependency_dict[OBSERVATION_PERIOD]
             # Add time_to_event
             cohort = cohort.join(
-                observation_period.select("person_id", "observation_period_end_date"), "person_id"
+                observation_period.select("person_id", "observation_period_end_date"),
+                cohort[person_id_column] == observation_period["person_id"]
+            ).select(
+                [cohort[c] for c in cohort.columns] +
+                [observation_period["observation_period_end_date"]]
             ).withColumn(
                 "study_end_date",
                 F.coalesce(F.col("outcome_date"), F.col("observation_period_end_date"))
-            ).drop("observation_period_end_date")
+            ).drop(
+                "observation_period_end_date"
+            )
         else:
             # Add time_to_event
             cohort = cohort.withColumn(
                 "study_end_date",
                 F.coalesce(
                     F.col("outcome_date"),
-                    F.date_add(
-                        cohort[index_date_column], self._prediction_window
-                    )
+                    F.expr(f"{index_date_column} + INTERVAL {self._prediction_window} DAYS")
                 )
             )
         cohort = cohort.withColumn("time_to_event", F.datediff("study_end_date", index_date_column))
@@ -578,11 +585,11 @@ class NestedCohortBuilder:
         # 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,
-                cohort[person_id_column] == patient_splits.person_id
+            cohort.alias("cohort").join(
+                patient_splits.alias("split"),
+                F.col(f"cohort.{person_id_column}") == F.col("split.person_id")
             ).select(
-                [cohort[c] for c in cohort.columns] + [patient_splits.split]
+                [F.col(f"cohort.{c}").alias(c) for c in cohort.columns] + [F.col("split.split").alias("split")]
             ).orderBy(person_id_column, index_date_column).write.mode(
                 "overwrite"
             ).parquet(
@@ -596,7 +603,9 @@ 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_column, index_date_column).write.mode("overwrite").parquet(self._output_data_folder)
+            cohort.orderBy(person_id_column, index_date_column).write.mode("overwrite").parquet(
+                os.path.join(self._output_data_folder, "data")
+            )
 
     def extract_ehr_records_for_cohort(self, cohort: DataFrame):
         """
@@ -618,6 +627,25 @@ class NestedCohortBuilder:
             aggregate_by_hour=self._aggregate_by_hour,
         )
 
+        if self._cache_events:
+            all_patient_events_dir = os.path.join(self._output_data_folder, "all_patient_events")
+            ehr_records.write.mode("overwrite").parquet(
+                all_patient_events_dir
+            )
+            ehr_records = self.spark.read.parquet(
+                all_patient_events_dir
+            )
+
+        if self._should_construct_artificial_visits:
+            ehr_records, visit_occurrence_with_artificial_visits = construct_artificial_visits(
+                ehr_records,
+                self._dependency_dict[VISIT_OCCURRENCE],
+                spark=self.spark if self._cache_events else None,
+                persistence_folder=self._output_data_folder if self._cache_events else None,
+            )
+            # Refresh the dependency
+            self._dependency_dict[VISIT_OCCURRENCE] = visit_occurrence_with_artificial_visits
+
         # Duplicate the records for cohorts that allow multiple entries
         ehr_records = ehr_records.alias("ehr").join(
             cohort.alias("cohort"), F.col("ehr.person_id") == F.col("cohort.person_id")
@@ -628,30 +656,29 @@ class NestedCohortBuilder:
         # 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 = F.col("ehr.date") <= F.current_date()
+                record_window_filter = F.col("ehr.datetime") <= F.current_timestamp()
             else:
-                record_window_filter = F.col("ehr.date") <= F.date_add(
-                    F.col("cohort.index_date"), self._prediction_window
+                record_window_filter = (
+                        F.col("ehr.datetime") <=
+                        F.expr(f"cohort.index_date - INTERVAL {self._hold_off_window} DAYS + INTERVAL 0.1 SECOND")
                 )
         else:
             # For patient level prediction, we remove all records post index date
             if self._is_observation_post_index:
-                record_window_filter = F.col("ehr.date").between(
+                record_window_filter = F.col("ehr.datetime").between(
                     F.col("cohort.index_date"),
-                    F.date_add(F.col("cohort.index_date"), self._observation_window),
+                    F.expr(f"cohort.index_date + INTERVAL {self._observation_window} DAYS + INTERVAL 0.1 SECOND")
                 )
             else:
                 if self._is_observation_window_unbounded:
-                    record_window_filter = F.col("ehr.date") <= F.date_sub(
-                        F.col("cohort.index_date"), self._hold_off_window
+                    record_window_filter = (
+                        F.col("ehr.datetime")
+                        <= F.expr(f"cohort.index_date - INTERVAL {self._hold_off_window} DAYS + INTERVAL 0.1 SECOND")
                     )
                 else:
-                    record_window_filter = F.col("ehr.date").between(
-                        F.date_sub(
-                            F.col("cohort.index_date"),
-                            self._observation_window + self._hold_off_window,
-                        ),
-                        F.date_sub(F.col("cohort.index_date"), self._hold_off_window),
+                    record_window_filter = F.col("ehr.datetime").between(
+                        F.expr(f"cohort.index_date - INTERVAL {self._observation_window + self._hold_off_window} DAYS"),
+                        F.expr(f"cohort.index_date - INTERVAL {self._hold_off_window} DAYS + INTERVAL 0.1 SECOND"),
                     )
 
         # Somehow the dataframe join does not work without using the alias
@@ -709,6 +736,7 @@ class NestedCohortBuilder:
                 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,
+                cohort_index=cohort.select("person_id", "cohort_member_id", "index_date")
             )
 
         return create_sequence_data(
@@ -832,4 +860,6 @@ def create_prediction_cohort(
         single_contribution=spark_args.single_contribution,
         exclude_features=spark_args.exclude_features,
         meds_format=spark_args.meds_format,
+        cache_events=spark_args.cache_events,
+        should_construct_artificial_visits=spark_args.should_construct_artificial_visits,
     ).build()

{cehrbert_data-0.0.6 → cehrbert_data-0.0.7}/src/cehrbert_data/cohorts/type_two_diabietes.py

@@ -20,8 +20,7 @@ person_ids_to_exclude_observation AS
     JOIN global_temp.{observation_exclusion_concepts} AS oec
         ON o.observation_concept_id = oec.concept_id
 )
-SELECT
-    distinct
+SELECT DISTINCT
     c.person_id,
     c.index_date,
     c.visit_occurrence_id
@@ -29,10 +28,18 @@ FROM
 (
     SELECT DISTINCT
         vo.person_id,
-        FIRST(DATE(vo.visit_start_date)) OVER (PARTITION BY co.person_id
-            ORDER BY DATE(vo.visit_start_date), vo.visit_occurrence_id) AS index_date,
-        FIRST(vo.visit_occurrence_id) OVER (PARTITION BY co.person_id
-            ORDER BY DATE(vo.visit_start_date), vo.visit_occurrence_id) AS visit_occurrence_id
+        vo.visit_occurrence_id,
+        CASE
+            WHEN co.condition_start_datetime < vo.visit_start_datetime AND datediff(vo.visit_start_datetime, co.condition_start_datetime) <= 1 
+                THEN to_timestamp(concat(date_format(vo.visit_start_date, 'yyyy-MM-dd'), ' 23:59:00'), 'yyyy-MM-dd HH:mm:ss')
+            ELSE to_timestamp(concat(date_format(co.condition_start_date, 'yyyy-MM-dd'), ' 23:59:00'), 'yyyy-MM-dd HH:mm:ss')
+        END AS index_date,
+        ROW_NUMBER() OVER(
+            PARTITION BY co.person_id
+            ORDER BY co.condition_start_datetime,
+                co.condition_start_date,
+                co.visit_occurrence_id
+        ) as r_number
     FROM global_temp.condition_occurrence AS co
     JOIN global_temp.{diabetes_inclusion_concepts} AS ie
         ON co.condition_concept_id = ie.concept_id
@@ -43,7 +50,9 @@ JOIN person_ids_to_include_drug AS d
     ON c.person_id = d.person_id
 LEFT JOIN person_ids_to_exclude_observation AS eo
     ON c.person_id = eo.person_id AND c.index_date > eo.observation_date
-WHERE eo.person_id IS NULL AND c.index_date >= '{date_lower_bound}'
+WHERE eo.person_id IS NULL
+    AND c.r_number = 1 
+    AND c.index_date >= '{date_lower_bound}'
 """
 
 DIABETES_INCLUSION = [443238, 201820, 442793, 4016045]

{cehrbert_data-0.0.6 → cehrbert_data-0.0.7}/src/cehrbert_data/decorators/__init__.py

@@ -2,4 +2,5 @@ from .artificial_time_token_decorator import AttEventDecorator
 from .death_event_decorator import DeathEventDecorator
 from .clinical_event_decorator import ClinicalEventDecorator
 from .demographic_event_decorator import DemographicEventDecorator
+from .prediction_token_decorator import PredictionEventDecorator
 from .patient_event_decorator_base import time_token_func, get_att_function, AttType