Pipeline setup step 2: Fill forecast.py

iup/eval.py

@@ -30,9 +30,9 @@ def check_date_match(data: IncidentUptakeData, pred: PointForecast):
     (data["time_end"] == pred["time_end"]).all()
 
     # 2. There should not be any duplicated date in either data or prediction.
-    assert not (any(data["time_end"].is_duplicated())), (
-        "Duplicated dates are found in data and prediction."
-    )
+    assert not (
+        any(data["time_end"].is_duplicated())
+    ), "Duplicated dates are found in data and prediction."
 
 
 def score(

scripts/eval.py

@@ -1,6 +1,8 @@
-                for score_fun in score_funs:
-                    score = eval.score(
-                        incident_test_data, incident_projections, score_fun
-                    )
-                    print(f"{model=} {forecast_date=} {score_fun=} {score=}")
-                    # save these scores somewhere
+import iup.eval
+
+
+def function_name(incident_test_data, incident_projections, score_funs):
+    for score_fun in score_funs:
+        score = iup.eval.score(incident_test_data, incident_projections, score_fun)
+        # save these scores somewhere
+        print(score)

scripts/forecast.py

@@ -3,7 +3,10 @@
 import polars as pl
 import yaml
 
-def run_all_forecasts() -> pl.DataFrame:
+import iup.models
+
+
+def run_all_forecasts(clean_data, config) -> pl.DataFrame:
     """Run all forecasts
 
     Returns:
@@ -18,53 +21,77 @@ def run_all_forecasts() -> pl.DataFrame:
     models = [getattr(iup.models, model_name) for model_name in config["models"]]
     assert all(issubclass(model, iup.models.UptakeModel) for model in models)
 
+    all_forecast = pl.DataFrame()
+
     for model in models:
-            for forecast_date in forecast_dates:
-                # Get data available as of the forecast date
+        for forecast_date in forecast_dates:
+            # Get data available as of the forecast date
+            forecast = run_forecast(
+                model,
+                clean_data,
+                grouping_factors=config["groups"],
+                forecast_start=forecast_date,
+                forecast_end=config["timeframe"]["end"],
+            )
+
+            forecast = forecast.with_columns(
+                forecast_start=forecast_date,
+                forecast_end=config["timeframe"]["end"],
+                model=pl.lit(model.__name__),
+            )
+
+            all_forecast = pl.concat([all_forecast, forecast])
+
+    return all_forecast
+
+
+def run_forecast(
+    model,
+    observed_data,
+    grouping_factors,
+    forecast_start,
+    forecast_end,
+) -> pl.DataFrame:
+    """Run a single model for a single forecast date"""
 
+    # preprocess.py returns cumulative data, need to convert to incidence for LinearIncidentUptakeModel #
+    incident_data = iup.CumulativeUptakeData(observed_data).to_incident(
+        grouping_factors
+    )
 
-def run_forecast() -> pl.DataFrame:
-    """Run a single model for a single forecast date"""
     incident_train_data = iup.IncidentUptakeData(
-        iup.IncidentUptakeData.split_train_test(
-            incident_data, config["timeframe"]["start"], "train"
-        )
+        iup.IncidentUptakeData.split_train_test(incident_data, forecast_start, "train")
     )
 
     # Fit models using the training data and make projections
     fit_model = model().fit(incident_train_data, grouping_factors)
 
     cumulative_projections = fit_model.predict(
-        config["timeframe"]["start"],
-        config["timeframe"]["end"],
+        forecast_start,
+        forecast_end,
         config["timeframe"]["interval"],
         grouping_factors,
     )
-    # save these projections somewhere
-
-    incident_projections = cumulative_projections.to_incident(
-        grouping_factors
-    )
-    # save these projections somewhere
 
-    # Evaluation / Post-processing --------------------------------------------
+    incident_projections = cumulative_projections.to_incident(grouping_factors)
 
-    incident_test_data = iup.IncidentUptakeData(
-        iup.IncidentUptakeData.split_train_test(
-            incident_data, config["timeframe"]["start"], "test"
-        )
-    ).filter(pl.col("date") <= config["timeframe"]["end"])
+    # Note that here returns incident projections only, for evaluation
+    return incident_projections
 
 
 if __name__ == "__main__":
     p = argparse.ArgumentParser()
     p.add_argument("--config", help="config file", default="scripts/config.yaml")
     p.add_argument("--input", help="input data")
+    p.add_argument("--output", help="output parquet file")
     args = p.parse_args()
 
     with open(args.config, "r") as f:
         config = yaml.safe_load(f)
 
-    input_data = pl.scan_parquet(args.input)
+    input_data = pl.scan_parquet(args.input).collect()
+
+    input_data = iup.CumulativeUptakeData(input_data)
 
-    run_all_forecasts(config=config, cache=args.cache)
+    all_forecast = run_all_forecasts(config=config, clean_data=input_data)
+    all_forecast.write_parquet(args.output)