Repository: mmcdermott/MEDS_Tabular_AutoML

Issue:
Updating path inputs and outputs to be consistent with MEDS v0.3 (`final_cohort` -> `data`, store all MEDS-Tab outputs in an output cohort directory instead of the input MEDS cohort directory so we don't overwrite MEDS cohort metadata).

Discussion:

Patch:
diff --git a/src/MEDS_tabular_automl/configs/default.yaml b/src/MEDS_tabular_automl/configs/default.yaml
index 8f8513c..82a2164 100644
--- a/src/MEDS_tabular_automl/configs/default.yaml
+++ b/src/MEDS_tabular_automl/configs/default.yaml
@@ -1,11 +1,13 @@
 MEDS_cohort_dir: ???
+output_cohort_dir: ???
 do_overwrite: False
 seed: 1
 tqdm: False
 worker: 0
 loguru_init: False
 
-log_dir: ${output_dir}/.logs/
+log_dir: ${output_cohort_dir}/.logs/
+cache_dir: ${output_cohort_dir}/.cache
 
 hydra:
   verbose: False
diff --git a/src/MEDS_tabular_automl/configs/describe_codes.yaml b/src/MEDS_tabular_automl/configs/describe_codes.yaml
index d171513..8d0aac3 100644
--- a/src/MEDS_tabular_automl/configs/describe_codes.yaml
+++ b/src/MEDS_tabular_automl/configs/describe_codes.yaml
@@ -2,13 +2,8 @@ defaults:
   - default
   - _self_
 
-# split we wish to get metadata for
-split: train
-# Raw data, must have a subdirectory "train" with the training data split
-input_dir: ${MEDS_cohort_dir}/final_cohort/${split}
+input_dir: ${MEDS_cohort_dir}/data
 # Where to store output code frequency data
-cache_dir: ${MEDS_cohort_dir}/.cache
-output_dir: ${MEDS_cohort_dir}
-output_filepath: ${output_dir}/code_metadata.parquet
+output_filepath: ${output_cohort_dir}/metadata/codes.parquet
 
 name: describe_codes
diff --git a/src/MEDS_tabular_automl/configs/launch_xgboost.yaml b/src/MEDS_tabular_automl/configs/launch_xgboost.yaml
index 3dce8bc..93118e5 100644
--- a/src/MEDS_tabular_automl/configs/launch_xgboost.yaml
+++ b/src/MEDS_tabular_automl/configs/launch_xgboost.yaml
@@ -9,13 +9,12 @@ defaults:
 task_name: task
 
 # Task cached data dir
-input_dir: ${MEDS_cohort_dir}/${task_name}/task_cache
+input_dir: ${output_cohort_dir}/${task_name}/task_cache
 # Directory with task labels
-input_label_dir: ${MEDS_cohort_dir}/${task_name}/labels/final_cohort
+input_label_dir: ${MEDS_cohort_dir}/${task_name}/labels/
 # Where to output the model and cached data
-output_dir: ${MEDS_cohort_dir}/model/model_${now:%Y-%m-%d_%H-%M-%S}
-output_filepath: ${output_dir}/model_metadata.parquet
-cache_dir: ${MEDS_cohort_dir}/.cache
+model_dir: ${output_cohort_dir}/model/model_${now:%Y-%m-%d_%H-%M-%S}
+output_filepath: ${model_dir}/model_metadata.json
 
 # Model parameters
 model_params:
@@ -31,13 +30,9 @@ model_params:
     keep_data_in_memory: True
     binarize_task: True
 
-hydra:
-  verbose: False
-  sweep:
-    dir: ${output_dir}/.logs/
-  run:
-    dir: ${output_dir}/.logs/
+log_dir: ${model_dir}/.logs/
 
+hydra:
   # Optuna Sweeper
   sweeper:
     sampler:
diff --git a/src/MEDS_tabular_automl/configs/tabularization.yaml b/src/MEDS_tabular_automl/configs/tabularization.yaml
index dd40e3f..ca2c4cb 100644
--- a/src/MEDS_tabular_automl/configs/tabularization.yaml
+++ b/src/MEDS_tabular_automl/configs/tabularization.yaml
@@ -5,8 +5,8 @@ defaults:
 
 # Raw data
 # Where the code metadata is stored
-input_code_metadata_fp: ${MEDS_cohort_dir}/code_metadata.parquet
-input_dir: ${MEDS_cohort_dir}/final_cohort
-output_dir: ${MEDS_cohort_dir}/tabularize
+input_code_metadata_fp: ${output_cohort_dir}/metadata/codes.parquet
+input_dir: ${MEDS_cohort_dir}/data
+output_dir: ${output_cohort_dir}/tabularize
 
 name: tabularization
diff --git a/src/MEDS_tabular_automl/configs/tabularization/default.yaml b/src/MEDS_tabular_automl/configs/tabularization/default.yaml
index d11dd62..3f8761c 100644
--- a/src/MEDS_tabular_automl/configs/tabularization/default.yaml
+++ b/src/MEDS_tabular_automl/configs/tabularization/default.yaml
@@ -1,7 +1,7 @@
 # User inputs
 allowed_codes: null
 min_code_inclusion_frequency: 10
-filtered_code_metadata_fp: ${MEDS_cohort_dir}/tabularized_code_metadata.parquet
+filtered_code_metadata_fp: ${output_cohort_dir}/tabularized_code_metadata.parquet
 window_sizes:
   - "1d"
   - "7d"
diff --git a/src/MEDS_tabular_automl/configs/task_specific_caching.yaml b/src/MEDS_tabular_automl/configs/task_specific_caching.yaml
index eb1c98e..ad8ed99 100644
--- a/src/MEDS_tabular_automl/configs/task_specific_caching.yaml
+++ b/src/MEDS_tabular_automl/configs/task_specific_caching.yaml
@@ -5,10 +5,10 @@ defaults:
 task_name: task
 
 # Tabularized Data
-input_dir: ${MEDS_cohort_dir}/tabularize
+input_dir: ${output_cohort_dir}/tabularize
 # Where the labels are stored, with columns patient_id, timestamp, label
-input_label_dir: ${MEDS_cohort_dir}/${task_name}/labels/final_cohort
+input_label_dir: ${MEDS_cohort_dir}/${task_name}/labels
 # Where to output the task specific tabularized data
-output_dir: ${MEDS_cohort_dir}/${task_name}/task_cache
+output_dir: ${output_cohort_dir}/${task_name}/task_cache
 
 name: task_specific_caching
diff --git a/src/MEDS_tabular_automl/scripts/describe_codes.py b/src/MEDS_tabular_automl/scripts/describe_codes.py
index fdee111..c29b542 100644
--- a/src/MEDS_tabular_automl/scripts/describe_codes.py
+++ b/src/MEDS_tabular_automl/scripts/describe_codes.py
@@ -8,7 +8,7 @@ import hydra
 import numpy as np
 import polars as pl
 from loguru import logger
-from omegaconf import DictConfig, OmegaConf
+from omegaconf import DictConfig
 
 from ..describe_codes import (
     compute_feature_frequencies,
@@ -36,15 +36,6 @@ def main(cfg: DictConfig):
     if not cfg.loguru_init:
         hydra_loguru_init()
 
-    # Store Config
-    output_dir = Path(cfg.output_dir)
-    output_dir.mkdir(exist_ok=True, parents=True)
-    OmegaConf.save(cfg, output_dir / "config.yaml")
-
-    # Create output dir
-    input_dir = Path(cfg.input_dir)
-    input_dir.mkdir(exist_ok=True, parents=True)
-
     # 0. Identify Output Columns and Frequencies
     logger.info("Iterating through shards and caching feature frequencies.")
 
diff --git a/src/MEDS_tabular_automl/scripts/launch_xgboost.py b/src/MEDS_tabular_automl/scripts/launch_xgboost.py
index 46233c2..fb88338 100644
--- a/src/MEDS_tabular_automl/scripts/launch_xgboost.py
+++ b/src/MEDS_tabular_automl/scripts/launch_xgboost.py
@@ -1,5 +1,4 @@
 from collections.abc import Callable, Mapping
-from datetime import datetime
 from importlib.resources import files
 from pathlib import Path
 
@@ -440,10 +439,10 @@ def main(cfg: DictConfig) -> float:
         # print("Held Out Iterator Time: \n", model.iheld_out._profile_durations())
 
         # save model
-        save_dir = Path(cfg.output_dir)
-        save_dir.mkdir(parents=True, exist_ok=True)
-        model_time = datetime.now().strftime("%H%M%S%f")
-        model.model.save_model(save_dir / f"{auc:.4f}_model_{model_time}.json")
+        output_fp = Path(cfg.output_filepath)
+        output_fp.parent.mkdir(parents=True, exist_ok=True)
+
+        model.model.save_model(output_fp)
     except Exception as e:
         logger.error(f"Error occurred: {e}")
         auc = 0.0


Test patch:
diff --git a/tests/test_integration.py b/tests/test_integration.py
index 3c0bee8..50174cd 100644
--- a/tests/test_integration.py
+++ b/tests/test_integration.py
@@ -45,10 +45,12 @@ def run_command(script: str, args: list[str], hydra_kwargs: dict[str, str], test
 def test_integration():
     # Step 0: Setup Environment
     with tempfile.TemporaryDirectory() as d:
-        MEDS_cohort_dir = Path(d) / "processed"
+        MEDS_cohort_dir = Path(d) / "MEDS_cohort_dir"
+        output_cohort_dir = Path(d) / "output_cohort_dir"
 
-        describe_codes_config = {
+        shared_config = {
             "MEDS_cohort_dir": str(MEDS_cohort_dir.resolve()),
+            "output_cohort_dir": str(output_cohort_dir.resolve()),
             "do_overwrite": False,
             "seed": 1,
             "hydra.verbose": True,
@@ -56,6 +58,8 @@ def test_integration():
             "loguru_init": True,
         }
 
+        describe_codes_config = {**shared_config}
+
         with initialize(
             version_base=None, config_path="../src/MEDS_tabular_automl/configs/"
         ):  # path to config.yaml
@@ -63,11 +67,11 @@ def test_integration():
             cfg = compose(config_name="describe_codes", overrides=overrides)  # config.yaml
 
         # Create the directories
-        (MEDS_cohort_dir / "final_cohort").mkdir(parents=True, exist_ok=True)
+        (MEDS_cohort_dir / "data").mkdir(parents=True, exist_ok=True)
 
         # Store MEDS outputs
         for split, data in MEDS_OUTPUTS.items():
-            file_path = MEDS_cohort_dir / "final_cohort" / f"{split}.parquet"
+            file_path = MEDS_cohort_dir / "data" / f"{split}.parquet"
             file_path.parent.mkdir(exist_ok=True)
             df = pl.read_csv(StringIO(data))
             df.with_columns(pl.col("timestamp").str.to_datetime("%Y-%m-%dT%H:%M:%S%.f")).write_parquet(
@@ -92,7 +96,6 @@ def test_integration():
             describe_codes_config,
             "describe_codes",
         )
-        assert (Path(cfg.output_dir) / "config.yaml").is_file()
         assert Path(cfg.output_filepath).is_file()
 
         feature_columns = get_feature_columns(cfg.output_filepath)
@@ -104,12 +107,7 @@ def test_integration():
 
         # Step 2: Run the static data tabularization script
         tabularize_config = {
-            "MEDS_cohort_dir": str(MEDS_cohort_dir.resolve()),
-            "do_overwrite": False,
-            "seed": 1,
-            "hydra.verbose": True,
-            "tqdm": False,
-            "loguru_init": True,
+            **shared_config,
             "tabularization.min_code_inclusion_frequency": 1,
             "tabularization.window_sizes": "[30d,365d,full]",
         }
@@ -158,12 +156,7 @@ def test_integration():
 
         # Step 3: Run the time series tabularization script
         tabularize_config = {
-            "MEDS_cohort_dir": str(MEDS_cohort_dir.resolve()),
-            "do_overwrite": False,
-            "seed": 1,
-            "hydra.verbose": True,
-            "tqdm": False,
-            "loguru_init": True,
+            **shared_config,
             "tabularization.min_code_inclusion_frequency": 1,
             "tabularization.window_sizes": "[30d,365d,full]",
         }
@@ -205,12 +198,7 @@ def test_integration():
             )
         # Step 4: Run the task_specific_caching script
         cache_config = {
-            "MEDS_cohort_dir": str(MEDS_cohort_dir.resolve()),
-            "do_overwrite": False,
-            "seed": 1,
-            "hydra.verbose": True,
-            "tqdm": False,
-            "loguru_init": True,
+            **shared_config,
             "tabularization.min_code_inclusion_frequency": 1,
             "tabularization.window_sizes": "[30d,365d,full]",
         }
@@ -220,7 +208,7 @@ def test_integration():
             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) / "final_cohort", "parquet"):
+        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]])
@@ -231,7 +219,7 @@ def test_integration():
             split = f.parent.stem
             shard_num = f.stem
             out_f = Path(cfg.input_label_dir) / Path(
-                get_shard_prefix(Path(cfg.MEDS_cohort_dir) / "final_cohort", f)
+                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)
diff --git a/tests/test_tabularize.py b/tests/test_tabularize.py
index 0a33409..0f48536 100644
--- a/tests/test_tabularize.py
+++ b/tests/test_tabularize.py
@@ -148,10 +148,12 @@ EXPECTED_STATIC_FILES = [
 
 def test_tabularize():
     with tempfile.TemporaryDirectory() as d:
-        MEDS_cohort_dir = Path(d) / "processed"
+        MEDS_cohort_dir = Path(d) / "MEDS_cohort_dir"
+        output_cohort_dir = Path(d) / "output_cohort_dir"
 
-        describe_codes_config = {
+        shared_config = {
             "MEDS_cohort_dir": str(MEDS_cohort_dir.resolve()),
+            "output_cohort_dir": str(output_cohort_dir.resolve()),
             "do_overwrite": False,
             "seed": 1,
             "hydra.verbose": True,
@@ -159,6 +161,8 @@ def test_tabularize():
             "loguru_init": True,
         }
 
+        describe_codes_config = {**shared_config}
+
         with initialize(
             version_base=None, config_path="../src/MEDS_tabular_automl/configs/"
         ):  # path to config.yaml
@@ -166,11 +170,11 @@ def test_tabularize():
             cfg = compose(config_name="describe_codes", overrides=overrides)  # config.yaml
 
         # Create the directories
-        (MEDS_cohort_dir / "final_cohort").mkdir(parents=True, exist_ok=True)
+        (MEDS_cohort_dir / "data").mkdir(parents=True, exist_ok=True)
 
         # Store MEDS outputs
         for split, data in MEDS_OUTPUTS.items():
-            file_path = MEDS_cohort_dir / "final_cohort" / f"{split}.parquet"
+            file_path = MEDS_cohort_dir / "data" / f"{split}.parquet"
             file_path.parent.mkdir(exist_ok=True)
             df = pl.read_csv(StringIO(data))
             df.with_columns(pl.col("timestamp").str.to_datetime("%Y-%m-%dT%H:%M:%S%.f")).write_parquet(
@@ -190,7 +194,6 @@ def test_tabularize():
         # Step 1: Describe Codes - compute code frequencies
         describe_codes.main(cfg)
 
-        assert (Path(cfg.output_dir) / "config.yaml").is_file()
         assert Path(cfg.output_filepath).is_file()
 
         feature_columns = get_feature_columns(cfg.output_filepath)
@@ -202,12 +205,7 @@ def test_tabularize():
 
         # Step 2: Tabularization
         tabularize_static_config = {
-            "MEDS_cohort_dir": str(MEDS_cohort_dir.resolve()),
-            "do_overwrite": False,
-            "seed": 1,
-            "hydra.verbose": True,
-            "tqdm": False,
-            "loguru_init": True,
+            **shared_config,
             "tabularization.min_code_inclusion_frequency": 1,
             "tabularization.window_sizes": "[30d,365d,full]",
         }
@@ -218,8 +216,11 @@ def test_tabularize():
             overrides = [f"{k}={v}" for k, v in tabularize_static_config.items()]
             cfg = compose(config_name="tabularization", overrides=overrides)  # config.yaml
         tabularize_static.main(cfg)
-        output_files = list(Path(cfg.output_dir).glob("**/static/**/*.npz"))
-        actual_files = [get_shard_prefix(Path(cfg.output_dir), each) + ".npz" for each in output_files]
+
+        output_dir = Path(cfg.output_cohort_dir) / "tabularize"
+
+        output_files = list(output_dir.glob("**/static/**/*.npz"))
+        actual_files = [get_shard_prefix(output_dir, each) + ".npz" for each in output_files]
         assert set(actual_files) == set(EXPECTED_STATIC_FILES)
         # Check the files are not empty
         for f in output_files:
@@ -252,9 +253,9 @@ def test_tabularize():
         tabularize_time_series.main(cfg)
 
         # confirm summary files exist:
-        output_files = list_subdir_files(cfg.output_dir, "npz")
+        output_files = list_subdir_files(str(output_dir.resolve()), "npz")
         actual_files = [
-            get_shard_prefix(Path(cfg.output_dir), each) + ".npz"
+            get_shard_prefix(output_dir, each) + ".npz"
             for each in output_files
             if "none/static" not in str(each)
         ]
@@ -282,12 +283,7 @@ def test_tabularize():
 
         # Step 3: Cache Task data
         cache_config = {
-            "MEDS_cohort_dir": str(MEDS_cohort_dir.resolve()),
-            "do_overwrite": False,
-            "seed": 1,
-            "hydra.verbose": True,
-            "tqdm": False,
-            "loguru_init": True,
+            **shared_config,
             "tabularization.min_code_inclusion_frequency": 1,
             "tabularization.window_sizes": "[30d,365d,full]",
         }
@@ -299,7 +295,7 @@ 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) / "final_cohort", "parquet"):
+        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]])
@@ -310,7 +306,7 @@ def test_tabularize():
             split = f.parent.stem
             shard_num = f.stem
             out_f = Path(cfg.input_label_dir) / Path(
-                get_shard_prefix(Path(cfg.MEDS_cohort_dir) / "final_cohort", f)
+                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)
@@ -318,12 +314,7 @@ def test_tabularize():
         cache_task.main(cfg)
 
         xgboost_config_kwargs = {
-            "MEDS_cohort_dir": str(MEDS_cohort_dir.resolve()),
-            "do_overwrite": False,
-            "seed": 1,
-            "hydra.verbose": True,
-            "tqdm": False,
-            "loguru_init": True,
+            **shared_config,
             "tabularization.min_code_inclusion_frequency": 1,
             "tabularization.window_sizes": "[30d,365d,full]",
         }
@@ -334,8 +325,10 @@ def test_tabularize():
             overrides = [f"{k}={v}" for k, v in xgboost_config_kwargs.items()]
             cfg = compose(config_name="launch_xgboost", overrides=overrides)  # config.yaml
 
+        output_dir = Path(cfg.output_cohort_dir) / "model"
+
         launch_xgboost.main(cfg)
-        output_files = list(Path(cfg.output_dir).glob("**/*.json"))
+        output_files = list(output_dir.glob("**/*.json"))
         assert len(output_files) == 1
 
 
@@ -359,6 +352,7 @@ def test_xgboost_config():
     )
     xgboost_config_kwargs = {
         "MEDS_cohort_dir": MEDS_cohort_dir,
+        "output_cohort_dir": "blah",
         "do_overwrite": False,
         "seed": 1,
         "hydra.verbose": True,
@@ -373,5 +367,4 @@ def test_xgboost_config():
     ):  # path to config.yaml
         overrides = [f"{k}={v}" for k, v in xgboost_config_kwargs.items()]
         cfg = compose(config_name="launch_xgboost", overrides=overrides)  # config.yaml
-    print(cfg.tabularization.window_sizes)
     assert cfg.tabularization.window_sizes