Forecasting API

tirex2.ForecastModel

High-level, batched forecasting interface around a :class:TiRex2 backbone.

The wrapper takes ownership of the model only as a delegate: it batches the :class:~tirex.model.types.TimeseriesType it is given (building them from a GluonTS dataset in :meth:forecast_gluon), feeds them to :meth:TiRex2.predict, and formats the per-series quantile forecasts into the requested output type. Attribute access falls through to the wrapped model, so the backbone's own methods (e.g. predict) remain reachable on the wrapper.

Parameters:

Name	Type	Description	Default
model	TiRex2	An instantiated, ready-for-inference backbone exposing predict(timeseries: list[TimeseriesType], prediction_length: int) -> list[Tensor] and a quantiles buffer holding the quantile levels it forecasts.	required

Source code in src/tirex2/api_adapter/forecast.py

class ForecastModel:
    """High-level, batched forecasting interface around a :class:`TiRex2` backbone.

    The wrapper takes ownership of the model only as a delegate: it batches the
    :class:`~tirex.model.types.TimeseriesType` it is given (building them from a GluonTS
    dataset in :meth:`forecast_gluon`), feeds them to :meth:`TiRex2.predict`, and formats the
    per-series quantile forecasts into the requested output type. Attribute access falls
    through to the wrapped model, so the backbone's own methods (e.g. ``predict``) remain
    reachable on the wrapper.

    Parameters
    ----------
    model : TiRex2
        An instantiated, ready-for-inference backbone exposing
        ``predict(timeseries: list[TimeseriesType], prediction_length: int) -> list[Tensor]``
        and a ``quantiles`` buffer holding the quantile levels it forecasts.
    """

    def __init__(self, model):
        self.model = model

    def _quantile_levels(self) -> list[float]:
        """Return the model's forecast quantile levels as clean Python floats (float32 noise rounded off)."""
        return [round(float(q), 6) for q in self.model.quantiles]

    def __getattr__(self, name):
        """Delegate unknown attribute lookups to the wrapped model."""
        try:
            model = object.__getattribute__(self, "model")
        except AttributeError:
            raise AttributeError(name)
        return getattr(model, name)

    def forecast(
        self,
        timeseries: list[TimeseriesType],
        prediction_length: int,
        output_type: ForecastOutputType = "torch",
        batch_size: int = 512,
        yield_per_batch: bool = False,
        **predict_kwargs,
    ):
        """Forecast a list of :class:`TimeseriesType`, each carrying a target and optional covariates.

        Extra ``predict_kwargs`` are forwarded verbatim to :meth:`TiRex2.predict`.
        In particular ``tta_sign_flip`` controls sign-flip test-time augmentation
        (roughly doubles inference cost), and ``tta_diff`` controls postprocessor
        differencing; when omitted, the checkpoint's configured defaults
        (``model-config.yaml``) are used. Pass ``True``/``False`` to override.

        Examples
        --------
        >>> import torch
        >>> from tirex2 import TimeseriesType, load_model
        >>> model = load_model("NX-AI/TiRex-2", device="cpu")
        >>> ts = TimeseriesType(target=torch.randn(1, 128), past_covariates=None, future_covariates=None)
        >>> forecasts = model.forecast([ts], prediction_length=32, output_type="numpy")
        >>> forecasts[0].shape
        (1, 9, 32)
        """
        assert batch_size >= 1, "Batch size must be >= 1"
        return _gen_forecast(
            self.model,
            list(timeseries),
            None,
            prediction_length,
            output_type,
            batch_size,
            yield_per_batch,
            self._quantile_levels(),
            **predict_kwargs,
        )

    def forecast_gluon(
        self,
        gluonDataset,
        prediction_length: int,
        output_type: ForecastOutputType = "torch",
        batch_size: int = 512,
        yield_per_batch: bool = False,
        multivariate: bool = False,
        data_kwargs: dict = {},
        **predict_kwargs,
    ):
        """Forecast every entry of a GluonTS dataset, carrying its covariates and metadata through.

        With ``multivariate=False`` (default) each target variate is rendered as its own
        univariate ``QuantileForecast``; with ``multivariate=True`` each series yields a single
        forecast retaining the variate axis, so a multivariate dataset is scored jointly rather
        than channel-by-channel. The flag only affects ``output_type="gluonts"`` formatting.

        Extra ``predict_kwargs`` are forwarded verbatim to :meth:`TiRex2.predict`.
        In particular ``tta_sign_flip`` controls sign-flip test-time augmentation
        (roughly doubles inference cost), and ``tta_diff`` controls postprocessor
        differencing; when omitted, the checkpoint's configured defaults
        (``model-config.yaml``) are used. Pass ``True``/``False`` to override.
        """
        assert batch_size >= 1, "Batch size must be >= 1"
        try:
            from .gluon import build_gluon_timeseries
        except ImportError:
            raise ValueError("forecast_gluon needs GluonTS but GluonTS is not available (not installed)!")

        timeseries, meta = build_gluon_timeseries(gluonDataset, multivariate=multivariate, **data_kwargs)
        return _gen_forecast(
            self.model,
            timeseries,
            meta,
            prediction_length,
            output_type,
            batch_size,
            yield_per_batch,
            self._quantile_levels(),
            **predict_kwargs,
        )

    def forecast_fev(
        self,
        window: "fev.EvaluationWindow",
        prediction_length: int,
        output_type: ForecastOutputType = "torch",
        batch_size: int = 512,
        yield_per_batch: bool = False,
        data_kwargs: dict | None = None,
        quantile_levels: list[float] | None = None,
        return_inference_time: bool = False,
        **predict_kwargs,
    ):
        """Forecast a single FEV evaluation window.

        The call mirrors :meth:`forecast_gluon`: convert the external dataset
        representation into :class:`TimeseriesType`, then delegate batching,
        prediction and output rendering to the common forecast path. Use
        ``output_type="fev"`` to return predictions in the format accepted by
        ``fev.Task.evaluation_summary``. Pass ``return_inference_time=True`` to
        also return the model-only prediction time, excluding FEV input
        conversion and final ``DatasetDict`` construction.
        """
        assert batch_size >= 1, "Batch size must be >= 1"
        try:
            import fev  # noqa: F401
        except ImportError:
            raise ValueError("forecast_fev needs fev but fev is not available (not installed)!")

        data_kwargs = data_kwargs or {}
        timeseries, meta = build_fev_timeseries(window, **data_kwargs)
        if output_type == "fev":
            requested_quantiles = quantile_levels if quantile_levels is not None else self._quantile_levels()
            for item in meta:
                item["quantile_levels"] = requested_quantiles

        return _gen_forecast(
            self.model,
            timeseries,
            meta,
            prediction_length,
            output_type,
            batch_size,
            yield_per_batch,
            self._quantile_levels(),
            return_inference_time=return_inference_time,
            **predict_kwargs,
        )

forecast

forecast(timeseries: list[TimeseriesType], prediction_length: int, output_type: ForecastOutputType = 'torch', batch_size: int = 512, yield_per_batch: bool = False, **predict_kwargs)

Forecast a list of :class:TimeseriesType, each carrying a target and optional covariates.

Extra predict_kwargs are forwarded verbatim to :meth:TiRex2.predict. In particular tta_sign_flip controls sign-flip test-time augmentation (roughly doubles inference cost), and tta_diff controls postprocessor differencing; when omitted, the checkpoint's configured defaults (model-config.yaml) are used. Pass True/False to override.

Examples:

>>> import torch
>>> from tirex2 import TimeseriesType, load_model
>>> model = load_model("NX-AI/TiRex-2", device="cpu")
>>> ts = TimeseriesType(target=torch.randn(1, 128), past_covariates=None, future_covariates=None)
>>> forecasts = model.forecast([ts], prediction_length=32, output_type="numpy")
>>> forecasts[0].shape
(1, 9, 32)

Source code in src/tirex2/api_adapter/forecast.py

def forecast(
    self,
    timeseries: list[TimeseriesType],
    prediction_length: int,
    output_type: ForecastOutputType = "torch",
    batch_size: int = 512,
    yield_per_batch: bool = False,
    **predict_kwargs,
):
    """Forecast a list of :class:`TimeseriesType`, each carrying a target and optional covariates.

    Extra ``predict_kwargs`` are forwarded verbatim to :meth:`TiRex2.predict`.
    In particular ``tta_sign_flip`` controls sign-flip test-time augmentation
    (roughly doubles inference cost), and ``tta_diff`` controls postprocessor
    differencing; when omitted, the checkpoint's configured defaults
    (``model-config.yaml``) are used. Pass ``True``/``False`` to override.

    Examples
    --------
    >>> import torch
    >>> from tirex2 import TimeseriesType, load_model
    >>> model = load_model("NX-AI/TiRex-2", device="cpu")
    >>> ts = TimeseriesType(target=torch.randn(1, 128), past_covariates=None, future_covariates=None)
    >>> forecasts = model.forecast([ts], prediction_length=32, output_type="numpy")
    >>> forecasts[0].shape
    (1, 9, 32)
    """
    assert batch_size >= 1, "Batch size must be >= 1"
    return _gen_forecast(
        self.model,
        list(timeseries),
        None,
        prediction_length,
        output_type,
        batch_size,
        yield_per_batch,
        self._quantile_levels(),
        **predict_kwargs,
    )

forecast_gluon

forecast_gluon(gluonDataset, prediction_length: int, output_type: ForecastOutputType = 'torch', batch_size: int = 512, yield_per_batch: bool = False, multivariate: bool = False, data_kwargs: dict = {}, **predict_kwargs)

Forecast every entry of a GluonTS dataset, carrying its covariates and metadata through.

With multivariate=False (default) each target variate is rendered as its own univariate QuantileForecast; with multivariate=True each series yields a single forecast retaining the variate axis, so a multivariate dataset is scored jointly rather than channel-by-channel. The flag only affects output_type="gluonts" formatting.

Extra predict_kwargs are forwarded verbatim to :meth:TiRex2.predict. In particular tta_sign_flip controls sign-flip test-time augmentation (roughly doubles inference cost), and tta_diff controls postprocessor differencing; when omitted, the checkpoint's configured defaults (model-config.yaml) are used. Pass True/False to override.

Source code in src/tirex2/api_adapter/forecast.py

def forecast_gluon(
    self,
    gluonDataset,
    prediction_length: int,
    output_type: ForecastOutputType = "torch",
    batch_size: int = 512,
    yield_per_batch: bool = False,
    multivariate: bool = False,
    data_kwargs: dict = {},
    **predict_kwargs,
):
    """Forecast every entry of a GluonTS dataset, carrying its covariates and metadata through.

    With ``multivariate=False`` (default) each target variate is rendered as its own
    univariate ``QuantileForecast``; with ``multivariate=True`` each series yields a single
    forecast retaining the variate axis, so a multivariate dataset is scored jointly rather
    than channel-by-channel. The flag only affects ``output_type="gluonts"`` formatting.

    Extra ``predict_kwargs`` are forwarded verbatim to :meth:`TiRex2.predict`.
    In particular ``tta_sign_flip`` controls sign-flip test-time augmentation
    (roughly doubles inference cost), and ``tta_diff`` controls postprocessor
    differencing; when omitted, the checkpoint's configured defaults
    (``model-config.yaml``) are used. Pass ``True``/``False`` to override.
    """
    assert batch_size >= 1, "Batch size must be >= 1"
    try:
        from .gluon import build_gluon_timeseries
    except ImportError:
        raise ValueError("forecast_gluon needs GluonTS but GluonTS is not available (not installed)!")

    timeseries, meta = build_gluon_timeseries(gluonDataset, multivariate=multivariate, **data_kwargs)
    return _gen_forecast(
        self.model,
        timeseries,
        meta,
        prediction_length,
        output_type,
        batch_size,
        yield_per_batch,
        self._quantile_levels(),
        **predict_kwargs,
    )

forecast_fev

forecast_fev(window: EvaluationWindow, prediction_length: int, output_type: ForecastOutputType = 'torch', batch_size: int = 512, yield_per_batch: bool = False, data_kwargs: dict | None = None, quantile_levels: list[float] | None = None, return_inference_time: bool = False, **predict_kwargs)

Forecast a single FEV evaluation window.

The call mirrors :meth:forecast_gluon: convert the external dataset representation into :class:TimeseriesType, then delegate batching, prediction and output rendering to the common forecast path. Use output_type="fev" to return predictions in the format accepted by fev.Task.evaluation_summary. Pass return_inference_time=True to also return the model-only prediction time, excluding FEV input conversion and final DatasetDict construction.

Source code in src/tirex2/api_adapter/forecast.py

def forecast_fev(
    self,
    window: "fev.EvaluationWindow",
    prediction_length: int,
    output_type: ForecastOutputType = "torch",
    batch_size: int = 512,
    yield_per_batch: bool = False,
    data_kwargs: dict | None = None,
    quantile_levels: list[float] | None = None,
    return_inference_time: bool = False,
    **predict_kwargs,
):
    """Forecast a single FEV evaluation window.

    The call mirrors :meth:`forecast_gluon`: convert the external dataset
    representation into :class:`TimeseriesType`, then delegate batching,
    prediction and output rendering to the common forecast path. Use
    ``output_type="fev"`` to return predictions in the format accepted by
    ``fev.Task.evaluation_summary``. Pass ``return_inference_time=True`` to
    also return the model-only prediction time, excluding FEV input
    conversion and final ``DatasetDict`` construction.
    """
    assert batch_size >= 1, "Batch size must be >= 1"
    try:
        import fev  # noqa: F401
    except ImportError:
        raise ValueError("forecast_fev needs fev but fev is not available (not installed)!")

    data_kwargs = data_kwargs or {}
    timeseries, meta = build_fev_timeseries(window, **data_kwargs)
    if output_type == "fev":
        requested_quantiles = quantile_levels if quantile_levels is not None else self._quantile_levels()
        for item in meta:
            item["quantile_levels"] = requested_quantiles

    return _gen_forecast(
        self.model,
        timeseries,
        meta,
        prediction_length,
        output_type,
        batch_size,
        yield_per_batch,
        self._quantile_levels(),
        return_inference_time=return_inference_time,
        **predict_kwargs,
    )