added observation_window to the meds data conversion

Author: ChaoPang

src/cehrbert/data_generators/hf_data_generator/meds_utils.py

@@ -54,6 +54,7 @@ def convert_one_patient(
     patient: meds_reader.Subject,
     conversion: MedsToCehrBertConversion,
     prediction_time: datetime = None,
+    observation_window: int = None,
     label: Union[int, float] = None,
 ) -> CehrBertPatient:
     """
@@ -80,6 +81,9 @@ def convert_one_patient(
         The cutoff time for processing events. Events occurring after this time are
         ignored.
 
+    observation_window: int, optional (default=None)
+        The observation window for extracting features
+
     label : Union[int, float], optional (default=None)
         The prediction label associated with this patient, which could represent a
         clinical outcome (e.g., survival or treatment response).
@@ -122,9 +126,7 @@ def convert_one_patient(
     )
     """
     demographics, patient_blocks = generate_demographics_and_patient_blocks(
-        conversion=conversion,
-        patient=patient,
-        prediction_time=prediction_time,
+        conversion=conversion, patient=patient, prediction_time=prediction_time, observation_window=observation_window
     )
 
     patient_block_dict = collections.defaultdict(list)
@@ -230,9 +232,9 @@ def _meds_to_cehrbert_generator(
     )
     with meds_reader.SubjectDatabase(path_to_db) as patient_database:
         for shard in shards:
-            for patient_id, prediction_time, label in shard:
+            for patient_id, prediction_time, observation_window, label in shard:
                 patient = patient_database[patient_id]
-                converted_patient = convert_one_patient(patient, conversion, prediction_time, label)
+                converted_patient = convert_one_patient(patient, conversion, prediction_time, observation_window, label)
                 # there are patients whose birthdate is none
                 visits = converted_patient["visits"]
                 if converted_patient["birth_datetime"] is None:
@@ -262,11 +264,11 @@ def _create_cehrbert_data_from_meds(
             subject_id = cohort_row.subject_id
             prediction_time = cohort_row.prediction_time
             label = int(cohort_row.boolean_value)
-            batches.append((subject_id, prediction_time, label))
+            batches.append((subject_id, prediction_time, data_args.observation_window, label))
     else:
         patient_split = get_subject_split(os.path.expanduser(data_args.data_folder))
         for subject_id in patient_split[split]:
-            batches.append((subject_id, None, None))
+            batches.append((subject_id, None, None, None))
 
     features = Features(
         {

src/cehrbert/data_generators/hf_data_generator/patient_block.py

@@ -1,7 +1,7 @@
 import re
 from collections import defaultdict
 from dataclasses import dataclass
-from datetime import datetime
+from datetime import datetime, timedelta
 from typing import Iterable, List, Optional, Tuple
 
 import meds_reader
@@ -321,19 +321,32 @@ def generate_demographics_and_patient_blocks(
     conversion: MedsToCehrBertConversion,
     patient: meds_reader.Subject,
     prediction_time: datetime = None,
+    observation_window: int = None,
 ) -> Tuple[PatientDemographics, List[PatientBlock]]:
     if isinstance(conversion, MedsToBertMimic4):
         return mimic_meds_generate_demographics_and_patient_blocks(
-            patient, conversion, prediction_time, conversion.default_visit_id
+            patient=patient,
+            conversion=conversion,
+            prediction_time=prediction_time,
+            observation_window=observation_window,
+            default_visit_id=conversion.default_visit_id,
         )
     elif isinstance(conversion, MedsToCehrbertOMOP):
-        return omop_meds_generate_demographics_and_patient_blocks(patient, conversion, prediction_time)
+        return omop_meds_generate_demographics_and_patient_blocks(
+            patient=patient,
+            conversion=conversion,
+            prediction_time=prediction_time,
+            observation_window=observation_window,
+        )
     else:
         raise RuntimeError(f"Unrecognized conversion: {conversion}")
 
 
 def omop_meds_generate_demographics_and_patient_blocks(
-    patient: meds_reader.Subject, conversion: MedsToCehrBertConversion, prediction_time: datetime = None
+    patient: meds_reader.Subject,
+    conversion: MedsToCehrBertConversion,
+    prediction_time: datetime = None,
+    observation_window: int = None,
 ) -> Tuple[PatientDemographics, List[PatientBlock]]:
     disconnect_problem_list_events = getattr(conversion, "disconnect_problem_list_events", False)
     birth_datetime = None
@@ -342,6 +355,9 @@ def omop_meds_generate_demographics_and_patient_blocks(
     ethnicity = None
     visit_events = defaultdict(list)
     unlinked_event_mapping = defaultdict(list)
+    observation_start_window: Optional[datetime] = None
+    if prediction_time and observation_window:
+        observation_start_window = prediction_time - timedelta(days=observation_window)
     for e in patient.events:
         # This indicates demographics features
         event_code_uppercase = e.code.upper()
@@ -355,10 +371,11 @@ def omop_meds_generate_demographics_and_patient_blocks(
             ethnicity = e.code
         elif e.time is not None:
             # Skip out of the loop if the events' time stamps are beyond the prediction time
-            if prediction_time is not None:
-                if e.time > prediction_time:
-                    break
-
+            if prediction_time is not None and e.time > prediction_time:
+                break
+            # Skip out of the loop if the events' time stamps are before the observation start window
+            if observation_start_window is not None and e.time < observation_start_window:
+                break
             if getattr(e, "visit_id", None):
                 visit_id = e.visit_id
                 visit_events[visit_id].append(e)
@@ -528,6 +545,7 @@ def mimic_meds_generate_demographics_and_patient_blocks(
     patient: meds_reader.Subject,
     conversion: MedsToCehrBertConversion,
     prediction_time: datetime = None,
+    observation_window: int = None,
     default_visit_id: int = 1,
 ) -> Tuple[PatientDemographics, List[PatientBlock]]:
     birth_datetime = None
@@ -539,13 +557,21 @@ def mimic_meds_generate_demographics_and_patient_blocks(
     current_date = None
     events_for_current_date = []
     patient_blocks = []
-    for e in patient.events:
+    observation_start_window: Optional[datetime] = None
+    if prediction_time and observation_window:
+        observation_start_window = prediction_time - timedelta(days=observation_window)
 
+    for e in patient.events:
         # Skip out of the loop if the events' time stamps are beyond the prediction time
         if prediction_time is not None and e.time is not None:
             if e.time > prediction_time:
                 break
 
+        # Skip out of the loop if the events' time stamps are before observation start window
+        if observation_start_window is not None and e.time is not None:
+            if e.time < observation_start_window:
+                break
+
         # This indicates demographics features
         event_code_uppercase = e.code.upper()
         if event_code_uppercase.startswith(birth_code):

src/cehrbert/runners/hf_runner_argument_dataclass.py

@@ -150,6 +150,12 @@ class DataTrainingArguments:
             "This depends on how the EHR records are generated, default to False"
         },
     )
+    observation_window: Optional[int] = dataclasses.field(
+        default=None,
+        metadata={
+            "help": "The observation window to use in the preprocessing.",
+        },
+    )
     streaming: Optional[bool] = dataclasses.field(
         default=False,
         metadata={"help": "The boolean indicator to indicate whether the data should be streamed"},