diff --git a/foundationforecast/core/__init__.py b/foundationforecast/core/__init__.py index 7e6a8c3..5a850a9 100644 --- a/foundationforecast/core/__init__.py +++ b/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", diff --git a/foundationforecast/core/forecaster.py b/foundationforecast/core/forecaster.py index 2a0b1b4..4b1b5e7 100644 --- a/foundationforecast/core/forecaster.py +++ b/foundationforecast/core/forecaster.py @@ -16,6 +16,7 @@ from tqdm import tqdm from utilsforecast.processing import ( backtest_splits, + counts_by_id, drop_index_if_pandas, join, maybe_compute_sort_indices, @@ -23,6 +24,8 @@ vertical_concat, ) +from .utils import PanelData, TimeSeriesDataset, grouped_std_by_id + def get_seasonality( freq: str, @@ -167,6 +170,33 @@ 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, @@ -174,6 +204,7 @@ def forecast( 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 diff --git a/foundationforecast/core/gluonts_forecaster.py b/foundationforecast/core/gluonts_forecaster.py index 25bc2d5..a85a359 100644 --- a/foundationforecast/core/gluonts_forecaster.py +++ b/foundationforecast/core/gluonts_forecaster.py @@ -2,6 +2,7 @@ from contextlib import contextmanager from typing import Any +import numpy as np import pandas as pd import torch import utilsforecast.processing as ufp @@ -9,9 +10,10 @@ 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( diff --git a/foundationforecast/core/multi_model.py b/foundationforecast/core/multi_model.py index 3d6635b..74e7b7d 100644 --- a/foundationforecast/core/multi_model.py +++ b/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( diff --git a/foundationforecast/core/utils.py b/foundationforecast/core/utils.py index 287be32..4308276 100644 --- a/foundationforecast/core/utils.py +++ b/foundationforecast/core/utils.py @@ -1,41 +1,94 @@ -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: + if not pd.api.types.is_datetime64_any_dtype(df[time_col]): + df = df.copy() + df[time_col] = pd.to_datetime(df[time_col]) + 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,7 +102,7 @@ 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): @@ -57,6 +110,6 @@ def __next__(self): 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 diff --git a/foundationforecast/models/chronos.py b/foundationforecast/models/chronos.py index 0691279..cda425a 100644 --- a/foundationforecast/models/chronos.py +++ b/foundationforecast/models/chronos.py @@ -15,7 +15,7 @@ from tqdm import tqdm from ..core.forecaster import Forecaster, QuantileConverter -from ..core.utils import TimeSeriesDataset +from ..core.utils import PanelData, TimeSeriesDataset, process_panel_from_df @dataclass @@ -152,18 +152,21 @@ def __init__( self.finetuning_config = finetuning_config @staticmethod - def _build_fit_inputs_from_df(df: pd.DataFrame) -> list[dict[str, Any]]: + def _build_fit_inputs_from_df( + df: pd.DataFrame, + panel: PanelData | None = None, + ) -> list[dict[str, Any]]: """Build list of fit inputs from a DataFrame (unique_id, ds, y).""" - df_sorted = df.sort_values(by=["unique_id", "ds"]) - return [ - {"target": group["y"].values} for _, group in df_sorted.groupby("unique_id") - ] + if panel is None: + panel = process_panel_from_df(df) + return [{"target": arr} for arr in panel.series_arrays] def _maybe_finetune( self, model: BaseChronosPipeline, df: pd.DataFrame, h: int, + panel: PanelData | None = None, ) -> BaseChronosPipeline: """If finetuning_config is set, finetune the model on df and return it.""" if self.finetuning_config is None: @@ -173,7 +176,7 @@ def _maybe_finetune( f"Finetuning is not supported for model {self.repo_id}; " "the loaded pipeline has no fit method." ) - train_inputs = self._build_fit_inputs_from_df(df) + train_inputs = self._build_fit_inputs_from_df(df, panel=panel) fit_kwargs: dict[str, Any] = { "inputs": train_inputs, "prediction_length": h, @@ -284,6 +287,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. @@ -336,12 +340,17 @@ def forecast( """ freq = self._maybe_infer_freq(df, freq) qc = QuantileConverter(level=level, quantiles=quantiles) - dataset = TimeSeriesDataset.from_df( - df, batch_size=self.batch_size, dtype=self.dtype + if panel is None: + panel = process_panel_from_df(df) + dataset = self._make_timeseries_dataset( + df, + batch_size=self.batch_size, + dtype=self.dtype, + panel=panel, ) fcst_df = dataset.make_future_dataframe(h=h, freq=freq) with self._get_model() as model: - model = self._maybe_finetune(model, df, h) + model = self._maybe_finetune(model, df, h, panel=panel) fcsts_mean_np, fcsts_quantiles_np = self._predict( model, dataset, @@ -350,10 +359,12 @@ def forecast( ) fcst_df[self.alias] = fcsts_mean_np.reshape(-1, 1) if qc.quantiles is not None and fcsts_quantiles_np is not None: - for i, q in enumerate(qc.quantiles): - fcst_df[f"{self.alias}-q-{int(q * 100)}"] = fcsts_quantiles_np[ - ..., i - ].reshape(-1, 1) + fcst_df = self._assign_quantile_forecasts( + fcst_df, + self.alias, + qc.quantiles, + fcsts_quantiles_np, + ) fcst_df = qc.maybe_convert_quantiles_to_level( fcst_df, models=[self.alias], diff --git a/foundationforecast/models/flowstate.py b/foundationforecast/models/flowstate.py index 6e92478..fb38350 100644 --- a/foundationforecast/models/flowstate.py +++ b/foundationforecast/models/flowstate.py @@ -14,7 +14,7 @@ from tsfm_public.models.flowstate.utils.utils import get_fixed_factor from ..core.forecaster import Forecaster, QuantileConverter, _DataProcessor -from ..core.utils import TimeSeriesDataset +from ..core.utils import PanelData, TimeSeriesDataset class FlowState(Forecaster, _DataProcessor): @@ -236,6 +236,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. @@ -285,10 +286,11 @@ def forecast( """ freq = self._maybe_infer_freq(df, freq) qc = QuantileConverter(level=level, quantiles=quantiles) - dataset = TimeSeriesDataset.from_df( + dataset = self._make_timeseries_dataset( df, batch_size=self.batch_size, dtype=self.dtype, + panel=panel, ) fcst_df = dataset.make_future_dataframe(h=h, freq=freq) scale_factor = self.scale_factor or get_fixed_factor(freq) @@ -314,10 +316,12 @@ def forecast( ) fcst_df[self.alias] = fcsts_mean_np.reshape(-1, 1) if qc.quantiles is not None and fcsts_quantiles_np is not None: - for i, q in enumerate(qc.quantiles): - fcst_df[f"{self.alias}-q-{int(q * 100)}"] = fcsts_quantiles_np[ - ..., i - ].reshape(-1, 1) + fcst_df = self._assign_quantile_forecasts( + fcst_df, + self.alias, + qc.quantiles, + fcsts_quantiles_np, + ) fcst_df = qc.maybe_convert_quantiles_to_level( fcst_df, models=[self.alias], diff --git a/foundationforecast/models/patchtst_fm.py b/foundationforecast/models/patchtst_fm.py index 45fed0f..9e969ae 100644 --- a/foundationforecast/models/patchtst_fm.py +++ b/foundationforecast/models/patchtst_fm.py @@ -13,7 +13,7 @@ from tsfm_public import PatchTSTFMForPrediction from ..core.forecaster import Forecaster, QuantileConverter, _DataProcessor -from ..core.utils import TimeSeriesDataset +from ..core.utils import PanelData, TimeSeriesDataset logger = logging.getLogger(__name__) @@ -255,6 +255,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. @@ -304,10 +305,11 @@ def forecast( """ freq = self._maybe_infer_freq(df, freq) qc = QuantileConverter(level=level, quantiles=quantiles) - dataset = TimeSeriesDataset.from_df( + dataset = self._make_timeseries_dataset( df, batch_size=self.batch_size, dtype=self.dtype, + panel=panel, ) fcst_df = dataset.make_future_dataframe(h=h, freq=freq) # scale_factor = self.scale_factor or get_fixed_factor(freq) @@ -336,10 +338,12 @@ def forecast( # should only enter when quantiles are used if qc.quantiles is not None and fcsts_quantiles_np is not None: - for i, q in enumerate(qc.quantiles): - fcst_df[f"{self.alias}-q-{int(q * 100)}"] = fcsts_quantiles_np[ - ..., i - ].reshape(-1) + fcst_df = self._assign_quantile_forecasts( + fcst_df, + self.alias, + qc.quantiles, + fcsts_quantiles_np, + ) fcst_df = qc.maybe_convert_quantiles_to_level( fcst_df, models=[self.alias], diff --git a/foundationforecast/models/sundial.py b/foundationforecast/models/sundial.py index 7cfd433..4498fd7 100644 --- a/foundationforecast/models/sundial.py +++ b/foundationforecast/models/sundial.py @@ -12,7 +12,7 @@ from transformers import AutoModelForCausalLM from ..core.forecaster import Forecaster, QuantileConverter, _DataProcessor -from ..core.utils import TimeSeriesDataset +from ..core.utils import PanelData, TimeSeriesDataset class Sundial(Forecaster, _DataProcessor): @@ -166,6 +166,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. @@ -215,7 +216,9 @@ def forecast( """ freq = self._maybe_infer_freq(df, freq) qc = QuantileConverter(level=level, quantiles=quantiles) - dataset = TimeSeriesDataset.from_df(df, batch_size=self.batch_size) + dataset = self._make_timeseries_dataset( + df, batch_size=self.batch_size, panel=panel + ) fcst_df = dataset.make_future_dataframe(h=h, freq=freq) with self._get_model() as model: fcsts_mean_np, fcsts_quantiles_np = self._predict( @@ -226,10 +229,12 @@ def forecast( ) fcst_df[self.alias] = fcsts_mean_np.reshape(-1, 1) if qc.quantiles is not None and fcsts_quantiles_np is not None: - for i, q in enumerate(qc.quantiles): - fcst_df[f"{self.alias}-q-{int(q * 100)}"] = fcsts_quantiles_np[ - ..., i - ].reshape(-1, 1) + fcst_df = self._assign_quantile_forecasts( + fcst_df, + self.alias, + qc.quantiles, + fcsts_quantiles_np, + ) fcst_df = qc.maybe_convert_quantiles_to_level( fcst_df, models=[self.alias], diff --git a/foundationforecast/models/t0.py b/foundationforecast/models/t0.py index 63c9661..9996dd3 100644 --- a/foundationforecast/models/t0.py +++ b/foundationforecast/models/t0.py @@ -14,7 +14,7 @@ from tqdm import tqdm from ..core.forecaster import Forecaster, QuantileConverter -from ..core.utils import TimeSeriesDataset +from ..core.utils import PanelData class T0(Forecaster): @@ -127,6 +127,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. @@ -177,7 +178,9 @@ def forecast( """ freq = self._maybe_infer_freq(df, freq) qc = QuantileConverter(level=level, quantiles=quantiles) - dataset = TimeSeriesDataset.from_df(df, batch_size=self.batch_size) + dataset = self._make_timeseries_dataset( + df, batch_size=self.batch_size, panel=panel + ) fcst_df = dataset.make_future_dataframe(h=h, freq=freq) # T0 interpolates arbitrary quantile levels from its trained knots, # so the median and any user-requested quantiles come from one pass. diff --git a/foundationforecast/models/tabpfn.py b/foundationforecast/models/tabpfn.py index f05755a..2da0312 100644 --- a/foundationforecast/models/tabpfn.py +++ b/foundationforecast/models/tabpfn.py @@ -26,6 +26,7 @@ ) from ..core.forecaster import Forecaster, QuantileConverter +from ..core.utils import PanelData class TabPFN(Forecaster): @@ -152,7 +153,9 @@ def forecast( freq: str | None = None, level: list[int | float] | None = None, quantiles: list[float] | None = None, + panel: PanelData | None = None, ) -> pd.DataFrame: + # NOTE: 'panel' is accepted for API compatibility; this model does not use it. """Generate forecasts for time series data using the model. This method produces point forecasts and, optionally, prediction diff --git a/foundationforecast/models/tafsut.py b/foundationforecast/models/tafsut.py index f088cc0..ad11ed8 100644 --- a/foundationforecast/models/tafsut.py +++ b/foundationforecast/models/tafsut.py @@ -8,7 +8,7 @@ from tqdm import tqdm from ..core.forecaster import Forecaster, QuantileConverter, _DataProcessor -from ..core.utils import TimeSeriesDataset +from ..core.utils import PanelData, TimeSeriesDataset class Tafsut(Forecaster, _DataProcessor): @@ -135,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. @@ -184,10 +185,11 @@ def forecast( """ freq = self._maybe_infer_freq(df, freq) qc = QuantileConverter(level=level, quantiles=quantiles) - dataset = TimeSeriesDataset.from_df( + dataset = self._make_timeseries_dataset( df, batch_size=self.batch_size, dtype=self.dtype, + panel=panel, ) fcst_df = dataset.make_future_dataframe(h=h, freq=freq) with self._get_model() as model: diff --git a/foundationforecast/models/timegpt.py b/foundationforecast/models/timegpt.py index 006e989..abd0aec 100644 --- a/foundationforecast/models/timegpt.py +++ b/foundationforecast/models/timegpt.py @@ -6,6 +6,7 @@ from nixtla import NixtlaClient from ..core.forecaster import Forecaster +from ..core.utils import PanelData @dataclass @@ -119,7 +120,9 @@ def forecast( freq: str | None = None, level: list[int | float] | None = None, quantiles: list[float] | None = None, + panel: PanelData | None = None, ) -> pd.DataFrame: + # NOTE: 'panel' is accepted for API compatibility; this model does not use it. """Generate forecasts for time series data using the model. This method produces point forecasts and, optionally, prediction diff --git a/foundationforecast/models/timesfm.py b/foundationforecast/models/timesfm.py index dbbd509..68b5162 100644 --- a/foundationforecast/models/timesfm.py +++ b/foundationforecast/models/timesfm.py @@ -14,7 +14,7 @@ from tqdm import tqdm from ..core.forecaster import Forecaster, QuantileConverter -from ..core.utils import TimeSeriesDataset +from ..core.utils import PanelData, TimeSeriesDataset # Legacy HF repo IDs from GIFT-Eval submissions without "pytorch" in the name. # Still loaded via timesfm_v1 PyTorch checkpoints (JAX is not supported). @@ -96,6 +96,7 @@ def forecast( freq: str | None = None, level: list[int | float] | None = None, quantiles: list[float] | None = None, + panel: PanelData | None = None, ) -> pd.DataFrame: freq = self._maybe_infer_freq(df, freq) qc = QuantileConverter(level=level, quantiles=quantiles) @@ -203,6 +204,7 @@ def forecast( freq: str | None = None, level: list[int | float] | None = None, quantiles: list[float] | None = None, + panel: PanelData | None = None, ) -> pd.DataFrame: freq = self._maybe_infer_freq(df, freq) qc = QuantileConverter(level=level, quantiles=quantiles) @@ -212,10 +214,11 @@ def forecast( "please use the default quantiles or default level, " "see https://github.com/google-research/timesfm/issues/286" ) - dataset = TimeSeriesDataset.from_df( + dataset = self._make_timeseries_dataset( df, batch_size=self.batch_size, dtype=torch.float32, + panel=panel, ) fcst_df = dataset.make_future_dataframe(h=h, freq=freq) with self._get_predictor(prediction_length=h) as model: @@ -311,6 +314,7 @@ def forecast( freq: str | None = None, level: list[int | float] | None = None, quantiles: list[float] | None = None, + panel: PanelData | None = None, ) -> pd.DataFrame: freq = self._maybe_infer_freq(df, freq) qc = QuantileConverter(level=level, quantiles=quantiles) @@ -320,10 +324,11 @@ def forecast( "please use the default quantiles or default level, " "see https://github.com/google-research/timesfm/issues/286" ) - dataset = TimeSeriesDataset.from_df( + dataset = self._make_timeseries_dataset( df, batch_size=self.batch_size, dtype=torch.float32, + panel=panel, ) fcst_df = dataset.make_future_dataframe(h=h, freq=freq) with self._get_predictor(prediction_length=h) as forecaster: diff --git a/foundationforecast/models/tirex.py b/foundationforecast/models/tirex.py index b3ac197..9db41d1 100644 --- a/foundationforecast/models/tirex.py +++ b/foundationforecast/models/tirex.py @@ -16,7 +16,7 @@ from tqdm import tqdm from ..core.forecaster import Forecaster, QuantileConverter -from ..core.utils import TimeSeriesDataset +from ..core.utils import PanelData, TimeSeriesDataset if TYPE_CHECKING: from tirex2.api_adapter import ForecastModel @@ -190,6 +190,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. @@ -246,9 +247,10 @@ def forecast( "TiRex only supports the default quantiles, " "please use the default quantiles or default level, " ) - dataset = TimeSeriesDataset.from_df( + dataset = self._make_timeseries_dataset( df, batch_size=self.batch_size, + panel=panel, ) fcst_df = dataset.make_future_dataframe(h=h, freq=freq) @@ -262,10 +264,12 @@ def forecast( ) fcst_df[self.alias] = fcsts_mean_np.reshape(-1, 1) if qc.quantiles is not None and fcsts_quantiles_np is not None: - for i, q in enumerate(qc.quantiles): - fcst_df[f"{self.alias}-q-{int(q * 100)}"] = fcsts_quantiles_np[ - ..., i - ].reshape(-1, 1) + fcst_df = self._assign_quantile_forecasts( + fcst_df, + self.alias, + qc.quantiles, + fcsts_quantiles_np, + ) fcst_df = qc.maybe_convert_quantiles_to_level( fcst_df, models=[self.alias], diff --git a/foundationforecast/models/toto.py b/foundationforecast/models/toto.py index 5578260..e365ee2 100644 --- a/foundationforecast/models/toto.py +++ b/foundationforecast/models/toto.py @@ -13,7 +13,7 @@ from tqdm import tqdm from ..core.forecaster import Forecaster, QuantileConverter -from ..core.utils import TimeSeriesDataset +from ..core.utils import PanelData, TimeSeriesDataset # Config key that only appears in Toto 2.0 checkpoints (a Toto2ModelConfig field). # Used to dispatch between the Toto 1.0 and Toto 2.0 backends. @@ -357,6 +357,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. @@ -407,7 +408,9 @@ def forecast( """ freq = self._maybe_infer_freq(df, freq) qc = QuantileConverter(level=level, quantiles=quantiles) - dataset = TimeSeriesDataset.from_df(df, batch_size=self.batch_size) + dataset = self._make_timeseries_dataset( + df, batch_size=self.batch_size, panel=panel + ) fcst_df = dataset.make_future_dataframe(h=h, freq=freq) forecast_fn = self._forecast_toto2 if self._is_toto2() else self._forecast with self._get_model() as model: @@ -419,10 +422,12 @@ def forecast( ) fcst_df[self.alias] = fcsts_mean_np.reshape(-1, 1) if qc.quantiles is not None and fcsts_quantiles_np is not None: - for i, q in enumerate(qc.quantiles): - fcst_df[f"{self.alias}-q-{int(q * 100)}"] = fcsts_quantiles_np[ - ..., i - ].reshape(-1, 1) + fcst_df = self._assign_quantile_forecasts( + fcst_df, + self.alias, + qc.quantiles, + fcsts_quantiles_np, + ) fcst_df = qc.maybe_convert_quantiles_to_level( fcst_df, models=[self.alias], diff --git a/tests/core/test_forecaster.py b/tests/core/test_forecaster.py index 07402c2..26bbff1 100644 --- a/tests/core/test_forecaster.py +++ b/tests/core/test_forecaster.py @@ -1,3 +1,4 @@ +import numpy as np import pandas as pd import pytest @@ -8,6 +9,7 @@ generate_series_with_anomalies, ) from foundationforecast.core.forecaster import ( + Forecaster, QuantileConverter, get_seasonality, maybe_infer_freq, @@ -83,9 +85,9 @@ def test_prepare_level_and_quantiles_with_levels(): @pytest.mark.parametrize( "quantiles,expected_level", [ - ([0.1, 0.5, 0.9], [0, 80]), - ([0.1, 0.5, 0.2, 0.9], [0, 60, 80]), - ([0.5], [0]), + ([0.1, 0.5, 0.9], [80]), + ([0.1, 0.5, 0.2, 0.9], [60, 80]), + ([0.5], None), ], ) def test_prepare_level_and_quantiles_with_quantiles(quantiles, expected_level): @@ -345,3 +347,20 @@ def test_validate_input(): def test_validate_input_errors(df, h, match): with pytest.raises(ValueError, match=match): DummyModel().validate_input(df, h) + + +def test_assign_quantile_forecasts(): + fcst_df = pd.DataFrame(index=range(4)) + quantiles = [0.1, 0.5, 0.9] + fcsts_quantiles_np = np.ones((2, 2, len(quantiles))) + out = Forecaster._assign_quantile_forecasts( + fcst_df, + "dummy", + quantiles, + fcsts_quantiles_np, + ) + assert list(out.columns) == [ + "dummy-q-10", + "dummy-q-50", + "dummy-q-90", + ] diff --git a/tests/core/test_utils.py b/tests/core/test_utils.py index 944f0d8..369824e 100644 --- a/tests/core/test_utils.py +++ b/tests/core/test_utils.py @@ -1,7 +1,11 @@ +import numpy as np import torch from tests.helpers import generate_series -from foundationforecast.core.utils import TimeSeriesDataset +from foundationforecast.core.utils import ( + TimeSeriesDataset, + process_panel_from_df, +) def test_timeseries_dataset_class_default_dtype_is_bfloat16(): @@ -16,3 +20,39 @@ def test_timeseries_dataset_respects_custom_dtype(): df = generate_series(1, "D", min_length=10, max_length=10) dataset = TimeSeriesDataset.from_df(df, batch_size=10, dtype=torch.float32) assert dataset.data[0].dtype == torch.float32 + + +def test_timeseries_dataset_multi_series_order_and_values(): + df = generate_series(3, "D", min_length=5, max_length=5, seed=0) + panel = process_panel_from_df(df) + dataset = TimeSeriesDataset.from_panel(panel, batch_size=2, dtype=torch.float32) + + assert len(panel.series_arrays) == 3 + assert len(dataset.uids) == 3 + for arr, tensor in zip(panel.series_arrays, dataset.data, strict=True): + np.testing.assert_allclose(arr, tensor.numpy()) + + batches = list(dataset) + assert len(batches) == 2 + assert len(batches[0]) == 2 + assert len(batches[1]) == 1 + + +def test_timeseries_dataset_from_panel_matches_from_df(): + df = generate_series(5, "D", min_length=8, max_length=12, seed=1) + panel = process_panel_from_df(df) + from_df = TimeSeriesDataset.from_df(df, batch_size=3, dtype=torch.float32) + from_panel = TimeSeriesDataset.from_panel(panel, batch_size=3, dtype=torch.float32) + + assert list(from_df.uids) == list(from_panel.uids) + for left, right in zip(from_df.data, from_panel.data, strict=True): + torch.testing.assert_close(left, right) + + +def test_lazy_batch_iteration(): + df = generate_series(4, "D", min_length=6, max_length=6, seed=2) + dataset = TimeSeriesDataset.from_df(df, batch_size=2, dtype=torch.float32) + assert dataset._tensors is None + first_batch = next(iter(dataset)) + assert dataset._tensors is None + assert len(first_batch) == 2 diff --git a/tests/helpers.py b/tests/helpers.py index a37b210..7631836 100644 --- a/tests/helpers.py +++ b/tests/helpers.py @@ -48,7 +48,8 @@ class DummyModel(Forecaster): def __init__(self, alias: str = "dummy"): self.alias = alias - def forecast(self, df, h, freq=None, level=None, quantiles=None): + def forecast(self, df, h, freq=None, level=None, quantiles=None, panel=None): + _ = panel freq = self._maybe_infer_freq(df, freq) qc = QuantileConverter(level=level, quantiles=quantiles) last_times = df.groupby("unique_id")["ds"].max() @@ -73,7 +74,8 @@ class SeasonalNaiveModel(Forecaster): def __init__(self, season_length: int | None = None): self.season_length = season_length - def forecast(self, df, h, freq=None, level=None, quantiles=None): + def forecast(self, df, h, freq=None, level=None, quantiles=None, panel=None): + _ = panel freq = self._maybe_infer_freq(df, freq) season_length = self._maybe_get_seasonality(freq) results = [] diff --git a/tests/test_foundation_forecast.py b/tests/test_foundation_forecast.py index 0cf16f7..aeb2b85 100644 --- a/tests/test_foundation_forecast.py +++ b/tests/test_foundation_forecast.py @@ -101,13 +101,15 @@ def test_foundation_forecast_fallback_model(): class FailingModel(Forecaster): alias = "FailingModel" - def forecast(self, df, h, freq=None, level=None, quantiles=None): + def forecast(self, df, h, freq=None, level=None, quantiles=None, panel=None): + _ = panel raise RuntimeError("Intentional failure") class FallbackModel(Forecaster): alias = "FallbackModel" - def forecast(self, df, h, freq=None, level=None, quantiles=None): + def forecast(self, df, h, freq=None, level=None, quantiles=None, panel=None): + _ = panel n = len(df["unique_id"].unique()) * h return pd.DataFrame( { @@ -132,7 +134,8 @@ def test_foundation_forecast_no_fallback_raises(): class FailingModel(Forecaster): alias = "FailingModel" - def forecast(self, df, h, freq=None, level=None, quantiles=None): + def forecast(self, df, h, freq=None, level=None, quantiles=None, panel=None): + _ = panel raise RuntimeError("Intentional failure") df = generate_series(n_series=1, freq="D", min_length=10)