Repository: mmcdermott/MEDS_Tabular_AutoML

Issue:
Updating how we handle patient splits to use `metadata/patient_splits.parquet`
Note that, given our usage of shard files to make XGBoost iteration easier and faster, the right way to handle this is to re-shard the data for tasks and MEDS data on input into split-separated shards. This is annoying, but should be an easy transformation to perform in a separate step which will make all the subsequent changes much easier to work with.
The way task dataframes are read in needs to be compliant with MEDS v0.3 datasets.
This entails:
  - [x] Updating column names.
  - [x] Updating tests to use proper names
  - [x] Enabling task dataframes to be differently sharded than data dataframes.
  - [x] Ensuring task dataframes do not need to have event IDs pre-computed.

Right now, event ID pre-computation and sharding matching happen in calling scripts, [here](https://github.com/mmcdermott/MEDS_TAB_MIMIC_IV/blob/76d9fc01f7f196dfeb35b8ff4be9d6529cfdc5a2/aces_task_extraction.py#L79):
```python
data_df = filter_parquet(in_fp, cfg.tabularization._resolved_codes)
data_df = get_unique_time_events_df(get_events_df(data_df, feature_columns))
data_df = data_df.drop(["code", "numerical_value"])
data_df = data_df.with_row_index("event_id")

output_df = label_df.lazy().join_asof(other=data_df, by="patient_id", on="timestamp")
```

Discussion:

Patch:
diff --git a/pyproject.toml b/pyproject.toml
index e82db50..b7972f9 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -16,7 +16,8 @@ classifiers = [
 ]
 dependencies = [
   "polars", "pyarrow", "loguru", "hydra-core", "numpy", "scipy<1.14.0", "pandas", "tqdm", "xgboost",
-  "scikit-learn", "hydra-optuna-sweeper", "hydra-joblib-launcher", "ml-mixins", "meds==0.3"
+  "scikit-learn", "hydra-optuna-sweeper", "hydra-joblib-launcher", "ml-mixins", "meds==0.3",
+  #"MEDS-transforms==0.0.4",
 ]
 
 [project.scripts]
diff --git a/src/MEDS_tabular_automl/configs/describe_codes.yaml b/src/MEDS_tabular_automl/configs/describe_codes.yaml
index 8d0aac3..ec980bf 100644
--- a/src/MEDS_tabular_automl/configs/describe_codes.yaml
+++ b/src/MEDS_tabular_automl/configs/describe_codes.yaml
@@ -2,7 +2,7 @@ defaults:
   - default
   - _self_
 
-input_dir: ${MEDS_cohort_dir}/data
+input_dir: ${output_cohort_dir}/data
 # Where to store output code frequency data
 output_filepath: ${output_cohort_dir}/metadata/codes.parquet
 
diff --git a/src/MEDS_tabular_automl/configs/launch_xgboost.yaml b/src/MEDS_tabular_automl/configs/launch_xgboost.yaml
index 93118e5..4fc1383 100644
--- a/src/MEDS_tabular_automl/configs/launch_xgboost.yaml
+++ b/src/MEDS_tabular_automl/configs/launch_xgboost.yaml
@@ -11,7 +11,7 @@ task_name: task
 # Task cached data dir
 input_dir: ${output_cohort_dir}/${task_name}/task_cache
 # Directory with task labels
-input_label_dir: ${MEDS_cohort_dir}/${task_name}/labels/
+input_label_dir: ${output_cohort_dir}/${task_name}/labels/
 # Where to output the model and cached data
 model_dir: ${output_cohort_dir}/model/model_${now:%Y-%m-%d_%H-%M-%S}
 output_filepath: ${model_dir}/model_metadata.json
diff --git a/src/MEDS_tabular_automl/configs/tabularization.yaml b/src/MEDS_tabular_automl/configs/tabularization.yaml
index ca2c4cb..cf03d63 100644
--- a/src/MEDS_tabular_automl/configs/tabularization.yaml
+++ b/src/MEDS_tabular_automl/configs/tabularization.yaml
@@ -6,7 +6,7 @@ defaults:
 # Raw data
 # Where the code metadata is stored
 input_code_metadata_fp: ${output_cohort_dir}/metadata/codes.parquet
-input_dir: ${MEDS_cohort_dir}/data
+input_dir: ${output_cohort_dir}/data
 output_dir: ${output_cohort_dir}/tabularize
 
 name: tabularization
diff --git a/src/MEDS_tabular_automl/configs/task_specific_caching.yaml b/src/MEDS_tabular_automl/configs/task_specific_caching.yaml
index ad8ed99..c002dfc 100644
--- a/src/MEDS_tabular_automl/configs/task_specific_caching.yaml
+++ b/src/MEDS_tabular_automl/configs/task_specific_caching.yaml
@@ -10,5 +10,8 @@ input_dir: ${output_cohort_dir}/tabularize
 input_label_dir: ${MEDS_cohort_dir}/${task_name}/labels
 # Where to output the task specific tabularized data
 output_dir: ${output_cohort_dir}/${task_name}/task_cache
+output_label_dir: ${output_cohort_dir}/${task_name}/labels
+
+label_column: "boolean_value"
 
 name: task_specific_caching
diff --git a/src/MEDS_tabular_automl/scripts/cache_task.py b/src/MEDS_tabular_automl/scripts/cache_task.py
index 62df9b8..15c194b 100644
--- a/src/MEDS_tabular_automl/scripts/cache_task.py
+++ b/src/MEDS_tabular_automl/scripts/cache_task.py
@@ -1,6 +1,7 @@
 #!/usr/bin/env python
 
 """Aggregates time-series data for feature columns across different window sizes."""
+from functools import partial
 from importlib.resources import files
 from pathlib import Path
 
@@ -10,6 +11,7 @@ import polars as pl
 import scipy.sparse as sp
 from omegaconf import DictConfig
 
+from ..describe_codes import filter_parquet, get_feature_columns
 from ..file_name import list_subdir_files
 from ..mapper import wrap as rwlock_wrap
 from ..utils import (
@@ -17,7 +19,9 @@ from ..utils import (
     STATIC_CODE_AGGREGATION,
     STATIC_VALUE_AGGREGATION,
     VALUE_AGGREGATIONS,
+    get_events_df,
     get_shard_prefix,
+    get_unique_time_events_df,
     hydra_loguru_init,
     load_matrix,
     load_tqdm,
@@ -37,6 +41,10 @@ VALID_AGGREGATIONS = [
 ]
 
 
+def write_lazyframe(df: pl.LazyFrame, fp: Path):
+    df.collect().write_parquet(fp, use_pyarrow=True)
+
+
 def generate_row_cached_matrix(matrix: sp.coo_array, label_df: pl.LazyFrame) -> sp.coo_array:
     """Generates row-cached matrix for a given matrix and label DataFrame.
 
@@ -80,31 +88,51 @@ def main(cfg: DictConfig):
     tabularization_tasks = list_subdir_files(cfg.input_dir, "npz")
     np.random.shuffle(tabularization_tasks)
 
+    label_dir = Path(cfg.input_label_dir)
+    label_df = pl.scan_parquet(label_dir / "**/*.parquet").rename(
+        {
+            "prediction_time": "time",
+            cfg.label_column: "label",
+        }
+    )
+
+    feature_columns = get_feature_columns(cfg.tabularization.filtered_code_metadata_fp)
+
     # iterate through them
     for data_fp in iter_wrapper(tabularization_tasks):
         # parse as time series agg
         split, shard_num, window_size, code_type, agg_name = Path(data_fp).with_suffix("").parts[-5:]
-        label_fp = Path(cfg.input_label_dir) / split / f"{shard_num}.parquet"
-        out_fp = (Path(cfg.output_dir) / get_shard_prefix(cfg.input_dir, data_fp)).with_suffix(".npz")
-        assert label_fp.exists(), f"Output file {label_fp} does not exist."
 
-        def read_fn(fps):
-            matrix_fp, label_fp = fps
-            return load_matrix(matrix_fp), pl.scan_parquet(label_fp)
+        raw_data_fp = Path(cfg.output_cohort_dir) / "data" / split / f"{shard_num}.parquet"
+        raw_data_df = filter_parquet(raw_data_fp, cfg.tabularization._resolved_codes)
+        raw_data_df = (
+            get_unique_time_events_df(get_events_df(raw_data_df, feature_columns))
+            .with_row_index("event_id")
+            .select("patient_id", "time", "event_id")
+        )
+        shard_label_df = label_df.join(
+            raw_data_df.select("patient_id").unique(), on="patient_id", how="inner"
+        ).join_asof(other=raw_data_df, by="patient_id", on="time")
 
-        def compute_fn(shard_dfs):
-            matrix, label_df = shard_dfs
-            cache_matrix = generate_row_cached_matrix(matrix, label_df)
-            return cache_matrix
+        shard_label_fp = Path(cfg.output_label_dir) / split / f"{shard_num}.parquet"
+        rwlock_wrap(
+            raw_data_fp,
+            shard_label_fp,
+            pl.scan_parquet,
+            write_lazyframe,
+            lambda df: shard_label_df,
+            do_overwrite=cfg.do_overwrite,
+            do_return=False,
+        )
 
-        def write_fn(cache_matrix, out_fp):
-            write_df(cache_matrix, out_fp, do_overwrite=cfg.do_overwrite)
+        out_fp = (Path(cfg.output_dir) / get_shard_prefix(cfg.input_dir, data_fp)).with_suffix(".npz")
+        compute_fn = partial(generate_row_cached_matrix, label_df=shard_label_df)
+        write_fn = partial(write_df, do_overwrite=cfg.do_overwrite)
 
-        in_fps = [data_fp, label_fp]
         rwlock_wrap(
-            in_fps,
+            data_fp,
             out_fp,
-            read_fn,
+            load_matrix,
             write_fn,
             compute_fn,
             do_overwrite=cfg.do_overwrite,


Test patch:
diff --git a/tests/test_integration.py b/tests/test_integration.py
index 7517903..d22eac5 100644
--- a/tests/test_integration.py
+++ b/tests/test_integration.py
@@ -67,14 +67,20 @@ def test_integration():
             cfg = compose(config_name="describe_codes", overrides=overrides)  # config.yaml
 
         # Create the directories
-        (MEDS_cohort_dir / "data").mkdir(parents=True, exist_ok=True)
+        (output_cohort_dir / "data").mkdir(parents=True, exist_ok=True)
 
         # Store MEDS outputs
+        all_data = []
         for split, data in MEDS_OUTPUTS.items():
-            file_path = MEDS_cohort_dir / "data" / f"{split}.parquet"
+            file_path = output_cohort_dir / "data" / f"{split}.parquet"
             file_path.parent.mkdir(exist_ok=True)
-            df = pl.read_csv(StringIO(data))
-            df.with_columns(pl.col("time").str.to_datetime("%Y-%m-%dT%H:%M:%S%.f")).write_parquet(file_path)
+            df = pl.read_csv(StringIO(data)).with_columns(
+                pl.col("time").str.to_datetime("%Y-%m-%dT%H:%M:%S%.f")
+            )
+            df.write_parquet(file_path)
+            all_data.append(df)
+
+        all_data = pl.concat(all_data, how="diagonal_relaxed").sort(by=["patient_id", "time"])
 
         # Check the files are not empty
         meds_files = list_subdir_files(Path(cfg.input_dir), "parquet")
@@ -84,7 +90,7 @@ def test_integration():
         for f in meds_files:
             assert pl.read_parquet(f).shape[0] > 0, "MEDS Data Tabular Dataframe Should not be Empty!"
         split_json = json.load(StringIO(SPLITS_JSON))
-        splits_fp = MEDS_cohort_dir / "splits.json"
+        splits_fp = output_cohort_dir / ".shards.json"
         json.dump(split_json, splits_fp.open("w"))
 
         # Step 1: Run the describe_codes script
@@ -205,22 +211,15 @@ def test_integration():
         ):  # path to config.yaml
             overrides = [f"{k}={v}" for k, v in cache_config.items()]
             cfg = compose(config_name="task_specific_caching", overrides=overrides)  # config.yaml
-        # Create fake labels
-        for f in list_subdir_files(Path(cfg.MEDS_cohort_dir) / "data", "parquet"):
-            df = pl.scan_parquet(f)
-            df = get_unique_time_events_df(get_events_df(df, feature_columns)).collect()
-            pseudo_labels = pl.Series(([0, 1] * df.shape[0])[: df.shape[0]])
-            df = df.with_columns(pl.Series(name="label", values=pseudo_labels))
-            df = df.select(pl.col(["patient_id", "time", "label"]))
-            df = df.with_row_index("event_id")
-
-            split = f.parent.stem
-            shard_num = f.stem
-            out_f = Path(cfg.input_label_dir) / Path(
-                get_shard_prefix(Path(cfg.MEDS_cohort_dir) / "data", f)
-            ).with_suffix(".parquet")
-            out_f.parent.mkdir(parents=True, exist_ok=True)
-            df.write_parquet(out_f)
+
+        df = get_unique_time_events_df(get_events_df(all_data.lazy(), feature_columns)).collect()
+        pseudo_labels = pl.Series(([0, 1] * df.shape[0])[: df.shape[0]])
+        df = df.with_columns(pl.Series(name="boolean_value", values=pseudo_labels))
+        df = df.select("patient_id", pl.col("time").alias("prediction_time"), "boolean_value")
+
+        out_fp = Path(cfg.input_label_dir) / "0.parquet"
+        out_fp.parent.mkdir(parents=True, exist_ok=True)
+        df.write_parquet(out_fp)
 
         stderr, stdout_ws = run_command("generate-subsets", ["[30d]"], {}, "generate-subsets window_sizes")
         stderr, stdout_agg = run_command(
diff --git a/tests/test_tabularize.py b/tests/test_tabularize.py
index c235f79..130721c 100644
--- a/tests/test_tabularize.py
+++ b/tests/test_tabularize.py
@@ -170,14 +170,20 @@ def test_tabularize():
             cfg = compose(config_name="describe_codes", overrides=overrides)  # config.yaml
 
         # Create the directories
-        (MEDS_cohort_dir / "data").mkdir(parents=True, exist_ok=True)
+        (output_cohort_dir / "data").mkdir(parents=True, exist_ok=True)
 
         # Store MEDS outputs
+        all_data = []
         for split, data in MEDS_OUTPUTS.items():
-            file_path = MEDS_cohort_dir / "data" / f"{split}.parquet"
+            file_path = output_cohort_dir / "data" / f"{split}.parquet"
             file_path.parent.mkdir(exist_ok=True)
-            df = pl.read_csv(StringIO(data))
-            df.with_columns(pl.col("time").str.to_datetime("%Y-%m-%dT%H:%M:%S%.f")).write_parquet(file_path)
+            df = pl.read_csv(StringIO(data)).with_columns(
+                pl.col("time").str.to_datetime("%Y-%m-%dT%H:%M:%S%.f")
+            )
+            df.write_parquet(file_path)
+            all_data.append(df)
+
+        all_data = pl.concat(all_data, how="diagonal_relaxed").sort(by=["patient_id", "time"])
 
         # Check the files are not empty
         meds_files = list_subdir_files(Path(cfg.input_dir), "parquet")
@@ -187,7 +193,7 @@ def test_tabularize():
         for f in meds_files:
             assert pl.read_parquet(f).shape[0] > 0, "MEDS Data Tabular Dataframe Should not be Empty!"
         split_json = json.load(StringIO(SPLITS_JSON))
-        splits_fp = MEDS_cohort_dir / "splits.json"
+        splits_fp = output_cohort_dir / ".shards.json"
         json.dump(split_json, splits_fp.open("w"))
         # Step 1: Describe Codes - compute code frequencies
         describe_codes.main(cfg)
@@ -293,21 +299,14 @@ def test_tabularize():
             cfg = compose(config_name="task_specific_caching", overrides=overrides)  # config.yaml
 
         # Create fake labels
-        for f in list_subdir_files(Path(cfg.MEDS_cohort_dir) / "data", "parquet"):
-            df = pl.scan_parquet(f)
-            df = get_unique_time_events_df(get_events_df(df, feature_columns)).collect()
-            pseudo_labels = pl.Series(([0, 1] * df.shape[0])[: df.shape[0]])
-            df = df.with_columns(pl.Series(name="label", values=pseudo_labels))
-            df = df.select(pl.col(["patient_id", "time", "label"]))
-            df = df.with_row_index("event_id")
-
-            split = f.parent.stem
-            shard_num = f.stem
-            out_f = Path(cfg.input_label_dir) / Path(
-                get_shard_prefix(Path(cfg.MEDS_cohort_dir) / "data", f)
-            ).with_suffix(".parquet")
-            out_f.parent.mkdir(parents=True, exist_ok=True)
-            df.write_parquet(out_f)
+        df = get_unique_time_events_df(get_events_df(all_data.lazy(), feature_columns)).collect()
+        pseudo_labels = pl.Series(([0, 1] * df.shape[0])[: df.shape[0]])
+        df = df.with_columns(pl.Series(name="boolean_value", values=pseudo_labels))
+        df = df.select("patient_id", pl.col("time").alias("prediction_time"), "boolean_value")
+
+        out_fp = Path(cfg.input_label_dir) / "0.parquet"
+        out_fp.parent.mkdir(parents=True, exist_ok=True)
+        df.write_parquet(out_fp)
 
         cache_task.main(cfg)