Test cases and (eventually) fixes for #114

EventStream/baseline/FT_task_baseline.py

@@ -47,6 +47,7 @@ def load_flat_rep(
     do_update_if_missing: bool = True,
     task_df_name: str | None = None,
     do_cache_filtered_task: bool = True,
+    overwrite_cache_filtered_task: bool = False,
     subjects_included: dict[str, set[int]] | None = None,
 ) -> dict[str, pl.LazyFrame]:
     """Loads a set of flat representations from a passed dataset that satisfy the given constraints.
@@ -68,14 +69,16 @@ def load_flat_rep(
         do_update_if_missing: If `True`, then if any window sizes or features are missing, the function will
             try to update the stored flat representations to reflect these. If `False`, if information is
             missing, it will raise a `FileNotFoundError` instead.
-        task_df_name: If specified, the flat representations loaded will be (inner) joined against the task
+        task_df_name: If specified, the flat representations loaded will be joined against the task
             dataframe of this name on the columns ``"subject_id"`` and ``"end_time"`` (which will be renamed
             to ``"timestamp"``). This is to avoid needing to load the full dataset in flattened form into
             memory. This is also used as a cache key; if a pre-filtered dataset is written to disk at a
             specified path for this task, then the data will be loaded from there, rather than from the base
             dataset.
         do_cache_filtered_task: If `True`, the flat representations will, after being filtered to just the
             relevant rows for the task, be cached to disk for faster re-use.
+        overwrite_cache_filtered_task: If `True`, the flat representations will be regenerated. If `False`,
+            the cached file will be loaded if exists.
         subjects_included: A dictionary by split of the subjects to include in the task. Omitted splits are
             used wholesale.
 
@@ -148,8 +151,9 @@ def load_flat_rep(
 
     by_split = {}
     for sp, all_sp_subjects in ESD.split_subjects.items():
+        all_sp_subjects = pl.Series(list(all_sp_subjects)).cast(ESD.subject_id_dtype)
         if task_df_name is not None:
-            sp_join_df = join_df.filter(pl.col("subject_id").is_in(list(all_sp_subjects)))
+            sp_join_df = join_df.filter(pl.col("subject_id").is_in(all_sp_subjects))
 
         static_df = pl.scan_parquet(flat_dir / "static" / sp / "*.parquet")
         if task_df_name is not None:
@@ -171,19 +175,28 @@ def load_flat_rep(
                 if task_df_name is not None:
                     fn = fp.parts[-1]
                     cached_fp = task_window_dir / fn
-                    if cached_fp.is_file():
+                    if cached_fp.is_file() and not overwrite_cache_filtered_task:
                         df = pl.scan_parquet(cached_fp).select("subject_id", "timestamp", *window_features)
                         if subjects_included.get(sp, None) is not None:
                             subjects = list(set(subjects).intersection(subjects_included[sp]))
-                            df = df.filter(pl.col("subject_id").is_in(subjects))
+                            df = df.filter(
+                                pl.col("subject_id").is_in(pl.Series(subjects).cast(ESD.subject_id_dtype))
+                            )
                         window_dfs.append(df)
                         continue
 
                 df = pl.scan_parquet(fp)
                 if task_df_name is not None:
-                    filter_join_df = sp_join_df.select(join_keys).filter(pl.col("subject_id").is_in(subjects))
+                    filter_join_df = sp_join_df.select(join_keys).filter(
+                        pl.col("subject_id").is_in(pl.Series(subjects).cast(ESD.subject_id_dtype))
+                    )
 
-                    df = df.join(filter_join_df, on=join_keys, how="inner")
+                    df = filter_join_df.join_asof(
+                        df,
+                        by="subject_id",
+                        on="timestamp",
+                        strategy="forward" if "-" in window_size else "backward",
+                    )
 
                     if do_cache_filtered_task:
                         cached_fp.parent.mkdir(exist_ok=True, parents=True)
@@ -193,7 +206,7 @@ def load_flat_rep(
                 df = df.select("subject_id", "timestamp", *window_features)
                 if subjects_included.get(sp, None) is not None:
                     subjects = list(set(subjects).intersection(subjects_included[sp]))
-                    df = df.filter(pl.col("subject_id").is_in(subjects))
+                    df = df.filter(pl.col("subject_id").is_in(pl.Series(subjects).cast(ESD.subject_id_dtype)))
 
                 window_dfs.append(df)
 

EventStream/data/dataset_base.py

@@ -1367,27 +1367,24 @@ def cache_deep_learning_representation(
         NRT_dir = self.config.save_dir / "NRT_reps"
 
         shards_fp = self.config.save_dir / "DL_shards.json"
-        if shards_fp.exists():
+        if shards_fp.exists() and not do_overwrite:
             shards = json.loads(shards_fp.read_text())
         else:
             shards = {}
 
             if subjects_per_output_file is None:
                 subject_chunks = [self.subject_ids]
             else:
-                subjects = np.random.permutation(list(self.subject_ids))
+                subjects = np.random.permutation(list(set(self.subject_ids)))
                 subject_chunks = np.array_split(
                     subjects,
                     np.arange(subjects_per_output_file, len(subjects), subjects_per_output_file),
                 )
 
-            subject_chunks = [[int(x) for x in c] for c in subject_chunks]
-
             for chunk_idx, subjects_list in enumerate(subject_chunks):
                 for split, subjects in self.split_subjects.items():
                     shard_key = f"{split}/{chunk_idx}"
-                    included_subjects = set(subjects_list).intersection({int(x) for x in subjects})
-                    shards[shard_key] = list(included_subjects)
+                    shards[shard_key] = list(set(subjects_list).intersection(subjects))
 
             shards_fp.write_text(json.dumps(shards))
 

EventStream/data/dataset_polars.py

@@ -690,6 +690,7 @@ def _agg_by_time(self):
         )
 
     def _update_subject_event_properties(self):
+        self.subject_id_dtype = self.events_df.schema["subject_id"]
         if self.events_df is not None:
             logger.debug("Collecting event types")
             self.event_types = (
@@ -699,15 +700,28 @@ def _update_subject_event_properties(self):
                 .to_list()
             )
 
-            n_events_pd = self.events_df.get_column("subject_id").value_counts(sort=False).to_pandas()
-            self.n_events_per_subject = n_events_pd.set_index("subject_id")["count"].to_dict()
+            logger.debug("Collecting subject event counts")
+            n_events = self.events_df.group_by("subject_id").agg(pl.len().alias("count"))
+            n_events = n_events.drop_nulls("subject_id")
+            # here we cast to str to avoid issues with the subject_id column being various other types as we
+            # will eventually JSON serialize it.
+            n_events = n_events.with_columns(pl.col("subject_id").cast(pl.Utf8))
+            self.n_events_per_subject = {
+                subject_id: count for subject_id, count in zip(n_events["subject_id"], n_events["count"])
+            }
             self.subject_ids = set(self.n_events_per_subject.keys())
 
         if self.subjects_df is not None:
-            logger.debug("Collecting subject event counts")
-            subjects_with_no_events = (
-                set(self.subjects_df.get_column("subject_id").to_list()) - self.subject_ids
+            subjects_df_subjects = (
+                self.subjects_df
+                .drop_nulls("subject_id")
+                .select(pl.col("subject_id").cast(pl.Utf8))
             )
+            subjects_df_subj_ids = set(subjects_df_subjects["subject_id"].to_list())
+            subj_no_in_df = self.subject_ids - subjects_df_subj_ids
+            if len(subj_no_in_df) > 0:
+                logger.warning(f"Found {len(subj_no_in_df)} subjects not in subject df!")
+            subjects_with_no_events = subjects_df_subj_ids - self.subject_ids
             for sid in subjects_with_no_events:
                 self.n_events_per_subject[sid] = 0
             self.subject_ids.update(subjects_with_no_events)
@@ -723,7 +737,20 @@ def _filter_col_inclusion(cls, df: DF_T, col_inclusion_targets: dict[str, bool |
                     filter_exprs.append(pl.col(col).is_null())
                 case _:
                     try:
-                        incl_list = pl.Series(list(incl_targets), dtype=df.schema[col])
+                        logger.debug(
+                            f"Converting inclusion targets of type {type(list(incl_targets)[0])} for "
+                            f"{col} to {df.schema[col]}"
+                        )
+                        if isinstance(list(incl_targets)[0], str):
+                            incl_list = pl.Series(list(incl_targets), dtype=pl.Utf8)
+                        else:
+                            incl_list = pl.Series(list(incl_targets), dtype=df.schema[col])
+
+                        incl_list = incl_list.cast(df.schema[col])
+
+                        logger.debug(
+                            f"Converted to Series of type {incl_list.dtype} with size {len(incl_list)}"
+                        )
                     except TypeError as e:
                         incl_targets_by_type = defaultdict(list)
                         for t in incl_targets:
@@ -1358,6 +1385,11 @@ def build_DL_cached_representation(
         # 1. Process subject data into the right format.
         if subject_ids:
             subjects_df = self._filter_col_inclusion(self.subjects_df, {"subject_id": subject_ids})
+            if len(subject_ids) != len(subjects_df):
+                raise ValueError(
+                    f"Size of given subject_ids are {len(subject_ids)}, but after _filter_col_inclusion "
+                    f"the size of subjects_df are {len(subjects_df)}"
+                )
         else:
             subjects_df = self.subjects_df
 
@@ -1369,6 +1401,7 @@ def build_DL_cached_representation(
                 pl.col("index").alias("static_indices"),
             )
         )
+        logger.debug(f"Size of static_data: {static_data.shape[0]}")
 
         # 2. Process event data into the right format.
         if subject_ids:
@@ -1378,6 +1411,7 @@ def build_DL_cached_representation(
             events_df = self.events_df
             event_ids = None
         event_data = self._melt_df(events_df, ["subject_id", "timestamp", "event_id"], event_measures)
+        logger.debug(f"Size of event_data: {event_data.shape[0]}")
 
         # 3. Process measurement data into the right base format:
         if event_ids:
@@ -1392,6 +1426,7 @@ def build_DL_cached_representation(
 
         if do_sort_outputs:
             dynamic_data = dynamic_data.sort("event_id", "measurement_id")
+        logger.debug(f"Size of dynamic_data: {dynamic_data.shape[0]}")
 
         # 4. Join dynamic and event data.
 

EventStream/data/pytorch_dataset.py

@@ -461,7 +461,16 @@ def _seeded_getitem(self, idx: int) -> dict[str, list[float]]:
 
         subject_id, st, end = self.index[idx]
 
-        shard = self.subj_map[subject_id]
+        if str(subject_id) not in self.subj_map:
+            err_str = [f"Subject {subject_id} ({type(subject_id)} -- as str) not found in the shard map!"]
+
+            if len(self.subj_map) < 10:
+                err_str.append("Subject IDs in map:")
+                err_str.extend(f"  * {k} ({type(k)}): {v}" for k, v in self.subj_map.items())
+
+            raise ValueError("\n".join(err_str))
+
+        shard = self.subj_map[str(subject_id)]
         subject_idx = self.subj_indices[subject_id]
         static_row = self.static_dfs[shard][subject_idx].to_dict()
 
@@ -471,7 +480,7 @@ def _seeded_getitem(self, idx: int) -> dict[str, list[float]]:
         }
 
         if self.config.do_include_subject_id:
-            out["subject_id"] = subject_id
+            out["subject_id"] = static_row["subject_id"].item()
 
         seq_len = end - st
         if seq_len > self.max_seq_len:

configs/dataset_base.yaml

@@ -19,7 +19,7 @@ center_and_scale: True
 
 hydra:
   job:
-    name: build_${cohort_name}
+    name: build_dataset
   run:
     dir: ${save_dir}/.logs
   sweep:

sample_data/dataset.yaml

@@ -10,7 +10,7 @@ subject_id_col: "MRN"
 raw_data_dir: "./sample_data/raw/"
 save_dir: "./sample_data/processed/${cohort_name}"
 
-DL_chunk_size: null
+DL_chunk_size: 25
 
 inputs:
   subjects:

sample_data/dataset_parquet.yaml

@@ -0,0 +1,69 @@
+defaults:
+  - dataset_base
+  - _self_
+
+# So that it can be run multiple times without issue.
+do_overwrite: True
+
+cohort_name: "sample"
+subject_id_col: "MRN"
+raw_data_dir: "./sample_data/raw_parquet"
+save_dir: "./sample_data/processed/${cohort_name}"
+
+DL_chunk_size: 25
+
+inputs:
+  subjects:
+    input_df: "${raw_data_dir}/subjects.parquet"
+  admissions:
+    input_df: "${raw_data_dir}/admit_vitals.parquet"
+    start_ts_col: "admit_date"
+    end_ts_col: "disch_date"
+    ts_format: "%m/%d/%Y, %H:%M:%S"
+    event_type: ["OUTPATIENT_VISIT", "ADMISSION", "DISCHARGE"]
+  vitals:
+    input_df: "${raw_data_dir}/admit_vitals.parquet"
+    ts_col: "vitals_date"
+    ts_format: "%m/%d/%Y, %H:%M:%S"
+  labs:
+    input_df: "${raw_data_dir}/labs.parquet"
+    ts_col: "timestamp"
+    ts_format: "%H:%M:%S-%Y-%m-%d"
+  medications:
+    input_df: "${raw_data_dir}/medications.parquet"
+    ts_col: "timestamp"
+    ts_format: "%H:%M:%S-%Y-%m-%d"
+    columns: {"name": "medication"}
+
+measurements:
+  static:
+    single_label_classification:
+      subjects: ["eye_color"]
+  functional_time_dependent:
+    age:
+      functor: AgeFunctor
+      necessary_static_measurements: { "dob": ["timestamp", "%m/%d/%Y"] }
+      kwargs: { dob_col: "dob" }
+  dynamic:
+    multi_label_classification:
+      admissions: ["department"]
+      medications:
+        - name: medication
+          modifiers: 
+            - [dose, "float"]
+            - [frequency, "categorical"]
+            - [duration, "categorical"]
+            - [generic_name, "categorical"]
+    univariate_regression:
+      vitals: ["HR", "temp"]
+    multivariate_regression:
+      labs: [["lab_name", "lab_value"]]
+
+outlier_detector_config:
+  stddev_cutoff: 1.5
+min_valid_vocab_element_observations: 5
+min_valid_column_observations: 5
+min_true_float_frequency: 0.1
+min_unique_numerical_observations: 20
+min_events_per_subject: 3
+agg_by_time_scale: "1h"

scripts/build_flat_reps.py

@@ -0,0 +1,45 @@
+#!/usr/bin/env python
+"""Builds a flat representation dataset given a hydra config file."""
+
+try:
+    import stackprinter
+
+    stackprinter.set_excepthook(style="darkbg2")
+except ImportError:
+    pass  # no need to fail because of missing dev dependency
+
+from pathlib import Path
+
+import hydra
+from omegaconf import DictConfig
+
+from EventStream.data.dataset_polars import Dataset
+
+
+@hydra.main(version_base=None, config_path="../configs", config_name="dataset_base")
+def main(cfg: DictConfig):
+    cfg = hydra.utils.instantiate(cfg, _convert_="all")
+    save_dir = Path(cfg.pop("save_dir"))
+    window_sizes = cfg.pop("window_sizes")
+    subjects_per_output_file = (
+        cfg.pop("subjects_per_output_file") if "subjects_per_output_file" in cfg else None
+    )
+
+    # Build flat reps for specified task and window sizes
+    ESD = Dataset.load(save_dir)
+    feature_inclusion_frequency, include_only_measurements = ESD._resolve_flat_rep_cache_params(
+        feature_inclusion_frequency=None, include_only_measurements=None
+    )
+    cache_kwargs = dict(
+        subjects_per_output_file=subjects_per_output_file,
+        feature_inclusion_frequency=feature_inclusion_frequency,  # 0.1
+        window_sizes=window_sizes,
+        include_only_measurements=include_only_measurements,
+        do_overwrite=False,
+        do_update=True,
+    )
+    ESD.cache_flat_representation(**cache_kwargs)
+
+
+if __name__ == "__main__":
+    main()

tests/data/test_pytorch_dataset.py

@@ -320,7 +320,7 @@ def setUp(self):
 
         shards_fp = self.path / "DL_shards.json"
         shards = {
-            f"{self.split}/0": list(set(DL_REP_DF["subject_id"].to_list())),
+            f"{self.split}/0": [str(x) for x in set(DL_REP_DF["subject_id"].to_list())],
         }
         shards_fp.write_text(json.dumps(shards))