fix handling of horizons for sarix and gbqr models

reichlab · Nov 19, 2024 · 06bcc51 · 06bcc51
1 parent a22b225
commit 06bcc51
Show file tree

Hide file tree

Showing 4 changed files with 956 additions and 956 deletions.
diff --git a/src/idmodels/gbqr.py b/src/idmodels/gbqr.py
@@ -247,7 +247,7 @@ def _format_as_flusight_output(self, preds_df, ref_date):
 
         preds_df["target_end_date"] = preds_df["wk_end_date"] + pd.to_timedelta(7*preds_df["horizon"], unit="days")
         preds_df["reference_date"] = ref_date
-        preds_df["horizon"] = preds_df["horizon"] - 2
+        preds_df["horizon"] = (pd.to_timedelta(preds_df["target_end_date"].dt.date - ref_date).dt.days / 7).astype(int)
         preds_df["target"] = "wk inc flu hosp"
 
         preds_df["output_type"] = "quantile"

diff --git a/src/idmodels/sarix.py b/src/idmodels/sarix.py
@@ -79,7 +79,7 @@ def run(self, run_config):
 
         preds_df["target_end_date"] = preds_df["wk_end_date"] + pd.to_timedelta(7*preds_df["horizon"], unit="days")
         preds_df["reference_date"] = run_config.ref_date
-        preds_df["horizon"] = preds_df["horizon"] - 2
+        preds_df["horizon"] = (pd.to_timedelta(preds_df["target_end_date"].dt.date - run_config.ref_date).dt.days / 7).astype(int)
         preds_df["output_type"] = "quantile"
         preds_df["target"] = "wk inc flu hosp"
         preds_df.drop(columns="wk_end_date", inplace=True)