From 8d1138fc6d988341332c8862c3f9c231ede6cf57 Mon Sep 17 00:00:00 2001 From: AzulGarza Date: Thu, 13 Aug 2026 14:22:17 -0600 Subject: [PATCH 1/3] feat: use foundationforecast --- docs/model-hub.md | 3 + mkdocs.yml | 2 + pyproject.toml | 12 +- tests/models/foundation/test_shims.py | 42 ++ tests/models/foundation/test_utils.py | 3 +- tests/test_forecaster.py | 2 +- timecopilot/forecaster.py | 116 +--- timecopilot/models/foundation/__init__.py | 30 + timecopilot/models/foundation/chronos.py | 362 +--------- timecopilot/models/foundation/flowstate.py | 261 +------ timecopilot/models/foundation/moirai.py | 127 +--- timecopilot/models/foundation/patchtst_fm.py | 274 +------- timecopilot/models/foundation/sundial.py | 237 +------ timecopilot/models/foundation/t0.py | 207 +----- timecopilot/models/foundation/tabpfn.py | 224 +----- timecopilot/models/foundation/timegpt.py | 190 +----- timecopilot/models/foundation/timesfm.py | 326 +-------- timecopilot/models/foundation/tirex.py | 273 +------- timecopilot/models/foundation/toto.py | 430 +----------- timecopilot/models/foundation/utils.py | 62 -- timecopilot/models/prophet.py | 3 + timecopilot/models/utils/forecaster.py | 683 +------------------ uv.lock | 54 +- 23 files changed, 199 insertions(+), 3724 deletions(-) create mode 100644 tests/models/foundation/test_shims.py delete mode 100644 timecopilot/models/foundation/utils.py diff --git a/docs/model-hub.md b/docs/model-hub.md index 62363103..b7d584f7 100644 --- a/docs/model-hub.md +++ b/docs/model-hub.md @@ -5,6 +5,9 @@ TimeCopilot provides a unified API for time series forecasting, integrating foun Here you'll find all the time series forecasting models available in TimeCopilot, organized by family. Click on any model name to jump to its detailed API documentation. +!!! note "Foundation model implementations" + Foundation model implementations live in the [`foundationforecast`](https://pypi.org/project/foundationforecast/) package on PyPI. TimeCopilot re-exports them at `timecopilot.models.foundation.*` for backward compatibility. Model docstrings are maintained in the [foundationforecast](https://github.com/TimeCopilot/foundationforecast) repository. + !!! note "Family example notebooks" Walkthrough notebooks for families of foundation models [here](examples/index.md#foundation-models). diff --git a/mkdocs.yml b/mkdocs.yml index 2afa7522..f87b42c7 100644 --- a/mkdocs.yml +++ b/mkdocs.yml @@ -169,6 +169,8 @@ plugins: python: paths: [timecopilot] options: + preload_modules: + - foundationforecast relative_crossrefs: true members_order: source separate_signature: true diff --git a/pyproject.toml b/pyproject.toml index a26b35a8..31312721 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -66,8 +66,8 @@ dependencies = [ "catboost>=1.2.10", "datasets>=4.1.1", "fire", + "foundationforecast>=0.1.1", "fsspec>=2025.9.0", - "gluonts[torch]", "huggingface-hub>=0.36.2,<2.0", "hydra-core>=1.3.2", "lightgbm>=4.6.0", @@ -92,17 +92,7 @@ dependencies = [ "pytorch-lightning==2.4.0", "scipy<=1.15.3", "statsforecast>=2.0.2", - "tabpfn-time-series==1.0.3 ; python_full_version < '3.13'", "tensorboard>=2.20.0", - "tfc-t0>=0.2.3 ; python_full_version >= '3.11' and python_full_version < '3.14'", - "timecopilot-chronos-forecasting>=0.2.2", - "timecopilot-granite-tsfm>=0.2.1 ; python_full_version >= '3.11' and python_full_version < '3.14'", - "timecopilot-timesfm>=0.3.0", - "timecopilot-tirex2>=0.1.0 ; python_full_version >= '3.11'", - "timecopilot-tirex>=0.1.1 ; python_full_version >= '3.11'", - "timecopilot-toto-2>=0.1.1", - "timecopilot-toto>=0.1.7", - "timecopilot-uni2ts>=0.1.3 ; python_full_version < '3.14'", "torchmetrics>=1.8.2", "transformers>=4.41,<6 ; python_full_version < '3.13'", "transformers>=4.48,<6 ; python_full_version >= '3.13'", diff --git a/tests/models/foundation/test_shims.py b/tests/models/foundation/test_shims.py new file mode 100644 index 00000000..b397ca3f --- /dev/null +++ b/tests/models/foundation/test_shims.py @@ -0,0 +1,42 @@ +import sys + +import pytest + +from timecopilot.models.foundation.chronos import Chronos +from timecopilot.models.foundation.moirai import Moirai +from timecopilot.models.foundation.timegpt import TimeGPT +from timecopilot.models.foundation.timesfm import TimesFM +from timecopilot.models.foundation.toto import Toto +from timecopilot.models.utils.forecaster import Forecaster + +SHIM_MODELS = [Chronos, Moirai, TimesFM, Toto, TimeGPT] + +if sys.version_info >= (3, 11): + from timecopilot.models.foundation.tirex import TiRex + + SHIM_MODELS.append(TiRex) + +if sys.version_info >= (3, 11) and sys.version_info < (3, 14): + from timecopilot.models.foundation.flowstate import FlowState + from timecopilot.models.foundation.patchtst_fm import PatchTSTFM + from timecopilot.models.foundation.t0 import T0 + + SHIM_MODELS.extend([FlowState, PatchTSTFM, T0]) + +if sys.version_info < (3, 13): + from timecopilot.models.foundation.sundial import Sundial + from timecopilot.models.foundation.tabpfn import TabPFN + + SHIM_MODELS.extend([Sundial, TabPFN]) + + +@pytest.mark.parametrize("model_cls", SHIM_MODELS) +def test_foundation_shim_is_forecaster(model_cls): + assert issubclass(model_cls, Forecaster) + + +@pytest.mark.parametrize("model_cls", [Chronos, Moirai, TimesFM, Toto]) +def test_foundation_shim_has_core_methods(model_cls): + assert hasattr(model_cls, "forecast") + assert hasattr(model_cls, "cross_validation") + assert hasattr(model_cls, "detect_anomalies") diff --git a/tests/models/foundation/test_utils.py b/tests/models/foundation/test_utils.py index c51e6574..86b163fa 100644 --- a/tests/models/foundation/test_utils.py +++ b/tests/models/foundation/test_utils.py @@ -1,6 +1,5 @@ import torch - -from timecopilot.models.foundation.utils import TimeSeriesDataset +from foundationforecast.core.utils import TimeSeriesDataset def test_timeseries_dataset_class_default_dtype_is_bfloat16(): diff --git a/tests/test_forecaster.py b/tests/test_forecaster.py index fae035f0..e5ee0de5 100644 --- a/tests/test_forecaster.py +++ b/tests/test_forecaster.py @@ -181,7 +181,7 @@ def test_clean_cache_runs_after_each_model(monkeypatch, models): monkeypatch.setattr( TimeCopilotForecaster, "_clean_model_cache", - staticmethod(lambda: calls.append("cleaned")), + lambda self: calls.append("cleaned"), ) df = generate_series(n_series=1, freq="D", min_length=10) diff --git a/timecopilot/forecaster.py b/timecopilot/forecaster.py index 93159af8..cf160464 100644 --- a/timecopilot/forecaster.py +++ b/timecopilot/forecaster.py @@ -3,6 +3,7 @@ from typing import TYPE_CHECKING, TypeVar import pandas as pd +from foundationforecast.core.multi_model import MultiModelForecasterMixin from .models.utils.forecaster import Forecaster @@ -29,7 +30,7 @@ ) -class TimeCopilotForecaster(Forecaster): +class TimeCopilotForecaster(MultiModelForecasterMixin, Forecaster): """ Unified forecaster for multiple time series models. @@ -72,41 +73,6 @@ def __init__( self.fallback_model = fallback_model self.clean_cache = clean_cache - def _validate_unique_aliases(self, models: list[Forecaster]) -> None: - """ - Validate that all models have unique aliases. - - Args: - models (list[Forecaster]): List of model instances to validate. - - Raises: - ValueError: If duplicate aliases are found. - """ - aliases = [model.alias for model in models] - duplicates = set([alias for alias in aliases if aliases.count(alias) > 1]) - - if duplicates: - raise ValueError( - f"Duplicate model aliases found: {sorted(duplicates)}. " - f"Each model must have a unique alias to avoid column name conflicts. " - f"Please provide different aliases when instantiating models of the " - f"same class." - ) - - @staticmethod - def _clean_model_cache() -> None: - """Release temporary Python and CUDA memory between model calls.""" - import gc - - gc.collect() - try: - import torch - - if torch.cuda.is_available(): - torch.cuda.empty_cache() - except ImportError: - pass - @staticmethod def _is_distributed_df(df: AnyDataFrame) -> bool: """ @@ -120,64 +86,6 @@ def _is_distributed_df(df: AnyDataFrame) -> bool: """ return not isinstance(df, pd.DataFrame) - def _call_models( - self, - attr: str, - merge_on: list[str], - df: pd.DataFrame, - h: int, - freq: str | None, - level: list[int | float] | None, - quantiles: list[float] | None, - **kwargs, - ) -> pd.DataFrame: - # infer just once to avoid multiple calls to _maybe_infer_freq - freq = self._maybe_infer_freq(df, freq) - res_df: pd.DataFrame | None = None - for model in self.models: - known_kwargs = { - "df": df, - "h": h, - "freq": freq, - "level": level, - } - if attr != "detect_anomalies": - known_kwargs["quantiles"] = quantiles - fn = getattr(model, attr) - try: - res_df_model = fn(**known_kwargs, **kwargs) - except (ValueError, RuntimeError) as e: - if self.fallback_model is None: - raise e - fn = getattr(self.fallback_model, attr) - try: - res_df_model = fn(**known_kwargs, **kwargs) - res_df_model = res_df_model.rename( - columns={ - col: ( - col.replace(self.fallback_model.alias, model.alias) - if col.startswith(self.fallback_model.alias) - else col - ) - for col in res_df_model.columns - } - ) - except (ValueError, RuntimeError) as e: - raise e - 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, how="left") - if self.clean_cache: - self._clean_model_cache() - return res_df - def _forecast_pandas( self, df: pd.DataFrame, @@ -192,9 +100,8 @@ def _forecast_pandas( This method is called directly for pandas DataFrames or by the distributed wrapper for each partition. """ - return self._call_models( - "forecast", - merge_on=["unique_id", "ds"], + return MultiModelForecasterMixin.forecast( + self, df=df, h=h, freq=freq, @@ -369,9 +276,8 @@ def _cross_validation_pandas( This method is called directly for pandas DataFrames or by the distributed wrapper for each partition. """ - return self._call_models( - "cross_validation", - merge_on=["unique_id", "ds", "cutoff"], + return MultiModelForecasterMixin.cross_validation( + self, df=df, h=h, freq=freq, @@ -564,15 +470,13 @@ def _detect_anomalies_pandas( This method is called directly for pandas DataFrames or by the distributed wrapper for each partition. """ - return self._call_models( - "detect_anomalies", - merge_on=["unique_id", "ds", "cutoff"], + return MultiModelForecasterMixin.detect_anomalies( + self, df=df, - h=h, # type: ignore + h=h, freq=freq, n_windows=n_windows, - level=level, # type: ignore - quantiles=None, + level=level, ) def _detect_anomalies_distributed( diff --git a/timecopilot/models/foundation/__init__.py b/timecopilot/models/foundation/__init__.py index 3b49b67d..6a377f56 100644 --- a/timecopilot/models/foundation/__init__.py +++ b/timecopilot/models/foundation/__init__.py @@ -1,5 +1,35 @@ +import sys + +from .chronos import Chronos, ChronosFinetuningConfig +from .moirai import Moirai +from .timegpt import TimeGPT, TimeGPTFinetuningConfig from .timesfm import TimesFM +from .toto import Toto __all__ = [ + "Chronos", + "ChronosFinetuningConfig", + "Moirai", + "TimeGPT", + "TimeGPTFinetuningConfig", "TimesFM", + "Toto", ] + +if sys.version_info >= (3, 11): + from .tirex import TiRex as TiRex + + __all__.append("TiRex") + +if sys.version_info >= (3, 11) and sys.version_info < (3, 14): + from .flowstate import FlowState as FlowState + from .patchtst_fm import PatchTSTFM as PatchTSTFM + from .t0 import T0 as T0 + + __all__.extend(["FlowState", "PatchTSTFM", "T0"]) + +if sys.version_info < (3, 13): + from .sundial import Sundial as Sundial + from .tabpfn import TabPFN as TabPFN + + __all__.extend(["Sundial", "TabPFN"]) diff --git a/timecopilot/models/foundation/chronos.py b/timecopilot/models/foundation/chronos.py index 50587d90..f83612cb 100644 --- a/timecopilot/models/foundation/chronos.py +++ b/timecopilot/models/foundation/chronos.py @@ -1,361 +1,11 @@ -from contextlib import contextmanager -from dataclasses import dataclass -from pathlib import Path -from typing import Any, Literal +from foundationforecast.models.chronos import Chronos as _Chronos +from foundationforecast.models.chronos import ChronosFinetuningConfig -import numpy as np -import pandas as pd -import torch -from chronos import ( - BaseChronosPipeline, - Chronos2Pipeline, - ChronosBoltPipeline, - ChronosPipeline, -) -from tqdm import tqdm +from ..utils.forecaster import Forecaster -from ..utils.forecaster import Forecaster, QuantileConverter -from .utils import TimeSeriesDataset +class Chronos(_Chronos, Forecaster): + pass -@dataclass -class ChronosFinetuningConfig: - """Configuration for finetuning a Chronos pipeline before forecasting. - Pass an instance to the ``Chronos`` constructor; when you call - ``forecast()``, the model is finetuned on the context data before - predicting. The forecast horizon ``h`` from ``forecast(df, h, ...)`` is - used as ``prediction_length`` for the internal ``fit()``. Parameters are - passed through to the chronos pipeline's ``fit()``, with ``finetune_steps`` - mapped to the library's ``num_steps``. - - Attributes: - finetune_steps: Number of training steps. Passed to the pipeline as - ``num_steps``. Defaults to 1000. - learning_rate: Optimizer learning rate. Defaults to None (chronos - uses 1e-6; for LoRA, 1e-5 is recommended). - batch_size: Training batch size for finetuning. Defaults to None - (chronos uses 256). The ``batch_size`` on ``Chronos`` is for - inference only. - finetune_mode: ``"full"`` (full parameter update) or ``"lora"`` - (low-rank adaptation). Defaults to None (chronos uses ``"full"``). - lora_config: LoRA configuration when ``finetune_mode="lora"``. Defaults - to None. See the Chronos-2 quickstart for details. - save_path: If set, the finetuned model is saved to this directory (path - or str). Use this same path as ``repo_id`` when creating - ``Chronos(repo_id=save_path, finetuning_config=None)`` for - subsequent forecasting without finetuning. - - Notes: - - Based on the [Chronos-2 quickstart](https://github.com/amazon-science/chronos-forecasting/blob/main/notebooks/chronos-2-quickstart.ipynb). - """ - - finetune_steps: int = 1000 - learning_rate: float | None = None - batch_size: int | None = None - finetune_mode: Literal["full", "lora"] | None = None - lora_config: Any = None - save_path: str | Path | None = None - - -class Chronos(Forecaster): - """ - Chronos models are large pre-trained models for time series forecasting, - supporting both probabilistic and point forecasts. See the - [official repo](https://github.com/amazon-science/chronos-forecasting) - for more details. - """ - - def __init__( - self, - repo_id: str = "amazon/chronos-t5-large", - batch_size: int = 16, - alias: str = "Chronos", - dtype: torch.dtype = torch.float32, - finetuning_config: ChronosFinetuningConfig | None = None, - ): - # ruff: noqa: E501 - """ - Args: - repo_id (str, optional): The Hugging Face Hub model ID or local - path to load the Chronos model from. Examples include - "amazon/chronos-t5-tiny", "amazon/chronos-t5-large", or a - local directory. You can also pass a path where a finetuned - model was saved (see ``finetuning_config.save_path``); use - that path as ``repo_id`` with ``finetuning_config=None`` to - reuse the saved model. Defaults to "amazon/chronos-t5-large". - See the full list of available models at - [Hugging Face](https://huggingface.co/collections/ - amazon/chronos-models-65f1791d630a8d57cb718444) - batch_size (int, optional): Batch size to use for inference only. - Larger models may require smaller batch sizes due to GPU - memory constraints. Defaults to 16. For Chronos-Bolt models, - higher batch sizes (e.g., 256) are possible. When finetuning, - use ``finetuning_config.batch_size`` to set the training - batch size (optional; library default when not set). - alias (str, optional): Name to use for the model in output - DataFrames and logs. Defaults to "Chronos". - dtype (torch.dtype, optional): Data type for model weights and - input tensors. Defaults to torch.float32 for numerical - precision. Use torch.bfloat16 for reduced memory usage on - supported hardware. - finetuning_config (ChronosFinetuningConfig | None, optional): If - provided, the model is finetuned on the forecast context - data before predicting. Set ``save_path`` on the config to - save the finetuned model; then use that path as ``repo_id`` - with ``finetuning_config=None`` for later forecasts. See - ChronosFinetuningConfig and the - [Chronos-2 quickstart](https://github.com/amazon-science/chronos-forecasting/blob/main/notebooks/chronos-2-quickstart.ipynb) - for parameter details. - - Notes: - **Available models:** - - | Model ID | Parameters | - | ---------------------------------------------------------------------- | ---------- | - | [`amazon/chronos-2`](https://huggingface.co/amazon/chronos-2) | 120M | - | [`autogluon/chronos-2-synth`](https://huggingface.co/autogluon/chronos-2-synth) | 120M | - | [`autogluon/chronos-2-small`](https://huggingface.co/autogluon/chronos-2-small) | 28M | - | [`amazon/chronos-bolt-tiny`](https://huggingface.co/amazon/chronos-bolt-tiny) | 9M | - | [`amazon/chronos-bolt-mini`](https://huggingface.co/amazon/chronos-bolt-mini) | 21M | - | [`amazon/chronos-bolt-small`](https://huggingface.co/amazon/chronos-bolt-small) | 48M | - | [`amazon/chronos-bolt-base`](https://huggingface.co/amazon/chronos-bolt-base) | 205M | - | [`amazon/chronos-t5-tiny`](https://huggingface.co/amazon/chronos-t5-tiny) | 8M | - | [`amazon/chronos-t5-mini`](https://huggingface.co/amazon/chronos-t5-mini) | 20M | - | [`amazon/chronos-t5-small`](https://huggingface.co/amazon/chronos-t5-small) | 46M | - | [`amazon/chronos-t5-base`](https://huggingface.co/amazon/chronos-t5-base) | 200M | - | [`amazon/chronos-t5-large`](https://huggingface.co/amazon/chronos-t5-large) | 710M | - - **Academic Reference:** - - - Paper: [Chronos: Learning the Language of Time Series](https://arxiv.org/abs/2403.07815) - - **Resources:** - - - GitHub: [amazon-science/chronos-forecasting](https://github.com/amazon-science/chronos-forecasting) - - HuggingFace: [amazon/chronos-models](https://huggingface.co/collections/amazon/chronos-models-65f1791d630a8d57cb718444) - - **Technical Details:** - - - The model is loaded onto the best available device (GPU if - available, otherwise CPU). - - For best performance with large models (e.g., "chronos-t5-large"), - a CUDA-compatible GPU is recommended. - - Model weights and input tensors use dtype (default: torch.float32) - for numerical precision. Can be overridden via the dtype parameter. - - """ - self.repo_id = repo_id - self.batch_size = batch_size - self.alias = alias - self.dtype = dtype - self.finetuning_config = finetuning_config - - @staticmethod - def _build_fit_inputs_from_df(df: pd.DataFrame) -> 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") - ] - - def _maybe_finetune( - self, - model: BaseChronosPipeline, - df: pd.DataFrame, - h: int, - ) -> BaseChronosPipeline: - """If finetuning_config is set, finetune the model on df and return it.""" - if self.finetuning_config is None: - return model - if not hasattr(model, "fit"): - raise ValueError( - 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) - fit_kwargs: dict[str, Any] = { - "inputs": train_inputs, - "prediction_length": h, - "num_steps": self.finetuning_config.finetune_steps, - } - if self.finetuning_config.learning_rate is not None: - fit_kwargs["learning_rate"] = self.finetuning_config.learning_rate - if self.finetuning_config.batch_size is not None: - fit_kwargs["batch_size"] = self.finetuning_config.batch_size - if self.finetuning_config.finetune_mode is not None: - fit_kwargs["finetune_mode"] = self.finetuning_config.finetune_mode - if self.finetuning_config.lora_config is not None: - fit_kwargs["lora_config"] = self.finetuning_config.lora_config - if self.finetuning_config.save_path is not None: - sp = Path(self.finetuning_config.save_path) - fit_kwargs["output_dir"] = str(sp.parent) - fit_kwargs["finetuned_ckpt_name"] = sp.name - return model.fit(**fit_kwargs) - - @contextmanager - def _get_model(self) -> BaseChronosPipeline: - device_map = "cuda:0" if torch.cuda.is_available() else "cpu" - repo_path = Path(self.repo_id) - # LoRA checkpoints save adapter_config.json; BaseChronosPipeline.from_pretrained - # uses AutoConfig and fails. Chronos2Pipeline.from_pretrained handles LoRA via PEFT. - if repo_path.is_dir() and (repo_path / "adapter_config.json").exists(): - cls = Chronos2Pipeline - else: - cls = BaseChronosPipeline - model = cls.from_pretrained( - self.repo_id, - device_map=device_map, - torch_dtype=self.dtype, - ) - try: - yield model - finally: - del model - torch.cuda.empty_cache() - - def _predict( - self, - model: BaseChronosPipeline, - dataset: TimeSeriesDataset, - h: int, - quantiles: list[float] | None, - ) -> tuple[np.ndarray, np.ndarray | None]: - """handles distinction between predict and predict_quantiles""" - if quantiles is not None: - fcsts = [ - model.predict_quantiles( - batch, - prediction_length=h, - quantile_levels=quantiles, - ) - for batch in tqdm(dataset) - ] # list of tuples - fcsts_quantiles, fcsts_mean = zip(*fcsts, strict=False) - if isinstance(model, Chronos2Pipeline): - fcsts_mean = [f_mean for fcst in fcsts_mean for f_mean in fcst] # type: ignore - fcsts_quantiles = [ - f_quantile - for fcst in fcsts_quantiles - for f_quantile in fcst # type: ignore - ] - fcsts_mean_np = torch.cat(fcsts_mean).numpy() - fcsts_quantiles_np = torch.cat(fcsts_quantiles).numpy() - else: - fcsts = [ - model.predict( - batch, - prediction_length=h, - ) - for batch in tqdm(dataset) - ] - if isinstance(model, Chronos2Pipeline): - fcsts = [f_fcst for fcst in fcsts for f_fcst in fcst] # type: ignore - fcsts = torch.cat(fcsts) - if isinstance(model, ChronosPipeline): - # for t5 models, `predict` returns a tensor of shape - # (batch_size, num_samples, prediction_length). - # notice that the method return samples. - # see https://github.com/amazon-science/chronos-forecasting/blob/6a9c8dadac04eb85befc935043e3e2cce914267f/src/chronos/chronos.py#L450-L537 - # also for these models, the following is how the mean is computed - # in the `predict_quantiles` method - # see https://github.com/amazon-science/chronos-forecasting/blob/6a9c8dadac04eb85befc935043e3e2cce914267f/src/chronos/chronos.py#L554 - fcsts_mean = fcsts.mean(dim=1) # type: ignore - elif isinstance(model, ChronosBoltPipeline | Chronos2Pipeline): - # for bolt models, `predict` returns a tensor of shape - # (batch_size, num_quantiles, prediction_length) - # notice that in this case, the method returns the default quantiles - # instead of samples - # see https://github.com/amazon-science/chronos-forecasting/blob/6a9c8dadac04eb85befc935043e3e2cce914267f/src/chronos/chronos_bolt.py#L479-L563 - # for these models, the median is prefered as mean forecasts - # as it can be seen in - # https://github.com/amazon-science/chronos-forecasting/blob/6a9c8dadac04eb85befc935043e3e2cce914267f/src/chronos/chronos_bolt.py#L615-L616 - fcsts_mean = fcsts[:, model.quantiles.index(0.5), :] # type: ignore - else: - raise ValueError(f"Unsupported model: {self.repo_id}") - fcsts_mean_np = fcsts_mean.numpy() # type: ignore - fcsts_quantiles_np = None - return fcsts_mean_np, fcsts_quantiles_np - - def forecast( - self, - df: pd.DataFrame, - h: int, - freq: str | None = None, - level: list[int | float] | None = None, - quantiles: list[float] | None = None, - ) -> pd.DataFrame: - """Generate forecasts for time series data using the model. - - This method produces point forecasts and, optionally, prediction - intervals or quantile forecasts. The input DataFrame can contain one - or multiple time series in stacked (long) format. - - Args: - df (pd.DataFrame): - DataFrame containing the time series to forecast. It must - include as columns: - - - "unique_id": an ID column to distinguish multiple series. - - "ds": a time column indicating timestamps or periods. - - "y": a target column with the observed values. - - h (int): - Forecast horizon specifying how many future steps to predict. - freq (str, optional): - Frequency of the time series (e.g. "D" for daily, "M" for - monthly). See [Pandas frequency aliases](https://pandas.pydata.org/ - pandas-docs/stable/user_guide/timeseries.html#offset-aliases) for - valid values. If not provided, the frequency will be inferred - from the data. - level (list[int | float], optional): - Confidence levels for prediction intervals, expressed as - percentages (e.g. [80, 95]). If provided, the returned - DataFrame will include lower and upper interval columns for - each specified level. - quantiles (list[float], optional): - List of quantiles to forecast, expressed as floats between 0 - and 1. Should not be used simultaneously with `level`. When - provided, the output DataFrame will contain additional columns - named in the format "model-q-{percentile}", where {percentile} - = 100 × quantile value. - - Returns: - pd.DataFrame: - DataFrame containing forecast results. Includes: - - - point forecasts for each timestamp and series. - - prediction intervals if `level` is specified. - - quantile forecasts if `quantiles` is specified. - - For multi-series data, the output retains the same unique - identifiers as the input DataFrame. - - When ``finetuning_config`` was set at construction, the model is - finetuned on ``df`` before predicting. - """ - 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 - ) - fcst_df = dataset.make_future_dataframe(h=h, freq=freq) - with self._get_model() as model: - model = self._maybe_finetune(model, df, h) - fcsts_mean_np, fcsts_quantiles_np = self._predict( - model, - dataset, - h, - quantiles=qc.quantiles, - ) - 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 = qc.maybe_convert_quantiles_to_level( - fcst_df, - models=[self.alias], - ) - return fcst_df +__all__ = ["Chronos", "ChronosFinetuningConfig"] diff --git a/timecopilot/models/foundation/flowstate.py b/timecopilot/models/foundation/flowstate.py index d8c41458..a728ffa5 100644 --- a/timecopilot/models/foundation/flowstate.py +++ b/timecopilot/models/foundation/flowstate.py @@ -1,261 +1,10 @@ -import sys -from contextlib import contextmanager +from foundationforecast.models.flowstate import FlowState as _FlowState -if sys.version_info < (3, 11) or sys.version_info >= (3, 14): - raise ImportError("FlowState requires Python >= 3.11 and < 3.14") +from ..utils.forecaster import Forecaster -import numpy as np -import pandas as pd -import torch -from tqdm import tqdm -from tsfm_public import FlowStateForPrediction -from tsfm_public.models.flowstate.utils.utils import get_fixed_factor -from ..utils.forecaster import Forecaster, QuantileConverter, _DataProcessor -from .utils import TimeSeriesDataset +class FlowState(_FlowState, Forecaster): + pass -class FlowState(Forecaster, _DataProcessor): - """ - FlowState is the first time-scale adjustable Time Series Foundation Model (TSFM), - open-sourced by IBM Research. Combining a State Space Model (SSM) Encoder with a - Functional Basis Decoder allows FlowState to transition into a timescale invariant - coefficient space and make a continuous forecast from this space. This allows - FlowState to seamlessly adjust to all possible sampling rates. - - See the [official repo](https://github.com/ibm-granite/granite-tsfm) and - [paper](https://arxiv.org/abs/2508.05287) for more details. - """ - - def __init__( - self, - repo_id: str = "ibm-research/flowstate", - scale_factor: float | None = None, - context_length: int = 2_048, - batch_size: int = 1_024, - alias: str = "FlowState", - ): - """ - Initialize FlowState time series foundation model. - - Args: - repo_id (str, optional): The Hugging Face Hub model ID or local path to - load the FlowState model from. Supported models: - - - `ibm-research/flowstate` (default) - - `ibm-granite/granite-timeseries-flowstate-r1`. - - scale_factor (float, optional): Scale factor for temporal adaptation. - If None, will be automatically determined based on the time series - frequency. The scale factor adjusts the model to different sampling - rates. For example, if your data has seasonality every N=96 time steps - (quarter hourly with daily cycle), scale_factor = 24/96 = 0.25. - context_length (int, optional): Maximum context length (input window size) - for the model. Controls how much history is used for each forecast. - Defaults to 2,048. The model supports flexible context lengths. - batch_size (int, optional): Batch size for inference. Defaults to 1,024. - Adjust based on available memory and model size. Larger batch sizes - can improve throughput but require more GPU memory. - alias (str, optional): Name to use for the model in output DataFrames and - logs. Defaults to "FlowState". - - Notes: - **Academic Reference:** - - - Paper: [FlowState: Sampling Rate Invariant Time Series Forecasting](https://arxiv.org/abs/2508.05287) - - **Resources:** - - - GitHub: [ibm-granite/granite-tsfm](https://github.com/ibm-granite/granite-tsfm) - - HuggingFace Models: [ibm-granite/granite-timeseries-flowstate-r1]( - https://huggingface.co/ibm-granite/granite-timeseries-flowstate-r1 - ), [ibm-research/flowstate](https://huggingface.co/ibm-research/flowstate). - - **Technical Details:** - - - The model is loaded onto the best available device (GPU if - available, otherwise CPU). - - FlowState uses State Space Model (SSM) encoder with Functional - Basis Decoder (FBD) for time-scale invariant forecasting. - - Recommended forecasting horizon: no more than 30 seasons. - - **Supported Models:** - - - `ibm-research/flowstate` (default) - - `ibm-granite/granite-timeseries-flowstate-r1`. - """ - self.repo_id = repo_id - self.scale_factor = scale_factor - self.context_length = context_length - self.batch_size = batch_size - self.alias = alias - self.device = "cuda" if torch.cuda.is_available() else "cpu" - self.dtype = torch.float32 - - @contextmanager - def _get_model(self) -> FlowStateForPrediction: - model = FlowStateForPrediction.from_pretrained(self.repo_id).to(self.device) - try: - model.eval() - yield model - finally: - del model - torch.cuda.empty_cache() - - def _predict_batch( - self, - model: FlowStateForPrediction, - batch: list[torch.Tensor], - h: int, - quantiles: list[float] | None, - supported_quantiles: list[float], - scale_factor: float, - ) -> tuple[np.ndarray, np.ndarray | None]: - context = self._prepare_and_validate_context(batch) - if context.shape[1] > self.context_length: - context = context[..., -self.context_length :] - context = self._maybe_impute_missing(context) - # context is (batch, context_length) - # then we convert it to (context_length, batch, 1) - context = context.unsqueeze(-1).transpose(0, 1) - context = context.to(self.device) - # (batch, quantiles, h, n_ch) - fcst = model( - context, - prediction_length=h, - scale_factor=scale_factor, - batch_first=False, - ).quantile_outputs - fcst = fcst.squeeze(-1).transpose(-1, -2) # now shape is (batch, h, quantiles) - fcst_mean = fcst[..., supported_quantiles.index(0.5)] - fcst_mean_np = fcst_mean.detach().numpy(force=True) - fcst_quantiles_np = ( - fcst.detach().numpy(force=True) if quantiles is not None else None - ) - return fcst_mean_np, fcst_quantiles_np - - def _predict( - self, - model: FlowStateForPrediction, - dataset: TimeSeriesDataset, - h: int, - quantiles: list[float] | None, - supported_quantiles: list[float], - scale_factor: float, - ) -> tuple[np.ndarray, np.ndarray | None]: - fcsts = [ - self._predict_batch( - model, - batch, - h, - quantiles, - supported_quantiles, - scale_factor, - ) - for batch in tqdm(dataset) - ] # list of tuples - fcsts_mean_tp, fcsts_quantiles_tp = zip(*fcsts, strict=False) - # handle single item forecast output - fcsts_mean_np = fcsts_mean_tp[0] - if fcsts_mean_tp[0].shape != tuple(): - fcsts_mean_np = np.concatenate(fcsts_mean_tp) - if quantiles is not None: - fcsts_quantiles_np = np.concatenate(fcsts_quantiles_tp) - else: - fcsts_quantiles_np = None - return fcsts_mean_np, fcsts_quantiles_np - - def forecast( - self, - df: pd.DataFrame, - h: int, - freq: str | None = None, - level: list[int | float] | None = None, - quantiles: list[float] | None = None, - ) -> pd.DataFrame: - """Generate forecasts for time series data using the model. - - This method produces point forecasts and, optionally, prediction - intervals or quantile forecasts. The input DataFrame can contain one - or multiple time series in stacked (long) format. - - Args: - df (pd.DataFrame): - DataFrame containing the time series to forecast. It must - include as columns: - - - "unique_id": an ID column to distinguish multiple series. - - "ds": a time column indicating timestamps or periods. - - "y": a target column with the observed values. - - h (int): - Forecast horizon specifying how many future steps to predict. - freq (str, optional): - Frequency of the time series (e.g. "D" for daily, "M" for - monthly). See [Pandas frequency aliases](https://pandas.pydata.org/ - pandas-docs/stable/user_guide/timeseries.html#offset-aliases) for - valid values. If not provided, the frequency will be inferred - from the data. - level (list[int | float], optional): - Confidence levels for prediction intervals, expressed as - percentages (e.g. [80, 95]). If provided, the returned - DataFrame will include lower and upper interval columns for - each specified level. - quantiles (list[float], optional): - List of quantiles to forecast, expressed as floats between 0 - and 1. Should not be used simultaneously with `level`. When - provided, the output DataFrame will contain additional columns - named in the format "model-q-{percentile}", where {percentile} - = 100 × quantile value. - - Returns: - pd.DataFrame: - DataFrame containing forecast results. Includes: - - - point forecasts for each timestamp and series. - - prediction intervals if `level` is specified. - - quantile forecasts if `quantiles` is specified. - - For multi-series data, the output retains the same unique - identifiers as the input DataFrame. - """ - 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, - ) - fcst_df = dataset.make_future_dataframe(h=h, freq=freq) - scale_factor = self.scale_factor or get_fixed_factor(freq) - with self._get_model() as model: - cfg = model.config - supported_quantiles = cfg.quantiles - if qc.quantiles is not None and not np.allclose( - qc.quantiles, - supported_quantiles, - ): - raise ValueError( - "FlowState only supports the default quantiles, " - f"supported quantiles are {supported_quantiles}, " - "please use the default quantiles or default level, " - ) - fcsts_mean_np, fcsts_quantiles_np = self._predict( - model, - dataset, - h, - quantiles=qc.quantiles, - supported_quantiles=supported_quantiles, - scale_factor=scale_factor, - ) - 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 = qc.maybe_convert_quantiles_to_level( - fcst_df, - models=[self.alias], - ) - return fcst_df +__all__ = ["FlowState"] diff --git a/timecopilot/models/foundation/moirai.py b/timecopilot/models/foundation/moirai.py index 76fec566..fc51093b 100644 --- a/timecopilot/models/foundation/moirai.py +++ b/timecopilot/models/foundation/moirai.py @@ -1,127 +1,10 @@ -from contextlib import contextmanager +from foundationforecast.models.moirai import Moirai as _Moirai -import torch -from gluonts.torch.model.predictor import PyTorchPredictor -from uni2ts.model.moirai import MoiraiForecast, MoiraiModule -from uni2ts.model.moirai2 import Moirai2Forecast, Moirai2Module -from uni2ts.model.moirai_moe import MoiraiMoEForecast, MoiraiMoEModule +from ..utils.forecaster import Forecaster -from ..utils.gluonts_forecaster import GluonTSForecaster +class Moirai(_Moirai, Forecaster): + pass -class Moirai(GluonTSForecaster): - """ - Moirai is a universal foundation model for time series forecasting, designed to - handle a wide range of frequencies, multivariate series, and covariates. It uses - a masked encoder-based transformer architecture with multi-patch size projection - layers and Any-Variate Attention, enabling zero-shot and probabilistic - forecasting. See the [official repo](https://github.com/ - SalesforceAIResearch/uni2ts), - [Hugging Face](https://huggingface.co/collections/ - Salesforce/moirai-r-models-65c8d3a94c51428c300e0742), and - [arXiv:2402.02592](https://arxiv.org/abs/2402.02592) for more details. - """ - def __init__( - self, - repo_id: str = "Salesforce/moirai-1.0-R-large", - filename: str = "model.ckpt", - context_length: int = 4096, - patch_size: int = 32, - num_samples: int = 100, - target_dim: int = 1, - feat_dynamic_real_dim: int = 0, - past_feat_dynamic_real_dim: int = 0, - batch_size: int = 32, - alias: str = "Moirai", - ): - """ - Args: - repo_id (str, optional): The Hugging Face Hub model ID or local path to - load the Moirai model from. Examples include - "Salesforce/moirai-1.0-R-large". Defaults to - "Salesforce/moirai-1.0-R-large". See the full list of models at - [Hugging Face](https://huggingface.co/collections/Salesforce/ - moirai-r-models-65c8d3a94c51428c300e0742). - filename (str, optional): Checkpoint filename for the model weights. - Defaults to "model.ckpt". - context_length (int, optional): Maximum context length (input window size) - for the model. Controls how much history is used for each forecast. - Defaults to 4096. - patch_size (int, optional): Patch size for patch-based input encoding. - Can be set to "auto" or a specific value (e.g., 8, 16, 32, 64, 128). - Defaults to 32. See the Moirai paper for recommended values by - frequency. Not used for Moirai-2.0. - num_samples (int, optional): Number of samples for probabilistic - forecasting. Controls the number of forecast samples drawn for - uncertainty estimation. Defaults to 100. - Not used for Moirai-2.0. - target_dim (int, optional): Number of target variables (for multivariate - forecasting). Defaults to 1. - feat_dynamic_real_dim (int, optional): Number of dynamic real covariates - known in the future. Defaults to 0. - past_feat_dynamic_real_dim (int, optional): Number of past dynamic real - covariates. Defaults to 0. - batch_size (int, optional): Batch size to use for inference. Defaults to - 32. Adjust based on available memory and model size. - alias (str, optional): Name to use for the model in output DataFrames and - logs. Defaults to "Moirai". - - Notes: - **Academic Reference:** - - - Paper: [Unified Training of Universal Time Series Forecasting Transformers](https://arxiv.org/abs/2402.02592) - - **Resources:** - - - GitHub: [SalesforceAIResearch/uni2ts](https://github.com/SalesforceAIResearch/uni2ts) - - HuggingFace: [Salesforce/moirai-r-models](https://huggingface.co/collections/Salesforce/moirai-r-models-65c8d3a94c51428c300e0742) - - **Technical Details:** - - - The model is loaded onto the best available device (GPU if available, - otherwise CPU). - """ - super().__init__( - repo_id=repo_id, - filename=filename, - alias=alias, - num_samples=num_samples, - ) - self.context_length = context_length - self.patch_size = patch_size - self.target_dim = target_dim - self.feat_dynamic_real_dim = feat_dynamic_real_dim - self.past_feat_dynamic_real_dim = past_feat_dynamic_real_dim - self.batch_size = batch_size - - @contextmanager - def get_predictor(self, prediction_length: int) -> PyTorchPredictor: - kwargs = { - "prediction_length": prediction_length, - "context_length": self.context_length, - "patch_size": self.patch_size, - "num_samples": self.num_samples, - "target_dim": self.target_dim, - "feat_dynamic_real_dim": self.feat_dynamic_real_dim, - "past_feat_dynamic_real_dim": self.past_feat_dynamic_real_dim, - } - if "moe" in self.repo_id: - model_cls, model_module = MoiraiMoEForecast, MoiraiMoEModule - elif "moirai-2.0" in self.repo_id: - model_cls, model_module = Moirai2Forecast, Moirai2Module - del kwargs["patch_size"] - del kwargs["num_samples"] - else: - model_cls, model_module = MoiraiForecast, MoiraiModule - model = model_cls( - module=model_module.from_pretrained(self.repo_id), - **kwargs, - ) - predictor = model.create_predictor(batch_size=self.batch_size) - - try: - yield predictor - finally: - del predictor, model - torch.cuda.empty_cache() +__all__ = ["Moirai"] diff --git a/timecopilot/models/foundation/patchtst_fm.py b/timecopilot/models/foundation/patchtst_fm.py index 6fb69a98..e27e8a81 100644 --- a/timecopilot/models/foundation/patchtst_fm.py +++ b/timecopilot/models/foundation/patchtst_fm.py @@ -1,274 +1,10 @@ -import sys -from contextlib import contextmanager +from foundationforecast.models.patchtst_fm import PatchTSTFM as _PatchTSTFM -if sys.version_info < (3, 11) or sys.version_info >= (3, 14): - raise ImportError("PatchTSTFM requires Python >= 3.11 and < 3.14") +from ..utils.forecaster import Forecaster -import numpy as np -import pandas as pd -import torch -from tqdm import tqdm -from tsfm_public import PatchTSTFMForPrediction -from ..utils.forecaster import Forecaster, QuantileConverter, _DataProcessor -from .utils import TimeSeriesDataset +class PatchTSTFM(_PatchTSTFM, Forecaster): + pass -# default to the median quantile -# PatchTST-FM supports quantiles from 0.01 to 0.99 -DEFAULT_QUANTILES = [0.5] - -class PatchTSTFM(Forecaster, _DataProcessor): - """ - PatchTST-FM is a Time Series Foundation Model (TSFM) from IBM Research based on a - standard patch Transformer. This generic architecture achieves state-of-the-art - zero-shot forecasting performance with a straightforward training protocol. The - work provides a transparent, reproducible baseline with comprehensive ablations - on model scaling, data composition, and training techniques. - - See the [official repo](https://github.com/ibm-granite/granite-tsfm) and - [paper](https://arxiv.org/abs/2602.06909) for more details. - """ - - # NOTE: may want to adjust default context_length, default on granite_tsfm is 8192 - def __init__( - self, - repo_id: str = "ibm-research/patchtst-fm-r1", - # scale_factor: float | None = None, - context_length: int = 8192, # default from granite-tsfm - batch_size: int = 2_048, - alias: str = "PatchTST-FM", - ): - """ - Initialize PatchTSTFM time series foundation model. - - Args: - repo_id (str, optional): The Hugging Face Hub model ID or local path to - load the PatchTST-FM model from. Supported models: - - - `ibm-research/patchtst-fm-r1` - - context_length (int, optional): Maximum context length (input window size) - for the model. Controls how much history is used for each forecast. - Defaults to 8,192. The model supports flexible context lengths. - batch_size (int, optional): Batch size for inference. Defaults to 2,048. - Adjust based on available memory and model size. Larger batch sizes - can improve throughput but require more GPU memory. - alias (str, optional): Name to use for the model in output DataFrames and - logs. Defaults to "PatchTST-FM". - - Notes: - **Academic Reference:** - - - Paper: [Revisiting the Generic Transformer: Deconstructing a - Strong Baseline for Time Series Foundation Models]( - https://arxiv.org/abs/2602.06909) - - **Resources:** - - - GitHub: [ibm-granite/granite-tsfm](https://github.com/ibm-granite/granite-tsfm) - - HuggingFace Models: [ibm-research/patchtst-fm-r1](https://huggingface.co/ibm-research/patchtst-fm-r1) - - **Technical Details:** - - - The model is loaded onto the best available device (GPU if - available, otherwise CPU). - - **Supported Models:** - - - `ibm-research/patchtst-fm-r1` (default) - """ - self.repo_id = repo_id - # self.scale_factor = scale_factor - self.context_length = context_length - self.batch_size = batch_size - # NOTE: 'mps' may not be 100% reliable, initial tests with the - # patchtst-fm gift_eval notebook resulted in predictions of 0 across - # the board. for now use mps when available, change if it becomes an issue. - self.device = "cuda" if torch.cuda.is_available() else "cpu" - # self.device = ( - # "cuda" - # if torch.cuda.is_available() - # else ("mps" if torch.mps.is_available() else "cpu") - # ) - self.alias = alias - self.dtype = torch.float32 - - @contextmanager - def _get_model(self) -> PatchTSTFMForPrediction: - model = PatchTSTFMForPrediction.from_pretrained(self.repo_id).to(self.device) - try: - model.eval() - yield model - finally: - del model - if self.device.startswith("cuda"): - torch.cuda.empty_cache() - elif self.device.startswith("mps"): - torch.mps.empty_cache() - - def _predict_batch( - self, - model: PatchTSTFMForPrediction, - batch: list[torch.Tensor] | torch.Tensor, - h: int, - quantiles: list[float] | None, - # scale_factor: float, - ) -> tuple[np.ndarray, np.ndarray | None]: - context = self._prepare_and_validate_context(batch) - if context.shape[1] > self.context_length: - context = context[..., -self.context_length :] - context = self._maybe_impute_missing(context) - # context is (batch, context_length) - - # input data is grouped by id - # input shape: (id_group/batch, data) - # output shape: (batch/id, quantiles, h) - quantile_levels = DEFAULT_QUANTILES if quantiles is None else quantiles - - fcst = model( - context, - prediction_length=h, - quantile_levels=quantile_levels, - # scale_factor=scale_factor, - # batch_first=False, - ).quantile_outputs - fcst = fcst.squeeze(-1).transpose(-1, -2) # now shape is (batch, h, quantiles) - - # may not be the ideal solution, but this should be more adaptable - # when quantiles can vary. - # there is no guarantee that 0.5 will be in the list of quantiles. - fcst_mean = fcst.mean(dim=-1).squeeze() if fcst.ndim >= 3 else fcst.squeeze() - # fcst_mean = fcst[..., quantile_levels.index(0.5)].squeeze() - fcst_mean_np = fcst_mean.detach().cpu().numpy() - fcst_quantiles_np = ( - fcst.detach().cpu().numpy() if quantiles is not None else None - ) - return fcst_mean_np, fcst_quantiles_np - - def _predict( - self, - model: PatchTSTFMForPrediction, - dataset: TimeSeriesDataset, - h: int, - quantiles: list[float] | None, - # scale_factor: float, - ) -> tuple[np.ndarray, np.ndarray | None]: - fcsts = [ - self._predict_batch( - model, - batch, - h, - quantiles, - # scale_factor, - ) - for batch in tqdm(dataset) - ] # list of tuples - fcsts_mean_tp, fcsts_quantiles_tp = zip(*fcsts, strict=False) - # handle single item forecast output - fcsts_mean_np = fcsts_mean_tp[0] - if fcsts_mean_tp[0].shape != tuple(): - fcsts_mean_np = np.concatenate(fcsts_mean_tp) - if quantiles is not None: - fcsts_quantiles_np = np.concatenate(fcsts_quantiles_tp) - else: - fcsts_quantiles_np = None - return fcsts_mean_np, fcsts_quantiles_np - - def forecast( - self, - df: pd.DataFrame, - h: int, - freq: str | None = None, - level: list[int | float] | None = None, - quantiles: list[float] | None = None, - ) -> pd.DataFrame: - """Generate forecasts for time series data using the model. - - This method produces point forecasts and, optionally, prediction - intervals or quantile forecasts. The input DataFrame can contain one - or multiple time series in stacked (long) format. - - Args: - df (pd.DataFrame): - DataFrame containing the time series to forecast. It must - include as columns: - - - "unique_id": an ID column to distinguish multiple series. - - "ds": a time column indicating timestamps or periods. - - "y": a target column with the observed values. - - h (int): - Forecast horizon specifying how many future steps to predict. - freq (str, optional): - Frequency of the time series (e.g. "D" for daily, "M" for - monthly). See [Pandas frequency aliases](https://pandas.pydata.org/ - pandas-docs/stable/user_guide/timeseries.html#offset-aliases) for - valid values. If not provided, the frequency will be inferred - from the data. - level (list[int | float], optional): - Confidence levels for prediction intervals, expressed as - percentages (e.g. [80, 95]). If provided, the returned - DataFrame will include lower and upper interval columns for - each specified level. - quantiles (list[float], optional): - List of quantiles to forecast, expressed as floats between 0 - and 1. Should not be used simultaneously with `level`. When - provided, the output DataFrame will contain additional columns - named in the format "model-q-{percentile}", where {percentile} - = 100 × quantile value. - - Returns: - pd.DataFrame: - DataFrame containing forecast results. Includes: - - - point forecasts for each timestamp and series. - - prediction intervals if `level` is specified. - - quantile forecasts if `quantiles` is specified. - - For multi-series data, the output retains the same unique - identifiers as the input DataFrame. - """ - 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, - ) - fcst_df = dataset.make_future_dataframe(h=h, freq=freq) - # scale_factor = self.scale_factor or get_fixed_factor(freq) - with self._get_model() as model: - cfg = model.config - supported_quantiles = cfg.quantile_levels - if qc.quantiles is not None and not set(qc.quantiles).issubset( - supported_quantiles - ): - raise ValueError( - "PatchTSTFM only supports the default quantiles, " - f"supported quantiles are {supported_quantiles}, " - f"quantiles provided are {qc.quantiles}, " - "please use the default quantiles or default level." - ) - - fcsts_mean_np, fcsts_quantiles_np = self._predict( - model, - dataset, - h, - quantiles=qc.quantiles, - # scale_factor=scale_factor, - ) - - fcst_df[self.alias] = fcsts_mean_np.reshape(-1) - - # 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 = qc.maybe_convert_quantiles_to_level( - fcst_df, - models=[self.alias], - ) - return fcst_df +__all__ = ["PatchTSTFM"] diff --git a/timecopilot/models/foundation/sundial.py b/timecopilot/models/foundation/sundial.py index daf0baf7..d047d8db 100644 --- a/timecopilot/models/foundation/sundial.py +++ b/timecopilot/models/foundation/sundial.py @@ -1,237 +1,10 @@ -import sys -from contextlib import contextmanager -from typing import TYPE_CHECKING +from foundationforecast.models.sundial import Sundial as _Sundial -if sys.version_info >= (3, 13) and not TYPE_CHECKING: - raise ImportError("Sundial requires Python < 3.13") +from ..utils.forecaster import Forecaster -import numpy as np -import pandas as pd -import torch -from tqdm import tqdm -from transformers import AutoModelForCausalLM -from ..utils.forecaster import Forecaster, QuantileConverter, _DataProcessor -from .utils import TimeSeriesDataset +class Sundial(_Sundial, Forecaster): + pass -class Sundial(Forecaster, _DataProcessor): - """ - Sundial is a family of generative time series foundation models, - pre-trained on TimeBench (10^12 time points). It uses the TimeFlow Loss to - predict next-patch distributions, allowing Transformers to be trained without - discrete tokenization and make non-deterministic predictions. The model supports - both point and probabilistic zero-shot forecasting. See the - [official repo](https://github.com/thuml/Sundial) for more details. - """ - - def __init__( - self, - repo_id: str = "thuml/sundial-base-128m", - num_samples: int = 100, - context_length: int = 2_880, - batch_size: int = 1_024, - alias: str = "Sundial", - ): - """ - Args: - repo_id (str, optional): The Hugging Face Hub model ID or local path to - load the Sundial model from. Examples include - "thuml/sundial-base-128m". Defaults to "thuml/sundial-base-128m". - See the full list of models at [Hugging Face](https://huggingface.co/ - thuml/sundial-base-128m). - num_samples (int, optional): Number of samples to generate for - probabilistic forecasting. More samples provide better distribution - estimates but increase computation time. Defaults to 100. - context_length (int, optional): Maximum context length (input window size) - for the model. Controls how much history is used for each forecast. - Defaults to 2,880. The model supports different lookback lengths. - batch_size (int, optional): Batch size for inference. Defaults to 1,024. - Adjust based on available memory and model size. Larger batch sizes - can improve throughput but require more GPU memory. - alias (str, optional): Name to use for the model in output DataFrames and - logs. Defaults to "Sundial". - - Notes: - **Academic Reference:** - - - Paper: [Sundial: A Family of Highly Capable Time Series Foundation Models](https://arxiv.org/abs/2502.00816) - - **Resources:** - - - GitHub: [thuml/Sundial](https://github.com/thuml/Sundial) - - HuggingFace: [thuml/sundial-base-128m](https://huggingface.co/thuml/sundial-base-128m) - - **Technical Details:** - - - The model is loaded onto the best available device (GPU if - available, otherwise CPU). - - The model weights are loaded with torch_dtype=torch.bfloat16 for - efficiency on supported hardware. - - The model is only available for Python < 3.13. - """ - self.repo_id = repo_id - self.num_samples = num_samples - self.context_length = context_length - self.batch_size = batch_size - self.alias = alias - self.device = "cuda" if torch.cuda.is_available() else "cpu" - self.dtype = torch.bfloat16 - - @contextmanager - def _get_model(self) -> AutoModelForCausalLM: - model = AutoModelForCausalLM.from_pretrained( - self.repo_id, - torch_dtype=self.dtype, - trust_remote_code=True, - ).to(self.device) - try: - yield model - finally: - del model - torch.cuda.empty_cache() - - def _predict_batch( - self, - model: AutoModelForCausalLM, - batch: list[torch.Tensor], - h: int, - quantiles: list[float] | None, - ) -> tuple[np.ndarray, np.ndarray | None]: - context = self._prepare_and_validate_context(batch) - if context.shape[1] > self.context_length: - context = context[..., -self.context_length :] - context = self._maybe_impute_missing(context) - context = context.to(self.device) - with torch.autocast(device_type=self.device, dtype=self.dtype): - # (batch_size, num_samples, h) - samples = model.generate( - context, - max_new_tokens=h, - revin=True, - num_samples=self.num_samples, - ) - q_median = torch.tensor( - 0.5, - device=self.device, - dtype=samples.dtype, - ) - fcst_mean = torch.quantile( - samples, - q=q_median, - dim=1, - ) - fcst_mean_np = fcst_mean.cpu().numpy() - if quantiles is not None: - quantiles_torch = torch.tensor( - quantiles, - device=self.device, - dtype=samples.dtype, - ) - # (num_quantiles, batch_size, h) - fcst_quantiles = torch.quantile( - samples, - q=quantiles_torch, - dim=1, - ) - fcst_quantiles_np = fcst_quantiles.cpu().numpy() - # (batch_size, h, num_quantiles) - fcst_quantiles_np = np.moveaxis(fcst_quantiles_np, 0, -1) - else: - fcst_quantiles_np = None - return fcst_mean_np, fcst_quantiles_np - - def _predict( - self, - model: AutoModelForCausalLM, - dataset: TimeSeriesDataset, - h: int, - quantiles: list[float] | None, - ) -> tuple[np.ndarray, np.ndarray | None]: - fcsts = [ - self._predict_batch(model, batch, h, quantiles) for batch in tqdm(dataset) - ] # list of tuples - fcsts_mean_tp, fcsts_quantiles_tp = zip(*fcsts, strict=False) - fcsts_mean_np = np.concatenate(fcsts_mean_tp) - if quantiles is not None: - fcsts_quantiles_np = np.concatenate(fcsts_quantiles_tp) - else: - fcsts_quantiles_np = None - return fcsts_mean_np, fcsts_quantiles_np - - def forecast( - self, - df: pd.DataFrame, - h: int, - freq: str | None = None, - level: list[int | float] | None = None, - quantiles: list[float] | None = None, - ) -> pd.DataFrame: - """Generate forecasts for time series data using the model. - - This method produces point forecasts and, optionally, prediction - intervals or quantile forecasts. The input DataFrame can contain one - or multiple time series in stacked (long) format. - - Args: - df (pd.DataFrame): - DataFrame containing the time series to forecast. It must - include as columns: - - - "unique_id": an ID column to distinguish multiple series. - - "ds": a time column indicating timestamps or periods. - - "y": a target column with the observed values. - - h (int): - Forecast horizon specifying how many future steps to predict. - freq (str, optional): - Frequency of the time series (e.g. "D" for daily, "M" for - monthly). See [Pandas frequency aliases](https://pandas.pydata.org/ - pandas-docs/stable/user_guide/timeseries.html#offset-aliases) for - valid values. If not provided, the frequency will be inferred - from the data. - level (list[int | float], optional): - Confidence levels for prediction intervals, expressed as - percentages (e.g. [80, 95]). If provided, the returned - DataFrame will include lower and upper interval columns for - each specified level. - quantiles (list[float], optional): - List of quantiles to forecast, expressed as floats between 0 - and 1. Should not be used simultaneously with `level`. When - provided, the output DataFrame will contain additional columns - named in the format "model-q-{percentile}", where {percentile} - = 100 × quantile value. - - Returns: - pd.DataFrame: - DataFrame containing forecast results. Includes: - - - point forecasts for each timestamp and series. - - prediction intervals if `level` is specified. - - quantile forecasts if `quantiles` is specified. - - For multi-series data, the output retains the same unique - identifiers as the input DataFrame. - """ - freq = self._maybe_infer_freq(df, freq) - qc = QuantileConverter(level=level, quantiles=quantiles) - dataset = TimeSeriesDataset.from_df(df, batch_size=self.batch_size) - fcst_df = dataset.make_future_dataframe(h=h, freq=freq) - with self._get_model() as model: - fcsts_mean_np, fcsts_quantiles_np = self._predict( - model, - dataset, - h, - quantiles=qc.quantiles, - ) - 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 = qc.maybe_convert_quantiles_to_level( - fcst_df, - models=[self.alias], - ) - return fcst_df +__all__ = ["Sundial"] diff --git a/timecopilot/models/foundation/t0.py b/timecopilot/models/foundation/t0.py index 490c2090..134395d7 100644 --- a/timecopilot/models/foundation/t0.py +++ b/timecopilot/models/foundation/t0.py @@ -1,207 +1,10 @@ -import json -import sys -from contextlib import contextmanager +from foundationforecast.models.t0 import T0 as _T0 -if sys.version_info < (3, 11) or sys.version_info >= (3, 14): - raise ImportError("T0 requires Python >= 3.11 and < 3.14") +from ..utils.forecaster import Forecaster -import numpy as np -import pandas as pd -import torch -from huggingface_hub import hf_hub_download -from huggingface_hub.constants import CONFIG_NAME -from t0 import T0Forecaster -from tqdm import tqdm -from ..utils.forecaster import Forecaster, QuantileConverter -from .utils import TimeSeriesDataset +class T0(_T0, Forecaster): + pass -class T0(Forecaster): - """ - T0 is an open-weights time series foundation model from - [The Forecasting Company](https://theforecastingcompany.com/). It is a - decoder-style patch transformer that alternates time and covariate - attention layers, producing probabilistic multi-horizon quantile - forecasts. It decodes up to 1,024 timesteps in a single forward pass and - falls back on autoregressive rollout for longer horizons. T0 natively - handles numerical covariates, both historical (known over the past) and - future (known over the forecast horizon). See the - [model card](https://huggingface.co/theforecastingcompany/t0-alpha) - for more details. - """ - - def __init__( - self, - repo_id: str = "theforecastingcompany/t0-alpha", - context_length: int = 4096, - batch_size: int = 16, - alias: str = "t0-alpha", - ): - # ruff: noqa: E501 - """ - Args: - repo_id (str, optional): The Hugging Face Hub model ID or local path to - load the T0 model from. Defaults to "theforecastingcompany/t0-alpha". - See the full list of models at - [Hugging Face](https://huggingface.co/theforecastingcompany). - context_length (int, optional): Maximum context length (input window - size) for the model. Series longer than this are truncated to the - most recent `context_length` observations. Defaults to 4096. - batch_size (int, optional): Batch size to use for inference. Defaults - to 16. Adjust based on available memory. - alias (str, optional): Name to use for the model in output DataFrames - and logs. Defaults to "t0-alpha". - - Notes: - **Requirements:** - - - T0 requires Python 3.11 to 3.13 (via the - [`tfc-t0`](https://pypi.org/project/tfc-t0/) package). - - **Available models:** - - | Model ID | Parameters | - | ------------------------------------------------------------------------------------------------- | ---------- | - | [`theforecastingcompany/t0-alpha`](https://huggingface.co/theforecastingcompany/t0-alpha) | ~102M | - - **Resources:** - - - HuggingFace: [theforecastingcompany/t0-alpha](https://huggingface.co/theforecastingcompany/t0-alpha) - - Platform: [Retrocast](https://app.retrocast.com/) - - **Technical Details:** - - - The model is loaded onto the best available device (GPU if - available, otherwise CPU). - - T0 predicts 5 quantile knots (0.1, 0.25, 0.5, 0.75, 0.9); the - median (0.5) is used as the point forecast and other requested - quantiles are obtained by linear interpolation across the knots. - - NaN values in the context are treated as missing observations. - - T0 natively supports past and known-future covariates through its - `predict` API; this integration currently exposes the univariate - path only. - """ - self.repo_id = repo_id - self.context_length = context_length - self.batch_size = batch_size - self.alias = alias - self.device = "cuda" if torch.cuda.is_available() else "cpu" - - @contextmanager - def _get_model(self) -> T0Forecaster: - # huggingface_hub may not inject config.json into model kwargs when the - # checkpoint repo is gated; pass the config explicitly. - config_path = hf_hub_download(self.repo_id, CONFIG_NAME) - with open(config_path, encoding="utf-8") as f: - config = json.load(f) - model = ( - T0Forecaster.from_pretrained(self.repo_id, **config).to(self.device).eval() - ) - try: - yield model - finally: - del model - torch.cuda.empty_cache() - - def _to_context(self, batch: list[torch.Tensor]) -> torch.Tensor: - """Left-pad a ragged batch with NaN (treated as missing by T0).""" - max_len = min( - max(len(ts) for ts in batch), - self.context_length, - ) - context = torch.full( - (len(batch), max_len), - float("nan"), - dtype=torch.float32, - ) - for idx, ts in enumerate(batch): - ts = ts[-max_len:] - context[idx, -len(ts) :] = ts.to(dtype=torch.float32) - return context - - def forecast( - self, - df: pd.DataFrame, - h: int, - freq: str | None = None, - level: list[int | float] | None = None, - quantiles: list[float] | None = None, - ) -> pd.DataFrame: - """Generate forecasts for time series data using the model. - - This method produces point forecasts and, optionally, prediction - intervals or quantile forecasts. The input DataFrame can contain one - or multiple time series in stacked (long) format. - - Args: - df (pd.DataFrame): - DataFrame containing the time series to forecast. It must - include as columns: - - - "unique_id": an ID column to distinguish multiple series. - - "ds": a time column indicating timestamps or periods. - - "y": a target column with the observed values. - - h (int): - Forecast horizon specifying how many future steps to predict. - freq (str, optional): - Frequency of the time series (e.g. "D" for daily, "M" for - monthly). See [Pandas frequency aliases](https://pandas.pydata.org/ - pandas-docs/stable/user_guide/timeseries.html#offset-aliases) for - valid values. If not provided, the frequency will be inferred - from the data. - level (list[int | float], optional): - Confidence levels for prediction intervals, expressed as - percentages (e.g. [80, 95]). If provided, the returned - DataFrame will include lower and upper interval columns for - each specified level. - quantiles (list[float], optional): - List of quantiles to forecast, expressed as floats between 0 - and 1. Should not be used simultaneously with `level`. When - provided, the output DataFrame will contain additional columns - named in the format "model-q-{percentile}", where {percentile} - = 100 × quantile value. Quantiles the model wasn't trained on - are linearly interpolated across its fixed knots. - - Returns: - pd.DataFrame: - DataFrame containing forecast results. Includes: - - - point forecasts for each timestamp and series. - - prediction intervals if `level` is specified. - - quantile forecasts if `quantiles` is specified. - - For multi-series data, the output retains the same unique - identifiers as the input DataFrame. - """ - freq = self._maybe_infer_freq(df, freq) - qc = QuantileConverter(level=level, quantiles=quantiles) - dataset = TimeSeriesDataset.from_df(df, batch_size=self.batch_size) - 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. - pred_quantiles = sorted(set(qc.quantiles or []) | {0.5}) - median_idx = pred_quantiles.index(0.5) - fcsts: list[np.ndarray] = [] - with self._get_model() as model: - for batch in tqdm(dataset): - out = model.predict( - self._to_context(batch), - horizon=h, - quantiles=pred_quantiles, - ) - # shape: (batch, h, n_quantiles) - fcsts.append(out.quantiles.cpu().numpy()) - fcsts_np = np.concatenate(fcsts, axis=0) - fcst_df[self.alias] = fcsts_np[..., median_idx].reshape(-1, 1) - if qc.quantiles is not None: - for q in qc.quantiles: - fcst_df[f"{self.alias}-q-{int(q * 100)}"] = fcsts_np[ - ..., pred_quantiles.index(q) - ].reshape(-1, 1) - fcst_df = qc.maybe_convert_quantiles_to_level( - fcst_df, - models=[self.alias], - ) - return fcst_df +__all__ = ["T0"] diff --git a/timecopilot/models/foundation/tabpfn.py b/timecopilot/models/foundation/tabpfn.py index 44b5cd02..0858ff37 100644 --- a/timecopilot/models/foundation/tabpfn.py +++ b/timecopilot/models/foundation/tabpfn.py @@ -1,224 +1,10 @@ -import sys -from contextlib import contextmanager +from foundationforecast.models.tabpfn import TabPFN as _TabPFN -if sys.version_info >= (3, 13): - raise ImportError("TabPFN requires Python < 3.13") +from ..utils.forecaster import Forecaster -import numpy as np -import pandas as pd -import torch -from tabpfn_client import set_access_token -from tabpfn_time_series import ( - TABPFN_TS_DEFAULT_QUANTILE_CONFIG, - FeatureTransformer, - TabPFNMode, - TabPFNTimeSeriesPredictor, - TimeSeriesDataFrame, -) -from tabpfn_time_series.data_preparation import generate_test_X -from tabpfn_time_series.features import ( - AutoSeasonalFeature, - CalendarFeature, - RunningIndexFeature, -) -from tabpfn_time_series.features.feature_generator_base import ( - FeatureGenerator, -) -from ..utils.forecaster import Forecaster, QuantileConverter +class TabPFN(_TabPFN, Forecaster): + pass -class TabPFN(Forecaster): - """ - TabPFN is a zero-shot time series forecasting model that frames univariate - forecasting as a tabular regression problem using TabPFNv2. It supports both - point and probabilistic forecasts, and can incorporate exogenous variables via - feature engineering. See the - [official repo](https://github.com/PriorLabs/tabpfn-time-series) for more details. - """ - - def __init__( - self, - features: list[FeatureGenerator] | None = None, - context_length: int = 4096, - mode: TabPFNMode | None = None, - api_key: str | None = None, - alias: str = "TabPFN", - ): - """ - Args: - features (list[FeatureGenerator], optional): List of TabPFN-TS feature - generators to use for feature engineering. If None, uses - `[RunningIndexFeature(), CalendarFeature(), AutoSeasonalFeature()]` - by default. - See - [TabPFN-TS features](https://github.com/PriorLabs/tabpfn-time-series/ - tree/main/tabpfn_time_series/features). - context_length (int, optional): Maximum context length (input window size) - for the model. Defaults to 4096. Controls how much history is used for - each forecast. - mode (TabPFNMode, optional): Inference mode for TabPFN. If None, uses LOCAL - (`"tabpfn-local"`) if a GPU is available, otherwise CLIENT (cloud - inference via `"tabpfn-client"`). See - [TabPFN-TS docs](https://github.com/PriorLabs/tabpfn-time-series/ - blob/3cd61ad556466de837edd1c6036744176145c024/tabpfn_time_series/ - predictor.py#L11) for available modes. - api_key (str, optional): API key for tabpfn-client cloud inference. Required - if using CLIENT mode and not already set in the environment. - alias (str, optional): Name to use for the model in output DataFrames and - logs. Defaults to "TabPFN". - - Notes: - **Academic Reference:** - - - Paper: [From Tables to Time: How TabPFN-v2 Outperforms - Specialized Time Series Forecasting Models](https://arxiv.org/abs/2501.02945) - - **Resources:** - - - GitHub: [PriorLabs/tabpfn-time-series](https://github.com/PriorLabs/tabpfn-time-series) - - **Technical Details:** - - - For LOCAL mode, a CUDA-capable GPU is recommended for best performance. - - The model is only available for Python < 3.13. - """ - if features is None: - features = [ - RunningIndexFeature(), - CalendarFeature(), - AutoSeasonalFeature(), - ] - self.feature_transformer = FeatureTransformer(features) - self.context_length = context_length - if mode is None: - mode = TabPFNMode.LOCAL if torch.cuda.is_available() else TabPFNMode.CLIENT - if mode == TabPFNMode.CLIENT and api_key is not None: - set_access_token(api_key) - self.mode = mode - self.alias = alias - - @contextmanager - def _get_model(self) -> TabPFNTimeSeriesPredictor: - model = TabPFNTimeSeriesPredictor(tabpfn_mode=self.mode) - try: - yield model - finally: - del model - torch.cuda.empty_cache() - - def _forecast( - self, - model: TabPFNTimeSeriesPredictor, - df: pd.DataFrame, - h: int, - quantiles: list[float] | None, - ) -> pd.DataFrame: - """handles distinction between quantiles and no quantiles""" - renamer = { - "unique_id": "item_id", - "ds": "timestamp", - "y": "target", - } - tsdf = df.rename(columns=renamer) - tsdf = TimeSeriesDataFrame(tsdf.set_index(["item_id", "timestamp"])) - if self.context_length > 0: - tsdf = tsdf.slice_by_timestep(-self.context_length, None) - future_tsdf = generate_test_X(tsdf, h) - tsdf, future_tsdf = self.feature_transformer.transform(tsdf, future_tsdf) - fcst_df = model.predict(tsdf, future_tsdf) - fcst_df = fcst_df.reset_index() - re_renamer = {v: k for k, v in renamer.items()} - re_renamer["target"] = self.alias - fcst_df = fcst_df.rename(columns=re_renamer) - if quantiles is None: - fcst_df = fcst_df[["unique_id", "ds", self.alias]] - else: - q_renamer = { - q_orig: f"{self.alias}-q-{int(100 * q_user)}" - for q_orig, q_user in zip( - TABPFN_TS_DEFAULT_QUANTILE_CONFIG, - quantiles, - strict=True, - ) - } - fcst_df = fcst_df.rename(columns=q_renamer) - return pd.DataFrame(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, - ) -> pd.DataFrame: - """Generate forecasts for time series data using the model. - - This method produces point forecasts and, optionally, prediction - intervals or quantile forecasts. The input DataFrame can contain one - or multiple time series in stacked (long) format. - - Args: - df (pd.DataFrame): - DataFrame containing the time series to forecast. It must - include as columns: - - - "unique_id": an ID column to distinguish multiple series. - - "ds": a time column indicating timestamps or periods. - - "y": a target column with the observed values. - - h (int): - Forecast horizon specifying how many future steps to predict. - freq (str, optional): - Frequency of the time series (e.g. "D" for daily, "M" for - monthly). See [Pandas frequency aliases](https://pandas.pydata.org/ - pandas-docs/stable/user_guide/timeseries.html#offset-aliases) for - valid values. If not provided, the frequency will be inferred - from the data. - level (list[int | float], optional): - Confidence levels for prediction intervals, expressed as - percentages (e.g. [80, 95]). If provided, the returned - DataFrame will include lower and upper interval columns for - each specified level. - quantiles (list[float], optional): - List of quantiles to forecast, expressed as floats between 0 - and 1. Should not be used simultaneously with `level`. When - provided, the output DataFrame will contain additional columns - named in the format "model-q-{percentile}", where {percentile} - = 100 × quantile value. - - Returns: - pd.DataFrame: - DataFrame containing forecast results. Includes: - - - point forecasts for each timestamp and series. - - prediction intervals if `level` is specified. - - quantile forecasts if `quantiles` is specified. - - For multi-series data, the output retains the same unique - identifiers as the input DataFrame. - """ - freq = self._maybe_infer_freq(df, freq) - qc = QuantileConverter(level=level, quantiles=quantiles) - if qc.quantiles is not None and not np.allclose( - qc.quantiles, - TABPFN_TS_DEFAULT_QUANTILE_CONFIG, - ): - raise ValueError( - "TabPFN only supports the default quantiles, " - "please use the default quantiles or default level, " - ) - with self._get_model() as model: - fcst_df = self._forecast( - model, - df, - h, - quantiles=qc.quantiles, - ) - if qc.quantiles is not None: - fcst_df = qc.maybe_convert_quantiles_to_level( - fcst_df, - models=[self.alias], - ) - return fcst_df +__all__ = ["TabPFN"] diff --git a/timecopilot/models/foundation/timegpt.py b/timecopilot/models/foundation/timegpt.py index 98984caf..03706285 100644 --- a/timecopilot/models/foundation/timegpt.py +++ b/timecopilot/models/foundation/timegpt.py @@ -1,191 +1,11 @@ -import os -from dataclasses import dataclass -from typing import Literal - -import pandas as pd -from nixtla import NixtlaClient +from foundationforecast.models.timegpt import TimeGPT as _TimeGPT +from foundationforecast.models.timegpt import TimeGPTFinetuningConfig from ..utils.forecaster import Forecaster -@dataclass -class TimeGPTFinetuningConfig: - """Configuration for finetuning TimeGPT before forecasting. - - Pass an instance to the ``TimeGPT`` constructor; when you call - ``forecast()``, the model is finetuned on the context data before - predicting. Parameters are passed through to ``NixtlaClient.forecast()``. - - Attributes: - finetune_steps: Number of training iterations. The model is trained - for this many steps on your data to minimize forecasting error. - Defaults to 10. - finetune_loss: Loss function used during finetuning. Options are - ``"default"``, ``"mae"``, ``"mse"``, ``"rmse"``, ``"mape"``, and - ``"smape"``. Defaults to ``"default"``. - finetune_depth: Controls how many model layers are finetuned, on a - scale from 1 (few parameters) to 5 (entire model). Higher values - increase training time and may overfit. Defaults to 1. - - Notes: - - Based on the [TimeGPT fine-tuning docs](https://docs.nixtla.io/forecasting/fine-tuning/steps). - """ - - finetune_steps: int = 10 - finetune_loss: Literal["default", "mae", "mse", "rmse", "mape", "smape"] = "default" - finetune_depth: Literal[1, 2, 3, 4, 5] = 1 - - -class TimeGPT(Forecaster): - """ - TimeGPT is a pre-trained foundation model for time series forecasting and anomaly - detection, developed by Nixtla. It is based on a large encoder-decoder transformer - architecture trained on over 100 billion data points from diverse domains. - See the [official repo](https://github.com/nixtla/nixtla), - [docs](https://www.nixtla.io/docs), - and [arXiv:2310.03589](https://arxiv.org/abs/2310.03589) for more details. - """ - - def __init__( - self, - api_key: str | None = None, - base_url: str | None = None, - max_retries: int = 1, - model: str = "timegpt-1", - alias: str = "TimeGPT", - finetuning_config: TimeGPTFinetuningConfig | None = None, - ): - """ - Args: - api_key (str, optional): API key for authenticating with the Nixtla TimeGPT - API. If not provided, will use the `NIXTLA_API_KEY` - environment variable. - base_url (str, optional): Base URL for the TimeGPT API. Defaults to the - official Nixtla endpoint. - max_retries (int, optional): Maximum number of retries for API requests. - Defaults to 1. - model (str, optional): Model name or version to use. Defaults to - "timegpt-1". See the [Nixtla docs](https://www.nixtla.io/docs) for - available models. - alias (str, optional): Name to use for the model in output DataFrames and - logs. Defaults to "TimeGPT". - finetuning_config (TimeGPTFinetuningConfig | None, optional): If - provided, the model is finetuned on the forecast context - data before predicting. See ``TimeGPTFinetuningConfig`` and - the [TimeGPT fine-tuning docs](https://docs.nixtla.io/forecasting/fine-tuning/steps) - for parameter details. - - Notes: - **Academic Reference:** - - - Paper: [TimeGPT-1](https://arxiv.org/abs/2310.03589) - - **Resources:** - - - GitHub: [Nixtla/nixtla](https://github.com/Nixtla/nixtla) - - **Technical Details:** - - - TimeGPT is a foundation model for time series forecasting designed for - production-ready forecasting with minimal setup. - - Provides zero-shot forecasting capabilities across various - domains and frequencies. - - Requires a valid API key from Nixtla to use. - - For more information, see the - [TimeGPT documentation](https://www.nixtla.io/docs). - """ - self.api_key = api_key - self.base_url = base_url - self.max_retries = max_retries - self.model = model - self.alias = alias - self.finetuning_config = finetuning_config - - def _get_client(self) -> NixtlaClient: - if self.api_key is None: # noqa: SIM108 - api_key = os.environ["NIXTLA_API_KEY"] - else: - api_key = self.api_key - return NixtlaClient( - api_key=api_key, - base_url=self.base_url, - max_retries=self.max_retries, - ) - - def forecast( - self, - df: pd.DataFrame, - h: int, - freq: str | None = None, - level: list[int | float] | None = None, - quantiles: list[float] | None = None, - ) -> pd.DataFrame: - """Generate forecasts for time series data using the model. - - This method produces point forecasts and, optionally, prediction - intervals or quantile forecasts. The input DataFrame can contain one - or multiple time series in stacked (long) format. - - Args: - df (pd.DataFrame): - DataFrame containing the time series to forecast. It must - include as columns: - - - "unique_id": an ID column to distinguish multiple series. - - "ds": a time column indicating timestamps or periods. - - "y": a target column with the observed values. - - h (int): - Forecast horizon specifying how many future steps to predict. - freq (str, optional): - Frequency of the time series (e.g. "D" for daily, "M" for - monthly). See [Pandas frequency aliases](https://pandas.pydata.org/ - pandas-docs/stable/user_guide/timeseries.html#offset-aliases) for - valid values. If not provided, the frequency will be inferred - from the data. - level (list[int | float], optional): - Confidence levels for prediction intervals, expressed as - percentages (e.g. [80, 95]). If provided, the returned - DataFrame will include lower and upper interval columns for - each specified level. - quantiles (list[float], optional): - List of quantiles to forecast, expressed as floats between 0 - and 1. Should not be used simultaneously with `level`. When - provided, the output DataFrame will contain additional columns - named in the format "model-q-{percentile}", where {percentile} - = 100 × quantile value. - - Returns: - pd.DataFrame: - DataFrame containing forecast results. Includes: - - - point forecasts for each timestamp and series. - - prediction intervals if `level` is specified. - - quantile forecasts if `quantiles` is specified. +class TimeGPT(_TimeGPT, Forecaster): + pass - For multi-series data, the output retains the same unique - identifiers as the input DataFrame. - When ``finetuning_config`` was set at construction, the model is - finetuned on ``df`` before predicting. - """ - freq = self._maybe_infer_freq(df, freq) - client = self._get_client() - finetune_kwargs: dict = {} - if self.finetuning_config is not None: - finetune_kwargs["finetune_steps"] = self.finetuning_config.finetune_steps - finetune_kwargs["finetune_loss"] = self.finetuning_config.finetune_loss - finetune_kwargs["finetune_depth"] = self.finetuning_config.finetune_depth - fcst_df = client.forecast( - df=df, - h=h, - freq=freq, - model=self.model, - level=level, - quantiles=quantiles, - **finetune_kwargs, - ) - fcst_df["ds"] = pd.to_datetime(fcst_df["ds"]) - cols = [col.replace("TimeGPT", self.alias) for col in fcst_df.columns] - fcst_df.columns = cols - return fcst_df +__all__ = ["TimeGPT", "TimeGPTFinetuningConfig"] diff --git a/timecopilot/models/foundation/timesfm.py b/timecopilot/models/foundation/timesfm.py index cc97b33f..cfc6f1c6 100644 --- a/timecopilot/models/foundation/timesfm.py +++ b/timecopilot/models/foundation/timesfm.py @@ -1,326 +1,10 @@ -import os -from contextlib import contextmanager +from foundationforecast.models.timesfm import TimesFM as _TimesFM -import numpy as np -import pandas as pd -import timesfm -import timesfm_v1 -import torch -from huggingface_hub import repo_exists -from timesfm import TimesFM_2p5_200M_torch -from timesfm_v1.timesfm_base import DEFAULT_QUANTILES as DEFAULT_QUANTILES_TFM -from tqdm import tqdm +from ..utils.forecaster import Forecaster -from ..utils.forecaster import Forecaster, QuantileConverter -from .utils import TimeSeriesDataset +class TimesFM(_TimesFM, Forecaster): + pass -class _TimesFMV1(Forecaster): - def __init__( - self, - repo_id: str, - context_length: int, - batch_size: int, - alias: str, - ): - self.repo_id = repo_id - self.context_length = context_length - self.batch_size = batch_size - self.alias = alias - @contextmanager - def _get_predictor( - self, - prediction_length: int, - quantiles: list[float] | None = None, - ) -> timesfm_v1.TimesFm: - backend = "gpu" if torch.cuda.is_available() else "cpu" - # these values are based on - # https://github.com/google-research/timesfm/blob/ba034ae71c2fc88eaf59f80b4a778cc2c0dca7d6/experiments/extended_benchmarks/run_timesfm.py#L91 - v2_version = "2.0" in self.repo_id - context_len = ( - min(self.context_length, 512) if not v2_version else self.context_length - ) - num_layers = 50 if v2_version else 20 - use_positional_embedding = not v2_version - - tfm_hparams = timesfm_v1.TimesFmHparams( - backend=backend, - horizon_len=prediction_length, - quantiles=quantiles, - context_len=context_len, - num_layers=num_layers, - use_positional_embedding=use_positional_embedding, - per_core_batch_size=self.batch_size, - ) - if os.path.exists(self.repo_id): - path = os.path.join(self.repo_id, "torch_model.ckpt") - tfm_checkpoint = timesfm_v1.TimesFmCheckpoint(path=path) - tfm = timesfm_v1.TimesFm( - hparams=tfm_hparams, - checkpoint=tfm_checkpoint, - ) - elif repo_exists(self.repo_id): - tfm_checkpoint = timesfm_v1.TimesFmCheckpoint( - huggingface_repo_id=self.repo_id - ) - tfm = timesfm_v1.TimesFm( - hparams=tfm_hparams, - checkpoint=tfm_checkpoint, - ) - else: - raise OSError( - f"Failed to load model. Searched for '{self.repo_id}' " - "as a local path to model directory and as a Hugging Face repo_id." - ) - - try: - yield tfm - finally: - del tfm - torch.cuda.empty_cache() - - def forecast( - self, - df: pd.DataFrame, - h: int, - freq: str | None = None, - level: list[int | float] | None = None, - quantiles: list[float] | None = None, - ) -> pd.DataFrame: - freq = self._maybe_infer_freq(df, freq) - qc = QuantileConverter(level=level, quantiles=quantiles) - if qc.quantiles is not None and len(qc.quantiles) != len(DEFAULT_QUANTILES_TFM): - raise ValueError( - "TimesFM only supports the default quantiles, " - "please use the default quantiles or default level, " - "see https://github.com/google-research/timesfm/issues/286" - ) - with self._get_predictor( - prediction_length=h, - quantiles=qc.quantiles or DEFAULT_QUANTILES_TFM, - ) as predictor: - fcst_df = predictor.forecast_on_df( - inputs=df, - freq=freq, - value_name="y", - model_name=self.alias, - num_jobs=1, - ) - if qc.quantiles is not None: - renamer = { - f"{self.alias}-q-{q}": f"{self.alias}-q-{int(q * 100)}" - for q in qc.quantiles - } - fcst_df = fcst_df.rename(columns=renamer) - fcst_df = qc.maybe_convert_quantiles_to_level( - fcst_df, - models=[self.alias], - ) - else: - fcst_df = fcst_df[["unique_id", "ds", self.alias]] - return fcst_df - - -class _TimesFMV2_p5(Forecaster): - def __init__( - self, - repo_id: str, - context_length: int, - batch_size: int, - alias: str, - **kwargs: dict, - ): - self.repo_id = repo_id - self.context_length = context_length - self.batch_size = batch_size - self.alias = alias - self.kwargs = kwargs - - @contextmanager - def _get_predictor( - self, - prediction_length: int, - ) -> TimesFM_2p5_200M_torch: - # `from_pretrained` handles both a local directory containing - # `model.safetensors` and a Hugging Face repo id, and the model picks - # the best available device on load. - if os.path.exists(self.repo_id) or repo_exists(self.repo_id): - tfm = TimesFM_2p5_200M_torch.from_pretrained(self.repo_id) - else: - raise OSError( - f"Failed to load model. Searched for '{self.repo_id}' " - "as a local path to model directory and as a Hugging Face repo_id." - ) - default_kwargs = { - "max_context": self.context_length, - "max_horizon": prediction_length, - "normalize_inputs": True, - "use_continuous_quantile_head": True, - "fix_quantile_crossing": True, - } - passed_kwargs = self.kwargs or {} - kwargs = {**default_kwargs, **passed_kwargs} - config = timesfm.ForecastConfig(**kwargs) - tfm.compile(config) - try: - yield tfm - finally: - del tfm - torch.cuda.empty_cache() - - def _predict( - self, - model: TimesFM_2p5_200M_torch, - dataset: TimeSeriesDataset, - h: int, - ) -> tuple[np.ndarray, np.ndarray]: - fcsts = [ - model.forecast( - inputs=batch, - horizon=h, - ) - for batch in tqdm(dataset) - ] - fcsts_mean, fcsts_quantiles = zip(*fcsts, strict=False) - fcsts_mean_np = np.concatenate(fcsts_mean) - fcsts_quantiles_np = np.concatenate(fcsts_quantiles) - return fcsts_mean_np, fcsts_quantiles_np - - def forecast( - self, - df: pd.DataFrame, - h: int, - freq: str | None = None, - level: list[int | float] | None = None, - quantiles: list[float] | None = None, - ) -> pd.DataFrame: - freq = self._maybe_infer_freq(df, freq) - qc = QuantileConverter(level=level, quantiles=quantiles) - if qc.quantiles is not None and len(qc.quantiles) != len(DEFAULT_QUANTILES_TFM): - raise ValueError( - "TimesFM only supports the default quantiles, " - "please use the default quantiles or default level, " - "see https://github.com/google-research/timesfm/issues/286" - ) - dataset = TimeSeriesDataset.from_df( - df, - batch_size=self.batch_size, - dtype=torch.float32, - ) - fcst_df = dataset.make_future_dataframe(h=h, freq=freq) - with self._get_predictor(prediction_length=h) as model: - fcsts_mean_np, fcsts_quantiles_np = self._predict( - model, - dataset, - h, - ) - fcst_df[self.alias] = fcsts_mean_np.reshape(-1, 1) - if qc.quantiles is not None: - for i, q in enumerate(qc.quantiles): - fcst_df[f"{self.alias}-q-{int(q * 100)}"] = fcsts_quantiles_np[ - ..., i + 1 # skip the first quantile (mean) - ].reshape(-1, 1) - fcst_df = qc.maybe_convert_quantiles_to_level( - fcst_df, - models=[self.alias], - ) - return fcst_df - - -class TimesFM(Forecaster): - """ - TimesFM is a large time series model for time series forecasting, supporting both - probabilistic and point forecasts. See the [official repo](https://github.com/ - google-research/timesfm) for more details. - """ - - def __new__( - cls, - repo_id: str = "google/timesfm-2.0-500m-pytorch", - context_length: int = 2048, - batch_size: int = 64, - alias: str = "TimesFM", - **kwargs: dict, - ): - if "pytorch" not in repo_id: - raise ValueError( - "TimesFM only supports pytorch models, " - "if you'd like to use jax, please open an issue" - ) - if "1.0" in repo_id or "2.0" in repo_id: - return _TimesFMV1( - repo_id=repo_id, - context_length=context_length, - batch_size=batch_size, - alias=alias, - ) - elif "2.5" in repo_id: - return _TimesFMV2_p5( - repo_id=repo_id, - context_length=context_length, - batch_size=batch_size, - alias=alias, - **kwargs, - ) - else: - raise ValueError( - "TimesFM only supports 1.0, 2.0 and 2.5 models, please use a " - "valid model id" - ) - - def __init__( - self, - repo_id: str = "google/timesfm-2.0-500m-pytorch", - context_length: int = 2048, - batch_size: int = 64, - alias: str = "TimesFM", - kwargs: dict | None = None, - ): - """ - Args: - repo_id (str, optional): The Hugging Face Hub model ID or local path to - load the TimesFM model from. Examples include - `google/timesfm-2.0-500m-pytorch`. Defaults to - `google/timesfm-2.0-500m-pytorch`. See the full list of models at - [Hugging Face](https://huggingface.co/collections/google/timesfm-release- - 66e4be5fdb56e960c1e482a6). Supported models: - - - `google/timesfm-1.0-200m-pytorch` - - `google/timesfm-2.0-500m-pytorch` - - `google/timesfm-2.5-200m-pytorch` - context_length (int, optional): Maximum context length (input window size) - for the model. Defaults to 2048. For TimesFM 2.0 models, max is 2048 - (must be a multiple of 32). For TimesFM 1.0 models, max is 512. See - [TimesFM docs](https://github.com/google-research/timesfm#loading-the- - model) for details. - batch_size (int, optional): Batch size for inference. Defaults to 64. - Adjust based on available memory and model size. - alias (str, optional): Name to use for the model in output DataFrames and - logs. Defaults to `TimesFM`. - kwargs (dict, optional): Additional keyword arguments to pass to the model. - Defaults to None. Only used for TimesFM 2.5 models. - - Notes: - **Academic Reference:** - - - Paper: [A decoder-only foundation model for time-series forecasting](https://arxiv.org/abs/2310.10688) - - **Resources:** - - - GitHub: [google-research/timesfm](https://github.com/google-research/timesfm) - - HuggingFace: [google/timesfm-release](https://huggingface.co/collections/google/timesfm-release-66e4be5fdb56e960c1e482a6) - - **Technical Details:** - - - Only PyTorch checkpoints are currently supported. JAX is not supported. - - The model is loaded onto the best available device (GPU if available, - otherwise CPU). - - **Supported Models:** - - - `google/timesfm-1.0-200m-pytorch` - - `google/timesfm-2.0-500m-pytorch` - - `google/timesfm-2.5-200m-pytorch` - """ - pass +__all__ = ["TimesFM"] diff --git a/timecopilot/models/foundation/tirex.py b/timecopilot/models/foundation/tirex.py index a529226b..0249e93e 100644 --- a/timecopilot/models/foundation/tirex.py +++ b/timecopilot/models/foundation/tirex.py @@ -1,273 +1,10 @@ -from __future__ import annotations +from foundationforecast.models.tirex import TiRex as _TiRex -import os -import sys -from contextlib import contextmanager -from typing import TYPE_CHECKING +from ..utils.forecaster import Forecaster -if sys.version_info < (3, 11): - raise ImportError("TiRex requires Python >= 3.11") -import numpy as np -import pandas as pd -import torch -from tirex import load_model -from tirex.base import PretrainedModel -from tqdm import tqdm +class TiRex(_TiRex, Forecaster): + pass -from ..utils.forecaster import Forecaster, QuantileConverter -from .utils import TimeSeriesDataset -if TYPE_CHECKING: - from tirex2.api_adapter import ForecastModel - -DEFAULT_QUANTILES_TIREX = [0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9] -_MEDIAN_QUANTILE_IDX = DEFAULT_QUANTILES_TIREX.index(0.5) - - -class TiRex(Forecaster): - """ - TiRex is a family of zero-shot time series forecasting models based on - xLSTM, supporting both point and quantile predictions. This class - transparently supports TiRex 1.0 and TiRex 2.0 checkpoints, dispatching - to the appropriate backend based on `repo_id`. See the - [TiRex repo](https://github.com/NX-AI/tirex) and - [TiRex-2 repo](https://github.com/NX-AI/tirex-2) for more details. - """ - - def __init__( - self, - repo_id: str = "NX-AI/TiRex", - batch_size: int = 16, - alias: str = "TiRex", - ): - """ - Args: - repo_id (str, optional): The Hugging Face Hub model ID or local path to load - the TiRex model from. Use "NX-AI/TiRex" for TiRex 1.0 or - "NX-AI/TiRex-2" for TiRex 2.0. Defaults to "NX-AI/TiRex". - See the full list of models at - [Hugging Face](https://huggingface.co/NX-AI). - batch_size (int, optional): Batch size to use for inference. Defaults to 16. - Adjust based on available memory and model size. - alias (str, optional): Name to use for the model in output DataFrames - and logs. Defaults to "TiRex". - - Notes: - **Academic References:** - - - TiRex 1.0: [TiRex: Zero-shot Time Series Forecasting with xLSTM](https://arxiv.org/abs/2505.23719) - - TiRex 2.0: [TiRex-2: Generalizing TiRex to Multivariate Data and Streaming](https://arxiv.org/abs/2607.01204) - - **Resources:** - - - GitHub: [NX-AI/tirex](https://github.com/NX-AI/tirex), - [NX-AI/tirex-2](https://github.com/NX-AI/tirex-2) - - HuggingFace: [NX-AI Models](https://huggingface.co/NX-AI) - - **Technical Details:** - - - TiRex 2.0 is loaded onto the best available device (CUDA, then MPS, - otherwise CPU). TiRex 1.0 uses CUDA when available, otherwise CPU. - - TiRex 1.0 on CPU disables CUDA kernels automatically. See the - [CUDA kernels section](https://github.com/NX-AI/tirex#cuda-kernels) - for details. - - TiRex 2.0 natively supports CPU, CUDA, and MPS devices. - - The model is only available for Python >= 3.11. - """ - self.repo_id = repo_id - self.batch_size = batch_size - self.alias = alias - - def _is_tirex2(self) -> bool: - repo = self.repo_id.rstrip("/") - return repo.endswith("TiRex-2") or repo.split("/")[-1] == "TiRex-2" - - @staticmethod - def _best_device_v2() -> str: - if torch.cuda.is_available(): - return "cuda" - if torch.backends.mps.is_available(): - return "mps" - return "cpu" - - @contextmanager - def _get_model(self) -> PretrainedModel | ForecastModel: - if self._is_tirex2(): - with self._get_model_v2() as model: - yield model - else: - with self._get_model_v1() as model: - yield model - - @contextmanager - def _get_model_v1(self) -> PretrainedModel: - device = "cuda" if torch.cuda.is_available() else "cpu" - if device == "cpu": - # see https://github.com/NX-AI/tirex/tree/main?tab=readme-ov-file#cuda-kernels - os.environ["TIREX_NO_CUDA"] = "1" - model = load_model(self.repo_id, device=device) - try: - yield model - finally: - del model - torch.cuda.empty_cache() - - @contextmanager - def _get_model_v2(self) -> ForecastModel: - from tirex2 import load_model as load_model_v2 - - device = self._best_device_v2() - model = load_model_v2(self.repo_id, device=device) - try: - yield model - finally: - del model - if device == "cuda": - torch.cuda.empty_cache() - elif device == "mps": - torch.mps.empty_cache() - - def _forecast_v1( - self, - model: PretrainedModel, - dataset: TimeSeriesDataset, - h: int, - quantiles: list[float] | None, - ) -> tuple[np.ndarray, np.ndarray | None]: - """handles distinction between quantiles and no quantiles""" - fcsts = [ - model.forecast( - batch, - prediction_length=h, - output_type="numpy", - ) - for batch in tqdm(dataset) - ] # list of tuples - fcsts_quantiles, fcsts_mean = zip(*fcsts, strict=False) - fcsts_mean_np = np.concatenate(fcsts_mean) - fcsts_quantiles_np = ( - None if quantiles is None else np.concatenate(fcsts_quantiles) - ) - - return fcsts_mean_np, fcsts_quantiles_np - - def _forecast_v2( - self, - model: ForecastModel, - dataset: TimeSeriesDataset, - h: int, - quantiles: list[float] | None, - ) -> tuple[np.ndarray, np.ndarray | None]: - from tirex2 import TimeseriesType - - timeseries = [ - TimeseriesType( - target=ts.float().unsqueeze(0), - past_covariates=None, - future_covariates=None, - ) - for ts in dataset.data - ] - forecasts = model.forecast( - timeseries=timeseries, - prediction_length=h, - output_type="numpy", - batch_size=self.batch_size, - ) - fcsts_mean_np = np.concatenate( - [f[0, _MEDIAN_QUANTILE_IDX, :] for f in forecasts], - ) - fcsts_quantiles_np = ( - None if quantiles is None else np.concatenate([f[0].T for f in forecasts]) - ) - return fcsts_mean_np, fcsts_quantiles_np - - def forecast( - self, - df: pd.DataFrame, - h: int, - freq: str | None = None, - level: list[int | float] | None = None, - quantiles: list[float] | None = None, - ) -> pd.DataFrame: - """Generate forecasts for time series data using the model. - - This method produces point forecasts and, optionally, prediction - intervals or quantile forecasts. The input DataFrame can contain one - or multiple time series in stacked (long) format. - - Args: - df (pd.DataFrame): - DataFrame containing the time series to forecast. It must - include as columns: - - - "unique_id": an ID column to distinguish multiple series. - - "ds": a time column indicating timestamps or periods. - - "y": a target column with the observed values. - - h (int): - Forecast horizon specifying how many future steps to predict. - freq (str, optional): - Frequency of the time series (e.g. "D" for daily, "M" for - monthly). See [Pandas frequency aliases](https://pandas.pydata.org/ - pandas-docs/stable/user_guide/timeseries.html#offset-aliases) for - valid values. If not provided, the frequency will be inferred - from the data. - level (list[int | float], optional): - Confidence levels for prediction intervals, expressed as - percentages (e.g. [80, 95]). If provided, the returned - DataFrame will include lower and upper interval columns for - each specified level. - quantiles (list[float], optional): - List of quantiles to forecast, expressed as floats between 0 - and 1. Should not be used simultaneously with `level`. When - provided, the output DataFrame will contain additional columns - named in the format "model-q-{percentile}", where {percentile} - = 100 × quantile value. - - Returns: - pd.DataFrame: - DataFrame containing forecast results. Includes: - - - point forecasts for each timestamp and series. - - prediction intervals if `level` is specified. - - quantile forecasts if `quantiles` is specified. - - For multi-series data, the output retains the - same unique identifiers as the input DataFrame. - """ - freq = self._maybe_infer_freq(df, freq) - qc = QuantileConverter(level=level, quantiles=quantiles) - if qc.quantiles is not None and len(qc.quantiles) != len( - DEFAULT_QUANTILES_TIREX - ): - raise ValueError( - "TiRex only supports the default quantiles, " - "please use the default quantiles or default level, " - ) - dataset = TimeSeriesDataset.from_df( - df, - batch_size=self.batch_size, - ) - - fcst_df = dataset.make_future_dataframe(h=h, freq=freq) - forecast_fn = self._forecast_v2 if self._is_tirex2() else self._forecast_v1 - with self._get_model() as model: - fcsts_mean_np, fcsts_quantiles_np = forecast_fn( - model, - dataset, - h, - quantiles=qc.quantiles, - ) - 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 = qc.maybe_convert_quantiles_to_level( - fcst_df, - models=[self.alias], - ) - return fcst_df +__all__ = ["TiRex"] diff --git a/timecopilot/models/foundation/toto.py b/timecopilot/models/foundation/toto.py index 059fe1b3..46bf03e9 100644 --- a/timecopilot/models/foundation/toto.py +++ b/timecopilot/models/foundation/toto.py @@ -1,430 +1,10 @@ -import json -from contextlib import contextmanager -from pathlib import Path +from foundationforecast.models.toto import Toto as _Toto -import numpy as np -import pandas as pd -import torch -from huggingface_hub import hf_hub_download -from toto.data.util.dataset import MaskedTimeseries -from toto.inference.forecaster import TotoForecaster -from toto.model.toto import Toto as TotoModel -from toto2 import Toto2Model -from tqdm import tqdm +from ..utils.forecaster import Forecaster -from ..utils.forecaster import Forecaster, QuantileConverter -from .utils import 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. -_TOTO2_CONFIG_KEY = "num_variate_layers_per_group" +class Toto(_Toto, Forecaster): + pass -class Toto(Forecaster): - """ - Toto is a family of foundation models for multivariate time series - forecasting, optimized for observability and high-dimensional data. This - class transparently supports both Toto 1.0 and Toto 2.0 checkpoints, - dispatching to the appropriate backend based on the loaded model. See the - [official repo](https://github.com/DataDog/toto) for more details. - """ - - def __init__( - self, - repo_id: str = "Datadog/Toto-Open-Base-1.0", - context_length: int = 4096, - batch_size: int = 16, - num_samples: int = 128, - samples_per_batch: int = 8, - decode_block_size: int | None = None, - alias: str = "Toto", - ): - # ruff: noqa: E501 - """ - Args: - repo_id (str, optional): The Hugging Face Hub model ID or local path to - load the Toto model from. This can be either a Toto 1.0 checkpoint - (e.g. "Datadog/Toto-Open-Base-1.0") or a Toto 2.0 checkpoint (e.g. - "Datadog/Toto-2.0-4m"). The model family is detected automatically - from the checkpoint configuration. Defaults to - "Datadog/Toto-Open-Base-1.0". See the full list of models at - [Hugging Face](https://huggingface.co/Datadog). - context_length (int, optional): Maximum context length (input window size) - for the model. Defaults to 4096. Should match the configuration of the - pretrained checkpoint. See [Toto docs](https://github.com/DataDog/toto# - toto-model) for details. - batch_size (int, optional): Batch size to use for inference. Defaults to 16. - Adjust based on available memory and model size. - num_samples (int, optional): Number of samples for probabilistic - forecasting. Controls the number of forecast samples drawn for - uncertainty estimation. Defaults to 128. Only used by Toto 1.0 - checkpoints; ignored by Toto 2.0, which predicts fixed quantile knots - directly. - samples_per_batch (int, optional): Number of samples processed per batch - during inference. Controls memory usage. Defaults to 8. Only used by - Toto 1.0 checkpoints; ignored by Toto 2.0. - decode_block_size (int | None, optional): Block size for Toto 2.0 block - decoding, expressed in time steps and divisible by the model patch - size. When None (default), Toto 2.0 forecasts in a single forward - pass, which is faster and better for short horizons. Larger values - (e.g. 768) improve long-term stability for very long horizons. Only - used by Toto 2.0 checkpoints; ignored by Toto 1.0. - alias (str, optional): Name to use for the model in output DataFrames and - logs. Defaults to "Toto". - - Notes: - **Available models:** - - *Toto 1.0 (sample-based):* - - | Model ID | Parameters | - | ----------------------------------------------------------------------------------- | ---------- | - | [`Datadog/Toto-Open-Base-1.0`](https://huggingface.co/Datadog/Toto-Open-Base-1.0) | 151M | - - *Toto 2.0 (quantile-knot based):* - - | Model ID | Parameters | - | --------------------------------------------------------------------- | ---------- | - | [`Datadog/Toto-2.0-4m`](https://huggingface.co/Datadog/Toto-2.0-4m) | 4M | - | [`Datadog/Toto-2.0-22m`](https://huggingface.co/Datadog/Toto-2.0-22m) | 22M | - | [`Datadog/Toto-2.0-313m`](https://huggingface.co/Datadog/Toto-2.0-313m) | 313M | - | [`Datadog/Toto-2.0-1B`](https://huggingface.co/Datadog/Toto-2.0-1B) | 1B | - | [`Datadog/Toto-2.0-2.5B`](https://huggingface.co/Datadog/Toto-2.0-2.5B) | 2.5B | - - **Academic Reference:** - - - Paper (Toto 1.0): [Building a Foundation Model for Time Series](https://arxiv.org/abs/2505.14766) - - Paper (Toto 2.0): [Toto 2.0: Time Series Forecasting Enters the Scaling Era](https://arxiv.org/abs/2605.20119) - - **Resources:** - - - GitHub: [DataDog/toto](https://github.com/DataDog/toto) - - HuggingFace: [Datadog Models](https://huggingface.co/Datadog) - - **Technical Details:** - - - The model is loaded onto the best available device (GPU if available, - otherwise CPU). - - For best performance, a CUDA-capable GPU is recommended. - - Toto 1.0 draws probabilistic samples and reports the sample mean as the - point forecast, with exact sample quantiles. Toto 2.0 predicts a fixed - set of quantile knots (0.1, 0.2, ..., 0.9); the median (0.5) is used as - the point forecast and requested quantiles are obtained by linear - interpolation across the knots. - """ - self.repo_id = repo_id - self.context_length = context_length - self.batch_size = batch_size - # Number of samples for probabilistic forecasting (Toto 1.0 only) - self.num_samples = num_samples - # Control memory usage during inference (Toto 1.0 only) - self.samples_per_batch = samples_per_batch - # Block decoding size (Toto 2.0 only) - self.decode_block_size = decode_block_size - self.alias = alias - self.device = "cuda" if torch.cuda.is_available() else "cpu" - self._is_toto2_cache: bool | None = None - - def _is_toto2(self) -> bool: - """Detect whether ``repo_id`` points to a Toto 2.0 checkpoint. - - Detection is based on the model ``config.json``: Toto 2.0 - configurations contain Toto2-specific fields that are absent from - Toto 1.0 checkpoints. The result is cached on the instance. - """ - if self._is_toto2_cache is not None: - return self._is_toto2_cache - repo_path = Path(self.repo_id) - if repo_path.is_dir(): - config_path = repo_path / "config.json" - else: - config_path = Path( - hf_hub_download(repo_id=self.repo_id, filename="config.json") - ) - config = json.loads(config_path.read_text()) - self._is_toto2_cache = _TOTO2_CONFIG_KEY in config - return self._is_toto2_cache - - @contextmanager - def _get_model(self) -> TotoForecaster | Toto2Model: - if self._is_toto2(): - model = Toto2Model.from_pretrained(self.repo_id).to(self.device).eval() - try: - yield model - finally: - del model - torch.cuda.empty_cache() - else: - model = TotoModel.from_pretrained(self.repo_id).to(self.device) - try: - yield TotoForecaster(model.model) - finally: - del model - torch.cuda.empty_cache() - - def _to_masked_timeseries(self, batch: list[torch.Tensor]) -> MaskedTimeseries: - batch_size = len(batch) - # using toch.float as stated in the docs - # https://github.com/DataDog/toto/blob/main/toto/notebooks/inference_tutorial.ipynb - padded_tensor = torch.zeros( - batch_size, - self.context_length, - dtype=torch.float, - device=self.device, - ) - padding_mask = torch.zeros( - batch_size, - self.context_length, - dtype=torch.float, - device=self.device, - ) - for idx, ts in enumerate(batch): - series_length = len(ts) - if series_length > self.context_length: - ts = ts[-self.context_length :] - series_length = self.context_length - padded_tensor[idx, -series_length:] = ts.to( - device=self.device, dtype=torch.float - ) - padding_mask[idx, -series_length:] = 1.0 - masked_ts = MaskedTimeseries( - series=padded_tensor, - padding_mask=padding_mask, - id_mask=torch.zeros_like(padded_tensor), - # Prepare timestamp information (optional, but expected by API; - # not used by the current model release) - timestamp_seconds=torch.zeros_like(padded_tensor), - time_interval_seconds=torch.full( - (batch_size,), - 1, - device=self.device, - ), - ) - return masked_ts - - def _to_toto2_inputs(self, batch: list[torch.Tensor]) -> dict[str, torch.Tensor]: - """Build the Toto 2.0 ``forecast`` inputs dict for a batch. - - Produces left-padded ``target``/``target_mask`` tensors of shape - ``(batch, n_var=1, context_length)`` together with zero - ``series_ids`` of shape ``(batch, n_var=1)``. - """ - batch_size = len(batch) - padded_tensor = torch.zeros( - batch_size, - self.context_length, - dtype=torch.float, - device=self.device, - ) - padding_mask = torch.zeros( - batch_size, - self.context_length, - dtype=torch.bool, - device=self.device, - ) - for idx, ts in enumerate(batch): - series_length = len(ts) - if series_length > self.context_length: - ts = ts[-self.context_length :] - series_length = self.context_length - padded_tensor[idx, -series_length:] = ts.to( - device=self.device, dtype=torch.float - ) - padding_mask[idx, -series_length:] = True - # add the variate dimension (n_var=1) - target = padded_tensor.unsqueeze(1) - target_mask = padding_mask.unsqueeze(1) - series_ids = torch.zeros( - batch_size, - 1, - dtype=torch.long, - device=self.device, - ) - return { - "target": target, - "target_mask": target_mask, - "series_ids": series_ids, - } - - def _forecast( - self, - model: TotoForecaster, - dataset: TimeSeriesDataset, - h: int, - quantiles: list[float] | None, - ) -> tuple[np.ndarray, np.ndarray | None]: - """handles distinction between quantiles and no quantiles""" - fcsts = [ - model.forecast( - self._to_masked_timeseries(batch), - prediction_length=h, - num_samples=self.num_samples, - samples_per_batch=self.samples_per_batch, - use_kv_cache=True, - ) - for batch in tqdm(dataset) - ] # list of fcsts objects - - fcsts_mean = [fcst.median.cpu().numpy() for fcst in fcsts] - fcsts_mean_np = np.concatenate(fcsts_mean, axis=1) - if fcsts_mean_np.shape[0] != 1: - raise ValueError( - f"fcsts_mean_np.shape[0] != 1: {fcsts_mean_np.shape[0]} != 1, " - "this is not expected, please open an issue on github" - ) - fcsts_mean_np = fcsts_mean_np.squeeze(axis=0) - if quantiles is not None: - quantiles_torch = torch.tensor( - quantiles, - device=self.device, - dtype=torch.float, - ) - fcsts_quantiles = [ - fcst.quantile(quantiles_torch).cpu().numpy() for fcst in fcsts - ] - fcsts_quantiles_np = np.concatenate(fcsts_quantiles, axis=2) - if fcsts_quantiles_np.shape[1] != 1: - raise ValueError( - "fcsts_quantiles_np.shape[1] != 1: " - f"{fcsts_quantiles_np.shape[1]} != 1, " - "this is not expected, please open an issue on github" - ) - fcsts_quantiles_np = np.moveaxis(fcsts_quantiles_np, 0, -1).squeeze(axis=0) - else: - fcsts_quantiles_np = None - return fcsts_mean_np, fcsts_quantiles_np - - def _forecast_toto2( - self, - model: Toto2Model, - dataset: TimeSeriesDataset, - h: int, - quantiles: list[float] | None, - ) -> tuple[np.ndarray, np.ndarray | None]: - """Forecast with a Toto 2.0 model. - - Toto 2.0 predicts a fixed set of quantile knots (0.1, ..., 0.9). The - median knot is used as the point forecast, and requested quantiles are - obtained by linear interpolation across the knots. - """ - knots = list(model.output_head.knots) - median_idx = knots.index(0.5) - # collect per-batch quantile knots of shape (n_knots, batch, h) - knot_fcsts: list[np.ndarray] = [] - for batch in tqdm(dataset): - inputs = self._to_toto2_inputs(batch) - with torch.no_grad(): - # shape: (n_knots, batch, n_var=1, h) - q = model.forecast( - inputs, - horizon=h, - decode_block_size=self.decode_block_size, - has_missing_values=True, - ) - q_np = q.float().cpu().numpy() - if q_np.shape[2] != 1: - raise ValueError( - f"toto2 forecast n_var != 1: {q_np.shape[2]} != 1, " - "this is not expected, please open an issue on github" - ) - knot_fcsts.append(q_np.squeeze(axis=2)) - # shape: (n_knots, n_series, h) - knots_np = np.concatenate(knot_fcsts, axis=1) - fcsts_mean_np = knots_np[median_idx] - if quantiles is not None: - # interpolate requested quantiles across the fixed knots, per - # (series, horizon) position. np.interp clamps at the edge knots. - knots_arr = np.asarray(knots) - fcsts_quantiles_np = np.stack( - [ - np.apply_along_axis( - lambda col, _q=q: np.interp(_q, knots_arr, col), - axis=0, - arr=knots_np, - ) - for q in quantiles - ], - axis=-1, - ) - else: - fcsts_quantiles_np = None - return fcsts_mean_np, fcsts_quantiles_np - - def forecast( - self, - df: pd.DataFrame, - h: int, - freq: str | None = None, - level: list[int | float] | None = None, - quantiles: list[float] | None = None, - ) -> pd.DataFrame: - """Generate forecasts for time series data using the model. - - This method produces point forecasts and, optionally, prediction - intervals or quantile forecasts. The input DataFrame can contain one - or multiple time series in stacked (long) format. - - Args: - df (pd.DataFrame): - DataFrame containing the time series to forecast. It must - include as columns: - - - "unique_id": an ID column to distinguish multiple series. - - "ds": a time column indicating timestamps or periods. - - "y": a target column with the observed values. - - h (int): - Forecast horizon specifying how many future steps to predict. - freq (str, optional): - Frequency of the time series (e.g. "D" for daily, "M" for - monthly). See [Pandas frequency aliases](https://pandas.pydata.org/ - pandas-docs/stable/user_guide/timeseries.html#offset-aliases) for - valid values. If not provided, the frequency will be inferred - from the data. - level (list[int | float], optional): - Confidence levels for prediction intervals, expressed as - percentages (e.g. [80, 95]). If provided, the returned - DataFrame will include lower and upper interval columns for - each specified level. - quantiles (list[float], optional): - List of quantiles to forecast, expressed as floats between 0 - and 1. Should not be used simultaneously with `level`. When - provided, the output DataFrame will contain additional columns - named in the format "model-q-{percentile}", where {percentile} - = 100 × quantile value. For Toto 2.0 checkpoints, quantiles are - linearly interpolated across the model's fixed knots. - - Returns: - pd.DataFrame: - DataFrame containing forecast results. Includes: - - - point forecasts for each timestamp and series. - - prediction intervals if `level` is specified. - - quantile forecasts if `quantiles` is specified. - - For multi-series data, the output retains the same unique - identifiers as the input DataFrame. - """ - freq = self._maybe_infer_freq(df, freq) - qc = QuantileConverter(level=level, quantiles=quantiles) - dataset = TimeSeriesDataset.from_df(df, batch_size=self.batch_size) - 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: - fcsts_mean_np, fcsts_quantiles_np = forecast_fn( - model, - dataset, - h, - quantiles=qc.quantiles, - ) - 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 = qc.maybe_convert_quantiles_to_level( - fcst_df, - models=[self.alias], - ) - return fcst_df +__all__ = ["Toto"] diff --git a/timecopilot/models/foundation/utils.py b/timecopilot/models/foundation/utils.py deleted file mode 100644 index 287be321..00000000 --- a/timecopilot/models/foundation/utils.py +++ /dev/null @@ -1,62 +0,0 @@ -from collections.abc import Iterable - -import pandas as pd -import torch -from utilsforecast.processing import make_future_dataframe - - -class TimeSeriesDataset: - def __init__( - self, - data: torch.Tensor, - uids: Iterable, - last_times: Iterable, - batch_size: int, - ): - self.data = data - 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.current_batch = 0 - - @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) - - def __len__(self): - return self.n_batches - - def make_future_dataframe(self, h: int, freq: str) -> pd.DataFrame: - return make_future_dataframe( - uids=self.uids, - last_times=pd.to_datetime(self.last_times), - h=h, - freq=freq, - ) # type: ignore - - def __iter__(self): - self.current_batch = 0 # Reset for new iteration - 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 diff --git a/timecopilot/models/prophet.py b/timecopilot/models/prophet.py index 84bd1e8f..3ee7745f 100644 --- a/timecopilot/models/prophet.py +++ b/timecopilot/models/prophet.py @@ -122,6 +122,9 @@ def _local_forecast_impl( fcst_df = qc.maybe_convert_quantiles_to_level(fcst_df, models=[self.alias]) return fcst_df + def _anomaly_min_series_length(self, h: int) -> int: + return h + 2 + def _local_forecast( self, df: pd.DataFrame, diff --git a/timecopilot/models/utils/forecaster.py b/timecopilot/models/utils/forecaster.py index 21205490..9fac3b1d 100644 --- a/timecopilot/models/utils/forecaster.py +++ b/timecopilot/models/utils/forecaster.py @@ -1,670 +1,23 @@ -import matplotlib.pyplot as plt -import numpy as np -import pandas as pd -import plotly.graph_objects -import torch -import utilsforecast.processing as ufp -from gluonts.time_feature.seasonality import ( - DEFAULT_SEASONALITIES, +from foundationforecast.core.forecaster import Forecaster as _Forecaster +from foundationforecast.core.forecaster import ( + QuantileConverter, + _DataProcessor, + get_seasonality, + maybe_convert_col_to_datetime, + maybe_infer_freq, ) -from gluonts.time_feature.seasonality import ( - get_seasonality as _get_seasonality, -) -from gluonts.transform import LastValueImputation -from prophet import Prophet as ProphetBase -from scipy import stats -from tqdm import tqdm -from utilsforecast.plotting import plot_series -from utilsforecast.processing import ( - backtest_splits, - drop_index_if_pandas, - join, - maybe_compute_sort_indices, - take_rows, - vertical_concat, -) -from utilsforecast.validation import ensure_time_dtype - - -def get_seasonality( - freq: str, - custom_seasonalities: dict[str, int] | None = None, -) -> int: - # fmt: off - """ - Get the seasonality of a frequency. - - Args: - freq (str): The frequency to get the seasonality of. - custom_seasonalities (dict[str, int] | None): Custom seasonalities to use. - If None, the default seasonalities are used. - - Returns: - int: The seasonality of the frequency. - - Example: - ```python - from timecopilot.models.utils.forecaster import get_seasonality - - get_seasonality("D", custom_seasonalities={"D": 7}) - # 7 - get_seasonality("D") # default seasonalities are used - # 1 - ``` - """ - # fmt: on - if custom_seasonalities is None: - custom_seasonalities = dict() - return _get_seasonality( - freq, - seasonalities=DEFAULT_SEASONALITIES | custom_seasonalities, - ) - - -def maybe_infer_freq(df: pd.DataFrame, freq: str | None) -> str: - """ - Infer the frequency of the time series data. - - Args: - df (pd.DataFrame): The time series data. - freq (str | None): The frequency of the time series data. If None, - the frequency will be inferred from the data. - - Returns: - str: The inferred frequency of the time series data. - """ - # based on https://github.com/Nixtla/nixtla/blob/bf67c76fd473a61c72b1f54725ffbcb51a3048c5/nixtla/nixtla_client.py#L208C1-L235C25 - if freq is not None: - return freq - sizes = df["unique_id"].value_counts(sort=True) - times = df.loc[df["unique_id"] == sizes.index[0], "ds"].sort_values() - if times.dt.tz is not None: - times = times.dt.tz_convert("UTC").dt.tz_localize(None) - inferred_freq = pd.infer_freq(times.values) - if inferred_freq is None: - raise RuntimeError( - "Could not infer the frequency of the time column. This could be due " - "to inconsistent intervals. Please check your data for missing, " - "duplicated or irregular timestamps" - ) - return inferred_freq - - -def maybe_convert_col_to_datetime(df: pd.DataFrame, col_name: str) -> pd.DataFrame: - if not pd.api.types.is_datetime64_any_dtype(df[col_name]): - df = df.copy() - df[col_name] = pd.to_datetime(df[col_name]) - return df - - -class Forecaster: - alias: str - - @staticmethod - def _maybe_infer_freq( - df: pd.DataFrame, - freq: str | None, - ) -> str: - return maybe_infer_freq(df, freq) - - def _maybe_get_seasonality(self, freq: str) -> int: - if hasattr(self, "season_length"): - if self.season_length is not None: - return self.season_length - else: - return get_seasonality(freq) - else: - return get_seasonality(freq) - - def forecast( - self, - df: pd.DataFrame, - h: int, - freq: str | None = None, - level: list[int | float] | None = None, - quantiles: list[float] | None = None, - ) -> pd.DataFrame: - """Generate forecasts for time series data using the model. - - This method produces point forecasts and, optionally, prediction - intervals or quantile forecasts. The input DataFrame can contain one - or multiple time series in stacked (long) format. - - Args: - df (pd.DataFrame): - DataFrame containing the time series to forecast. It must - include as columns: - - - "unique_id": an ID column to distinguish multiple series. - - "ds": a time column indicating timestamps or periods. - - "y": a target column with the observed values. - - h (int): - Forecast horizon specifying how many future steps to predict. - freq (str, optional): - Frequency of the time series (e.g. "D" for daily, "M" for - monthly). See [Pandas frequency aliases](https://pandas.pydata.org/ - pandas-docs/stable/user_guide/timeseries.html#offset-aliases) for - valid values. If not provided, the frequency will be inferred - from the data. - level (list[int | float], optional): - Confidence levels for prediction intervals, expressed as - percentages (e.g. [80, 95]). If provided, the returned - DataFrame will include lower and upper interval columns for - each specified level. - quantiles (list[float], optional): - List of quantiles to forecast, expressed as floats between 0 - and 1. Should not be used simultaneously with `level`. When - provided, the output DataFrame will contain additional columns - named in the format "model-q-{percentile}", where {percentile} - = 100 × quantile value. - - Returns: - pd.DataFrame: - DataFrame containing forecast results. Includes: - - - point forecasts for each timestamp and series. - - prediction intervals if `level` is specified. - - quantile forecasts if `quantiles` is specified. - - For multi-series data, the output retains the same unique - identifiers as the input DataFrame. - """ - raise NotImplementedError("This method must be implemented in a subclass.") - - def cross_validation( - self, - df: pd.DataFrame, - h: int, - freq: str | None = None, - n_windows: int = 1, - step_size: int | None = None, - level: list[int | float] | None = None, - quantiles: list[float] | None = None, - ) -> pd.DataFrame: - """ - Perform cross-validation on time series data. - - This method splits the time series into multiple training and testing - windows and generates forecasts for each window. It enables evaluating - forecast accuracy over different historical periods. Supports point - forecasts and, optionally, prediction intervals or quantile forecasts. - - Args: - df (pd.DataFrame): - DataFrame containing the time series to forecast. It must - include as columns: - - - "unique_id": an ID column to distinguish multiple series. - - "ds": a time column indicating timestamps or periods. - - "y": a target column with the observed values. - - h (int): - Forecast horizon specifying how many future steps to predict in - each window. - freq (str, optional): - Frequency of the time series (e.g. "D" for daily, "M" for - monthly). See [Pandas frequency aliases](https://pandas.pydata. - org/pandas-docs/stable/user_guide/timeseries.html#offset-aliases) - for valid values. If not provided, the frequency will be inferred - from the data. - n_windows (int, optional): - Number of cross-validation windows to generate. Defaults to 1. - step_size (int, optional): - Step size between the start of consecutive windows. If None, it - defaults to `h`. - level (list[int | float], optional): - Confidence levels for prediction intervals, expressed as - percentages (e.g. [80, 95]). When specified, the output - DataFrame includes lower and upper interval columns for each - level. - quantiles (list[float], optional): - Quantiles to forecast, expressed as floats between 0 and 1. - Should not be used simultaneously with `level`. If provided, - additional columns named "model-q-{percentile}" will appear in - the output, where {percentile} is 100 × quantile value. - - Returns: - pd.DataFrame: - DataFrame containing the forecasts for each cross-validation - window. The output includes: - - - "unique_id" column to indicate the series. - - "ds" column to indicate the timestamp. - - "y" column to indicate the target. - - "cutoff" column to indicate which window each forecast - belongs to. - - point forecasts for each timestamp and series. - - prediction intervals if `level` is specified. - - quantile forecasts if `quantiles` is specified. - """ - freq = self._maybe_infer_freq(df, freq) - df = maybe_convert_col_to_datetime(df, "ds") - # mlforecast cv code - results = [] - sort_idxs = maybe_compute_sort_indices(df, "unique_id", "ds") - if sort_idxs is not None: - df = take_rows(df, sort_idxs) - splits = backtest_splits( - df, - n_windows=n_windows, - h=h, - id_col="unique_id", - time_col="ds", - freq=pd.tseries.frequencies.to_offset(freq), - step_size=h if step_size is None else step_size, - ) - for _, (cutoffs, train, valid) in tqdm(enumerate(splits)): - if len(valid.columns) > 3: - raise NotImplementedError( - "Cross validation with exogenous variables is not yet supported." - ) - y_pred = self.forecast( - df=train, - h=h, - freq=freq, - level=level, - quantiles=quantiles, - ) - y_pred = join(y_pred, cutoffs, on="unique_id", how="left") - result = join( - valid[["unique_id", "ds", "y"]], - y_pred, - on=["unique_id", "ds"], - ) - if result.shape[0] < valid.shape[0]: - raise ValueError( - "Cross validation result produced less results than expected. " - "Please verify that the frequency parameter (freq) " - "matches your series' " - "and that there aren't any missing periods." - ) - results.append(result) - out = vertical_concat(results) - out = drop_index_if_pandas(out) - first_out_cols = ["unique_id", "ds", "cutoff", "y"] - remaining_cols = [c for c in out.columns if c not in first_out_cols] - fcst_cv_df = out[first_out_cols + remaining_cols] - return fcst_cv_df - - def detect_anomalies( - self, - df: pd.DataFrame, - h: int | None = None, - freq: str | None = None, - n_windows: int | None = None, - level: int | float = 99, - ) -> pd.DataFrame: - """ - Detect anomalies in time-series using a cross-validated z-score test. - - This method uses rolling-origin cross-validation to (1) produce - adjusted (out-of-sample) predictions and (2) estimate the - standard deviation of forecast errors. It then computes a per-point z-score, - flags values outside a two-sided prediction interval (with confidence `level`), - and returns a DataFrame with results. - - Args: - df (pd.DataFrame): - DataFrame containing the time series to detect anomalies. - Minimum series length is `h + 1` for most models, for Prophet models - it is `h + 2`. - h (int, optional): - Forecast horizon specifying how many future steps to predict. - In each cross validation window. If not provided, the seasonality - of the data (inferred from the frequency) is used. - freq (str, optional): - Frequency of the time series (e.g. "D" for daily, "M" for - monthly). See [Pandas frequency aliases](https://pandas.pydata. - org/pandas-docs/stable/user_guide/timeseries.html#offset-aliases) - for valid values. If not provided, the frequency will be inferred - from the data. - n_windows (int, optional): - Number of cross-validation windows to generate. - If not provided, the maximum number of windows - (computed by the shortest time series) is used. - If provided, the number of windows is the minimum - between the maximum number of windows - (computed by the shortest time series) - and the number of windows provided. - level (int | float): - Confidence levels for z-score, expressed as - percentages (e.g. 80, 95). Default is 99. - - Returns: - pd.DataFrame: - DataFrame containing the forecasts for each cross-validation - window. The output includes: - - - "unique_id" column to indicate the series. - - "ds" column to indicate the timestamp. - - "y" column to indicate the target. - - model column to indicate the model. - - lower prediction interval. - - upper prediction interval. - - anomaly column to indicate if the value is an anomaly. - an anomaly is defined as a value that is outside of the - prediction interval (True or False). - """ - freq = self._maybe_infer_freq(df, freq) - 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() - # we require at least one observation before the first forecast - max_possible_windows = (min_series_length - 1) // h - # Prophet needs a slightly different calculation to guarantee 2 input rows - if isinstance(self, ProphetBase): - max_possible_windows = (min_series_length - 2) // h - if n_windows is None: - _n_windows = max_possible_windows - else: - _n_windows = min(n_windows, max_possible_windows) - if _n_windows < 1: - # min series length should be 1 higher for Prophet than other models - exp_min_series_length = h + 2 if isinstance(self, ProphetBase) else h + 1 - raise ValueError( - f"Cannot perform anomaly detection: series too short. " - f"Minimum series length required: {exp_min_series_length}, " - f"actual minimum length: {min_series_length}" - ) - cv_results = self.cross_validation( - df=df, - h=h, - freq=freq, - n_windows=_n_windows, - step_size=h, # this is the default but who knows, anxiety - ) - cv_results["residuals"] = cv_results["y"] - cv_results[self.alias] - residual_stats = ( - cv_results.groupby("unique_id")["residuals"].std().reset_index() - ) - residual_stats.columns = ["unique_id", "residual_std"] - cv_results = cv_results.merge(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) - an_col = f"{self.alias}-anomaly" - cv_results[an_col] = np.abs(cv_results["z_score"]) > critical_z - lo_col = f"{self.alias}-lo-{int(level)}" - hi_col = f"{self.alias}-hi-{int(level)}" - margin = critical_z * cv_results["residual_std"] - cv_results[lo_col] = cv_results[self.alias] - margin - cv_results[hi_col] = cv_results[self.alias] + margin - output_cols = [ - "unique_id", - "ds", - "cutoff", - "y", - self.alias, - lo_col, - hi_col, - an_col, - ] - result = cv_results[output_cols].copy() - result = drop_index_if_pandas(result) - return result - - @staticmethod - def plot( - df: pd.DataFrame | None = None, - forecasts_df: pd.DataFrame | None = None, - ids: list[str] | None = None, - plot_random: bool = True, - max_ids: int | None = 8, - models: list[str] | None = None, - level: list[float] | None = None, - max_insample_length: int | None = None, - plot_anomalies: bool = False, - engine: str = "matplotlib", - palette: str | None = None, - seed: int | None = None, - resampler_kwargs: dict | None = None, - ax: plt.Axes | np.ndarray | plotly.graph_objects.Figure | None = None, - ): - """Plot forecasts and insample values. - - Args: - df (pd.DataFrame, optional): DataFrame with columns - [`unique_id`, `ds`, `y`]. Defaults to None. - forecasts_df (pd.DataFrame, optional): DataFrame with - columns [`unique_id`, `ds`] and models. Defaults to None. - ids (list[str], optional): Time Series to plot. If None, time series - are selected randomly. Defaults to None. - plot_random (bool, optional): Select time series to plot randomly. - Defaults to True. - max_ids (int, optional): Maximum number of ids to plot. Defaults to 8. - models (list[str], optional): Models to plot. Defaults to None. - level (list[float], optional): Prediction intervals to plot. - Defaults to None. - max_insample_length (int, optional): Maximum number of train/insample - observations to be plotted. Defaults to None. - plot_anomalies (bool, optional): Plot anomalies for each prediction - interval. Defaults to False. - engine (str, optional): Library used to plot. 'plotly', 'plotly-resampler' - or 'matplotlib'. Defaults to 'matplotlib'. - palette (str, optional): Name of the matplotlib colormap to use for the - plots. If None, uses the current style. Defaults to None. - seed (int, optional): Seed used for the random number generator. Only - used if plot_random is True. Defaults to 0. - resampler_kwargs (dict, optional): Keyword arguments to be passed to - plotly-resampler constructor. For further custumization ("show_dash") - call the method, store the plotting object and add the extra arguments - to its `show_dash` method. Defaults to None. - ax (matplotlib axes, array of matplotlib axes or plotly Figure, optional): - Object where plots will be added. Defaults to None. - """ - df = ensure_time_dtype(df, time_col="ds") - if forecasts_df is not None: - forecasts_df = ensure_time_dtype(forecasts_df, time_col="ds") - if any("anomaly" in col for col in forecasts_df.columns): - df = None - models = [ - col.split("-")[0] - for col in forecasts_df.columns - if col.endswith("-anomaly") - ] - forecasts_df = ufp.drop_columns( - forecasts_df, - [f"{model}-anomaly" for model in models], - ) - lv_cols = [ - c.replace(f"{model}-lo-", "") - for model in models - for c in forecasts_df.columns - if f"{model}-lo-" in c - ] - level = [float(c) if "." in c else int(c) for c in lv_cols] - level = list(set(level)) - plot_anomalies = True - return plot_series( - df=df, - forecasts_df=forecasts_df, - ids=ids, - plot_random=plot_random, - max_ids=max_ids, - models=models, - level=level, - max_insample_length=max_insample_length, - plot_anomalies=plot_anomalies, - engine=engine, - resampler_kwargs=resampler_kwargs, - palette=palette, - seed=seed, - id_col="unique_id", - time_col="ds", - target_col="y", - ax=ax, - ) - - -class QuantileConverter: - """Handles inputs and outputs for probabilistic forecasts.""" - - def __init__( - self, - level: list[int | float] | None = None, - quantiles: list[float] | None = None, - ): - level, quantiles, level_was_provided = self._prepare_level_and_quantiles( - level, quantiles - ) - self.level = level - self.quantiles = quantiles - # this is used to determine whether to return the level or the quantiles - self.level_was_provided = level_was_provided - - @staticmethod - def _prepare_level_and_quantiles( - level: list[int | float] | None, - quantiles: list[float] | None, - ) -> tuple[list[int | float] | None, list[float] | None, bool]: - # based on https://github.com/Nixtla/nixtla/blob/e74d98d9346a055153f84801cac94715c2342946/nixtla/nixtla_client.py#L444 - if level is not None and quantiles is not None: - raise ValueError( - "You must not provide both `level` and `quantiles` simultaneously." - ) - if quantiles is None and level is not None: - _quantiles = [] - for lv in level: - q_lo, q_hi = QuantileConverter._level_to_quantiles(lv) - _quantiles.append(q_lo) - _quantiles.append(q_hi) - quantiles = sorted(set(_quantiles)) - level_was_provided = True - return level, quantiles, level_was_provided - if level is None and quantiles is not None: - # we recover level from quantiles - 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] - level_was_provided = False - return sorted(set(level)), quantiles, level_was_provided - else: - return None, None, False - - @staticmethod - def _level_to_quantiles(level: int | float) -> tuple[float, float]: - """ - Given a prediction interval level (e.g. 80) return the lower & upper - quantiles that delimit the central interval (e.g. 0.10, 0.90). - """ - # handle trailing 9s, can occur with level == 80 - alpha = round(1 - level / 100, 2) - q_lo = alpha / 2 - q_hi = 1 - q_lo - return q_lo, q_hi - - def maybe_convert_level_to_quantiles( - self, - df: pd.DataFrame, - models: list[str], - ) -> pd.DataFrame: - """ - Receives a DataFrame with levels and returns - a DataFrame with quantiles if level was provided - """ - if self.level_was_provided or self.level is None: - return df - if self.quantiles is None: - raise ValueError("No quantiles were provided.") - out_cols = [c for c in df.columns if "-lo-" not in c and "-hi-" not in c] - df = ufp.copy_if_pandas(df, deep=False) - for model in models: - for q in sorted(self.quantiles): - if q == 0.5: - col = model - else: - lv = int(100 - 200 * q) - hi_or_lo = "lo" if lv > 0 else "hi" - lv = abs(lv) - col = f"{model}-{hi_or_lo}-{lv}" - q_col = f"{model}-q-{int(q * 100)}" - df = ufp.assign_columns(df, q_col, df[col]) - out_cols.append(q_col) - return df[out_cols] - - def maybe_convert_quantiles_to_level( - self, - df: pd.DataFrame, - models: list[str], - ) -> pd.DataFrame: - """ - Receives a DataFrame with quantiles and returns - a DataFrame with levels if quantiles were provided - """ - if not self.level_was_provided or self.quantiles is None: - return df - if self.level is None: - raise ValueError("No levels were provided.") - out_cols = [c for c in df.columns if "-q-" not in c] - df = ufp.copy_if_pandas(df, deep=False) - for model in models: - if 0 in self.level: - mid_col = f"{model}-q-50" - if mid_col in df: - df = ufp.assign_columns(df, model, df[mid_col]) - if model not in out_cols: - out_cols.append(model) - for lv in self.level: - q_lo, q_hi = self._level_to_quantiles(lv) - lo_src = f"{model}-q-{int(q_lo * 100)}" - hi_src = f"{model}-q-{int(q_hi * 100)}" - lo_tgt = f"{model}-lo-{lv}" - hi_tgt = f"{model}-hi-{lv}" - if lo_src in df and hi_src in df: - df = ufp.assign_columns(df, lo_tgt, df[lo_src]) - df = ufp.assign_columns(df, hi_tgt, df[hi_src]) - out_cols.extend([lo_tgt, hi_tgt]) - return df[out_cols] - -class _DataProcessor: - def __init__(self, dtype: torch.dtype, device: torch.device) -> None: - self.dtype = dtype - self.device = device +__all__ = [ + "Forecaster", + "QuantileConverter", + "_DataProcessor", + "get_seasonality", + "maybe_convert_col_to_datetime", + "maybe_infer_freq", +] - def _left_pad_and_stack_1D(self, tensors: list[torch.Tensor]) -> torch.Tensor: - max_len = max(len(c) for c in tensors) - padded = [] - for c in tensors: - assert isinstance(c, torch.Tensor) - assert c.ndim == 1 - padding = torch.full( - size=(max_len - len(c),), - fill_value=torch.nan, - device=c.device, - dtype=c.dtype, - ) - padded.append(torch.concat((padding, c), dim=-1)) - return torch.stack(padded) - def _prepare_and_validate_context( - self, - context: list[torch.Tensor] | torch.Tensor, - ) -> torch.Tensor: - if isinstance(context, list): - context = self._left_pad_and_stack_1D(context) - assert isinstance(context, torch.Tensor) - if context.ndim == 1: - context = context.unsqueeze(0) - assert context.ndim == 2 - return context +class Forecaster(_Forecaster): + """TimeCopilot Forecaster base (extends foundationforecast).""" - def _maybe_impute_missing( - self, batch: torch.Tensor, dtype=torch.float32 - ) -> torch.Tensor: - if torch.isnan(batch).any(): - batch = batch.to(dtype=dtype).detach().cpu().numpy() - imputed_rows = [] - for i in range(batch.shape[0]): - row = batch[i] - imputed_row = LastValueImputation()(row) - imputed_rows.append(imputed_row) - batch = np.vstack(imputed_rows) - batch = torch.tensor( - batch, - dtype=self.dtype, - device=self.device, - ) - return batch + pass diff --git a/uv.lock b/uv.lock index e4add748..948dc8be 100644 --- a/uv.lock +++ b/uv.lock @@ -1751,6 +1751,36 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/c7/93/0dd45cd283c32dea1545151d8c3637b4b8c53cdb3a625aeb2885b184d74d/fonttools-4.60.1-py3-none-any.whl", hash = "sha256:906306ac7afe2156fcf0042173d6ebbb05416af70f6b370967b47f8f00103bbb", size = 1143175, upload-time = "2025-09-29T21:13:24.134Z" }, ] +[[package]] +name = "foundationforecast" +version = "0.1.1" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "gluonts", extra = ["torch"] }, + { name = "huggingface-hub" }, + { name = "nixtla" }, + { name = "pandas", version = "2.3.3", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.13'" }, + { name = "scipy" }, + { name = "tabpfn-time-series", marker = "python_full_version < '3.13'" }, + { name = "tfc-t0", marker = "python_full_version >= '3.11' and python_full_version < '3.14'" }, + { name = "timecopilot-chronos-forecasting" }, + { name = "timecopilot-granite-tsfm", marker = "python_full_version >= '3.11' and python_full_version < '3.14'" }, + { name = "timecopilot-timesfm" }, + { name = "timecopilot-tirex", marker = "python_full_version >= '3.11'" }, + { name = "timecopilot-tirex2", marker = "python_full_version >= '3.11'" }, + { name = "timecopilot-toto" }, + { name = "timecopilot-toto-2" }, + { name = "timecopilot-uni2ts", marker = "python_full_version < '3.14'" }, + { name = "torch", version = "2.8.0", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.14'" }, + { name = "torch", version = "2.13.0", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.14'" }, + { name = "transformers" }, + { name = "utilsforecast" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/a3/93/c3307da9f4a0293930cad98060132893efb85a12ff8250cb8440c1aefba2/foundationforecast-0.1.1.tar.gz", hash = "sha256:5b619c6bf5bc46c9b81964827973d9d4025faef65b69f85c957556c89cd3f9bc", size = 2340749, upload-time = "2026-08-13T20:10:28.393Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/ca/d4/a81f0379a40d7289399cf906d4cb9339c4cb9242df04564bc1e15691ad21/foundationforecast-0.1.1-py3-none-any.whl", hash = "sha256:598fc883089f636aed8a152777295e9db838693852d4ebe9c78cc8f161e15d96", size = 52972, upload-time = "2026-08-13T20:10:26.925Z" }, +] + [[package]] name = "frozenlist" version = "1.8.0" @@ -7501,8 +7531,8 @@ dependencies = [ { name = "catboost" }, { name = "datasets" }, { name = "fire" }, + { name = "foundationforecast" }, { name = "fsspec" }, - { name = "gluonts", extra = ["torch"] }, { name = "huggingface-hub" }, { name = "hydra-core" }, { name = "lightgbm" }, @@ -7527,17 +7557,7 @@ dependencies = [ { name = "pytorch-lightning" }, { name = "scipy" }, { name = "statsforecast" }, - { name = "tabpfn-time-series", marker = "python_full_version < '3.13'" }, { name = "tensorboard" }, - { name = "tfc-t0", marker = "python_full_version >= '3.11' and python_full_version < '3.14'" }, - { name = "timecopilot-chronos-forecasting" }, - { name = "timecopilot-granite-tsfm", marker = "python_full_version >= '3.11' and python_full_version < '3.14'" }, - { name = "timecopilot-timesfm" }, - { name = "timecopilot-tirex", marker = "python_full_version >= '3.11'" }, - { name = "timecopilot-tirex2", marker = "python_full_version >= '3.11'" }, - { name = "timecopilot-toto" }, - { name = "timecopilot-toto-2" }, - { name = "timecopilot-uni2ts", marker = "python_full_version < '3.14'" }, { name = "torchmetrics" }, { name = "transformers" }, { name = "tsfeatures" }, @@ -7588,9 +7608,9 @@ requires-dist = [ { name = "dask", marker = "extra == 'distributed'", specifier = "<=2024.12.1" }, { name = "datasets", specifier = ">=4.1.1" }, { name = "fire" }, + { name = "foundationforecast", specifier = ">=0.1.1" }, { name = "fsspec", specifier = ">=2025.9.0" }, { name = "fugue", extras = ["dask", "ray", "spark"], marker = "extra == 'distributed'", specifier = ">=0.9.0" }, - { name = "gluonts", extras = ["torch"] }, { name = "huggingface-hub", specifier = ">=0.36.2,<2.0" }, { name = "hydra-core", specifier = ">=1.3.2" }, { name = "lightgbm", specifier = ">=4.6.0" }, @@ -7617,17 +7637,7 @@ requires-dist = [ { name = "ray", marker = "extra == 'distributed'", specifier = "==2.48" }, { name = "scipy", specifier = "<=1.15.3" }, { name = "statsforecast", specifier = ">=2.0.2" }, - { name = "tabpfn-time-series", marker = "python_full_version < '3.13'", specifier = "==1.0.3" }, { name = "tensorboard", specifier = ">=2.20.0" }, - { name = "tfc-t0", marker = "python_full_version >= '3.11' and python_full_version < '3.14'", specifier = ">=0.2.3" }, - { name = "timecopilot-chronos-forecasting", specifier = ">=0.2.2" }, - { name = "timecopilot-granite-tsfm", marker = "python_full_version >= '3.11' and python_full_version < '3.14'", specifier = ">=0.2.1" }, - { name = "timecopilot-timesfm", specifier = ">=0.3.0" }, - { name = "timecopilot-tirex", marker = "python_full_version >= '3.11'", specifier = ">=0.1.1" }, - { name = "timecopilot-tirex2", marker = "python_full_version >= '3.11'", specifier = ">=0.1.0" }, - { name = "timecopilot-toto", specifier = ">=0.1.7" }, - { name = "timecopilot-toto-2", specifier = ">=0.1.1" }, - { name = "timecopilot-uni2ts", marker = "python_full_version < '3.14'", specifier = ">=0.1.3" }, { name = "torchmetrics", specifier = ">=1.8.2" }, { name = "transformers", marker = "python_full_version < '3.13'", specifier = ">=4.41,<6" }, { name = "transformers", marker = "python_full_version >= '3.13'", specifier = ">=4.48,<6" }, From 4683cc8048503ffa557ee368fbdd929e8a88329b Mon Sep 17 00:00:00 2001 From: AzulGarza Date: Thu, 13 Aug 2026 14:31:48 -0600 Subject: [PATCH 2/3] fix: add correct timesfm imports and missing method --- tests/models/foundation/test_timesfm.py | 15 +++++++-------- timecopilot/models/utils/forecaster.py | 5 ++++- 2 files changed, 11 insertions(+), 9 deletions(-) diff --git a/tests/models/foundation/test_timesfm.py b/tests/models/foundation/test_timesfm.py index 75955c72..d0a3e19c 100644 --- a/tests/models/foundation/test_timesfm.py +++ b/tests/models/foundation/test_timesfm.py @@ -1,21 +1,20 @@ import os import pytest - -from timecopilot.models.foundation.timesfm import _TimesFMV1, _TimesFMV2_p5 +from foundationforecast.models.timesfm import _TimesFMV1, _TimesFMV2_p5 MODEL_PARAMS = [ ( _TimesFMV1, [ - "timecopilot.models.foundation.timesfm.timesfm_v1.TimesFmCheckpoint", - "timecopilot.models.foundation.timesfm.timesfm_v1.TimesFm", + "foundationforecast.models.timesfm.timesfm_v1.TimesFmCheckpoint", + "foundationforecast.models.timesfm.timesfm_v1.TimesFm", ], ), ( _TimesFMV2_p5, [ - "timecopilot.models.foundation.timesfm.TimesFM_2p5_200M_torch", + "foundationforecast.models.timesfm.TimesFM_2p5_200M_torch", ], ), ] @@ -24,7 +23,7 @@ @pytest.mark.parametrize("model_class, mock_paths", MODEL_PARAMS) def test_load_model_from_local_path(mocker, model_class, mock_paths): """Tests loading from a local path.""" - module_path = "timecopilot.models.foundation.timesfm" + module_path = "foundationforecast.models.timesfm" mock_os_exists = mocker.patch(f"{module_path}.os.path.exists", return_value=True) mock_loader = [mocker.patch(i) for i in mock_paths] @@ -55,7 +54,7 @@ def test_load_model_from_local_path(mocker, model_class, mock_paths): @pytest.mark.parametrize("model_class, mock_paths", MODEL_PARAMS) def test_load_model_from_hf_repo(mocker, model_class, mock_paths): """Tests loading from a Hugging Face repo.""" - module_path = "timecopilot.models.foundation.timesfm" + module_path = "foundationforecast.models.timesfm" mock_os_exists = mocker.patch(f"{module_path}.os.path.exists", return_value=False) mock_repo_exists = mocker.patch(f"{module_path}.repo_exists", return_value=True) mock_loader = [mocker.patch(i) for i in mock_paths] @@ -87,7 +86,7 @@ def test_load_model_from_hf_repo(mocker, model_class, mock_paths): def test_model_raises_OSError_on_failed_load(mocker, model_class, _): """Tests that an OSError is raised on a failed load attempt.""" - module_path = "timecopilot.models.foundation.timesfm" + module_path = "foundationforecast.models.timesfm" mock_os_exists = mocker.patch(f"{module_path}.os.path.exists", return_value=False) mock_repo_exists = mocker.patch(f"{module_path}.repo_exists", return_value=False) diff --git a/timecopilot/models/utils/forecaster.py b/timecopilot/models/utils/forecaster.py index 9fac3b1d..c3e2da35 100644 --- a/timecopilot/models/utils/forecaster.py +++ b/timecopilot/models/utils/forecaster.py @@ -1,3 +1,4 @@ +import pandas as pd from foundationforecast.core.forecaster import Forecaster as _Forecaster from foundationforecast.core.forecaster import ( QuantileConverter, @@ -20,4 +21,6 @@ class Forecaster(_Forecaster): """TimeCopilot Forecaster base (extends foundationforecast).""" - pass + @staticmethod + def _maybe_infer_freq(df: pd.DataFrame, freq: str | None) -> str: + return maybe_infer_freq(df, freq) From d00cf9183a7ef5ddf4b9be8593762f3c340f5540 Mon Sep 17 00:00:00 2001 From: AzulGarza Date: Thu, 13 Aug 2026 14:35:45 -0600 Subject: [PATCH 3/3] fix: add correct patch --- tests/models/foundation/test_chronos.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/tests/models/foundation/test_chronos.py b/tests/models/foundation/test_chronos.py index b1277332..59f799c2 100644 --- a/tests/models/foundation/test_chronos.py +++ b/tests/models/foundation/test_chronos.py @@ -12,7 +12,7 @@ def test_chronos_default_dtype_is_float32(): def test_chronos_model_uses_configured_dtype(mocker): """Ensure Chronos loads models with the configured dtype.""" mock_pipeline = mocker.patch( - "timecopilot.models.foundation.chronos.BaseChronosPipeline.from_pretrained" + "foundationforecast.models.chronos.BaseChronosPipeline.from_pretrained" ) mocker.patch("torch.cuda.is_available", return_value=False) @@ -43,12 +43,12 @@ def test_chronos_forecast_uses_configured_dtype(mocker): # Patch dataset creation to capture dtype argument mock_from_df = mocker.patch( - "timecopilot.models.foundation.chronos.TimeSeriesDataset.from_df" + "foundationforecast.models.chronos.TimeSeriesDataset.from_df" ) # Avoid real model loading and CUDA branching mocker.patch( - "timecopilot.models.foundation.chronos.BaseChronosPipeline.from_pretrained" + "foundationforecast.models.chronos.BaseChronosPipeline.from_pretrained" ) mocker.patch("torch.cuda.is_available", return_value=False)