feat: use common syntaxis and improve speed

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,
@@ -141,6 +144,31 @@ 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))]
+        return ufp.assign_columns(fcst_df, q_cols, q_vals)
+
     def forecast(
         self,
         df: pd.DataFrame,
@@ -225,7 +253,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
@@ -247,11 +275,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)
@@ -311,8 +340,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,9 @@ 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]
+            fcst_df = ufp.assign_columns(fcst_df, q_cols, q_vals)
         return fcst_df
 
     def gluonts_fcsts_to_df(
@@ -98,17 +97,34 @@ 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
+            ]
+            fcst_df = ufp.assign_columns(fcst_df, q_cols, q_vals)
+        return fcst_df
 
     def forecast(
         self,
@@ -117,6 +133,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 +183,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 +202,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:
@@ -49,6 +51,9 @@ def _call_models(
         **kwargs,
     ) -> pd.DataFrame:
         freq = maybe_infer_freq(df, freq)
+        panel: PanelData | None = None
+        if attr == "forecast":
+            panel = process_panel_from_df(df)
         res_df: pd.DataFrame | None = None
         for model in self.models:
             known_kwargs = {
@@ -59,6 +64,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)
@@ -82,7 +89,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:

foundationforecast/core/utils.py

@@ -1,41 +1,91 @@
-from collections.abc import Iterable
+from __future__ import annotations
 
+from typing import NamedTuple
+
+import numpy as np
 import pandas as pd
 import torch
-from utilsforecast.processing import make_future_dataframe
+from utilsforecast.processing import group_by_agg, make_future_dataframe, process_df
+
+
+class PanelData(NamedTuple):
+    uids: pd.Series | np.ndarray
+    last_times: np.ndarray
+    series_arrays: list[np.ndarray]
+
+
+def process_panel_from_df(
+    df: pd.DataFrame,
+    id_col: str = "unique_id",
+    time_col: str = "ds",
+    target_col: str = "y",
+) -> PanelData:
+    processed = process_df(df, id_col, time_col, target_col)
+    series_arrays = [
+        processed.data[s:e, 0]
+        for s, e in zip(processed.indptr[:-1], processed.indptr[1:], strict=True)
+    ]
+    return PanelData(processed.uids, processed.last_times, series_arrays)
+
+
+def grouped_std_by_id(
+    df: pd.DataFrame,
+    id_col: str,
+    value_col: str,
+) -> pd.DataFrame:
+    out = group_by_agg(df, id_col, {value_col: "std"})
+    std_col = "residual_std" if value_col == "residuals" else f"{value_col}_std"
+    return out.rename(columns={value_col: std_col})
 
 
 class TimeSeriesDataset:
     def __init__(
         self,
-        data: torch.Tensor,
-        uids: Iterable,
-        last_times: Iterable,
+        series_arrays: list[np.ndarray],
+        uids: pd.Series | np.ndarray,
+        last_times: np.ndarray,
         batch_size: int,
+        dtype: torch.dtype = torch.bfloat16,
     ):
-        self.data = data
+        self._series_arrays = series_arrays
         self.uids = uids
         self.last_times = last_times
         self.batch_size = batch_size
-        self.n_batches = len(data) // self.batch_size + (
-            0 if len(data) % self.batch_size == 0 else 1
+        self.dtype = dtype
+        self._tensors: list[torch.Tensor] | None = None
+        self.n_batches = len(series_arrays) // self.batch_size + (
+            0 if len(series_arrays) % self.batch_size == 0 else 1
         )
         self.current_batch = 0
 
+    @property
+    def data(self) -> list[torch.Tensor]:
+        if self._tensors is None:
+            self._tensors = [
+                torch.as_tensor(arr, dtype=self.dtype) for arr in self._series_arrays
+            ]
+        return self._tensors
+
+    @classmethod
+    def from_panel(
+        cls,
+        panel: PanelData,
+        batch_size: int,
+        dtype: torch.dtype = torch.bfloat16,
+    ) -> TimeSeriesDataset:
+        return cls(panel.series_arrays, panel.uids, panel.last_times, batch_size, dtype)
+
     @classmethod
     def from_df(
         cls,
         df: pd.DataFrame,
         batch_size: int,
         dtype: torch.dtype = torch.bfloat16,
-    ):
-        tensors = []
-        df_sorted = df.sort_values(by=["unique_id", "ds"])
-        for _, group in df_sorted.groupby("unique_id"):
-            tensors.append(torch.tensor(group["y"].values, dtype=dtype))
-        uids = df_sorted["unique_id"].unique()
-        last_times = df_sorted.groupby("unique_id")["ds"].tail(1)
-        return cls(tensors, uids, last_times, batch_size)
+        panel: PanelData | None = None,
+    ) -> TimeSeriesDataset:
+        if panel is None:
+            panel = process_panel_from_df(df)
+        return cls.from_panel(panel, batch_size, dtype)
 
     def __len__(self):
         return self.n_batches
@@ -49,14 +99,14 @@ def make_future_dataframe(self, h: int, freq: str) -> pd.DataFrame:
         )  # type: ignore
 
     def __iter__(self):
-        self.current_batch = 0  # Reset for new iteration
+        self.current_batch = 0
         return self
 
     def __next__(self):
         if self.current_batch < self.n_batches:
             start_idx = self.current_batch * self.batch_size
             end_idx = start_idx + self.batch_size
             self.current_batch += 1
-            return self.data[start_idx:end_idx]
-        else:
-            raise StopIteration
+            batch_arrays = self._series_arrays[start_idx:end_idx]
+            return [torch.as_tensor(arr, dtype=self.dtype) for arr in batch_arrays]
+        raise StopIteration