From 4dd3cadefb0d081f732b7a6999ceb8dfc8f4120f Mon Sep 17 00:00:00 2001 From: Matthew McDermott Date: Mon, 27 May 2024 09:02:32 -0400 Subject: [PATCH] Update src/MEDS_tabular_automl/generate_summarized_reps.py Co-authored-by: coderabbitai[bot] <136622811+coderabbitai[bot]@users.noreply.github.com> --- src/MEDS_tabular_automl/generate_summarized_reps.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/src/MEDS_tabular_automl/generate_summarized_reps.py b/src/MEDS_tabular_automl/generate_summarized_reps.py index 9888aa3..23db3dc 100644 --- a/src/MEDS_tabular_automl/generate_summarized_reps.py +++ b/src/MEDS_tabular_automl/generate_summarized_reps.py @@ -145,10 +145,10 @@ def _generate_summary(df: DF_T, window_size: str, agg: str) -> pl.LazyFrame: │ 2 ┆ 2021-01-04 ┆ 0 ┆ 1 │ └────────────┴────────────┴──────────────────┴──────────────────┘ """ - assert agg in VALID_AGGREGATIONS, f"Invalid aggregation: {agg}" - assert agg.split("/")[0] in [ - c.split("/")[0] for c in df.columns - ], f"df is invalid, no column with prefix: `{agg.split('/')[0]}`" + if agg not in VALID_AGGREGATIONS: + raise ValueError(f"Invalid aggregation: {agg}. Valid options are: {VALID_AGGREGATIONS}") + if agg.split("/")[0] not in [c.split("/")[0] for c in df.columns]: + raise ValueError(f"DataFrame is invalid, no column with prefix: `{agg.split('/')[0]}`") if window_size == "full": out_df = df.groupby("patient_id").agg(