Adds a multi-stage integration test for pre-processing.

src/MEDS_transforms/configs/stage_configs/fit_vocabulary_indices.yaml

@@ -1,4 +1,3 @@
 fit_vocabulary_indices:
   is_metadata: true
   ordering_method: "lexicographic"
-  output_dir: "${cohort_dir}"

src/MEDS_transforms/transforms/tensorization.py

@@ -5,6 +5,7 @@
 
 import hydra
 import polars as pl
+from loguru import logger
 from nested_ragged_tensors.ragged_numpy import JointNestedRaggedTensorDict
 from omegaconf import DictConfig
 
@@ -87,9 +88,17 @@ def convert_to_NRT(df: pl.LazyFrame) -> JointNestedRaggedTensorDict:
 
     time_delta_col = time_delta_cols[0]
 
-    return JointNestedRaggedTensorDict(
-        df.select(time_delta_col, "code", "numeric_value").collect().to_dict(as_series=False)
-    )
+    tensors_dict = df.select(time_delta_col, "code", "numeric_value").collect().to_dict(as_series=False)
+
+    if all((not v) for v in tensors_dict.values()):
+        logger.warning("All columns are empty. Returning an empty tensor dict.")
+        return JointNestedRaggedTensorDict({})
+
+    for k, v in tensors_dict.items():
+        if not v:
+            raise ValueError(f"Column {k} is empty")
+
+    return JointNestedRaggedTensorDict(tensors_dict)
 
 
 @hydra.main(

src/MEDS_transforms/transforms/tokenization.py

@@ -184,27 +184,27 @@ def extract_seq_of_patient_events(df: pl.LazyFrame) -> pl.LazyFrame:
         ...         None, datetime(2021, 1, 1), datetime(2021, 1, 1), datetime(2021, 1, 13),
         ...         None, datetime(2021, 1, 2), datetime(2021, 1, 2)],
         ...     "code": [100, 101, 102, 103, 200, 201, 202],
-        ...     "numeric_value": [1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0]
+        ...     "numeric_value": pl.Series([1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0], dtype=pl.Float32)
         ... }).lazy()
         >>> extract_seq_of_patient_events(df).collect()
         shape: (2, 4)
-        ┌────────────┬─────────────────┬───────────────────────────┬─────────────────────┐
-        │ patient_id ┆ time_delta_days ┆ code                      ┆ numeric_value       │
-        │ ---        ┆ ---             ┆ ---                       ┆ ---                 │
-        │ i64        ┆ list[f64]       ┆ list[list[f64]]           ┆ list[list[f64]]     │
-        ╞════════════╪═════════════════╪═══════════════════════════╪═════════════════════╡
-        │ 1          ┆ [NaN, 12.0]     ┆ [[101.0, 102.0], [103.0]] ┆ [[2.0, 3.0], [4.0]] │
-        │ 2          ┆ [NaN]           ┆ [[201.0, 202.0]]          ┆ [[6.0, 7.0]]        │
-        └────────────┴─────────────────┴───────────────────────────┴─────────────────────┘
+        ┌────────────┬─────────────────┬─────────────────────┬─────────────────────┐
+        │ patient_id ┆ time_delta_days ┆ code                ┆ numeric_value       │
+        │ ---        ┆ ---             ┆ ---                 ┆ ---                 │
+        │ i64        ┆ list[f32]       ┆ list[list[i64]]     ┆ list[list[f32]]     │
+        ╞════════════╪═════════════════╪═════════════════════╪═════════════════════╡
+        │ 1          ┆ [NaN, 12.0]     ┆ [[101, 102], [103]] ┆ [[2.0, 3.0], [4.0]] │
+        │ 2          ┆ [NaN]           ┆ [[201, 202]]        ┆ [[6.0, 7.0]]        │
+        └────────────┴─────────────────┴─────────────────────┴─────────────────────┘
     """
 
     _, dynamic = split_static_and_dynamic(df)
 
-    time_delta_days_expr = (pl.col("time").diff().dt.total_seconds() / SECONDS_PER_DAY).cast(pl.Float64)
+    time_delta_days_expr = (pl.col("time").diff().dt.total_seconds() / SECONDS_PER_DAY).cast(pl.Float32)
 
     return (
         dynamic.group_by("patient_id", "time", maintain_order=True)
-        .agg(fill_to_nans("code").name.keep(), fill_to_nans("numeric_value").name.keep())
+        .agg(pl.col("code").name.keep(), fill_to_nans("numeric_value").name.keep())
         .group_by("patient_id", maintain_order=True)
         .agg(
             fill_to_nans(time_delta_days_expr).alias("time_delta_days"),

src/MEDS_transforms/utils.py

@@ -244,7 +244,7 @@ def populate_stage(
          'reducer_output_dir': '/c/d/metadata'}
         >>> populate_stage("stage6", *args) # doctest: +NORMALIZE_WHITESPACE
         {'is_metadata': False, 'data_input_dir': '/c/d/stage4',
-         'metadata_input_dir': '/c/d/stage5', 'output_dir': '/c/d/data', 'reducer_output_dir': None}
+         'metadata_input_dir': '/c/d/metadata', 'output_dir': '/c/d/data', 'reducer_output_dir': None}
         >>> populate_stage("stage7", *args) # doctest: +NORMALIZE_WHITESPACE
         Traceback (most recent call last):
             ...
@@ -304,7 +304,7 @@ def populate_stage(
     if is_first_metadata_stage:
         default_metadata_input_dir = pipeline_input_metadata_dir
     else:
-        default_metadata_input_dir = prior_metadata_stage["output_dir"]
+        default_metadata_input_dir = prior_metadata_stage["reducer_output_dir"]
 
     # Now, we need to set output directories. The output directory for the stage will either be a stage
     # specific output directory, or, for the last data or metadata stages, respectively, will be the global

tests/test_add_time_derived_measurements.py

@@ -240,5 +240,5 @@ def test_add_time_derived_measurements():
         transform_script=ADD_TIME_DERIVED_MEASUREMENTS_SCRIPT,
         stage_name="add_time_derived_measurements",
         transform_stage_kwargs={"age": {"DOB_code": "DOB"}},
-        want_outputs=WANT_SHARDS,
+        want_data=WANT_SHARDS,
     )

tests/test_aggregate_code_metadata.py

@@ -180,7 +180,7 @@ def test_aggregate_code_metadata():
         transform_script=AGGREGATE_CODE_METADATA_SCRIPT,
         stage_name="aggregate_code_metadata",
         transform_stage_kwargs={"aggregations": AGGREGATIONS, "do_summarize_over_all_codes": True},
-        want_outputs=WANT_OUTPUT_CODE_METADATA_FILE,
-        code_metadata=MEDS_CODE_METADATA_FILE,
+        want_metadata=WANT_OUTPUT_CODE_METADATA_FILE,
+        input_code_metadata=MEDS_CODE_METADATA_FILE,
         do_use_config_yaml=True,
     )

tests/test_filter_measurements.py

@@ -113,7 +113,7 @@ def test_filter_measurements():
         transform_script=FILTER_MEASUREMENTS_SCRIPT,
         stage_name="filter_measurements",
         transform_stage_kwargs={"min_patients_per_code": 2},
-        want_outputs=WANT_SHARDS,
+        want_data=WANT_SHARDS,
     )
 
 
@@ -223,6 +223,6 @@ def test_match_revise_filter_measurements():
                 {"_matcher": {"code": "EYE_COLOR//HAZEL"}, "min_patients_per_code": 4},
             ],
         },
-        want_outputs=MR_WANT_SHARDS,
+        want_data=MR_WANT_SHARDS,
         do_use_config_yaml=True,
     )

tests/test_filter_patients.py

@@ -82,5 +82,5 @@ def test_filter_patients():
         transform_script=FILTER_PATIENTS_SCRIPT,
         stage_name="filter_patients",
         transform_stage_kwargs={"min_events_per_patient": 5},
-        want_outputs=WANT_SHARDS,
+        want_data=WANT_SHARDS,
     )

tests/test_fit_vocabulary_indices.py

@@ -33,5 +33,5 @@ def test_fit_vocabulary_indices_with_default_stage_config():
         transform_script=FIT_VOCABULARY_INDICES_SCRIPT,
         stage_name="fit_vocabulary_indices",
         transform_stage_kwargs=None,
-        want_outputs=parse_code_metadata_csv(WANT_CSV),
+        want_metadata=parse_code_metadata_csv(WANT_CSV),
     )