feat: reuse shared panel processing across forecasters (refresh of #2)

foundationforecast/core/__init__.py

@@ -8,14 +8,22 @@
 )
 from .gluonts_forecaster import GluonTSForecaster
 from .multi_model import MultiModelForecasterMixin
-from .utils import TimeSeriesDataset
+from .utils import (
+    PanelData,
+    TimeSeriesDataset,
+    grouped_std_by_id,
+    process_panel_from_df,
+)
 
 __all__ = [
     "Forecaster",
     "GluonTSForecaster",
     "MultiModelForecasterMixin",
+    "PanelData",
     "QuantileConverter",
     "TimeSeriesDataset",
+    "grouped_std_by_id",
+    "process_panel_from_df",
     "_DataProcessor",
     "get_seasonality",
     "maybe_convert_col_to_datetime",

foundationforecast/core/forecaster.py

@@ -16,13 +16,16 @@
 from tqdm import tqdm
 from utilsforecast.processing import (
     backtest_splits,
+    counts_by_id,
     drop_index_if_pandas,
     join,
     maybe_compute_sort_indices,
     take_rows,
     vertical_concat,
 )
 
+from .utils import PanelData, TimeSeriesDataset, grouped_std_by_id
+
 
 def get_seasonality(
     freq: str,
@@ -167,13 +170,41 @@ def _maybe_get_seasonality(self, freq: str) -> int:
             return get_seasonality(freq)
         return get_seasonality(freq)
 
+    @staticmethod
+    def _make_timeseries_dataset(
+        df: pd.DataFrame,
+        batch_size: int,
+        dtype: torch.dtype = torch.bfloat16,
+        panel: PanelData | None = None,
+    ) -> TimeSeriesDataset:
+        return TimeSeriesDataset.from_df(
+            df,
+            batch_size=batch_size,
+            dtype=dtype,
+            panel=panel,
+        )
+
+    @staticmethod
+    def _assign_quantile_forecasts(
+        fcst_df: pd.DataFrame,
+        alias: str,
+        quantiles: list[float],
+        fcsts_quantiles_np: np.ndarray,
+    ) -> pd.DataFrame:
+        q_cols = [f"{alias}-q-{int(q * 100)}" for q in quantiles]
+        q_vals = [fcsts_quantiles_np[..., i].reshape(-1) for i in range(len(quantiles))]
+        for q_col, q_val in zip(q_cols, q_vals, strict=True):
+            fcst_df = ufp.assign_columns(fcst_df, q_col, q_val)
+        return fcst_df
+
     def forecast(
         self,
         df: pd.DataFrame,
         h: int,
         freq: str | None = None,
         level: list[int | float] | None = None,
         quantiles: list[float] | None = None,
+        panel: PanelData | None = None,
     ) -> pd.DataFrame:
         raise NotImplementedError("This method must be implemented in a subclass.")
 
@@ -252,7 +283,7 @@ def detect_anomalies(
         df = maybe_convert_col_to_datetime(df, "ds")
         if h is None:
             h = self._maybe_get_seasonality(freq)
-        min_series_length = df.groupby("unique_id").size().min()
+        min_series_length = counts_by_id(df, "unique_id")["counts"].min()
         min_required = self._anomaly_min_series_length(h)
         reserved = min_required - h
         max_possible_windows = (min_series_length - reserved) // h
@@ -274,11 +305,12 @@ def detect_anomalies(
             step_size=h,
         )
         cv_results["residuals"] = cv_results["y"] - cv_results[self.alias]
-        residual_stats = (
-            cv_results.groupby("unique_id")["residuals"].std().reset_index()
+        residual_stats = grouped_std_by_id(
+            cv_results,
+            "unique_id",
+            "residuals",
         )
-        residual_stats.columns = ["unique_id", "residual_std"]
-        cv_results = cv_results.merge(residual_stats, on="unique_id", how="left")
+        cv_results = join(cv_results, residual_stats, on="unique_id", how="left")
         cv_results["z_score"] = cv_results["residuals"] / cv_results["residual_std"]
         alpha = 1 - level / 100
         critical_z = stats.norm.ppf(1 - alpha / 2)
@@ -338,8 +370,8 @@ def _prepare_level_and_quantiles(
         if level is None and quantiles is not None:
             if not all(0 < q < 1 for q in quantiles):
                 raise ValueError("`quantiles` should be floats between 0 and 1.")
-            level = [abs(int(100 - 200 * q)) for q in quantiles]
-            return sorted(set(level)), quantiles, False
+            level = sorted({abs(int(100 - 200 * q)) for q in quantiles if q != 0.5})
+            return level or None, quantiles, False
         return None, None, False
 
     @staticmethod

foundationforecast/core/gluonts_forecaster.py

@@ -2,16 +2,18 @@
 from contextlib import contextmanager
 from typing import Any
 
+import numpy as np
 import pandas as pd
 import torch
 import utilsforecast.processing as ufp
 from gluonts.dataset.pandas import PandasDataset
 from gluonts.model.forecast import Forecast
 from gluonts.torch.model.predictor import PyTorchPredictor
 from huggingface_hub import hf_hub_download
-from tqdm import tqdm
+from utilsforecast.processing import make_future_dataframe
 
 from .forecaster import Forecaster, QuantileConverter
+from .utils import PanelData, process_panel_from_df
 
 
 def fix_freq(freq: str) -> str:
@@ -84,12 +86,10 @@ def gluonts_instance_fcst_to_df(
             }
         )
         if quantiles is not None:
-            for q in quantiles:
-                fcst_df = ufp.assign_columns(
-                    fcst_df,
-                    f"{model_name}-q-{int(q * 100)}",
-                    fcst.quantile(q),
-                )
+            q_cols = [f"{model_name}-q-{int(q * 100)}" for q in quantiles]
+            q_vals = [fcst.quantile(q) for q in quantiles]
+            for q_col, q_val in zip(q_cols, q_vals, strict=True):
+                fcst_df = ufp.assign_columns(fcst_df, q_col, q_val)
         return fcst_df
 
     def gluonts_fcsts_to_df(
@@ -98,17 +98,35 @@ def gluonts_fcsts_to_df(
         freq: str,
         model_name: str,
         quantiles: list[float] | None,
+        h: int,
+        panel: PanelData,
     ) -> pd.DataFrame:
-        df = []
-        for fcst in tqdm(fcsts):
-            fcst_df = self.gluonts_instance_fcst_to_df(
-                fcst=fcst,
-                freq=freq,
-                model_name=model_name,
-                quantiles=quantiles,
-            )
-            df.append(fcst_df)
-        return pd.concat(df).reset_index(drop=True)
+        fcsts_list = list(fcsts)
+        if not fcsts_list:
+            return pd.DataFrame()
+        fcst_by_id = {fcst.item_id: fcst for fcst in fcsts_list}
+        ordered_fcsts = [fcst_by_id[uid] for uid in panel.uids]
+        point_fcsts = np.stack([fcst.median for fcst in ordered_fcsts])
+        fcst_df = make_future_dataframe(
+            uids=panel.uids,
+            last_times=pd.to_datetime(panel.last_times),
+            h=h,
+            freq=freq,
+        )
+        fcst_df = ufp.assign_columns(
+            fcst_df,
+            model_name,
+            point_fcsts.reshape(-1),
+        )
+        if quantiles is not None:
+            q_cols = [f"{model_name}-q-{int(q * 100)}" for q in quantiles]
+            q_vals = [
+                np.stack([fcst.quantile(q) for fcst in ordered_fcsts]).reshape(-1)
+                for q in quantiles
+            ]
+            for q_col, q_val in zip(q_cols, q_vals, strict=True):
+                fcst_df = ufp.assign_columns(fcst_df, q_col, q_val)
+        return fcst_df
 
     def forecast(
         self,
@@ -117,6 +135,7 @@ def forecast(
         freq: str | None = None,
         level: list[int | float] | None = None,
         quantiles: list[float] | None = None,
+        panel: PanelData | None = None,
     ) -> pd.DataFrame:
         """Generate forecasts for time series data using the model.
 
@@ -166,6 +185,8 @@ def forecast(
         df = maybe_convert_col_to_float32(df, "y")
         freq = self._maybe_infer_freq(df, freq)
         qc = QuantileConverter(level=level, quantiles=quantiles)
+        if panel is None:
+            panel = process_panel_from_df(df)
         gluonts_dataset = PandasDataset.from_long_dataframe(
             df.copy(deep=False),
             target="y",
@@ -183,6 +204,8 @@ def forecast(
             freq=freq,
             model_name=self.alias,
             quantiles=qc.quantiles,
+            h=h,
+            panel=panel,
         )
         if qc.quantiles is not None:
             fcst_df = qc.maybe_convert_quantiles_to_level(

foundationforecast/core/multi_model.py

@@ -1,8 +1,10 @@
 from __future__ import annotations
 
 import pandas as pd
+import utilsforecast.processing as ufp
 
 from .forecaster import Forecaster, maybe_infer_freq
+from .utils import PanelData, process_panel_from_df
 
 
 class MultiModelForecasterMixin:
@@ -46,10 +48,13 @@ def _call_models(
         freq: str | None,
         level: list[int | float] | None,
         quantiles: list[float] | None,
+        panel: PanelData | None = None,
         **kwargs,
     ) -> pd.DataFrame:
         Forecaster.validate_input(df, h)
         freq = maybe_infer_freq(df, freq)
+        if panel is None and attr == "forecast":
+            panel = process_panel_from_df(df)
         res_df: pd.DataFrame | None = None
         for model in self.models:
             known_kwargs = {
@@ -60,6 +65,8 @@ def _call_models(
             }
             if attr != "detect_anomalies":
                 known_kwargs["quantiles"] = quantiles
+            if panel is not None:
+                known_kwargs["panel"] = panel
             fn = getattr(model, attr)
             try:
                 res_df_model = fn(**known_kwargs, **kwargs)
@@ -83,7 +90,7 @@ def _call_models(
             else:
                 if "y" in res_df_model:
                     res_df_model = res_df_model.drop(columns=["y"])
-                res_df = res_df.merge(res_df_model, on=merge_on, how="left")
+                res_df = ufp.join(res_df, res_df_model, on=merge_on, how="left")
             if self.clean_cache:
                 self._clean_model_cache()
         if res_df is None:
@@ -97,6 +104,7 @@ def forecast(
         freq: str | None = None,
         level: list[int | float] | None = None,
         quantiles: list[float] | None = None,
+        panel: PanelData | None = None,
     ) -> pd.DataFrame:
         return self._call_models(
             "forecast",
@@ -106,6 +114,7 @@ def forecast(
             freq=freq,
             level=level,
             quantiles=quantiles,
+            panel=panel,
         )
 
     def cross_validation(