feat: use timecopilot forecaster class in agent (#65)

Author: AzulGarza

tests/test_forecaster.py

@@ -22,6 +22,7 @@ def test_forecaster_forecast(models, freq, h):
     df = generate_series(n_series=n_uids, freq=freq, min_length=30)
     forecaster = TimeCopilotForecaster(models=models)
     fcst_df = forecaster.forecast(df=df, h=h, freq=freq)
+    assert len(fcst_df.columns) == 2 + len(models)
     assert len(fcst_df) == h * n_uids
     for model in models:
         assert model.alias in fcst_df.columns
@@ -45,10 +46,40 @@ def test_forecaster_cross_validation(models, freq, h, n_windows, step_size):
         n_windows=n_windows,
         step_size=step_size,
     )
+    assert len(fcst_df.columns) == 4 + len(models)
     uids = df["unique_id"].unique()
     for uid in uids:  # noqa: B007
         fcst_df_uid = fcst_df.query("unique_id == @uid")
         assert fcst_df_uid["cutoff"].nunique() == n_windows
         assert len(fcst_df_uid) == n_windows * h
     for model in models:
         assert model.alias in fcst_df.columns
+
+
+def test_forecaster_forecast_with_level(models):
+    n_uids = 3
+    level = [80, 90]
+    df = generate_series(n_series=n_uids, freq="D", min_length=30)
+    forecaster = TimeCopilotForecaster(models=models)
+    fcst_df = forecaster.forecast(df=df, h=2, freq="D", level=level)
+    assert len(fcst_df) == 2 * n_uids
+    assert len(fcst_df.columns) == 2 + len(models) * (1 + 2 * len(level))
+    for model in models:
+        assert model.alias in fcst_df.columns
+        for lv in level:
+            assert f"{model.alias}-lo-{lv}" in fcst_df.columns
+            assert f"{model.alias}-hi-{lv}" in fcst_df.columns
+
+
+def test_forecaster_forecast_with_quantiles(models):
+    n_uids = 3
+    quantiles = [0.1, 0.9]
+    df = generate_series(n_series=n_uids, freq="D", min_length=30)
+    forecaster = TimeCopilotForecaster(models=models)
+    fcst_df = forecaster.forecast(df=df, h=2, freq="D", quantiles=quantiles)
+    assert len(fcst_df) == 2 * n_uids
+    assert len(fcst_df.columns) == 2 + len(models) * (1 + len(quantiles))
+    for model in models:
+        assert model.alias in fcst_df.columns
+        for q in quantiles:
+            assert f"{model.alias}-q-{int(100 * q)}" in fcst_df.columns

timecopilot/agent.py

@@ -30,6 +30,7 @@
 )
 from tsfeatures.tsfeatures import _get_feats
 
+from .forecaster import TimeCopilotForecaster
 from .models.benchmarks import (
     ADIDA,
     IMAPA,
@@ -400,7 +401,6 @@ async def cross_validation_tool(
             ctx: RunContext[ExperimentDataset],
             models: list[str],
         ) -> str:
-            models_fcst_cv = None
             callable_models = []
             for str_model in models:
                 if str_model not in MODELS:
@@ -409,21 +409,14 @@ async def cross_validation_tool(
                         f"{', '.join(MODELS.keys())}"
                     )
                 callable_models.append(MODELS[str_model])
-            for model in callable_models:
-                fcst_cv = model.cross_validation(
-                    df=ctx.deps.df,
-                    h=ctx.deps.h,
-                    freq=ctx.deps.freq,
-                )
-                if models_fcst_cv is None:
-                    models_fcst_cv = fcst_cv
-                else:
-                    models_fcst_cv = models_fcst_cv.merge(
-                        fcst_cv.drop(columns=["y"]),
-                        on=["unique_id", "cutoff", "ds"],
-                    )
+            forecaster = TimeCopilotForecaster(models=callable_models)
+            fcst_cv = forecaster.cross_validation(
+                df=ctx.deps.df,
+                h=ctx.deps.h,
+                freq=ctx.deps.freq,
+            )
             eval_df = ctx.deps.evaluate_forecast_df(
-                forecast_df=models_fcst_cv,
+                forecast_df=fcst_cv,
                 models=[model.alias for model in callable_models],
             )
             eval_df = eval_df.groupby(
@@ -444,7 +437,8 @@ async def forecast_tool(
             model: str,
         ) -> str:
             callable_model = MODELS[model]
-            fcst_df = callable_model.forecast(
+            forecaster = TimeCopilotForecaster(models=[callable_model])
+            fcst_df = forecaster.forecast(
                 df=ctx.deps.df,
                 h=ctx.deps.h,
                 freq=ctx.deps.freq,

timecopilot/forecaster.py

@@ -14,14 +14,29 @@ def _call_models(
         df: pd.DataFrame,
         h: int,
         freq: str,
+        level: list[int] | None = None,
+        quantiles: list[float] | None = None,
         **kwargs,
     ) -> pd.DataFrame:
         res_df: pd.DataFrame | None = None
         for model in self.models:
-            res_df_model = getattr(model, attr)(df=df, h=h, freq=freq, **kwargs)
+            res_df_model = getattr(model, attr)(
+                df=df,
+                h=h,
+                freq=freq,
+                level=level,
+                quantiles=quantiles,
+                **kwargs,
+            )
             if res_df is None:
                 res_df = res_df_model
             else:
+                if "y" in res_df_model:
+                    # drop y to avoid duplicate columns
+                    # y was added by the previous condition
+                    # to cross validation
+                    # (the initial model)
+                    res_df_model = res_df_model.drop(columns=["y"])
                 res_df = res_df.merge(
                     res_df_model,
                     on=merge_on,
@@ -33,13 +48,17 @@ def forecast(
         df: pd.DataFrame,
         h: int,
         freq: str,
+        level: list[int] | None = None,
+        quantiles: list[float] | None = None,
     ) -> pd.DataFrame:
         return self._call_models(
             "forecast",
             merge_on=["unique_id", "ds"],
             df=df,
             h=h,
             freq=freq,
+            level=level,
+            quantiles=quantiles,
         )
 
     def cross_validation(
@@ -49,6 +68,8 @@ def cross_validation(
         freq: str,
         n_windows: int = 1,
         step_size: int | None = None,
+        level: list[int] | None = None,
+        quantiles: list[float] | None = None,
     ) -> pd.DataFrame:
         return self._call_models(
             "cross_validation",
@@ -58,4 +79,6 @@ def cross_validation(
             freq=freq,
             n_windows=n_windows,
             step_size=step_size,
+            level=level,
+            quantiles=quantiles,
         )