diff --git a/.github/workflows/test_foundation_zoo.yml b/.github/workflows/test_foundation_zoo.yml new file mode 100644 index 000000000..fbadf8cbc --- /dev/null +++ b/.github/workflows/test_foundation_zoo.yml @@ -0,0 +1,65 @@ +name: Test foundation model zoo + +# Exercises tslearn.foundation against real pre-trained checkpoints from the +# Hugging Face Hub (see tests/test_foundation_zoo.py and the +# plot_foundation_model_zoo.py gallery example it backs). Each model pulls in +# its own package, sometimes with conflicting version pins (torch, +# transformers...), so every model runs as its own isolated job rather than +# being part of the regular test suite. + +on: + workflow_dispatch: + schedule: + - cron: "30 3 * * 0" # weekly, downloads are too heavy to run nightly + push: + branches: + - "foundation-module" + +jobs: + test_foundation_zoo: + runs-on: ubuntu-latest + strategy: + fail-fast: false + matrix: + include: + - name: chronos_bolt + deps: | + python -m pip install "chronos-forecasting>=2.0" + - name: timesfm + deps: | + python -m pip install "timesfm[torch]" + - name: moirai + deps: | + python -m pip install uni2ts + - name: ttm + deps: | + python -m pip install granite-tsfm + - name: moment + deps: | + python -m pip install momentfm + - name: time_moe + deps: | + python -m pip install "transformers==4.40.1" + env: + TSLEARN_RUN_FOUNDATION_ZOO: "1" + steps: + - uses: actions/checkout@v4 + + - name: Set up Python + uses: actions/setup-python@v5 + with: + python-version: "3.10" + + - name: Install tslearn + run: | + python -m pip install --upgrade pip + python -m pip install --extra-index-url https://download.pytorch.org/whl/cpu .[tests] torch + shell: bash + + - name: Install ${{ matrix.name }} dependencies + run: ${{ matrix.deps }} + shell: bash + + - name: Test ${{ matrix.name }} + run: python -m pytest -v tests/test_foundation_zoo.py -k "${{ matrix.name }}" + shell: bash diff --git a/CHANGELOG.md b/CHANGELOG.md index 3736cca7a..99f88c199 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -13,6 +13,15 @@ Changelogs for this project are recorded in this file since v0.2.0. ### Added +* New `tslearn.foundation` module to re-use pre-trained time series models, such as + the ones published on the Hugging Face Hub, behind the usual tslearn API. It provides + `ZeroShotForecaster`, `LinearProbeForecaster` and the underlying + `TimeSeriesFoundationEmbedder` feature extractor, which allows one to choose + the probed layer, the tokens taken into account and the pooling applied to their + representations. With `pooling=None`, the latter acts as a time series to time series + transform, returning one representation per token, and being a regular scikit-learn + transformer, it composes with a `sklearn.pipeline.Pipeline` for tasks such as linear + probing for classification. Requires PyTorch. * soft-dtw related tools now support Itakura and Sakoe-Chiba global constraints. ([#189](https://github.com/tslearn-team/tslearn/issues/189)) * Multithreading support added to `cdist_soft_dtw` and `softdtw_barycenter`. ([#310](https://github.com/tslearn-team/tslearn/issues/310)) diff --git a/docs/_static/img/foundation_tokens.svg b/docs/_static/img/foundation_tokens.svg new file mode 100644 index 000000000..cac9a3fd7 --- /dev/null +++ b/docs/_static/img/foundation_tokens.svg @@ -0,0 +1,54 @@ + + + Token layout of a patch-based foundation model, such as Chronos-2 + + + + 0 + 1 + N-1 + -2 + -1 + + + + + + + + + … + + + Token 1 + Token 2 + Token N + + + + + + + Register + token + Forecast + token + + + + + + tokens=(0, -2) + kept, and averaged over + + + + + + excluded + do not represent the series + + diff --git a/docs/examples/classification/plot_foundation_linear_probe.py b/docs/examples/classification/plot_foundation_linear_probe.py new file mode 100644 index 000000000..fe0a98a40 --- /dev/null +++ b/docs/examples/classification/plot_foundation_linear_probe.py @@ -0,0 +1,216 @@ +""" +Linear probing a pre-trained model for classification +===================================================== + +Time series foundation models are pre-trained for forecasting, yet the +representations they build along the way carry a lot of information about the +shape of a series, which makes them useful for other tasks as well. Linear +probing [1]_ is a cheap way to implement that principle: the pre-trained model +is kept frozen and used as a feature extractor, and a plain classifier is fitted +on top of its representations. Because the head is linear and the backbone is +never updated, the accuracy reached tells us how linearly separable the classes +already are in the representation space. + +This example applies that idea to a UCR dataset, composing +:class:`~tslearn.foundation.TimeSeriesFoundationEmbedder` with a classifier +inside a :class:`~sklearn.pipeline.Pipeline`, and using Chronos-2 [2]_ as the +frozen backbone. Running it requires the ``chronos-forecasting`` package:: + + pip install "chronos-forecasting>=2.0" + +References +---------- +.. [1] G. Alain and Y. Bengio. Understanding intermediate layers using linear + classifier probes. ICLR Workshop, 2017. +.. [2] A. F. Ansari, O. Shchur, J. Küken, et al. Chronos-2: From Univariate to + Universal Forecasting. arXiv:2510.15821, 2025. +""" + +# Author: Romain Tavenard +# License: BSD 3 clause +# sphinx_gallery_thumbnail_number = 2 + +############################################################################## +# Data +# ---- +# +# We use the ``Trace`` dataset from the UCR archive, which gathers four classes +# of transient signals recorded in a nuclear power plant, with only 100 training +# series in total. + +import numpy as np + +from tslearn.datasets import CachedDatasets +from tslearn.preprocessing import TimeSeriesScalerMeanVariance + +X_train, y_train, X_test, y_test = CachedDatasets().load_dataset("Trace") + +scaler = TimeSeriesScalerMeanVariance() +X_train = scaler.fit_transform(X_train) +X_test = scaler.transform(X_test) + +print(f"{X_train.shape=}, {X_test.shape=}, {len(np.unique(y_train))} classes") + +############################################################################## +# Probing the pre-trained model +# ----------------------------- +# +# The embedder only needs the pre-trained model and, since Chronos-2 returns +# forecasts rather than hidden states, an explicit layer to read +# representations from. A forward hook is placed on the selected block, so no +# modification of the model is needed. Being a regular scikit-learn +# transformer, it composes with a classifier inside a +# :class:`~sklearn.pipeline.Pipeline`, which is what "linear probing" amounts +# to here. +# +# Fitting is fast: each series goes through the model exactly once, and the +# only thing actually trained is a logistic regression over a few hundred +# features. + +from chronos import Chronos2Pipeline +from sklearn.linear_model import LogisticRegression +from sklearn.pipeline import Pipeline + +from tslearn.foundation import TimeSeriesFoundationEmbedder + +pipeline = Chronos2Pipeline.from_pretrained("autogluon/chronos-2-small") + +embedder = TimeSeriesFoundationEmbedder( + pipeline.model, + layer=-2, + pooling="mean", + # Chronos-2 appends a register token and a forecast token after its context + # tokens; neither represents the series, so they are left out of the average + tokens=(0, -2), +) +clf = Pipeline([("embed", embedder), ("classify", LogisticRegression(max_iter=1000))]) +clf.fit(X_train, y_train) + +print(f"Embedding size: {embedder.embedding_size_}") +print(f"Test accuracy: {clf.score(X_test, y_test):.3f}") + +############################################################################## +# Which layer, which pooling? +# --------------------------- +# +# The representations of a pre-trained model change a lot from one layer to the +# next, and the last ones are usually the most specialized towards the +# pre-training objective, here forecasting. Sweeping over layers is therefore +# worthwhile, and cheap: nothing is being trained beyond the linear head. +# +# ``pooling`` matters just as much. Chronos-2 emits one token per patch of the +# context plus a register token, so we compare averaging those tokens against +# picking the register token alone, the latter playing the role a ``[CLS]`` +# token plays in a text encoder. + +n_layers = len(pipeline.model.encoder.block) +poolings = ["mean", "max", "token"] +# Every other layer is enough to see the trend, and halves the build time +layers = sorted({*range(0, n_layers, 2), n_layers - 1}) + +accuracies = {} +for layer in layers: + for pooling in poolings: + model = Pipeline([ + ("embed", TimeSeriesFoundationEmbedder( + pipeline.model, + layer=layer, + pooling=pooling, + # Chronos-2 places its register token after the context tokens + token_index=-2 if pooling == "token" else 0, + tokens=(0, -2), + )), + ("classify", LogisticRegression(max_iter=1000)), + ]).fit(X_train, y_train) + accuracies[layer, pooling] = model.score(X_test, y_test) + +header = "layer " + " ".join(f"{pooling:>6}" for pooling in poolings) +print(header) +for layer in layers: + row = " ".join(f"{accuracies[layer, pooling]:>6.3f}" for pooling in poolings) + print(f"{layer:>5} {row}") + +############################################################################## + +import matplotlib.pyplot as plt + +fig, ax = plt.subplots(figsize=(8, 4), layout="constrained") +for pooling in poolings: + ax.plot( + layers, + [accuracies[layer, pooling] for layer in layers], + marker="o", + label=f"pooling={pooling}", + ) +ax.set_xlabel("probed layer") +ax.set_ylabel("test accuracy") +ax.set_title("Linear probe accuracy across the layers of Chronos-2") +ax.legend() +plt.show() + +############################################################################## +# Using the representations elsewhere +# ----------------------------------- +# +# :class:`~tslearn.foundation.TimeSeriesFoundationEmbedder` is a regular +# scikit-learn transformer, so it composes with the rest of the ecosystem the +# same way it did with the classifier above: here it is chained with a +# :class:`~sklearn.decomposition.PCA` inside a +# :class:`~sklearn.pipeline.Pipeline`, to project the frozen representations +# of the test set onto two dimensions and see whether the classes separate. + +from sklearn.decomposition import PCA + +best_layer, best_pooling = max(accuracies, key=accuracies.get) +projector = Pipeline([ + ("embed", TimeSeriesFoundationEmbedder( + pipeline.model, layer=best_layer, tokens=(0, -2) + )), + ("pca", PCA(n_components=2)), +]) +projected = projector.fit_transform(X_test) + +fig, ax = plt.subplots(figsize=(6, 5), layout="constrained") +for label in np.unique(y_test): + mask = y_test == label + ax.scatter(projected[mask, 0], projected[mask, 1], label=f"class {label}", s=20) +ax.set_xlabel("first principal component") +ax.set_ylabel("second principal component") +ax.set_title("Frozen representations of the Trace test set") +ax.legend() +plt.show() + +############################################################################## +# Representations as time series +# ------------------------------ +# +# Setting ``pooling=None`` skips the aggregation step altogether. Instead of +# one vector per series, the transformer then returns one vector per token, +# that is, a time series dataset of shape ``(n_ts, n_tokens, dim)``. Combined +# with ``tokens``, which drops the tokens that do not represent the input, this +# turns the pre-trained model into a time series to time series transform whose +# output can be fed to any other tslearn estimator. +# +# Chronos-2 emits one token per patch of 16 timesteps, so the representations +# below are a 16-times subsampled, high-dimensional view of the input series. +# Clustering those sequences under DTW recovers the classes reasonably well, +# without the labels ever being used. + +from tslearn.clustering import TimeSeriesKMeans + +from sklearn.metrics import adjusted_rand_score + +embedder = TimeSeriesFoundationEmbedder( + pipeline.model, layer=best_layer, pooling=None, tokens=(0, -2) +) +clusterer = TimeSeriesKMeans( + n_clusters=len(np.unique(y_test)), metric="dtw", max_iter=5, random_state=0 +) +pipeline = Pipeline(steps=[ + ("embed", embedder), + ("cluster", clusterer) +]) +labels = pipeline.fit_predict(X_test) + +print(f"Adjusted Rand index against the true classes: " + f"{adjusted_rand_score(y_test, labels):.3f}") diff --git a/docs/examples/forecasting/plot_foundation_forecasting.py b/docs/examples/forecasting/plot_foundation_forecasting.py new file mode 100644 index 000000000..a4ac76abb --- /dev/null +++ b/docs/examples/forecasting/plot_foundation_forecasting.py @@ -0,0 +1,270 @@ +""" +Forecasting with a pre-trained model +==================================== + +This example showcases the :mod:`tslearn.foundation` module, which allows one +to re-use a pre-trained time series model behind the usual tslearn API. Two +adaptation strategies are compared on the same data: + +* zero-shot forecasting, with :class:`~tslearn.foundation.ZeroShotForecaster`, + where the pre-trained model is used exactly as it is; +* linear probing, with :class:`~tslearn.foundation.LinearProbeForecaster`, + where the pre-trained model is kept frozen and only a linear head mapping its + representations to future values is fitted. + +Both are compared to the classical :class:`~tslearn.forecasting.AutoVARIMA` +baseline. + +The pre-trained model used here is Chronos-2 [1]_, an encoder-only model +pre-trained on a large corpus of real and synthetic series. Running this +example therefore requires the ``chronos-forecasting`` package:: + + pip install "chronos-forecasting>=2.0" + +Note that nothing in :mod:`tslearn.foundation` is specific to Chronos: any +PyTorch model exposing similar conventions can be plugged in the same way. + +References +---------- +.. [1] A. F. Ansari, O. Shchur, J. Küken, et al. Chronos-2: From Univariate to + Universal Forecasting. arXiv:2510.15821, 2025. +""" + +############################################################################## +# Data +# ---- +# +# We use a set of sine waves with varying frequencies, phases and noise levels. + +import numpy as np + +from tslearn.utils import to_time_series_dataset + +rng = np.random.RandomState(0) + +n_ts, sz, horizon, context_length = 32, 512, 48, 128 +t = np.arange(sz + horizon) + +periods = rng.uniform(30, 50, size=n_ts) +phases = rng.uniform(0, 2 * np.pi, size=n_ts) +trends = rng.uniform(-2e-3, 2e-3, size=n_ts) + +full_series = ( + np.sin(2 * np.pi * t[None, :] / periods[:, None] + phases[:, None]) + + trends[:, None] * t[None, :] + + 0.1 * rng.randn(n_ts, sz + horizon) +) +full_series = to_time_series_dataset(full_series) + +X_train, X_test = full_series[:, :sz], full_series[:, sz:] +print(f"{X_train.shape=}, {X_test.shape=}") + +############################################################################## +# Zero-shot forecasting +# --------------------- +# +# A time series foundation model is pre-trained to forecast unseen series out +# of the box, so no training is needed at all. Wrapping the inference pipeline +# in a :class:`~tslearn.foundation.ZeroShotForecaster` is enough to obtain a +# tslearn estimator, whose ``predict`` method returns an array of shape +# ``(n_ts, n, d)`` like every other forecaster of the library. +# +# Calling ``fit`` is not needed here since the purpose of a +# :class:`~tslearn.foundation.ZeroShotForecaster` is to reuse a pre-trained +# model without any training. +# +# Note that the raw series are passed as they are, with no prior scaling: +# :class:`~tslearn.foundation.ZeroShotForecaster` wraps Chronos-2's full +# inference pipeline, which normalizes every context internally before +# forecasting. + +from chronos import Chronos2Pipeline + +from tslearn.foundation import ZeroShotForecaster + +pipeline = Chronos2Pipeline.from_pretrained("autogluon/chronos-2-small") + +zero_shot = ZeroShotForecaster(pipeline, horizon_axis=-1) +y_zero_shot = zero_shot.predict(X_train, n=horizon) +print(f"{y_zero_shot.shape=}") + +############################################################################## +# ``horizon_axis`` says which axis of the model's output holds the forecast +# horizon. There is no shared convention across implementations, Chronos-2 +# returns ``(n_series, n_quantiles, horizon)`` from ``predict`` so ``-1`` is +# used. Left to its default of ``"auto"``, the estimator would +# infer it from the returned shape. +# +# The remaining axes are understood as holding quantile levels or sample paths +# and are reduced to a point forecast according to the ``quantile`` parameter, +# which defaults to the median. + +############################################################################## +# Linear probing +# -------------- +# +# Zero-shot use ignores the fact that we do have training data at hand. Linear +# probing exploits it, while leaving the pre-trained weights untouched: the +# model is only used to embed the context window, and a linear map from those +# representations to the next ``horizon`` values is fitted. +# +# Training pairs are cut out of the series with a sliding window. A larger +# ``stride`` yields fewer, less redundant windows and a faster fit, which +# matters because every window has to go through the pre-trained model once. +# +# The probing estimator needs the model itself rather than the inference +# pipeline, since it reads hidden states rather than forecasts. This means it +# bypasses the normalization the pipeline applies internally, unlike +# :class:`~tslearn.foundation.ZeroShotForecaster` above, so it has to be +# restored explicitly, by placing a +# :class:`~tslearn.preprocessing.TimeSeriesScalerMeanVariance` ahead of it in +# a :class:`~sklearn.pipeline.Pipeline`. +# +# Three options drive which representations are used: +# +# * ``layer`` selects the block to probe: an integer places a forward hook on +# the corresponding block, so ``layer=-2`` probes the penultimate one. The +# last layers of a pre-trained model tend to specialize towards its +# pre-training objective, so probing slightly earlier is often beneficial. +# ``layer=None`` reads the model's own output instead, which only works for +# models returning their hidden states; Chronos-2 returns quantile forecasts +# only, so an explicit layer is required here. +# * ``pooling`` says how the token representations are aggregated into a single +# vector. Chronos-2 emits one token per patch of the context, plus a register +# token and a forecast token, so both averaging (``pooling="mean"``) and +# picking a single token (``pooling="token"``) are sensible. +# Chronos-2's forecast token is a natural choice +# here: ``token_index=-1`` reads it directly, rather than averaging over +# representations that were never meant to summarize the series for +# forecasting. +# * ``tokens`` restricts which tokens take part in that aggregation. Chronos-2 +# appends a register token and a forecast token after its context tokens, and +# neither represents the input series, so ``tokens=(0, -2)`` focuses on the +# slice 0:-2 and keeps the average clean. +# +# .. image:: /_static/img/foundation_tokens.svg +# :width: 700 +# :align: center +# :alt: Chronos-2 appends a register token and a forecast token after its +# context tokens; tokens=(0, -2) keeps only the context tokens. + +from sklearn.pipeline import Pipeline + +from tslearn.foundation import LinearProbeForecaster +from tslearn.preprocessing import TimeSeriesScalerMeanVariance + +forecaster = LinearProbeForecaster( + pipeline.model, + context_length=context_length, + horizon=horizon, + stride=16, + layer=-2, + pooling="mean", + tokens=(0, -2), +) +scaler = TimeSeriesScalerMeanVariance(per_timeseries=False) +probe = Pipeline([("scale", scaler), ("probe", forecaster)]) +probe.fit(X_train) + +print(f"{forecaster.n_windows_} training windows, " + f"{forecaster.embedder_.embedding_size_}-dimensional embeddings") + +############################################################################## +# The head fitted on top of the frozen representations defaults to a +# :class:`sklearn.linear_model.RidgeCV`, which picks its regularization +# strength by cross-validation. Any scikit-learn regressor can be passed +# instead through the ``probe`` parameter. +# +# The scaler also normalizes the training targets, since they are cut from +# the same series as the context, so the forecasts come out on that +# normalized scale. They are put back in the original units using the +# ``mean_`` and ``std_`` the scaler computed when it was fitted. + +y_probe = probe.predict(X_train) * scaler.std_ + scaler.mean_ + +############################################################################## +# Comparison +# ---------- +# +# Both strategies are compared to AutoVARIMA on the mean absolute error over +# the held-out horizon. + +from tslearn.forecasting import AutoVARIMA +from tslearn.metrics.performance import mae + +varima = AutoVARIMA(seasonal_period=40).fit(X_train) +y_varima = varima.predict(n=horizon) + +scores = { + "AutoVARIMA (with seasonal differentiation)": mae(X_test, y_varima), + "Chronos-2 (zero-shot)": mae(X_test, y_zero_shot), + "Chronos-2 (linear probe)": mae(X_test, y_probe), +} +for name, score in scores.items(): + print(f"{name:>42}: MAE = {score:.4f}") + +############################################################################## +# Let us look at a few series to see where the differences come from. + +import matplotlib.pyplot as plt + +shown = [0, 1, 2] +fig, axes = plt.subplots( + len(shown), 1, sharex=True, figsize=(11, 8), layout="constrained" +) +context = 150 +for ax, i in zip(axes, shown): + ax.plot( + np.arange(sz - context, sz), + X_train[i, -context:, 0], + color="0.4", + label="context", + ) + ax.plot( + np.arange(sz, sz + horizon), + X_test[i, :, 0], + color="k", + label="ground truth", + ) + for values, label in [ + (y_varima, "AutoVARIMA"), + (y_zero_shot, "zero-shot"), + (y_probe, "linear probe"), + ]: + ax.plot(np.arange(sz, sz + horizon), values[i, :, 0], label=label, alpha=0.8) + ax.axvline(sz, color="0.8", linestyle="--") + ax.set_ylabel(f"series {i}") +axes[0].legend(loc="upper left", ncol=5, fontsize="small") +fig.suptitle("Forecasting a held-out horizon", fontsize=14) +plt.show() + +############################################################################## +# Choosing what to probe +# ---------------------- +# +# Which layer and which pooling work best is data-dependent, and cheap enough +# to explore: the frozen model is the expensive part, and it is never trained. +# Below we score a few configurations on the held-out horizon. + +results = {} +for layer in [-1, -2, -4]: + for pooling in ["mean", "token"]: + model = Pipeline([ + ("scale", scaler), + ("probe", LinearProbeForecaster( + pipeline.model, + context_length=context_length, + horizon=horizon, + stride=48, + layer=layer, + pooling=pooling, + # Chronos-2's forecast token is its last one + token_index=-1, + tokens=(0, -2), + )), + ]).fit(X_train) + y = model.predict(X_train) * scaler.std_ + scaler.mean_ + results[(layer, pooling)] = mae(X_test, y) + +for (layer, pooling), score in sorted(results.items(), key=lambda kv: kv[1]): + print(f"layer={str(layer):>5}, pooling={pooling:>4}: MAE = {score:.4f}") diff --git a/docs/examples/forecasting/plot_foundation_model_zoo.py b/docs/examples/forecasting/plot_foundation_model_zoo.py new file mode 100644 index 000000000..dc4e1b733 --- /dev/null +++ b/docs/examples/forecasting/plot_foundation_model_zoo.py @@ -0,0 +1,429 @@ +""" +A tour of the pre-trained forecasting model zoo +================================================= + +:mod:`tslearn.foundation` does not target a specific pre-trained model: it +only assumes that a model exposes a forecasting method, or a PyTorch +``forward``, following one of a handful of widespread conventions (see +:class:`~tslearn.foundation.ZeroShotForecaster` and +:class:`~tslearn.foundation.LinearProbeForecaster`). This example surveys +seven models for time series +forecasting, and shows, for each of them, what it takes to plug it in as a +:class:`~tslearn.foundation.ZeroShotForecaster`, a +:class:`~tslearn.foundation.LinearProbeForecaster`, or both. The smallest +checkpoint published for every model is used throughout, to keep downloads +light. + +For a line-by-line walkthrough of the module's parameters, see the +:doc:`Chronos-2 example `, which this one +complements rather than repeats. + +================= ======================= ========= ============ +Model Package Zero-shot Linear probe +================= ======================= ========= ============ +Chronos-2 [1]_ ``chronos-forecasting`` yes yes +Chronos-Bolt [2]_ ``chronos-forecasting`` yes yes +TimesFM [3]_ ``timesfm`` yes -- +Moirai [4]_ ``uni2ts`` yes -- +TTM [5]_ ``granite-tsfm`` yes -- +MOMENT [6]_ ``momentfm`` -- yes +Time-MoE [7]_ ``transformers`` yes yes +================= ======================= ========= ============ + +The right-hand column is not a limitation of tslearn: it reflects how each +model exposes its internals. Models built around a clean +``forward(context) -> hidden_states`` (the two Chronos variants, Time-MoE) +probe as easily as they forecast zero-shot. Models whose low-level module +expects an already-patchified input, an explicit padding mask, or a packed +multi-series format (TimesFM, Moirai, TTM) are naturally used zero-shot, and +would need a bespoke wrapper to expose clean per-token representations for +probing. MOMENT is the mirror case: only its reconstruction head is +pre-trained, so its zero-shot *forecasts* would come out of a head that has +never been trained, while its embeddings, precisely what linear probing +needs, are excellent. + +References +---------- +.. [1] A. F. Ansari, O. Shchur, J. Küken, et al. Chronos-2: From Univariate + to Universal Forecasting. arXiv:2510.15821, 2025. +.. [2] A. F. Ansari, L. Stella, C. Turkmen, et al. Chronos: Learning the + Language of Time Series. TMLR, 2024. +.. [3] A. Das, W. Kong, R. Sen, Y. Zhou. A decoder-only foundation model for + time-series forecasting. ICML, 2024. +.. [4] G. Woo, C. Liu, A. Kumar, et al. Unified Training of Universal Time + Series Forecasting Transformers. ICML, 2024. +.. [5] V. Ekambaram, A. Jati, P. Dayama, et al. Tiny Time Mixers (TTMs): Fast + Pre-trained Models for Enhanced Zero/Few-Shot Forecasting of Multivariate + Time Series. NeurIPS, 2024. +.. [6] M. Goswami, K. Szafer, A. Choudhry, et al. MOMENT: A Family of Open + Time-series Foundation Models. ICML, 2024. +.. [7] X. Shi, S. Wang, Y. Nie, et al. Time-MoE: Billion-Scale Time Series + Foundation Models with Mixture of Experts. ICLR, 2025. +""" + +############################################################################## +# Data +# ---- +# +# The same small set of noisy sine waves used by the Chronos-2 example is +# reused here, so that Chronos-Bolt's forecasts below can be read against +# the same yardstick. + +import numpy as np + +from tslearn.utils import to_time_series_dataset + +rng = np.random.RandomState(0) + +n_ts, sz, horizon, context_length = 16, 256, 24, 96 +t = np.arange(sz + horizon) + +periods = rng.uniform(30, 50, size=n_ts) +phases = rng.uniform(0, 2 * np.pi, size=n_ts) + +full_series = np.sin( + 2 * np.pi * t[None, :] / periods[:, None] + phases[:, None] +) + 0.1 * rng.randn(n_ts, sz + horizon) +full_series = to_time_series_dataset(full_series) + +X_train, X_test = full_series[:, :sz], full_series[:, sz:] + +############################################################################## +# Chronos-2 +# --------- +# +# Chronos-2 [1]_ is covered in depth by the :doc:`plot_foundation_forecasting` +# example: zero-shot forecasting with :class:`~tslearn.foundation.ZeroShotForecaster` +# needs nothing beyond the pipeline itself, and linear probing with +# :class:`~tslearn.foundation.LinearProbeForecaster` reads its forecast token +# (``layer=-2, pooling="token", token_index=-1, tokens=(0, -2)``). It is not +# repeated here. + +############################################################################## +# Chronos-Bolt +# ------------ +# +# Chronos-Bolt [2]_ is a faster, encoder-decoder T5 model, distilled to +# predict all quantiles in a single forward pass rather than +# autoregressively, from the same package as Chronos-2:: +# +# pip install "chronos-forecasting>=2.0" +# +# A details set it apart from Chronos-2: its ``forward`` takes a raw +# ``context`` and exposes an encoder made of a +# plain stack of T5 blocks, at ``encoder.block``, with no register or +# forecast token to filter out, so probing needs no ``tokens`` argument. +# +# As with Chronos-2, the raw series are passed to +# :class:`~tslearn.foundation.ZeroShotForecaster` unscaled, since it wraps +# the full pipeline. The linear probe, on the other hand, reads +# ``pipeline.model`` directly and so needs scaling restored explicitly through +# a :class:`~sklearn.pipeline.Pipeline`, as detailed in the +# :doc:`plot_foundation_forecasting` example. + +import torch +from chronos import BaseChronosPipeline +from sklearn.pipeline import Pipeline + +from tslearn.foundation import LinearProbeForecaster, ZeroShotForecaster +from tslearn.preprocessing import TimeSeriesScalerMeanVariance + +pipeline = BaseChronosPipeline.from_pretrained( + "amazon/chronos-bolt-small", device_map="cpu" +) + +zero_shot = ZeroShotForecaster(pipeline) +y_zero_shot = zero_shot.predict(X_train, n=horizon) + +forecaster = LinearProbeForecaster( + pipeline.model, + context_length=context_length, + horizon=horizon, + stride=8, + layer=-1, + layers_path="encoder.block", + pooling="mean", +) +scaler = TimeSeriesScalerMeanVariance(per_timeseries=False) +probe = Pipeline([("scale", scaler), ("probe", forecaster)]) +probe.fit(X_train) + +# The scaler also normalizes the training targets, so forecasts come out on +# that scale and are put back in the original units, using the ``mean_`` and +# ``std_`` the scaler computed when it was fitted, before comparison. +y_probe = probe.predict(X_train) * scaler.std_ + scaler.mean_ + +from tslearn.metrics.performance import mae + +print(f"{'Chronos-Bolt (zero-shot)':>28}: MAE = {mae(X_test, y_zero_shot):.4f}") +print(f"{'Chronos-Bolt (linear probe)':>28}: MAE = {mae(X_test, y_probe):.4f}") + +############################################################################## +import matplotlib.pyplot as plt + +fig, ax = plt.subplots(figsize=(11, 3), layout="constrained") +context = 80 +ax.plot(np.arange(sz - context, sz), X_train[0, -context:, 0], color="0.4", label="context") +ax.plot(np.arange(sz, sz + horizon), X_test[0, :, 0], color="k", label="ground truth") +ax.plot(np.arange(sz, sz + horizon), y_zero_shot[0, :, 0], label="zero-shot", alpha=0.8) +ax.plot(np.arange(sz, sz + horizon), y_probe[0, :, 0], label="linear probe", alpha=0.8) +ax.axvline(sz, color="0.8", linestyle="--") +ax.legend(loc="upper left", ncol=4, fontsize="small") +ax.set_title("Chronos-Bolt") +plt.show() + +############################################################################## +# TimesFM +# ------- +# +# TimesFM [3]_ is a decoder-only model, pre-trained to predict a whole patch +# of future values at once rather than one step at a time:: +# +# pip install "timesfm[torch]" +# +# Its ``forecast`` method takes the horizon *before* the context +# (``forecast(horizon, inputs)``), which conflicts with the positional +# calling convention :class:`~tslearn.foundation.ZeroShotForecaster` uses +# when auto-detecting a method, so ``predict_fn`` is required here too. The +# horizon and context length are fixed once and for all through a +# :class:`~timesfm.ForecastConfig` passed to ``compile``, which builds the +# model's static computation graph. +# +# .. code-block:: python +# +# import timesfm +# +# from tslearn.foundation import ZeroShotForecaster +# +# model = timesfm.TimesFM_2p5_200M_torch.from_pretrained( +# "google/timesfm-2.5-200m-pytorch" +# ) +# model.compile(timesfm.ForecastConfig( +# max_context=context_length, max_horizon=horizon, normalize_inputs=True +# )) +# +# zero_shot = ZeroShotForecaster( +# model, +# predict_fn=lambda model, context, horizon: model.forecast( +# horizon=horizon, inputs=list(context) +# )[0], +# context_length=context_length, +# ) +# y_zero_shot = zero_shot.predict(X_train, n=horizon) +# +# TimesFM's own module (``model.model``) does not lend itself to linear +# probing out of the box: its ``forward`` expects the context already split +# into patches and concatenated with an explicit padding mask, a +# transformation that :class:`~tslearn.foundation.TimeSeriesFoundationEmbedder` +# does not perform, so probing it would require wrapping that patching logic +# in a small ``torch.nn.Module`` of one's own. + +############################################################################## +# Moirai +# ------ +# +# Moirai [4]_ is a masked-encoder model, natively multivariate and trained +# across many frequencies at once:: +# +# pip install "uni2ts @ git+https://github.com/SalesforceAIResearch/uni2ts.git" +# +# The forecasting horizon, and the number of sample paths drawn from the +# predictive distribution, are baked into the model at construction time +# rather than passed to a ``predict``-like method, so, again, a +# ``predict_fn`` builds the ``past_target`` / ``past_observed_target`` / +# ``past_is_pad`` triplet that :class:`~uni2ts.model.moirai.MoiraiForecast` +# expects. +# +# .. code-block:: python +# +# import torch +# +# from uni2ts.model.moirai import MoiraiForecast, MoiraiModule +# +# from tslearn.foundation import ZeroShotForecaster +# +# module = MoiraiModule.from_pretrained("Salesforce/moirai-1.1-R-small") +# forecast_model = MoiraiForecast( +# module=module, +# prediction_length=horizon, +# context_length=context_length, +# patch_size=32, +# target_dim=1, +# feat_dynamic_real_dim=0, +# past_feat_dynamic_real_dim=0, +# ) +# +# def predict_fn(model, context, horizon): +# past_target = context.unsqueeze(-1) +# past_observed = torch.ones_like(past_target, dtype=torch.bool) +# past_is_pad = torch.zeros(past_target.shape[:2], dtype=torch.bool) +# return model(past_target, past_observed, past_is_pad, num_samples=20) +# +# zero_shot = ZeroShotForecaster( +# forecast_model, predict_fn=predict_fn, context_length=context_length +# ) +# y_zero_shot = zero_shot.predict(X_train, n=horizon) +# +# Its backbone (``MoiraiModule``) consumes a packed representation of the +# whole batch, with explicit ``sample_id``, ``variate_id`` and ``time_id`` +# tensors describing which timestep and which series each row belongs to. +# Reproducing that packing outside of ``uni2ts`` itself is enough work that +# zero-shot use is, in practice, the natural way to reach for Moirai. + +############################################################################## +# TTM (Tiny Time Mixers) +# ----------------------- +# +# TTM [5]_ departs from the transformer architecture used by every other +# model in this list: it is a light-weight MLP-Mixer, which is also why it +# is small enough to not need a dedicated "small" checkpoint:: +# +# pip install "granite-tsfm[notebooks] @ git+https://github.com/ibm-granite/granite-tsfm.git@v0.2.22" +# +# Its context length and horizon are fixed by the checkpoint (512 and 96 for +# ``granite-timeseries-ttm-r2``) rather than adjustable at call time, and, +# since the raw model is called directly, ``predict_fn`` again does the +# work that a ``predict`` method would otherwise do. +# +# .. code-block:: python +# +# import torch +# +# from tsfm_public.models.tinytimemixer import TinyTimeMixerForPrediction +# +# from tslearn.foundation import ZeroShotForecaster +# +# model = TinyTimeMixerForPrediction.from_pretrained( +# "ibm-granite/granite-timeseries-ttm-r2", num_input_channels=1 +# ) +# +# def predict_fn(model, context, horizon): +# past_values = context.unsqueeze(-1) +# return model(past_values=past_values).prediction_outputs +# +# zero_shot = ZeroShotForecaster( +# model, predict_fn=predict_fn, context_length=model.config.context_length +# ) +# y_zero_shot = zero_shot.predict(X_train, n=model.config.prediction_length) +# +# TTM's mixer blocks operate on a ``(batch, channels, patches, dim)`` tensor +# rather than the ``(batch, tokens, dim)`` shape +# :class:`~tslearn.foundation.TimeSeriesFoundationEmbedder` expects, so +# probing it directly would need a small reshaping wrapper too. + +############################################################################## +# MOMENT +# ------ +# +# MOMENT [6]_ is pre-trained as a masked autoencoder: reconstructing masked +# patches is what gives it strong general-purpose embeddings, but it means +# its forecasting head, unlike its embeddings, has never actually been +# trained:: +# +# pip install momentfm +# +# Loaded with ``task_name="embedding"``, a forward pass already returns a +# single pooled vector per series (``output.embeddings``), so this is the +# rare model that needs no ``layer``/``pooling`` gymnastics to reach a +# feature matrix, only a hook on its encoder to read the token-level +# representations that :class:`~tslearn.foundation.LinearProbeForecaster` +# expects. +# +# As with the linear probes above, ``model`` is the bare backbone rather than +# a pipeline, so scaling has to be restored the same way, through a +# :class:`~sklearn.pipeline.Pipeline`. +# +# .. code-block:: python +# +# from momentfm import MOMENTPipeline +# from sklearn.pipeline import Pipeline +# +# from tslearn.foundation import LinearProbeForecaster +# from tslearn.preprocessing import TimeSeriesScalerMeanVariance +# +# model = MOMENTPipeline.from_pretrained( +# "AutonLab/MOMENT-1-small", model_kwargs={"task_name": "embedding"} +# ) +# model.init() +# +# scaler = TimeSeriesScalerMeanVariance(per_timeseries=False) +# probe = Pipeline([("scale", scaler), ("probe", LinearProbeForecaster( +# model, +# context_length=context_length, +# horizon=horizon, +# stride=32, +# layer=-1, +# layers_path="encoder.block", +# pooling="mean", +# input_layout="channels_first", +# ))]) +# probe.fit(X_train) +# y_probe = probe.predict(X_train) * scaler.std_ + scaler.mean_ +# +# ``input_layout="channels_first"`` matters here: MOMENT's ``forward`` +# expects an explicit channel axis (``x_enc`` of shape +# ``(batch, channels, sz)``), unlike most other models in this list which +# take a plain ``(batch, sz)`` batch of univariate contexts. Calling +# :meth:`~tslearn.foundation.LinearProbeForecaster.predict` (a zero-shot +# forecast) would technically run, but on a randomly initialized head, so it +# is deliberately left out. + +############################################################################## +# Time-MoE +# -------- +# +# Time-MoE [7]_ is a decoder-only, GPT-style model trained with next-value +# prediction and a sparse mixture-of-experts feed-forward block, distributed +# as ``transformers`` "remote code" rather than through its own package:: +# +# pip install transformers +# +# .. code-block:: python +# +# from transformers import AutoModelForCausalLM +# +# from tslearn.foundation import LinearProbeForecaster, ZeroShotForecaster +# +# model = AutoModelForCausalLM.from_pretrained( +# "Maple728/TimeMoE-50M", trust_remote_code=True +# ) +# +# Forecasting several steps ahead means generating token by token, through +# ``generate`` rather than a single forward pass; ``max_new_tokens`` is not +# one of the horizon argument names +# :class:`~tslearn.foundation.ZeroShotForecaster` looks for, and the output +# needs slicing to drop the echoed context, so a ``predict_fn`` is used once +# more: +# +# .. code-block:: python +# +# def predict_fn(model, context, horizon): +# out = model.generate(input_ids=context, max_new_tokens=horizon) +# return out[:, -horizon:] +# +# zero_shot = ZeroShotForecaster(model, predict_fn=predict_fn) +# y_zero_shot = zero_shot.predict(X_train, n=horizon) +# +# Being decoder-only, its natural pooling is ``"last"``, the representation +# of the final context token, which has attended to every earlier one. Here +# too, ``model`` is the bare backbone rather than a normalizing pipeline, so +# scaling is restored through a :class:`~sklearn.pipeline.Pipeline`: +# +# .. code-block:: python +# +# from sklearn.pipeline import Pipeline +# +# from tslearn.preprocessing import TimeSeriesScalerMeanVariance +# +# scaler = TimeSeriesScalerMeanVariance(per_timeseries=False) +# probe = Pipeline([("scale", scaler), ("probe", LinearProbeForecaster( +# model, +# context_length=context_length, +# horizon=horizon, +# stride=8, +# layer=-1, +# layers_path="model.layers", +# pooling="last", +# ))]) +# probe.fit(X_train) +# y_probe = probe.predict(X_train) * scaler.std_ + scaler.mean_ diff --git a/docs/gen_modules/tslearn.foundation.rst b/docs/gen_modules/tslearn.foundation.rst new file mode 100644 index 000000000..4a81b52f3 --- /dev/null +++ b/docs/gen_modules/tslearn.foundation.rst @@ -0,0 +1,18 @@ +.. _mod-tslearn.foundation: + +tslearn.foundation +================== + +.. automodule:: tslearn.foundation + + .. rubric:: Classes + + .. autosummary:: + :toctree: foundation + :template: class.rst + + ZeroShotForecaster + LinearProbeForecaster + TimeSeriesFoundationEmbedder + + \ No newline at end of file diff --git a/docs/reference.rst b/docs/reference.rst index 3dededac0..09e8b81fd 100644 --- a/docs/reference.rst +++ b/docs/reference.rst @@ -16,6 +16,7 @@ The complete ``tslearn`` project is automatically documented for every module. datasets early_classification forecasting + foundation generators matrix_profile metrics diff --git a/docs/requirements_rtd.txt b/docs/requirements_rtd.txt index e872cc434..55b73d6c2 100644 --- a/docs/requirements_rtd.txt +++ b/docs/requirements_rtd.txt @@ -4,10 +4,16 @@ # # pip-compile --extra=docs --output-file=docs/requirements_rtd.txt # +accelerate==1.14.0 + # via chronos-forecasting accessible-pygments==0.0.5 # via pydata-sphinx-theme alabaster==1.0.0 # via sphinx +annotated-doc==0.0.5 + # via typer +anyio==4.14.2 + # via httpx babel==2.18.0 # via # pydata-sphinx-theme @@ -15,9 +21,16 @@ babel==2.18.0 beautifulsoup4==4.14.3 # via pydata-sphinx-theme certifi==2026.2.25 - # via requests + # via + # httpcore + # httpx + # requests charset-normalizer==3.4.7 # via requests +chronos-forecasting==2.3.1 + # via tslearn (pyproject.toml) +click==8.4.2 + # via huggingface-hub contourpy==1.3.3 # via matplotlib cycler==0.12.1 @@ -26,14 +39,42 @@ docutils==0.22.4 # via # pydata-sphinx-theme # sphinx +einops==0.8.2 + # via chronos-forecasting +filelock==3.32.2 + # via + # huggingface-hub + # torch fonttools==4.62.1 # via matplotlib +fsspec==2026.7.0 + # via + # huggingface-hub + # torch +h11==0.16.0 + # via httpcore +hf-xet==1.6.0 + # via huggingface-hub +httpcore==1.0.9 + # via httpx +httpx==0.28.1 + # via huggingface-hub +huggingface-hub==1.27.0 + # via + # accelerate + # tokenizers + # transformers idna==3.15 - # via requests + # via + # anyio + # httpx + # requests imagesize==2.0.0 # via sphinx jinja2==3.1.6 - # via sphinx + # via + # sphinx + # torch joblib==1.5.3 # via # scikit-learn @@ -42,53 +83,98 @@ kiwisolver==1.5.0 # via matplotlib llvmlite==0.47.0 # via numba +markdown-it-py==4.2.0 + # via rich markupsafe==3.0.3 # via jinja2 matplotlib==3.10.8 # via tslearn (pyproject.toml) +mdurl==0.1.2 + # via markdown-it-py +mpmath==1.3.0 + # via sympy +networkx==3.6.1 + # via torch numba==0.65.0 # via tslearn (pyproject.toml) numpy==2.4.4 # via + # accelerate + # chronos-forecasting # contourpy # matplotlib # numba + # pandas + # patsy # scikit-learn # scipy + # statsmodels + # transformers # tslearn (pyproject.toml) numpydoc==1.10.0 # via tslearn (pyproject.toml) packaging==26.1 # via + # accelerate + # huggingface-hub # matplotlib # sphinx + # statsmodels + # transformers +pandas==3.0.5 + # via + # chronos-forecasting + # statsmodels +patsy==1.0.2 + # via statsmodels pillow==12.3.0 # via # matplotlib # sphinx-gallery +psutil==7.2.2 + # via accelerate pydata-sphinx-theme==0.17.0 # via tslearn (pyproject.toml) pygments==2.20.0 # via # accessible-pygments # pydata-sphinx-theme + # rich # sphinx pypandoc==1.17 # via tslearn (pyproject.toml) pyparsing==3.3.2 # via matplotlib python-dateutil==2.9.0.post0 - # via matplotlib + # via + # matplotlib + # pandas +pyyaml==6.0.3 + # via + # accelerate + # huggingface-hub + # transformers +regex==2026.7.19 + # via transformers requests==2.33.1 # via sphinx +rich==15.0.0 + # via typer roman-numerals==4.1.0 # via sphinx +safetensors==0.8.0 + # via + # accelerate + # transformers scikit-learn==1.8.0 # via tslearn (pyproject.toml) scipy==1.17.1 # via # scikit-learn + # statsmodels # tslearn (pyproject.toml) +shellingham==1.5.4 + # via typer six==1.17.0 # via python-dateutil snowballstemmer==3.0.1 @@ -118,11 +204,36 @@ sphinxcontrib-qthelp==2.0.0 # via sphinx sphinxcontrib-serializinghtml==2.0.0 # via sphinx +statsmodels==0.14.6 + # via tslearn (pyproject.toml) +sympy==1.14.0 + # via torch threadpoolctl==3.6.0 # via scikit-learn +tokenizers==0.22.2 + # via transformers +torch==2.13.0 + # via + # accelerate + # chronos-forecasting + # tslearn (pyproject.toml) +tqdm==4.70.0 + # via + # huggingface-hub + # transformers +transformers==5.15.0 + # via chronos-forecasting +typer==0.27.1 + # via transformers typing-extensions==4.15.0 # via + # anyio # beautifulsoup4 + # huggingface-hub # pydata-sphinx-theme + # torch urllib3==2.7.0 # via requests + +# The following packages are considered to be unsafe in a requirements file: +# setuptools diff --git a/pyproject.toml b/pyproject.toml index c1fe364ae..af9508da8 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -42,8 +42,10 @@ changelog = "https://github.com/tslearn-team/tslearn/CHANGELOG.md" [project.optional-dependencies] pytorch = ['torch'] +# Re-use of pre-trained models through the tslearn.foundation module +foundation = ['torch'] tests = [ - "pytest", + "pytest", "h5py" ] docs = [ @@ -54,6 +56,9 @@ docs = [ "numpydoc", "matplotlib", "pypandoc", + # Used by the tslearn.foundation examples of the gallery + "torch", + "chronos-forecasting>=2.0", ] all_features = [ "torch", @@ -76,3 +81,6 @@ testpaths = [ "tests", "tslearn" ] + +[tool.coverage.run] +omit = ["tests/test_foundation_zoo.py"] \ No newline at end of file diff --git a/tests/test_estimators.py b/tests/test_estimators.py index 5d78dc994..0cd0f2856 100644 --- a/tests/test_estimators.py +++ b/tests/test_estimators.py @@ -91,6 +91,14 @@ def get_estimators(type_filter='all'): # only keep those that are from tslearn all_classes = filter(lambda c: not is_sklearn(c), all_classes) + # tslearn.foundation estimators wrap an externally provided pre-trained + # model, so they cannot be instantiated without arguments, which is what + # the common checks below rely on. They are covered by tests/test_foundation.py. + all_classes = filter( + lambda c: not inspect.getmodule(c).__name__.startswith('tslearn.foundation'), + all_classes + ) + # Now filter out the estimators that are not of the specified type filters = { 'all': [ClassifierMixin, RegressorMixin, diff --git a/tests/test_foundation.py b/tests/test_foundation.py new file mode 100644 index 000000000..84cdffd3f --- /dev/null +++ b/tests/test_foundation.py @@ -0,0 +1,1240 @@ +"""Tests for the :mod:`tslearn.foundation` module. + +These tests run entirely offline, against small hand-written models that +reproduce the calling conventions of the pre-trained models published on the +Hugging Face Hub, so that no download is ever needed. +""" + +import numpy as np + +import pytest + +from sklearn.base import clone +from sklearn.linear_model import Ridge +from sklearn.pipeline import make_pipeline +from sklearn.svm import LinearSVC + +try: + import torch +except ImportError: + torch = None + +from tslearn.generators import random_walks +from tslearn.foundation import ( # noqa: E402 + LinearProbeForecaster, + TimeSeriesFoundationEmbedder, + ZeroShotForecaster, +) + + +PATCH_SIZE = 4 +D_MODEL = 8 + + +@pytest.mark.skipif(torch is not None, reason="PyTorch is installed") +def test_no_torch(): + with pytest.raises(ValueError, match="torch is not installed"): + TimeSeriesFoundationEmbedder("whatever") + with pytest.raises(ValueError, match="torch is not installed"): + LinearProbeForecaster("whatever") + with pytest.raises(ValueError, match="torch is not installed"): + ZeroShotForecaster("whatever") + + +if torch: + class _Block(torch.nn.Module): + """A single, deliberately simple, encoder block.""" + + def __init__(self, d_model=D_MODEL): + super().__init__() + self.linear = torch.nn.Linear(d_model, d_model) + + def forward(self, hidden_states): + return hidden_states + torch.tanh(self.linear(hidden_states)) + + + class _Encoder(torch.nn.Module): + def __init__(self, n_layers=3, d_model=D_MODEL): + super().__init__() + self.block = torch.nn.ModuleList([_Block(d_model) for _ in range(n_layers)]) + self.final_layer_norm = torch.nn.LayerNorm(d_model) + + def forward(self, hidden_states): + for block in self.block: + hidden_states = block(hidden_states) + return self.final_layer_norm(hidden_states) + + + class _Output: + """Stands in for a ``transformers`` ``ModelOutput``.""" + + def __init__(self, last_hidden_state): + self.last_hidden_state = last_hidden_state + + + class _DummyBackbone(torch.nn.Module): + """A patch-based encoder taking univariate series, like Chronos-2. + + A register token is appended after the context tokens, so that the + ``pooling="token"`` code path can be exercised on a non-zero ``token_index``. + """ + + def __init__(self, n_layers=3, d_model=D_MODEL, patch_size=PATCH_SIZE, seed=0): + super().__init__() + # Seeded so that results do not depend on which test ran before + torch.manual_seed(seed) + self.patch_size = patch_size + self.embedding = torch.nn.Linear(patch_size, d_model) + self.register_token = torch.nn.Parameter(torch.randn(1, 1, d_model)) + self.encoder = _Encoder(n_layers, d_model) + + def forward(self, context): + batch_size, sz = context.shape + n_patches = sz // self.patch_size + patches = context[:, : n_patches * self.patch_size] + patches = patches.reshape(batch_size, n_patches, self.patch_size) + hidden_states = self.embedding(patches) + register = self.register_token.expand(batch_size, -1, -1) + hidden_states = torch.cat([hidden_states, register], dim=1) + return _Output(self.encoder(hidden_states)) + + + class _MultivariateBackbone(torch.nn.Module): + """An encoder taking (batch, sz, d) arrays, like ``transformers`` models.""" + + def __init__(self, n_channels, d_model=D_MODEL, seed=0): + super().__init__() + torch.manual_seed(seed) + self.embedding = torch.nn.Linear(n_channels, d_model) + self.encoder = _Encoder(2, d_model) + + def forward(self, past_values): + return _Output(self.encoder(self.embedding(past_values))) + + + class _DummyPipeline: + """A zero-shot forecaster mimicking ``Chronos2Pipeline``. + + It returns one tensor per series, of shape + ``(n_variates, prediction_length, n_quantiles)``. + """ + + n_quantiles = 9 + + def predict(self, inputs, prediction_length=1, batch_size=256): + inputs = np.asarray(inputs, dtype=np.float64) + # Naive forecast: repeat the last observed value + last = inputs[:, -1] + quantile_offsets = np.linspace(-1.0, 1.0, self.n_quantiles) + forecast = ( + last[:, None, None] + + np.zeros((1, prediction_length, 1)) + + quantile_offsets[None, None, :] + ) + return [torch.as_tensor(row[None]) for row in forecast] + + + def _dataset(n_ts=6, sz=32, d=1, seed=0): + return random_walks(n_ts=n_ts, sz=sz, d=d, random_state=seed) + + +# --------------------------------------------------------------------------- +# TimeSeriesFoundationEmbedder +# --------------------------------------------------------------------------- + +@pytest.mark.skipif(torch is None, reason="PyTorch not installed") +def test_embedder_shapes_and_layouts(): + X = _dataset(n_ts=5, sz=32, d=1) + embedder = TimeSeriesFoundationEmbedder(_DummyBackbone()) + embeddings = embedder.fit_transform(X) + assert embeddings.shape == (5, D_MODEL) + assert embedder.embedding_size_ == D_MODEL + + # Multivariate series are embedded channel per channel and concatenated + X_multi = _dataset(n_ts=5, sz=32, d=3) + embeddings = TimeSeriesFoundationEmbedder(_DummyBackbone()).fit_transform(X_multi) + assert embeddings.shape == (5, 3 * D_MODEL) + + # Natively multivariate models get the whole (n_ts, sz, d) array + embedder = TimeSeriesFoundationEmbedder( + _MultivariateBackbone(n_channels=3), input_layout="channels_last" + ) + assert embedder.fit_transform(X_multi).shape == (5, D_MODEL) + + +@pytest.mark.skipif(torch is None, reason="PyTorch not installed") +def test_embedder_channel_stacking_is_order_preserving(): + # The embedding of a multivariate series must be the concatenation of the + # embeddings of its channels, in channel order. + X = _dataset(n_ts=4, sz=32, d=2) + backbone = _DummyBackbone() + joint = TimeSeriesFoundationEmbedder(backbone).fit_transform(X) + per_channel = [ + TimeSeriesFoundationEmbedder(backbone).fit_transform(X[:, :, k : k + 1]) + for k in range(2) + ] + np.testing.assert_allclose( + joint, np.concatenate(per_channel, axis=1), rtol=1e-5, atol=1e-6 + ) + + +@pytest.mark.skipif(torch is None, reason="PyTorch not installed") +@pytest.mark.parametrize("pooling", ["mean", "max", "token", "last", "flatten"]) +def test_embedder_poolings(pooling): + X = _dataset(n_ts=4, sz=32, d=1) + embedder = TimeSeriesFoundationEmbedder(_DummyBackbone(), pooling=pooling) + embeddings = embedder.fit_transform(X) + n_tokens = 32 // PATCH_SIZE + 1 # context patches + register token + expected = n_tokens * D_MODEL if pooling == "flatten" else D_MODEL + assert embeddings.shape == (4, expected) + assert embedder.n_tokens_ is None + + +@pytest.mark.skipif(torch is None, reason="PyTorch not installed") +@pytest.mark.parametrize("pooling", [None, "none"]) +def test_embedder_without_pooling_is_a_series_to_series_transform(pooling): + X = _dataset(n_ts=4, sz=32, d=1) + n_tokens = 32 // PATCH_SIZE + 1 # context patches + register token + + embedder = TimeSeriesFoundationEmbedder(_DummyBackbone(), pooling=pooling) + embeddings = embedder.fit_transform(X) + assert embeddings.shape == (4, n_tokens, D_MODEL) + assert embedder.n_tokens_ == n_tokens + assert embedder.embedding_size_ == D_MODEL + + # The output is a valid tslearn dataset, usable by other estimators + from tslearn.clustering import TimeSeriesKMeans + + labels = TimeSeriesKMeans(n_clusters=2, max_iter=2, random_state=0).fit_predict( + embeddings + ) + assert labels.shape == (4,) + + +@pytest.mark.skipif(torch is None, reason="PyTorch not installed") +def test_embedder_token_selection(): + X = _dataset(n_ts=4, sz=32, d=1) + n_context = 32 // PATCH_SIZE # tokens that represent the series + backbone = _DummyBackbone() + + # Dropping the trailing register token leaves the context tokens + embeddings = TimeSeriesFoundationEmbedder( + backbone, pooling=None, tokens=(0, -1) + ).fit_transform(X) + assert embeddings.shape == (4, n_context, D_MODEL) + + # Equivalent spellings + np.testing.assert_allclose( + embeddings, + TimeSeriesFoundationEmbedder( + backbone, pooling=None, tokens=slice(0, -1) + ).fit_transform(X), + ) + np.testing.assert_allclose( + embeddings, + TimeSeriesFoundationEmbedder( + backbone, pooling=None, tokens=(None, n_context) + ).fit_transform(X), + ) + + # Token selection also applies before pooling, and changes the result: + # averaging over the context tokens is not averaging over all of them + without_register = TimeSeriesFoundationEmbedder( + backbone, pooling="mean", tokens=(0, -1) + ).fit_transform(X) + over_everything = TimeSeriesFoundationEmbedder( + backbone, pooling="mean" + ).fit_transform(X) + np.testing.assert_allclose( + without_register, embeddings.mean(axis=1), rtol=1e-5, atol=1e-6 + ) + assert not np.allclose(without_register, over_everything) + + # ... but not to "token", which is meant to reach an excluded token + np.testing.assert_allclose( + TimeSeriesFoundationEmbedder( + backbone, pooling="token", token_index=-1, tokens=(0, -1) + ).fit_transform(X), + TimeSeriesFoundationEmbedder( + backbone, pooling="token", token_index=-1 + ).fit_transform(X), + ) + + +@pytest.mark.skipif(torch is None, reason="PyTorch not installed") +def test_embedder_token_selection_multivariate(): + X = _dataset(n_ts=4, sz=32, d=3) + n_context = 32 // PATCH_SIZE + embeddings = TimeSeriesFoundationEmbedder( + _DummyBackbone(), pooling=None, tokens=(0, -1) + ).fit_transform(X) + # Channels are concatenated along the feature axis, one row per token + assert embeddings.shape == (4, n_context, 3 * D_MODEL) + + # Channel k of the output must be the univariate embedding of channel k + backbone = _DummyBackbone() + joint = TimeSeriesFoundationEmbedder( + backbone, pooling=None, tokens=(0, -1) + ).fit_transform(X) + for k in range(3): + alone = TimeSeriesFoundationEmbedder( + backbone, pooling=None, tokens=(0, -1) + ).fit_transform(X[:, :, k : k + 1]) + np.testing.assert_allclose( + joint[:, :, k * D_MODEL : (k + 1) * D_MODEL], + alone, + rtol=1e-5, + atol=1e-6, + ) + + +@pytest.mark.skipif(torch is None, reason="PyTorch not installed") +def test_embedder_token_selection_errors(): + X = _dataset(n_ts=4, sz=32, d=1) + with pytest.raises(ValueError, match="`tokens` must be"): + TimeSeriesFoundationEmbedder(_DummyBackbone(), tokens=3).fit(X) + with pytest.raises(ValueError, match="`tokens` must be"): + TimeSeriesFoundationEmbedder(_DummyBackbone(), tokens=(1, 2, 3)).fit(X) + with pytest.raises(ValueError, match="selects no token"): + TimeSeriesFoundationEmbedder(_DummyBackbone(), tokens=(0, 0)).fit(X) + + +@pytest.mark.skipif(torch is None, reason="PyTorch not installed") +def test_probes_reject_unpooled_representations(): + X = _dataset(n_ts=4, sz=32, d=1) + for pooling in (None, "none"): + with pytest.raises(ValueError, match="pooling=None"): + LinearProbeForecaster( + _DummyBackbone(), pooling=pooling, context_length=16, horizon=2 + ).fit(X) + + +@pytest.mark.skipif(torch is None, reason="PyTorch not installed") +def test_probes_accept_token_selection(): + X_fc = _dataset(n_ts=6, sz=48, d=1) + model = LinearProbeForecaster( + _DummyBackbone(), context_length=16, horizon=2, stride=8, tokens=(0, -1) + ).fit(X_fc) + assert model.predict(X_fc).shape == (6, 2, 1) + + +@pytest.mark.skipif(torch is None, reason="PyTorch not installed") +def test_embedder_token_index_selects_the_register_token(): + X = _dataset(n_ts=3, sz=32, d=1) + # The register token is appended last and does not depend on the input, so + # selecting it must yield identical embeddings for all series. + embeddings = TimeSeriesFoundationEmbedder( + _DummyBackbone(), pooling="token", token_index=-1, layer=0 + ).fit_transform(X) + np.testing.assert_allclose( + embeddings, np.repeat(embeddings[:1], 3, axis=0), rtol=1e-5, atol=1e-6 + ) + + # Whereas a context token does depend on the input + embeddings = TimeSeriesFoundationEmbedder( + _DummyBackbone(), pooling="token", token_index=0, layer=0 + ).fit_transform(X) + assert not np.allclose(embeddings, embeddings[:1]) + + +@pytest.mark.skipif(torch is None, reason="PyTorch not installed") +def test_embedder_layer_selection(): + X = _dataset(n_ts=4, sz=32, d=1) + backbone = _DummyBackbone(n_layers=3) + embeddings = [ + TimeSeriesFoundationEmbedder(backbone, layer=layer).fit_transform(X) + for layer in (0, 1, 2, None) + ] + # Every layer yields a different representation of the same data + for i in range(len(embeddings)): + for j in range(i + 1, len(embeddings)): + assert not np.allclose(embeddings[i], embeddings[j]) + + # Negative indices address the stack from the end + np.testing.assert_allclose( + TimeSeriesFoundationEmbedder(backbone, layer=-1).fit_transform(X), + embeddings[2], + rtol=1e-5, + atol=1e-6, + ) + + # An explicit path to the layer stack gives the same result as autodetection + np.testing.assert_allclose( + TimeSeriesFoundationEmbedder( + backbone, layer=1, layers_path="encoder.block" + ).fit_transform(X), + embeddings[1], + rtol=1e-5, + atol=1e-6, + ) + + +@pytest.mark.skipif(torch is None, reason="PyTorch not installed") +def test_embedder_leaves_the_model_frozen(): + X = _dataset(n_ts=4, sz=32, d=1) + backbone = _DummyBackbone() + before = [p.detach().clone() for p in backbone.parameters()] + TimeSeriesFoundationEmbedder(backbone).fit_transform(X) + for parameter, reference in zip(backbone.parameters(), before): + assert parameter.grad is None + torch.testing.assert_close(parameter.detach(), reference) + + +@pytest.mark.skipif(torch is None, reason="PyTorch not installed") +def test_embedder_errors(): + X = _dataset(n_ts=4, sz=32, d=1) + with pytest.raises(ValueError, match="pooling"): + TimeSeriesFoundationEmbedder(_DummyBackbone(), pooling="median").fit(X) + with pytest.raises(ValueError, match="input_layout"): + TimeSeriesFoundationEmbedder(_DummyBackbone(), input_layout="nchw").fit(X) + with pytest.raises(TypeError, match="torch.nn.Module"): + TimeSeriesFoundationEmbedder(_DummyPipeline()).fit(X) + with pytest.raises(ValueError, match="out of range"): + TimeSeriesFoundationEmbedder(_DummyBackbone(n_layers=3), layer=7).fit(X) + with pytest.raises(ValueError, match="token_index"): + TimeSeriesFoundationEmbedder( + _DummyBackbone(), pooling="token", token_index=999 + ).fit(X) + with pytest.raises(ValueError, match="features"): + TimeSeriesFoundationEmbedder(_DummyBackbone()).fit(X).transform( + _dataset(n_ts=4, sz=32, d=2) + ) + with pytest.raises(RuntimeError, match="Could not determine which argument"): + class _InvalidForwardBackbone(torch.nn.Module): + def forward(self, *args, **kwargs): pass + TimeSeriesFoundationEmbedder(_InvalidForwardBackbone()).fit(X) + with pytest.raises(RuntimeError, match="forward hook was never triggered"): + class _InvalidLayerBackbone(torch.nn.Module): + def __init__(self): + super().__init__() + self.linears = torch.nn.ModuleList([ + torch.nn.Linear(32, 32), + torch.nn.Linear(D_MODEL, D_MODEL) + ]) + def forward(self, data): + self.linears[0](data) + TimeSeriesFoundationEmbedder(_InvalidLayerBackbone(), layer=1).fit(X) + + +@pytest.mark.skipif(torch is None, reason="PyTorch not installed") +def test_embedder_in_a_sklearn_pipeline(): + X = _dataset(n_ts=20, sz=32, d=1) + y = np.arange(20) % 2 + pipeline = make_pipeline( + TimeSeriesFoundationEmbedder(_DummyBackbone()), LinearSVC() + ) + assert pipeline.fit(X, y).predict(X).shape == (20,) + assert clone(pipeline) is not pipeline + + +@pytest.mark.skipif(torch is None, reason="PyTorch not installed") +def test_embedder_channels_first_layout(): + X = _dataset(n_ts=4, sz=32, d=3) + seen = {} + + class _ChannelsFirstBackbone(torch.nn.Module): + """Expects a (batch, d, sz) input, like some ``transformers`` models.""" + + def __init__(self): + super().__init__() + self.linear = torch.nn.Linear(32, D_MODEL) + + def forward(self, series): + seen["shape"] = tuple(series.shape) + return _Output(self.linear(series)) + + embedder = TimeSeriesFoundationEmbedder( + _ChannelsFirstBackbone(), input_layout="channels_first", pooling="mean" + ) + embeddings = embedder.fit_transform(X) + assert seen["shape"] == (4, 3, 32) + assert embeddings.shape == (4, D_MODEL) + + +@pytest.mark.skipif(torch is None, reason="PyTorch not installed") +def test_embedder_layers_path_with_numeric_segment(): + X = _dataset(n_ts=3, sz=32, d=1) + + class _WrappedBackbone(torch.nn.Module): + """Nests the encoder in a top-level ``ModuleList``, as some models do, + requiring a numeric index to reach the stack of layers.""" + + def __init__(self, n_layers=3, d_model=D_MODEL, patch_size=PATCH_SIZE, seed=0): + super().__init__() + torch.manual_seed(seed) + self.patch_size = patch_size + self.embedding = torch.nn.Linear(patch_size, d_model) + self.register_token = torch.nn.Parameter(torch.randn(1, 1, d_model)) + self.stages = torch.nn.ModuleList([_Encoder(n_layers, d_model)]) + + def forward(self, context): + batch_size, sz = context.shape + n_patches = sz // self.patch_size + patches = context[:, : n_patches * self.patch_size] + patches = patches.reshape(batch_size, n_patches, self.patch_size) + hidden_states = self.embedding(patches) + register = self.register_token.expand(batch_size, -1, -1) + hidden_states = torch.cat([hidden_states, register], dim=1) + return _Output(self.stages[0](hidden_states)) + + backbone = _WrappedBackbone() + embeddings = TimeSeriesFoundationEmbedder( + backbone, layer=1, layers_path="stages.0.block" + ).fit_transform(X) + assert embeddings.shape == (3, D_MODEL) + + with pytest.raises(AttributeError, match="Could not resolve"): + TimeSeriesFoundationEmbedder( + backbone, layer=0, layers_path="stages.0.bogus" + ).fit(X) + + +@pytest.mark.skipif(torch is None, reason="PyTorch not installed") +def test_embedder_autodetect_layers_failure(): + X = _dataset(n_ts=3, sz=32, d=1) + + class _NoLayerStack(torch.nn.Module): + def __init__(self): + super().__init__() + self.linear = torch.nn.Linear(32, D_MODEL) + + def forward(self, x): pass + + with pytest.raises(ValueError, match="Could not automatically locate"): + TimeSeriesFoundationEmbedder(_NoLayerStack(), layer=0).fit(X) + + class _HeterogeneousLayersBackbone(torch.nn.Module): + def __init__(self, d_model=D_MODEL): + super().__init__() + self.bloc = torch.nn.ModuleList([ + torch.nn.Linear(d_model, d_model), + torch.nn.LayerNorm(d_model) + ]) + + def forward(self, x): pass + + with pytest.raises(ValueError, match="Could not automatically locate"): + TimeSeriesFoundationEmbedder(_HeterogeneousLayersBackbone(), layer=0).fit(X) + + +@pytest.mark.skipif(torch is None, reason="PyTorch not installed") +def test_embedder_accepts_dict_model_outputs(): + X = _dataset(n_ts=3, sz=32, d=1) + + class _DictOutputBackbone(torch.nn.Module): + """Returns a plain dict, as some pipelines do instead of a ModelOutput.""" + + def __init__(self): + super().__init__() + self.linear = torch.nn.Linear(32, D_MODEL) + + def forward(self, x): + return {"hidden_states": self.linear(x).unsqueeze(1)} + + embeddings = TimeSeriesFoundationEmbedder(_DictOutputBackbone()).fit_transform(X) + assert embeddings.shape == (3, D_MODEL) + + +@pytest.mark.skipif(torch is None, reason="PyTorch not installed") +def test_embedder_accepts_hidden_states_tuple_attribute(): + X = _dataset(n_ts=3, sz=32, d=1) + + class _LayerStackOutput: + """Stands in for a ``transformers`` output exposing all-layer states.""" + + def __init__(self, layers): + self.hidden_states = layers + + class _AllLayersBackbone(torch.nn.Module): + def __init__(self): + super().__init__() + torch.manual_seed(0) + self.linear = torch.nn.Linear(32, D_MODEL) + + def forward(self, x): + hidden = self.linear(x).unsqueeze(1) + return _LayerStackOutput((hidden, 2 * hidden)) + + backbone = _AllLayersBackbone() + embeddings = TimeSeriesFoundationEmbedder(backbone).fit_transform(X) + with torch.no_grad(): + raw = backbone.linear(torch.as_tensor(X[:, :, 0], dtype=torch.float32)) + # Only the last element of the `hidden_states` tuple is read, matching + # transformers' convention that it holds one tensor per layer + np.testing.assert_allclose(embeddings, 2 * raw.numpy(), rtol=1e-5, atol=1e-6) + + +@pytest.mark.skipif(torch is None, reason="PyTorch not installed") +def test_embedder_accepts_plain_tuple_outputs(): + X = _dataset(n_ts=3, sz=32, d=1) + + class _PlainTupleBackbone(torch.nn.Module): + """Returns a raw tuple, with no ``ModelOutput`` wrapper at all.""" + + def __init__(self): + super().__init__() + self.linear = torch.nn.Linear(32, D_MODEL) + + def forward(self, x): + hidden = self.linear(x) + # The 3d sequence of hidden states must be preferred over the + # pooled 2d tensor that precedes it in the tuple + return (hidden, hidden.unsqueeze(1)) + + embeddings = TimeSeriesFoundationEmbedder(_PlainTupleBackbone()).fit_transform(X) + assert embeddings.shape == (3, D_MODEL) + + +@pytest.mark.skipif(torch is None, reason="PyTorch not installed") +def test_embedder_plain_tuple_falls_back_to_any_tensor(): + X = _dataset(n_ts=3, sz=32, d=1) + + class _NoSequenceBackbone(torch.nn.Module): + """Returns an already-pooled 2d tensor wrapped in a 1-tuple.""" + + def __init__(self): + super().__init__() + self.linear = torch.nn.Linear(32, D_MODEL) + + def forward(self, x): + return (self.linear(x),) + + embeddings = TimeSeriesFoundationEmbedder(_NoSequenceBackbone()).fit_transform(X) + assert embeddings.shape == (3, D_MODEL) + + +@pytest.mark.skipif(torch is None, reason="PyTorch not installed") +def test_embedder_rejects_unrecognized_model_output(): + X = _dataset(n_ts=3, sz=32, d=1) + + class _BogusOutputBackbone(torch.nn.Module): + def forward(self, x): + return "not a tensor" + + with pytest.raises(TypeError, match="Could not extract a hidden state"): + TimeSeriesFoundationEmbedder(_BogusOutputBackbone()).fit(X) + + +@pytest.mark.skipif(torch is None, reason="PyTorch not installed") +def test_embedder_rejects_invalid_hidden_state_rank(): + X = _dataset(n_ts=3, sz=32, d=1) + + class _WeirdRankBackbone(torch.nn.Module): + def forward(self, x): + return torch.tensor(x[:, None, None, :]) # 4d, not a valid hidden-state rank + + with pytest.raises(ValueError, match="Expected hidden states of shape"): + TimeSeriesFoundationEmbedder(_WeirdRankBackbone()).fit(X) + + +@pytest.mark.skipif(torch is None, reason="PyTorch not installed") +def test_embedder_explicit_device(): + X = _dataset(n_ts=3, sz=32, d=1) + embeddings = TimeSeriesFoundationEmbedder( + _DummyBackbone(), device="cpu" + ).fit_transform(X) + assert embeddings.shape == (3, D_MODEL) + + +@pytest.mark.skipif(torch is None, reason="PyTorch not installed") +def test_embedder_explicit_input_name(): + X = _dataset(n_ts=3, sz=32, d=1) + + class _NamedArgBackbone(torch.nn.Module): + def __init__(self): + super().__init__() + self.linear = torch.nn.Linear(32, D_MODEL) + + def forward(self, my_custom_arg): + return _Output(self.linear(my_custom_arg).unsqueeze(1)) + + embedder = TimeSeriesFoundationEmbedder( + _NamedArgBackbone(), input_name="my_custom_arg" + ) + embeddings = embedder.fit_transform(X) + assert embeddings.shape == (3, D_MODEL) + + +@pytest.mark.skipif(torch is None, reason="PyTorch not installed") +def test_embedder_resolves_input_name_by_position_when_unrecognized(): + X = _dataset(n_ts=3, sz=32, d=1) + + class _PositionalArgBackbone(torch.nn.Module): + def __init__(self): + super().__init__() + self.linear = torch.nn.Linear(32, D_MODEL) + + # `data` is not one of CANDIDATE_INPUT_NAMES, forcing the fallback to + # the first positional parameter of `forward` + def forward(self, data): + return _Output(self.linear(data).unsqueeze(1)) + + embedder = TimeSeriesFoundationEmbedder(_PositionalArgBackbone()) + embeddings = embedder.fit_transform(X) + assert embedder._input_name == "data" + assert embeddings.shape == (3, D_MODEL) + + +@pytest.mark.skipif(torch is None, reason="PyTorch not installed") +def test_embedder_verbose_prints_progress(capsys): + X = _dataset(n_ts=5, sz=32, d=1) + TimeSeriesFoundationEmbedder( + _DummyBackbone(), batch_size=2, verbose=1 + ).fit_transform(X) + captured = capsys.readouterr() + assert "Embedded 5/5 series" in captured.out + + +@pytest.mark.skipif(torch is None, reason="PyTorch not installed") +def test_embedder_flatten_pooling_warns_on_length_mismatch(): + X_fit = _dataset(n_ts=3, sz=32, d=1) + X_other = _dataset(n_ts=3, sz=48, d=1, seed=1) + embedder = TimeSeriesFoundationEmbedder( + _DummyBackbone(), pooling="flatten" + ).fit(X_fit) + with pytest.warns(UserWarning, match="pooling='flatten' produces context-length"): + embedder.transform(X_other) + + +# --------------------------------------------------------------------------- +# ZeroShotForecaster +# --------------------------------------------------------------------------- + + +@pytest.mark.skipif(torch is None, reason="PyTorch not installed") +def test_zero_shot_forecaster(): + X = _dataset(n_ts=5, sz=32, d=1) + model = ZeroShotForecaster(_DummyPipeline()) + + with pytest.warns(UserWarning, match="does not train anything"): + model.fit(X) + + predicted = model.predict(X, n=7) + assert predicted.shape == (5, 7, 1) + # The dummy pipeline repeats the last value, and the median quantile of its + # symmetric output is that value exactly + np.testing.assert_allclose( + predicted, np.repeat(X[:, -1:], 7, axis=1), rtol=1e-5, atol=1e-6 + ) + + +@pytest.mark.skipif(torch is None, reason="PyTorch not installed") +def test_zero_shot_forecaster_multivariate(): + X = _dataset(n_ts=5, sz=32, d=3) + predicted = ZeroShotForecaster(_DummyPipeline()).predict(X, n=4) + assert predicted.shape == (5, 4, 3) + # Channels must not get mixed up while being folded in and out of the batch + np.testing.assert_allclose( + predicted, np.repeat(X[:, -1:], 4, axis=1), rtol=1e-5, atol=1e-6 + ) + + +@pytest.mark.skipif(torch is None, reason="PyTorch not installed") +def test_zero_shot_forecaster_quantile_selection(): + X = _dataset(n_ts=3, sz=32, d=1) + low = ZeroShotForecaster(_DummyPipeline(), quantile=0.1).predict(X, n=3) + high = ZeroShotForecaster(_DummyPipeline(), quantile=0.9).predict(X, n=3) + assert np.all(low < high) + + +@pytest.mark.skipif(torch is None, reason="PyTorch not installed") +def test_zero_shot_forecaster_custom_predict_fn(): + X = _dataset(n_ts=4, sz=32, d=1) + + def predict_fn(model, context, horizon): + assert context.shape == (4, 32) + return np.zeros((context.shape[0], horizon)) + + predicted = ZeroShotForecaster(object(), predict_fn=predict_fn).predict(X, n=5) + np.testing.assert_array_equal(predicted, np.zeros((4, 5, 1))) + + +@pytest.mark.skipif(torch is None, reason="PyTorch not installed") +def test_zero_shot_forecaster_context_length(): + X = _dataset(n_ts=4, sz=64, d=1) + seen = {} + + def predict_fn(model, context, horizon): + seen["shape"] = context.shape + return np.zeros((context.shape[0], horizon)) + + ZeroShotForecaster( + object(), predict_fn=predict_fn, context_length=16 + ).predict(X, n=2) + assert seen["shape"] == (4, 16) + + +@pytest.mark.skipif(torch is None, reason="PyTorch not installed") +def test_zero_shot_forecaster_negotiates_the_context_format(): + X = _dataset(n_ts=4, sz=32, d=1) + + class _NeedsVariateAxis: + """Rejects 2d contexts, like ``Chronos2Pipeline`` does.""" + + def predict(self, inputs, prediction_length=1): + inputs = np.asarray(inputs) + if inputs.ndim != 3: + raise ValueError( + "Expected 3-d tensor with shape " + "(n_series, n_variates, history_length)." + ) + return np.zeros((inputs.shape[0], prediction_length)) + + model = ZeroShotForecaster(_NeedsVariateAxis()) + assert model.predict(X, n=3).shape == (4, 3, 1) + assert model.context_format_ == "3d" + # The negotiated format is reused rather than re-discovered + assert model.predict(X, n=3).shape == (4, 3, 1) + + class _NeedsList: + def predict(self, inputs, prediction_length=1): + if not isinstance(inputs, list): + raise TypeError("A list of series is expected.") + return np.zeros((len(inputs), prediction_length)) + + model = ZeroShotForecaster(_NeedsList()) + assert model.predict(X, n=3).shape == (4, 3, 1) + assert model.context_format_ == "list" + + class _RejectsEverything: + def predict(self, inputs, prediction_length=1): + raise ValueError("nope") + + with pytest.raises(ValueError, match="rejected every supported way"): + ZeroShotForecaster(_RejectsEverything()).predict(X, n=3) + + +@pytest.mark.skipif(torch is None, reason="PyTorch not installed") +def test_zero_shot_forecaster_explicit_horizon_axis(): + X = _dataset(n_ts=4, sz=32, d=1) + + class _HorizonFirst: + """Returns (n_series, horizon, n_quantiles), like predict_quantiles.""" + + def predict(self, inputs, prediction_length=1): + n_rows = len(inputs) + levels = np.arange(5)[None, None, :] + return np.zeros((n_rows, prediction_length, 5)) + levels + + class _HorizonLast: + """Returns (n_series, n_quantiles, horizon), like Chronos-2.""" + + def predict(self, inputs, prediction_length=1): + n_rows = len(inputs) + levels = np.arange(5)[None, :, None] + return np.zeros((n_rows, 5, prediction_length)) + levels + + # Both conventions are handled once the horizon axis is stated + for model, axis in [(_HorizonFirst(), 1), (_HorizonLast(), 2)]: + predicted = ZeroShotForecaster(model, horizon_axis=axis).predict(X, n=5) + np.testing.assert_allclose(predicted, np.full((4, 5, 1), 2.0)) + + # Negative indices are accepted + np.testing.assert_allclose( + ZeroShotForecaster(_HorizonLast(), horizon_axis=-1).predict(X, n=5), + np.full((4, 5, 1), 2.0), + ) + + # A wrong axis is reported rather than silently mis-slicing + with pytest.raises(ValueError, match="cannot hold a horizon"): + ZeroShotForecaster(_HorizonLast(), horizon_axis=1).predict(X, n=7) + with pytest.raises(ValueError, match="out of range"): + ZeroShotForecaster(_HorizonLast(), horizon_axis=5).predict(X, n=3) + with pytest.raises(ValueError, match="out of range"): + ZeroShotForecaster(_HorizonLast(), horizon_axis=0).predict(X, n=3) + with pytest.raises(ValueError, match="must be an integer or 'auto'"): + ZeroShotForecaster(_HorizonLast(), horizon_axis="last").predict(X, n=3) + + +@pytest.mark.skipif(torch is None, reason="PyTorch not installed") +def test_zero_shot_forecaster_explicit_horizon_axis_lifts_ambiguity(): + X = _dataset(n_ts=4, sz=32, d=1) + + class _Ambiguous: + def predict(self, inputs, prediction_length=1): + n_rows = len(inputs) + # (n_series, horizon, horizon), the second axis being the horizon + steps = np.arange(prediction_length)[None, :, None] + return np.zeros((n_rows, prediction_length, prediction_length)) + steps + + import warnings + + with warnings.catch_warnings(): + warnings.simplefilter("error") + predicted = ZeroShotForecaster(_Ambiguous(), horizon_axis=1).predict(X, n=4) + np.testing.assert_allclose( + predicted, np.tile(np.arange(4.0)[None, :, None], (4, 1, 1)) + ) + + +@pytest.mark.skipif(torch is None, reason="PyTorch not installed") +def test_zero_shot_forecaster_warns_on_ambiguous_output_shape(): + X = _dataset(n_ts=4, sz=32, d=1) + + class _Ambiguous: + """Returns as many quantiles as there are forecast steps.""" + + def predict(self, inputs, prediction_length=1): + n_rows = len(inputs) + return np.zeros((n_rows, prediction_length, prediction_length)) + + with pytest.warns(UserWarning, match="Set `horizon_axis` explicitly"): + ZeroShotForecaster(_Ambiguous()).predict(X, n=5) + + # A horizon of one is not ambiguous, singleton axes are simply reduced + import warnings + + with warnings.catch_warnings(): + warnings.simplefilter("error") + assert ZeroShotForecaster(_Ambiguous()).predict(X, n=1).shape == (4, 1, 1) + + +@pytest.mark.skipif(torch is None, reason="PyTorch not installed") +def test_zero_shot_forecaster_reduces_quantiles_to_a_point_forecast(): + X = _dataset(n_ts=4, sz=32, d=1) + + class _Quantiles: + """Returns (n_series, n_quantiles, horizon), like Chronos-2.""" + + def predict(self, inputs, prediction_length=1): + n_rows = len(inputs) + levels = np.arange(9)[None, :, None] + return np.zeros((n_rows, 9, prediction_length)) + levels + + predicted = ZeroShotForecaster(_Quantiles()).predict(X, n=4) + # The median of 0..8 is 4, whatever the horizon + np.testing.assert_allclose(predicted, np.full((4, 4, 1), 4.0)) + np.testing.assert_allclose( + ZeroShotForecaster(_Quantiles(), quantile=None).predict(X, n=4), + np.full((4, 4, 1), 4.0), + ) + np.testing.assert_allclose( + ZeroShotForecaster(_Quantiles(), quantile=0.0).predict(X, n=4), + np.zeros((4, 4, 1)), + ) + + +@pytest.mark.skipif(torch is None, reason="PyTorch not installed") +def test_zero_shot_forecaster_errors(): + X = _dataset(n_ts=4, sz=32, d=1) + with pytest.raises(ValueError, match="`X` is required"): + ZeroShotForecaster(_DummyPipeline()).predict(None) + with pytest.raises(ValueError, match="positive integer"): + ZeroShotForecaster(_DummyPipeline()).predict(X, n=0) + with pytest.raises(ValueError, match="forecasting method"): + ZeroShotForecaster(object()).predict(X, n=1) + with pytest.raises(ValueError, match="input_layout"): + ZeroShotForecaster(_DummyPipeline(), input_layout="bogus").predict(X) + + +@pytest.mark.skipif(torch is None, reason="PyTorch not installed") +def test_zero_shot_forecaster_skips_methods_without_a_horizon_argument(): + X = _dataset(n_ts=4, sz=32, d=1) + calls = [] + + class _PartialModel: + """Exposes a `predict` with no forecast-horizon argument, which must + be skipped in favor of `forecast`.""" + + def predict(self, inputs): # pragma: no cover + calls.append("predict") + raise AssertionError("predict should never be called") + + def forecast(self, inputs, prediction_length=1): + calls.append("forecast") + n_rows = len(inputs) + return np.zeros((n_rows, prediction_length)) + + predicted = ZeroShotForecaster(_PartialModel()).predict(X, n=3) + assert predicted.shape == (4, 3, 1) + assert calls == ["forecast"] + + +@pytest.mark.skipif(torch is None, reason="PyTorch not installed") +def test_zero_shot_forecaster_accepts_object_attribute_output(): + X = _dataset(n_ts=4, sz=32, d=1) + + class _AttributeOutput: + """Stands in for an output object exposing the forecast as an attribute.""" + + def __init__(self, mean): + self.mean = mean + + class _ObjectOutputModel: + def predict(self, inputs, prediction_length=1): + n_rows = len(inputs) + return _AttributeOutput(np.zeros((n_rows, prediction_length))) + + predicted = ZeroShotForecaster(_ObjectOutputModel()).predict(X, n=3) + assert predicted.shape == (4, 3, 1) + + +@pytest.mark.skipif(torch is None, reason="PyTorch not installed") +def test_zero_shot_forecaster_univariate_row_mismatch_error(): + X = _dataset(n_ts=4, sz=32, d=1) + + class _WrongRowCountModel: + def predict(self, inputs, prediction_length=1): + return np.zeros((len(inputs) - 1, prediction_length)) + + with pytest.raises( + ValueError, match="forecasts for 3 series while 4 were provided" + ): + ZeroShotForecaster(_WrongRowCountModel()).predict(X, n=2) + + +@pytest.mark.skipif(torch is None, reason="PyTorch not installed") +def test_zero_shot_forecaster_univariate_forecast_too_short_error(): + X = _dataset(n_ts=4, sz=32, d=1) + + class _TooShortModel: + def predict(self, inputs, prediction_length=1): + n_rows = len(inputs) + return np.zeros((n_rows, prediction_length - 1)) # one step short + + with pytest.raises(ValueError, match="shorter than the requested horizon"): + ZeroShotForecaster(_TooShortModel()).predict(X, n=3) + + +@pytest.mark.skipif(torch is None, reason="PyTorch not installed") +def test_zero_shot_forecaster_accepts_dict_output(): + X = _dataset(n_ts=4, sz=32, d=1) + + class _DictOutputModel: + """Returns a plain dict, as some pipelines do instead of a ModelOutput.""" + + def predict(self, inputs, prediction_length=1): + n_rows = len(inputs) + return {"predictions": np.zeros((n_rows, prediction_length))} + + predicted = ZeroShotForecaster(_DictOutputModel()).predict(X, n=3) + assert predicted.shape == (4, 3, 1) + + +@pytest.mark.skipif(torch is None, reason="PyTorch not installed") +def test_zero_shot_forecaster_empty_output_error(): + X = _dataset(n_ts=4, sz=32, d=1) + + class _EmptyOutputModel: + def predict(self, inputs, prediction_length=1): + return [] + + with pytest.raises(ValueError, match="empty forecast"): + ZeroShotForecaster(_EmptyOutputModel()).predict(X, n=3) + + +@pytest.mark.skipif(torch is None, reason="PyTorch not installed") +def test_zero_shot_forecaster_heterogeneous_output_keeps_first_entry(): + X = _dataset(n_ts=4, sz=32, d=1) + + class _HeterogeneousOutputModel: + """Returns (forecast, extra_diagnostics), as some pipelines do.""" + + def predict(self, inputs, prediction_length=1): + n_rows = len(inputs) + forecast = np.zeros((n_rows, prediction_length)) + diagnostics = np.zeros((n_rows, 2)) # a different shape + return (forecast, diagnostics) + + predicted = ZeroShotForecaster(_HeterogeneousOutputModel()).predict(X, n=3) + assert predicted.shape == (4, 3, 1) + + +@pytest.mark.skipif(torch is None, reason="PyTorch not installed") +def test_zero_shot_forecaster_horizon_axis_prefers_smallest_longer_axis(): + X = _dataset(n_ts=4, sz=32, d=1) + + class _PaddedForecastModel: + """Pads the forecast length to a fixed block size larger than asked, + with no axis exactly matching the horizon.""" + + def predict(self, inputs, prediction_length=1): + n_rows = len(inputs) + n_samples = 3 # sample paths, always shorter than the horizon here + padded_length = prediction_length + 5 + return np.zeros((n_rows, n_samples, padded_length)) + + predicted = ZeroShotForecaster(_PaddedForecastModel()).predict(X, n=4) + np.testing.assert_allclose(predicted, np.zeros((4, 4, 1))) + + +@pytest.mark.skipif(torch is None, reason="PyTorch not installed") +def test_zero_shot_forecaster_horizon_axis_cannot_be_identified(): + X = _dataset(n_ts=4, sz=32, d=1) + + class _TooShortEverywhereModel: + def predict(self, inputs, prediction_length=1): + n_rows = len(inputs) + return np.zeros((n_rows, 2, 2)) + + with pytest.raises( + ValueError, match="Could not identify the forecast horizon axis" + ): + ZeroShotForecaster(_TooShortEverywhereModel()).predict(X, n=4) + + +@pytest.mark.skipif(torch is None, reason="PyTorch not installed") +def test_zero_shot_forecaster_natively_multivariate_channels_last(): + X = _dataset(n_ts=4, sz=32, d=3) + + class _NativelyMultivariateModel: + """Returns per-channel forecasts directly, given a (n_ts, sz, d) input.""" + + def predict(self, inputs, prediction_length=1): + inputs = np.asarray(inputs) + last = inputs[:, -1, :] # (n_ts, d) + return np.repeat(last[:, None, :], prediction_length, axis=1) + + predicted = ZeroShotForecaster( + _NativelyMultivariateModel(), input_layout="channels_last" + ).predict(X, n=4) + assert predicted.shape == (4, 4, 3) + np.testing.assert_allclose(predicted, np.repeat(X[:, -1:], 4, axis=1)) + + +@pytest.mark.skipif(torch is None, reason="PyTorch not installed") +def test_zero_shot_forecaster_natively_multivariate_channels_first(): + X = _dataset(n_ts=4, sz=32, d=3) + + class _ChannelsFirstModel: + """Expects a (n_ts, d, sz) input and returns (n_ts, d, horizon).""" + + def predict(self, inputs, prediction_length=1): + inputs = np.asarray(inputs) + last = inputs[:, :, -1] # (n_ts, d) + return np.repeat(last[:, :, None], prediction_length, axis=2) + + predicted = ZeroShotForecaster( + _ChannelsFirstModel(), input_layout="channels_first" + ).predict(X, n=4) + assert predicted.shape == (4, 4, 3) + np.testing.assert_allclose(predicted, np.repeat(X[:, -1:], 4, axis=1)) + + +@pytest.mark.skipif(torch is None, reason="PyTorch not installed") +def test_zero_shot_forecaster_natively_multivariate_errors(): + X = _dataset(n_ts=4, sz=32, d=3) + + class _WrongBatchModel: + def predict(self, inputs, prediction_length=1): + return np.zeros((len(inputs) - 1, prediction_length, 3)) + + with pytest.raises( + ValueError, match="forecasts for 3 series while 4 were provided" + ): + ZeroShotForecaster( + _WrongBatchModel(), input_layout="channels_last" + ).predict(X, n=2) + + class _WrongChannelsModel: + def predict(self, inputs, prediction_length=1): + n_rows = len(inputs) + return np.zeros((n_rows, prediction_length, 5)) # wrong d + + with pytest.raises(ValueError, match="Expected forecasts of shape"): + ZeroShotForecaster( + _WrongChannelsModel(), input_layout="channels_last" + ).predict(X, n=2) + + +# --------------------------------------------------------------------------- +# LinearProbeForecaster +# --------------------------------------------------------------------------- + + +@pytest.mark.skipif(torch is None, reason="PyTorch not installed") +def test_linear_probe_forecaster(): + X = _dataset(n_ts=8, sz=64, d=1) + model = LinearProbeForecaster( + _DummyBackbone(), context_length=16, horizon=4, stride=4 + ) + model.fit(X) + # 8 series x windows of span 20 taken every 4 timestamps + assert model.n_windows_ == 8 * len(range(0, 64 - 20 + 1, 4)) + predicted = model.predict(X) + assert predicted.shape == (8, 4, 1) + # A shorter horizon truncates the forecast rather than refitting + np.testing.assert_allclose(model.predict(X, n=2), predicted[:, :2]) + + +@pytest.mark.skipif(torch is None, reason="PyTorch not installed") +def test_linear_probe_forecaster_multivariate(): + X = _dataset(n_ts=6, sz=48, d=2) + model = LinearProbeForecaster( + _DummyBackbone(), context_length=16, horizon=3, stride=8 + ).fit(X) + assert model.predict(X).shape == (6, 3, 2) + assert model.embedder_.embedding_size_ == 2 * D_MODEL + + +@pytest.mark.skipif(torch is None, reason="PyTorch not installed") +def test_linear_probe_forecaster_learns_something(): + # On a dataset of pure sine waves with varying phases, a probe on top of a + # frozen encoder should do clearly better than forecasting the last value. + rng = np.random.RandomState(0) + t = np.linspace(0, 8 * np.pi, 96) + phases = rng.uniform(0, 2 * np.pi, size=40) + X = np.sin(t[None, :] + phases[:, None])[:, :, None] + + model = LinearProbeForecaster( + _DummyBackbone(), context_length=16, horizon=4, stride=2 + ).fit(X[:30]) + assert model.score(X[30:]) > 0.5 + + +@pytest.mark.skipif(torch is None, reason="PyTorch not installed") +def test_linear_probe_forecaster_accepts_any_probe(): + X = _dataset(n_ts=6, sz=48, d=1) + model = LinearProbeForecaster( + _DummyBackbone(), + probe=Ridge(alpha=1.0), + context_length=16, + horizon=2, + stride=8, + ).fit(X) + assert isinstance(model.probe_, Ridge) + # The passed estimator is cloned, not fitted in place + assert model.probe_ is not model.probe + assert not hasattr(model.probe, "coef_") + + +@pytest.mark.skipif(torch is None, reason="PyTorch not installed") +def test_linear_probe_forecaster_errors(): + X = _dataset(n_ts=4, sz=48, d=1) + with pytest.raises(ValueError, match="context_length \\+ horizon"): + LinearProbeForecaster( + _DummyBackbone(), context_length=40, horizon=16 + ).fit(X) + with pytest.raises(ValueError, match="positive integer"): + LinearProbeForecaster(_DummyBackbone(), horizon=0).fit(X) + + model = LinearProbeForecaster( + _DummyBackbone(), context_length=16, horizon=4, stride=8 + ).fit(X) + with pytest.raises(ValueError, match="horizon of 4"): + model.predict(X, n=8) + with pytest.raises(ValueError, match="context_length"): + model.predict(_dataset(n_ts=4, sz=8, d=1)) + with pytest.raises(ValueError, match="features"): + model.predict(_dataset(n_ts=4, sz=48, d=3)) + + +@pytest.mark.skipif(torch is None, reason="PyTorch not installed") +def test_linear_probe_forecaster_fit_predict(): + X = _dataset(n_ts=6, sz=48, d=1) + model = LinearProbeForecaster( + _DummyBackbone(), context_length=16, horizon=2, stride=8 + ) + np.testing.assert_allclose(model.fit_predict(X, n=2), model.predict(X, n=2)) + + +@pytest.mark.skipif(torch is None, reason="PyTorch not installed") +def test_estimators_are_clonable(): + for estimator in ( + TimeSeriesFoundationEmbedder(_DummyBackbone()), + ZeroShotForecaster(_DummyPipeline()), + LinearProbeForecaster(_DummyBackbone()), + ): + cloned = clone(estimator) + assert cloned is not estimator + assert cloned.get_params().keys() == estimator.get_params().keys() diff --git a/tests/test_foundation_zoo.py b/tests/test_foundation_zoo.py new file mode 100644 index 000000000..3098c6ee5 --- /dev/null +++ b/tests/test_foundation_zoo.py @@ -0,0 +1,223 @@ +"""Tests exercising :mod:`tslearn.foundation` against real pre-trained models. + +Unlike :mod:`test_foundation`, these tests download actual checkpoints from +the Hugging Face Hub and require model-specific packages (``timesfm``, +``uni2ts``, ``granite-tsfm``, ``momentfm``, a recent ``transformers``...) +that are *not* part of tslearn's own dependencies and that, in some cases, +pin conflicting versions of ``torch`` or ``transformers`` against one +another. Each test is therefore independent, skipped unless both its own +package is importable and the ``TSLEARN_RUN_FOUNDATION_ZOO`` environment +variable is set, and meant to be run in its own dedicated environment (see +``.github/workflows/test_foundation_zoo.yml``) rather than as part of the +regular test suite. + +The code below mirrors the snippets given in the +``plot_foundation_model_zoo.py`` gallery example: this file is what backs +the claim, made there, that each recipe actually works. +""" + +import os + +import numpy as np + +import pytest + +from tslearn.generators import random_walks + +pytestmark = pytest.mark.skipif( + not os.environ.get("TSLEARN_RUN_FOUNDATION_ZOO"), + reason="Set TSLEARN_RUN_FOUNDATION_ZOO=1 to run tests that download " + "real pre-trained models from the Hugging Face Hub.", +) +torch = pytest.importorskip("torch", reason="torch not installed") + +N_TS, SZ, D, CONTEXT_LENGTH, HORIZON = 5, 200, 1, 64, 12 + + +@pytest.mark.parametrize("data", [ + random_walks(n_ts=N_TS, sz=SZ, random_state=0).astype(np.float64), + torch.rand(N_TS, SZ, D, dtype=torch.float64), +]) +def test_chronos_bolt(data): + chronos = pytest.importorskip("chronos") + + from tslearn.foundation import LinearProbeForecaster, ZeroShotForecaster + + pipeline = chronos.BaseChronosPipeline.from_pretrained( + "amazon/chronos-bolt-small", device_map="cpu" + ) + + zero_shot = ZeroShotForecaster(pipeline) + y_zero_shot = zero_shot.predict(data, n=HORIZON) + assert y_zero_shot.shape == (N_TS, HORIZON, 1) + + probe = LinearProbeForecaster( + pipeline.model, + context_length=CONTEXT_LENGTH, + horizon=HORIZON, + stride=8, + layer=-1, + layers_path="encoder.block", + pooling="mean", + ) + probe.fit(data) + y_probe = probe.predict(data) + assert y_probe.shape == (N_TS, HORIZON, 1) + + +@pytest.mark.parametrize("data", [ + random_walks(n_ts=N_TS, sz=SZ, random_state=0), + torch.rand(N_TS, SZ, D, dtype=torch.float64), +]) +def test_timesfm(data): + timesfm = pytest.importorskip("timesfm") + + from tslearn.foundation import ZeroShotForecaster + + model = timesfm.TimesFM_2p5_200M_torch.from_pretrained( + "google/timesfm-2.5-200m-pytorch" + ) + model.compile( + timesfm.ForecastConfig( + max_context=CONTEXT_LENGTH, max_horizon=HORIZON, normalize_inputs=True + ) + ) + + zero_shot = ZeroShotForecaster( + model, + predict_fn=lambda model, context, horizon: model.forecast( + horizon=horizon, inputs=list(context) + )[0], + context_length=CONTEXT_LENGTH, + ) + y_zero_shot = zero_shot.predict(data, n=HORIZON) + assert y_zero_shot.shape == (N_TS, HORIZON, 1) + + +@pytest.mark.parametrize("data", [ + random_walks(n_ts=N_TS, sz=SZ, random_state=0), + torch.rand(N_TS, SZ, D, dtype=torch.float64), +]) +def test_moirai(data): + pytest.importorskip("uni2ts") + from uni2ts.model.moirai import MoiraiForecast, MoiraiModule + + from tslearn.foundation import ZeroShotForecaster + + module = MoiraiModule.from_pretrained("Salesforce/moirai-1.1-R-small") + forecast_model = MoiraiForecast( + module=module, + prediction_length=HORIZON, + context_length=CONTEXT_LENGTH, + patch_size=32, + target_dim=1, + feat_dynamic_real_dim=0, + past_feat_dynamic_real_dim=0, + ) + + def predict_fn(model, context, horizon): + past_target = context.unsqueeze(-1) + past_observed = torch.ones_like(past_target, dtype=torch.bool) + past_is_pad = torch.zeros(past_target.shape[:2], dtype=torch.bool) + with torch.no_grad(): + return model(past_target, past_observed, past_is_pad, num_samples=20) + + zero_shot = ZeroShotForecaster( + forecast_model, predict_fn=predict_fn, context_length=CONTEXT_LENGTH + ) + y_zero_shot = zero_shot.predict(data, n=HORIZON) + assert y_zero_shot.shape == (N_TS, HORIZON, 1) + + +@pytest.mark.parametrize("data", [ + random_walks(n_ts=N_TS, sz=SZ, random_state=0), + torch.rand(N_TS, SZ, D, dtype=torch.float64), +]) +def test_ttm(data): + pytest.importorskip("tsfm_public") + from tsfm_public.models.tinytimemixer import TinyTimeMixerForPrediction + + from tslearn.foundation import ZeroShotForecaster + + model = TinyTimeMixerForPrediction.from_pretrained( + "ibm-granite/granite-timeseries-ttm-r2", num_input_channels=1, revision="main" + ) + + def predict_fn(model, context, horizon): + past_values = context.unsqueeze(-1) + with torch.no_grad(): + return model(past_values=past_values).prediction_outputs + + zero_shot = ZeroShotForecaster( + model, predict_fn=predict_fn, context_length=model.config.context_length + ) + X = random_walks( + n_ts=N_TS, sz=model.config.context_length + 50, random_state=0 + ).astype(np.float32) + y_zero_shot = zero_shot.predict(X, n=model.config.prediction_length) + assert y_zero_shot.shape == (N_TS, model.config.prediction_length, 1) + + +@pytest.mark.parametrize("data", [ + random_walks(n_ts=N_TS, sz=SZ, random_state=0), + torch.rand(N_TS, SZ, D, dtype=torch.float64), +]) +def test_moment(data): + pytest.importorskip("momentfm") + from momentfm import MOMENTPipeline + + from tslearn.foundation import LinearProbeForecaster + + model = MOMENTPipeline.from_pretrained( + "AutonLab/MOMENT-1-small", model_kwargs={"task_name": "embedding"} + ) + model.init() + + probe = LinearProbeForecaster( + model, + context_length=CONTEXT_LENGTH, + horizon=HORIZON, + stride=32, + layer=-1, + layers_path="encoder.block", + pooling="mean", + input_layout="channels_first", + ) + probe.fit(data) + y_probe = probe.predict(data) + assert y_probe.shape == (N_TS, HORIZON, 1) + + +@pytest.mark.parametrize("data", [ + random_walks(n_ts=N_TS, sz=SZ, random_state=0), + torch.rand(N_TS, SZ, D, dtype=torch.float64), +]) +def test_time_moe(data): + transformers = pytest.importorskip("transformers") + + from tslearn.foundation import LinearProbeForecaster, ZeroShotForecaster + + model = transformers.AutoModelForCausalLM.from_pretrained( + "Maple728/TimeMoE-50M", trust_remote_code=True + ) + + def predict_fn(model, context, horizon): + out = model.generate(input_ids=context, max_new_tokens=horizon) + return out[:, -horizon:] + + zero_shot = ZeroShotForecaster(model, predict_fn=predict_fn) + y_zero_shot = zero_shot.predict(data, n=HORIZON) + assert y_zero_shot.shape == (N_TS, HORIZON, 1) + + probe = LinearProbeForecaster( + model, + context_length=CONTEXT_LENGTH, + horizon=HORIZON, + stride=8, + layer=-1, + layers_path="model.layers", + pooling="last", + ) + probe.fit(data) + y_probe = probe.predict(data) + assert y_probe.shape == (N_TS, HORIZON, 1) diff --git a/tslearn/foundation/__init__.py b/tslearn/foundation/__init__.py new file mode 100644 index 000000000..dbf7f78b6 --- /dev/null +++ b/tslearn/foundation/__init__.py @@ -0,0 +1,37 @@ +""" +The :mod:`tslearn.foundation` module gathers estimators that re-use pre-trained +time series models, such as the ones published on the Hugging Face Hub, behind +the usual tslearn API. + +Two adaptation strategies are covered: + +* zero-shot forecasting, with :class:`ZeroShotForecaster`, which uses a + pre-trained forecaster as-is; +* linear probing for forecasting, with :class:`LinearProbeForecaster`. + +Linear probing relies on :class:`TimeSeriesFoundationEmbedder`, which turns any frozen +PyTorch model into a feature extractor and lets one choose which layer to +read representations from and how to pool them. Being a regular +scikit-learn transformer, it also covers linear probing for classification, +or any other downstream task, by composing with a +:class:`sklearn.pipeline.Pipeline`. + +These estimators are deliberately agnostic to any specific model +implementation: they duck-type the wrapped object and offer escape hatches +(``predict_fn``, ``input_name``, ``layers_path``, ``input_layout``) for models +that depart from the most widespread conventions. + +Notes +----- + This module requires PyTorch, which is an optional dependency of tslearn. + +""" + +from ._embedding import TimeSeriesFoundationEmbedder +from ._forecasting import LinearProbeForecaster, ZeroShotForecaster + +__all__ = [ + "LinearProbeForecaster", + "TimeSeriesFoundationEmbedder", + "ZeroShotForecaster", +] diff --git a/tslearn/foundation/_embedding.py b/tslearn/foundation/_embedding.py new file mode 100644 index 000000000..8e011d864 --- /dev/null +++ b/tslearn/foundation/_embedding.py @@ -0,0 +1,573 @@ +"""Model-agnostic extraction of representations from pre-trained time series models.""" + +import inspect +import warnings + +import numpy as np + +from sklearn.base import BaseEstimator, TransformerMixin +from sklearn.utils.validation import check_is_fitted + +from tslearn.bases import TimeSeriesMixin +from tslearn.utils import check_array, to_time_series_dataset + +try: + import torch +except ImportError: + torch = None + + +#: Names of ``forward`` arguments that are commonly used by pre-trained time +#: series models to receive the raw context values. +CANDIDATE_INPUT_NAMES = ( + "context", + "past_values", + "past_target", + "input_values", + "inputs", + "x_enc", + "series", + "x", +) + +#: Attribute names under which model outputs commonly expose hidden states. +CANDIDATE_HIDDEN_STATE_NAMES = ( + "last_hidden_state", + "hidden_states", + "encoder_last_hidden_state", +) + +#: Supported ways of aggregating token representations. ``None``, for which +#: ``"none"`` is accepted as an alias, keeps them all and turns the estimator +#: into a time series to time series transform. +POOLINGS = ("mean", "max", "token", "last", "flatten", "none", None) + +#: Supported layouts for the array handed over to the wrapped model. +LAYOUTS = ("univariate", "channels_last", "channels_first") + + +def _normalize_pooling(pooling): + return None if pooling == "none" else pooling + + +def _token_slice(tokens): + """Turn the ``tokens`` parameter into a slice over the token axis.""" + if tokens is None: + return slice(None) + if isinstance(tokens, slice): + return tokens + if isinstance(tokens, (tuple, list)) and len(tokens) == 2: + start, stop = tokens + if all(bound is None or isinstance(bound, (int, np.integer)) + for bound in (start, stop)): + return slice(start, stop) + raise ValueError( + "`tokens` must be None, a slice, or a (start, stop) pair of integers, " + f"got {tokens!r}." + ) + + +def _layout_to_model_input(X, layout): + """Lay a ``(n_ts, sz, d)`` dataset out as expected by the wrapped model. + + With the ``"univariate"`` layout, channels are unfolded along the batch + axis so that the returned array has shape ``(n_ts * d, sz)``, series + ``i`` and channel ``k`` sitting at row ``i * d + k``. + """ + n_ts, sz, d = X.shape + if layout == "univariate": + return torch.swapaxes(X, 1, 2).reshape(n_ts * d, sz).contiguous() + if layout == "channels_first": + return torch.swapaxes(X, 1, 2).contiguous() + if layout == "channels_last": + return X.contiguous() + raise ValueError(f"`input_layout` must be one of {LAYOUTS}, got '{layout}'.") # pragma: no cover + +def _resolve_attribute_path(model, path): + """Return the sub-module of ``model`` designated by a dotted ``path``.""" + module = model + for attribute in path.split("."): + if attribute.isdigit(): + module = module[int(attribute)] + else: + if not hasattr(module, attribute): + raise AttributeError( + f"Could not resolve `layers_path='{path}'`: the model has no " + f"attribute '{attribute}'." + ) + module = getattr(module, attribute) + return module + + +def _autodetect_layers(model): + """Find the container of repeated blocks of a transformer-like model. + + The heuristic looks for the longest :class:`torch.nn.ModuleList` whose + children all share the same type, which is how stacks of identical + encoder/decoder blocks are declared in essentially every implementation + published on the Hugging Face Hub. + """ + best_modules = None + for module in model.modules(): + if not isinstance(module, torch.nn.ModuleList) or len(module) < 2: + continue + types = {type(block) for block in module} + if len(types) > 1: + continue + if best_modules is None or len(module) > len(best_modules): + best_modules = module + if best_modules is None: + raise ValueError( + "Could not automatically locate the stack of layers of the provided " + "model. Pass an explicit `layers_path` (e.g. 'encoder.block') to " + "select the layers to probe." + ) + return best_modules + + +def _as_tensor(output): + """Extract the hidden state tensor out of an arbitrary module output.""" + if isinstance(output, torch.Tensor): + return output + for name in CANDIDATE_HIDDEN_STATE_NAMES: + value = getattr(output, name, None) + if isinstance(value, torch.Tensor): + return value + if isinstance(value, (tuple, list)) and len(value) > 0: + return value[-1] + if isinstance(output, dict): + for name in CANDIDATE_HIDDEN_STATE_NAMES: + value = output.get(name) + if isinstance(value, torch.Tensor): + return value + if isinstance(output, (tuple, list)): + for item in output: + if isinstance(item, torch.Tensor) and item.ndim == 3: + return item + for item in output: + if isinstance(item, torch.Tensor): + return item + raise TypeError( + f"Could not extract a hidden state tensor out of an output of type " + f"{type(output).__name__}. Models that only return their task-specific " + "outputs, such as forecasts, expose no hidden state to pool: select an " + "intermediate layer with the `layer` parameter, so that " + "representations are read through a forward hook instead." + ) + + +class TimeSeriesFoundationEmbedder(TimeSeriesMixin, TransformerMixin, BaseEstimator): + """Turn a frozen pre-trained time series model into a feature extractor. + + This transformer runs a (typically pre-trained) PyTorch model in inference + mode and pools its hidden states into a single vector per time series. It + is the building block used by + :class:`~tslearn.foundation.LinearProbeForecaster`, and can also be used on + its own, for instance inside a :class:`sklearn.pipeline.Pipeline` ahead of + a classifier, to implement linear probing for classification. + + The transformer is deliberately agnostic to any particular model + implementation. It only assumes that: + + * ``model`` is a :class:`torch.nn.Module`; + * its ``forward`` method accepts the raw context values through one + argument (auto-detected among common names, see ``input_name``); + * it returns hidden states of shape + ``(batch, n_tokens, dim)`` if ``pooling=None`` and ``(batch, dim)`` + otherwise. + + Parameters + ---------- + model : torch.nn.Module + A pre-trained model used as a frozen feature extractor. Its + parameters are never updated by this class. + layer : int or None (default: None) + Which layer to read representations from. When None, the output of + the model itself is used, which is usually the last hidden state + after the final normalization layer. When an integer is given, a + forward hook is placed on the corresponding block of the model's + layer stack, so that ``layer=0`` probes the first block and + ``layer=-1`` the last one. Probing intermediate layers is often + preferable, as the last layers of a pre-trained model tend to + specialize towards its pre-training objective [1]_. + layers_path : str or None (default: None) + Dotted path to the module holding the stack of layers, e.g. + ``"encoder.block"``. Only used when ``layer`` is an integer. When + None, the stack is auto-detected as the longest list of identical + sub-modules. + pooling : {"mean", "max", "token", "last", "flatten", None} (default: "mean") + How to aggregate the ``n_tokens`` representations of a series into a + single vector. ``"token"`` selects the single token at index + ``token_index``, which is how one reaches a token that carries a + meaning of its own, such as a class token for + classification, or a forecast token for forecasting. + ``"last"`` selects the last token and + ``"flatten"`` concatenates them all, which yields a much larger, and + context-length dependent, feature vector. + + ``None``, for which the string ``"none"`` is accepted as an alias, + applies no pooling at all: :meth:`transform` then returns a time + series dataset of shape ``(n_ts, n_tokens, dim)`` rather than a + flat feature matrix, which turns this estimator into a time series to + time series transform. Combined with ``tokens``, this is a way to + obtain one representation per timestep, or per patch of timesteps, + that can be fed to any other tslearn estimator. Note that the + resulting series are not accepted by the linear probes of this + module, which need a flat feature matrix. + tokens : slice, (start, stop) pair or None (default: None) + Which tokens to keep before pooling. Models often emit tokens that do + not represent the input series, such as class, register or forecast + tokens; excluding them is usually what one wants, both to build clean + per-timestep representations and to avoid polluting an average. + When None, all tokens are kept. This parameter does not affect + ``pooling="token"``. + token_index : int (default: 0) + Index of the token to select when ``pooling="token"``. + input_layout : {"univariate", "channels_last", "channels_first"} (default: "univariate") + How the context is laid out when handed over to the model. + ``"univariate"`` forecasts every channel of every series + independently. The other two layouts feed + a ``(n_ts, sz, d)`` or ``(n_ts, d, sz)`` array respectively, for + models that are natively multivariate. + input_name : str or None (default: None) + Name of the ``forward`` argument receiving the context values. + model_kwargs : dict or None (default: None) + Extra keyword arguments passed to every ``forward`` call. + batch_size : int (default: 32) + Number of series embedded at once. + verbose : int (default: 0) + When positive, prints progress information. + + Attributes + ---------- + n_features_in_ : int + Number of features (channels) of the series seen during fit. + embedding_size_ : int + Dimension of the produced embeddings. When ``pooling`` is None, this + is the number of features of each timestep of the produced series. + n_tokens_ : int or None + Number of tokens kept per series when ``pooling`` is None, that is, + the length of the produced series. None otherwise. + + Notes + ----- + This estimator does not support variable length time series, as the + NaN padding used by tslearn to represent them would propagate through + most models. + + Examples + -------- + >>> from chronos import Chronos2Model # doctest: +SKIP + >>> model = Chronos2Model.from_pretrained("amazon/chronos-2") # doctest: +SKIP + >>> embedder = TimeSeriesFoundationEmbedder(model, layer=-2) # doctest: +SKIP + >>> embedder.fit_transform(X).shape # doctest: +SKIP + (10, 512) + + Keeping one representation per token, as a time series dataset: + + >>> embedder = TimeSeriesFoundationEmbedder( # doctest: +SKIP + ... model, layer=-2, pooling=None, tokens=(0, -2)) + >>> embedder.fit_transform(X).shape # doctest: +SKIP + (10, 16, 512) + + References + ---------- + .. [1] G. Alain and Y. Bengio. Understanding intermediate layers using + linear classifier probes. ICLR Workshop, 2017. + + """ + + def __init__( + self, + model, + layer=None, + layers_path=None, + pooling="mean", + tokens=None, + token_index=0, + input_layout="univariate", + input_name=None, + model_kwargs=None, + batch_size=32, + device=None, + verbose=0, + ): + if torch is None: + raise ValueError( + "Could not use TimeSeriesFoundationEmbedder since torch is not installed" + ) + self.model = model + self.layer = layer + self.layers_path = layers_path + self.pooling = pooling + self.tokens = tokens + self.token_index = token_index + self.input_layout = input_layout + self.input_name = input_name + self.model_kwargs = model_kwargs + self.batch_size = batch_size + self.device = device + self.verbose = verbose + + def _validate_params_(self): + if self.pooling not in POOLINGS: + raise ValueError( + f"`pooling` must be one of {POOLINGS}, got '{self.pooling}'." + ) + if self.input_layout not in LAYOUTS: + raise ValueError( + f"`input_layout` must be one of {LAYOUTS}, got " + f"'{self.input_layout}'." + ) + if not isinstance(self.model, torch.nn.Module): + raise TypeError( + "`model` must be a torch.nn.Module, got " + f"{type(self.model)}. Pre-trained models are typically obtained " + "through a `from_pretrained` call." + ) + + @property + def _device(self): + try: + return next(self.model.parameters()).device + except StopIteration: + return torch.device("cpu") + + @property + def _dtype(self): + try: + return next(self.model.parameters()).dtype + except StopIteration: + return torch.get_default_dtype() + + def _resolve_input_name(self): + if self.input_name is not None: + return self.input_name + + parameters = inspect.signature(self.model.forward).parameters + # First found among candidates + for name in CANDIDATE_INPUT_NAMES: + if name in parameters: + return name + # First found among signature parameters + for name, parameter in parameters.items(): + if parameter.kind in ( + parameter.POSITIONAL_ONLY, + parameter.POSITIONAL_OR_KEYWORD, + ): + return name + + raise RuntimeError( + "Could not determine which argument of the model's `forward` method " + "should receive the time series. Pass an explicit `input_name`." + ) + + def _resolve_layers(self): + if self.layers_path is not None: + return _resolve_attribute_path( + self.model, self.layers_path + ) + return _autodetect_layers(self.model) + + def _forward_hidden_states(self, batch): + """Run the model on a batch and return hidden states of shape + (batch, n_tokens, dim).""" + kwargs = dict(self.model_kwargs or {}) + if self._input_name: + kwargs[self._input_name] = batch + + if self.layer is None: + with torch.no_grad(): + output = self.model(**kwargs) + return _as_tensor(output) + + layers = self._resolve_layers() + try: + layer_module = layers[self.layer] + except IndexError: + raise ValueError( + f"`layer={self.layer}` is out of range: the model's layer stack " + f"holds {len(layers)} layers." + ) + + captured = {} + + def hook(_module, _inputs, output): + captured["hidden_states"] = _as_tensor(output) + + handle = layer_module.register_forward_hook(hook) + try: + with torch.no_grad(): + self.model(**kwargs) + finally: + handle.remove() + + if "hidden_states" not in captured: + raise RuntimeError( + "The forward hook was never triggered; the selected layer does " + "not seem to take part in the model's forward pass." + ) + return captured["hidden_states"] + + def _pool(self, hidden_states): + if hidden_states.ndim == 2: + # Already pooled by the model itself + return hidden_states + if hidden_states.ndim != 3: + raise ValueError( + "Expected hidden states of shape (batch, n_tokens, dim), got " + f"shape {tuple(hidden_states.shape)}." + ) + pooling = _normalize_pooling(self.pooling) + + # The selected token is deliberately read before any token selection, + # as it is usually one of the tokens `tokens` is meant to filter out. + if pooling == "token": + n_tokens = hidden_states.shape[1] + if not -n_tokens <= self.token_index < n_tokens: + raise ValueError( + f"`token_index={self.token_index}` is out of range: the " + f"model produced {n_tokens} tokens." + ) + return hidden_states[:, self.token_index] + + hidden_states = hidden_states[:, _token_slice(self.tokens)] + if hidden_states.shape[1] == 0: + raise ValueError( + f"`tokens={self.tokens!r}` selects no token at all out of the " + "sequence produced by the model." + ) + if pooling is None: + return hidden_states + if pooling == "mean": + return hidden_states.mean(dim=1) + if pooling == "max": + return hidden_states.max(dim=1).values + if pooling == "last": + return hidden_states[:, -1] + return hidden_states.reshape(hidden_states.shape[0], -1) + + def _check_input(self, X): + X = check_array(X, allow_nd=True, force_all_finite=True) + X = to_time_series_dataset(X, dtype=self._dtype, be="torch") + return X + + def _embed(self, X): + """Embed a checked (n_ts, sz, d) dataset. + + Returns a ``(n_ts, D)`` feature matrix, or a ``(n_ts, n_tokens, D)`` + time series dataset when ``pooling`` is None. + """ + n_ts, sz, d = X.shape + flat = _layout_to_model_input(X, self.input_layout) + + embeddings = [] + for start in range(0, flat.shape[0], self.batch_size): + batch = flat[start : start + self.batch_size].to(self._device) + hidden_states = self._forward_hidden_states(batch) + embeddings.append(self._pool(hidden_states).to(torch.float32)) + if self.verbose: + print( + f"Embedded {min(start + self.batch_size, flat.shape[0])}" + f"/{flat.shape[0]} series" + ) + embeddings = torch.concatenate(embeddings, axis=0) + + if self.input_layout == "univariate": + if embeddings.ndim == 3: + # (n_ts * d, n_tokens, dim) -> (n_ts, n_tokens, d * dim), so + # that the per-channel representations of a same timestep sit + # side by side, as in any other tslearn time series dataset. + n_tokens, dim = embeddings.shape[1:] + embeddings = embeddings.reshape(n_ts, d, n_tokens, dim) + embeddings = torch.swapaxes(embeddings, 1, 2) + embeddings = embeddings.reshape(n_ts, n_tokens, d * dim) + else: + # Concatenate the per-channel embeddings of a same series + embeddings = embeddings.reshape(n_ts, -1) + return embeddings + + def fit(self, X, y=None): + """Check the model and the input data. No parameter is learnt. + + Parameters + ---------- + X : array-like of shape=(n_ts, sz, d) + Time series dataset. + y : Ignored + + Returns + ------- + self + The fitted estimator + + """ + self._validate_params_() + X = self._check_input(X) + + self.model.eval() + self._input_name = self._resolve_input_name() + + self.n_features_in_ = X.shape[2] + self._sz_fit_ = X.shape[1] + + # A single series is enough to know the shape of the representations + embedded = self._embed(X[:1]) + self.embedding_size_ = embedded.shape[-1] + self.n_tokens_ = embedded.shape[1] if embedded.ndim == 3 else None + + return self + + def transform(self, X): + """Embed a time series dataset. + + Parameters + ---------- + X : array-like of shape=(n_ts, sz, d) + Time series dataset. + + Returns + ------- + array of shape=(n_ts, embedding_size_), or + (n_ts, n_tokens_, embedding_size_) when ``pooling`` is None + Frozen representations of the input series. + + """ + check_is_fitted(self, "embedding_size_") + X = self._check_input(X) + if X.shape[2] != self.n_features_in_: + raise ValueError( + f"Series with {self.n_features_in_} features were expected, got " + f"{X.shape[2]}." + ) + if self.pooling == "flatten" and X.shape[1] != self._sz_fit_: + warnings.warn( + "pooling='flatten' produces context-length dependent features; " + f"series of length {self._sz_fit_} were seen at fit time and " + f"series of length {X.shape[1]} are being transformed." + ) + return self._embed(X) + + def fit_transform(self, X, y=None, **fit_params): + """Fit the estimator and embed a time series dataset. + + Parameters + ---------- + X : array-like of shape=(n_ts, sz, d) + Time series dataset. + y : Ignored + + Returns + ------- + array of shape=(n_ts, embedding_size_), or + (n_ts, n_tokens_, embedding_size_) when ``pooling`` is None + Frozen representations of the input series. + + """ + return self.fit(X, y).transform(X) + + def __sklearn_tags__(self): + tags = super().__sklearn_tags__() + tags.target_tags.required = False + tags.input_tags.allow_nan = False + return tags diff --git a/tslearn/foundation/_forecasting.py b/tslearn/foundation/_forecasting.py new file mode 100644 index 000000000..7267639eb --- /dev/null +++ b/tslearn/foundation/_forecasting.py @@ -0,0 +1,823 @@ +"""Re-use of pre-trained time series models for forecasting.""" + +import inspect +import warnings + +import numpy as np + +from sklearn.base import BaseEstimator, clone +from sklearn.linear_model import RidgeCV +from sklearn.utils.validation import check_is_fitted + +from tslearn.bases import TimeSeriesMixin +from tslearn.utils import check_array, to_time_series_dataset + +from ._embedding import ( + LAYOUTS, + TimeSeriesFoundationEmbedder, + _layout_to_model_input, + _normalize_pooling, +) + +try: + import torch +except ImportError: + torch = None + + +#: Method names commonly exposed by pre-trained forecasters, by decreasing +#: order of preference. +CANDIDATE_PREDICT_METHODS = ( + "predict", + "forecast", + "predict_quantiles", + "generate", +) + +#: Names of the argument through which a forecast horizon is commonly passed. +CANDIDATE_HORIZON_NAMES = ( + "prediction_length", + "horizon", + "forecast_horizon", + "prediction_horizon", + "num_steps", + "n_steps", + "steps", + "h", + "n", +) + +#: Attribute names under which forecasts are commonly returned. +CANDIDATE_FORECAST_NAMES = ( + "quantile_preds", + "prediction_outputs", + "predictions", + "sequences", + "mean", +) + + +#: Ways of presenting a batch of univariate contexts to a model, by decreasing +#: order of preference. Implementations disagree on this: Chronos-Bolt expects a +#: ``(batch, length)`` array while Chronos-2 requires an explicit variate axis, +#: and GluonTS-derived pipelines take a list of series. The first representation +#: the model accepts is used, and remembered in the ``context_format_`` +#: attribute. +CONTEXT_FORMATS = ("2d", "3d", "list") + + +def _format_context(context, context_format): + """Present a ``(n_rows, sz)`` batch of univariate contexts to a model.""" + if context_format == "2d": + return context + if context_format == "3d": + # (n_rows, n_variates=1, sz) + return context[:, None, :] + if context_format == "list": + return [row for row in context] + raise ValueError( # pragma: no cover + f"`context_format` must be one of {CONTEXT_FORMATS}, got " + f"'{context_format}'." + ) + + +def _to_numpy(value): + if torch is not None and isinstance(value, torch.Tensor): + return value.detach().to(torch.float32).cpu().numpy() + return np.asarray(value, dtype=np.float64) + + +def _unwrap_forecast(output): + """Extract a numeric array out of an arbitrary forecasting output.""" + for name in CANDIDATE_FORECAST_NAMES: + value = getattr(output, name, None) + if value is not None and not callable(value): + return _to_numpy(value) + if isinstance(output, dict): + for name in CANDIDATE_FORECAST_NAMES: + if name in output: + return _to_numpy(output[name]) + if isinstance(output, (list, tuple)): + if len(output) == 0: + raise ValueError("The model returned an empty forecast.") + items = [_to_numpy(item) for item in output] + shapes = {item.shape for item in items} + if len(shapes) == 1: + # A list with one entry per series, as returned e.g. by pipelines + return np.stack(items, axis=0) + # A tuple of heterogeneous outputs: keep the first one + return items[0] + return _to_numpy(output) + + +def _check_probe_pooling(pooling): + """Reject the poolings that do not yield a flat feature matrix.""" + if _normalize_pooling(pooling) is None: + raise ValueError( + "A linear probe needs one feature vector per series, so " + "`pooling=None` is not supported here. Use another pooling, or " + "TimeSeriesFoundationEmbedder directly if you want to keep one " + "representation per token." + ) + + +def _resolve_horizon_axis(shape, horizon, horizon_axis): + """Locate the axis of a forecast array that holds the forecast horizon.""" + ndim = len(shape) + if horizon_axis != "auto": + axis = horizon_axis + ndim if horizon_axis < 0 else horizon_axis + if not 1 <= axis < ndim: + raise ValueError( + f"`horizon_axis={horizon_axis}` is out of range for a forecast " + f"of shape {shape}; it must designate one of the axes 1 to " + f"{ndim - 1}, the axis 0 being the series." + ) + if shape[axis] < horizon: + raise ValueError( + f"`horizon_axis={horizon_axis}` designates an axis of size " + f"{shape[axis]} in a forecast of shape {shape}, which cannot " + f"hold a horizon of {horizon}." + ) + return axis + + # An axis whose size is exactly the horizon is a much safer guess than one + # that is merely large enough; among those, the last one is retained, as + # forecasts are most often returned with time as their trailing axis. + exact_axes = [axis for axis in range(1, ndim) if shape[axis] == horizon] + if exact_axes: + ambiguous = [axis for axis in exact_axes if shape[axis] > 1] + if len(ambiguous) > 1: + warnings.warn( + f"The forecast returned by the model has shape {shape}, in " + f"which several axes could be the horizon axis of length " + f"{horizon}; axis {exact_axes[-1]} was assumed. Set " + "`horizon_axis` explicitly to remove the ambiguity.", + UserWarning, + stacklevel=3, + ) + return exact_axes[-1] + + longer_axes = [axis for axis in range(1, ndim) if shape[axis] > horizon] + if not longer_axes: + raise ValueError( + f"Could not identify the forecast horizon axis in an output of " + f"shape {shape} for a horizon of {horizon}. Set `horizon_axis` " + "explicitly, or pass a `predict_fn` to control how the model is " + "called." + ) + return min(longer_axes, key=lambda axis: shape[axis]) + + +def _reduce_to_point_forecast(forecast, n_rows, horizon, quantile, horizon_axis): + """Reduce an arbitrarily shaped univariate forecast to (n_rows, horizon). + + Any axis that is neither the batch axis nor the horizon axis is understood + as holding quantile levels or sample paths, and is therefore reduced. + """ + forecast = np.asarray(forecast, dtype=np.float64) + if forecast.shape[0] != n_rows: + raise ValueError( + f"The model returned forecasts for {forecast.shape[0]} series while " + f"{n_rows} were provided." + ) + if forecast.ndim == 1: # pragma: no cover + forecast = forecast[:, None] + if forecast.ndim == 2: + if forecast.shape[1] < horizon: + raise ValueError( + f"The model returned forecasts of length {forecast.shape[1]}, " + f"which is shorter than the requested horizon {horizon}." + ) + return forecast[:, :horizon] + + axis = _resolve_horizon_axis(forecast.shape, horizon, horizon_axis) + forecast = np.moveaxis(forecast, axis, -1) + forecast = forecast.reshape(n_rows, -1, forecast.shape[-1]) + if forecast.shape[1] == 1: + forecast = forecast[:, 0] + elif quantile is None: + forecast = forecast.mean(axis=1) + else: + forecast = np.quantile(forecast, quantile, axis=1) + return forecast[:, :horizon] + + +class _BaseFoundationForecaster(TimeSeriesMixin, BaseEstimator): + """Shared input handling for pre-trained forecasters.""" + + def _check_input(self, X): + X = check_array(X, allow_nd=True, force_all_finite=True) + return to_time_series_dataset(X, self._dtype, be="torch") + + def _check_layout(self): + if self.input_layout not in LAYOUTS: + raise ValueError( + f"`input_layout` must be one of {LAYOUTS}, got " + f"'{self.input_layout}'." + ) + horizon_axis = getattr(self, "horizon_axis", "auto") + if horizon_axis != "auto" and not isinstance( + horizon_axis, (int, np.integer) + ): + raise ValueError( + f"`horizon_axis` must be an integer or 'auto', got " + f"{horizon_axis!r}." + ) + + def fit_predict(self, X, y=None, n=1): + """Fit the estimator and forecast ``n`` timestamps for the given data. + + Parameters + ---------- + X : array-like of shape=(n_ts, sz, d) + Time series dataset. + y : Ignored + n : int (default: 1) + The number of timestamps to forecast, a.k.a. the horizon. + + Returns + ------- + array of shape=(n_ts, n, d) + Array of forecasted timestamps + + """ + return self.fit(X, y).predict(X, n=n) + + def __sklearn_tags__(self): + tags = super().__sklearn_tags__() + tags.target_tags.required = False + tags.input_tags.allow_nan = False + return tags + + +class ZeroShotForecaster(_BaseFoundationForecaster): + """Forecast with a pre-trained model, without any training. + + Time series foundation models are pre-trained on large corpora of series + and are meant to forecast unseen series out of the box. This estimator + wraps such a model behind the usual tslearn forecasting API, so that it can + be compared to, or combined with, the other estimators of the library. + Calling :meth:`fit` is therefore not required and does not learn anything. + + The estimator makes no assumption about a particular implementation. It + only assumes that ``model`` exposes a forecasting method (auto-detected + among ``predict_quantiles``, ``predict``, ``forecast`` and ``generate``) + accepting the context values as its first argument and a forecast horizon + as a keyword argument. When this is not the case, or when finer control is + needed, pass an explicit ``predict_fn``. + + Parameters + ---------- + model : object + A pre-trained forecasting model or inference pipeline. + predict_fn : callable or None (default: None) + Called as ``predict_fn(model, context, horizon)``, where ``context`` + is a NumPy array laid out as prescribed by ``input_layout``, and + expected to return forecasts that can be broadcast to + ``(n_rows, horizon)``. When None, the model's own forecasting method + is auto-detected. + input_layout : {"univariate", "channels_last", "channels_first"} (default: "univariate") + How the context is laid out when handed over to the model. + ``"univariate"`` forecasts every channel of every series + independently. The other two layouts feed + a ``(n_ts, sz, d)`` or ``(n_ts, d, sz)`` array respectively, for + models that are natively multivariate. + context_length : int or None (default: None) + When set, only the last ``context_length`` timestamps of each series + are used as context. + horizon_axis : int or "auto" (default: "auto") + Which axis of the array returned by the model holds the forecast + horizon, axis 0 being the series. + When ``"auto"``, the axis is inferred from the returned + shape, preferring an axis whose size is exactly the horizon and, among + those, the last one. A warning is raised if the shape leaves the + choice ambiguous, which happens when the model returns as many + quantiles as there are forecast steps. All the axes that are neither + the series axis nor the horizon axis are taken to hold quantile levels + or sample paths, and are reduced according to ``quantile``. + quantile : float or None (default: 0.5) + Probabilistic models return several quantiles or sample paths; this + is the quantile used to derive a point forecast from them. When None, + the mean is used instead. Note that the quantile is taken over the + values the model returned. + model_kwargs : dict or None (default: None) + Extra keyword arguments passed to the model's forecasting method. + + Attributes + ---------- + context_format_ : str + How the batch of contexts ended up being presented to the model, + among ``"2d"``, ``"3d"`` and ``"list"``. Set after the first call to + :meth:`predict` with the ``"univariate"`` layout. + + Notes + ----- + Unlike :class:`~tslearn.forecasting.VARIMA`, this estimator holds no + state about the fitted data, so :meth:`predict` requires ``X``. + + See Also + -------- + LinearProbeForecaster: Fit a linear head on top of a frozen pre-trained model. + + Examples + -------- + >>> from chronos import Chronos2Pipeline # doctest: +SKIP + >>> pipeline = Chronos2Pipeline.from_pretrained("amazon/chronos-2") # doctest: +SKIP + >>> model = ZeroShotForecaster(pipeline) # doctest: +SKIP + >>> model.predict(X, n=24).shape # doctest: +SKIP + (10, 24, 1) + + """ + + def __init__( + self, + model, + predict_fn=None, + input_layout="univariate", + context_length=None, + horizon_axis="auto", + quantile=0.5, + model_kwargs=None, + ): + if torch is None: + raise ValueError( + "Could not use ZeroShotForecaster since torch is not installed" + ) + + self.model = model + self.predict_fn = predict_fn + self.input_layout = input_layout + self.context_length = context_length + self.horizon_axis = horizon_axis + self.quantile = quantile + self.model_kwargs = model_kwargs + + @property + def _dtype(self): + try: + return next(self.model.parameters()).dtype + except (StopIteration, AttributeError): + return torch.get_default_dtype() + + def _resolve_predict_fn(self): + """Build a ``(context, horizon) -> forecast`` callable from the model.""" + if self.predict_fn is not None: + return lambda context, horizon: self.predict_fn( + self.model, context, horizon + ) + + for method_name in CANDIDATE_PREDICT_METHODS: + method = getattr(self.model, method_name, None) + if not callable(method): + continue + try: + parameters = inspect.signature(method).parameters + except (TypeError, ValueError): # pragma: no cover + continue + horizon_name = next( + (name for name in CANDIDATE_HORIZON_NAMES if name in parameters), + None, + ) + if horizon_name is None: + continue + + def call(context, horizon, method=method, horizon_name=horizon_name): + kwargs = dict(self.model_kwargs or {}) + kwargs[horizon_name] = horizon + return method(context, **kwargs) + + return call + + raise ValueError( + "Could not find a forecasting method on the provided model. The " + f"model is expected to expose one of {CANDIDATE_PREDICT_METHODS} " + "accepting a forecast horizon. Pass an explicit `predict_fn` " + "otherwise." + ) + + def _call_univariate(self, call, context, horizon): + """Call the model, negotiating how the context batch is presented. + + Implementations disagree on whether a batch of univariate series should + be handed over as a 2d array, as a 3d array with an explicit variate + axis, or as a list of series. Rather than requiring the user to know, + the candidate representations are tried in turn and the first one the + model accepts is remembered for subsequent calls. + """ + if self.predict_fn is not None: + return call(context, horizon) + + known_format = getattr(self, "context_format_", None) + formats = ( + (known_format,) if known_format is not None else CONTEXT_FORMATS + ) + errors = {} + for context_format in formats: + try: + forecast = call(_format_context(context, context_format), horizon) + except (ValueError, TypeError, IndexError, RuntimeError) as error: + errors[context_format] = f"{type(error).__name__}: {error}" + continue + self.context_format_ = context_format + return forecast + + details = "\n".join(f" - as a {key} input: {value}" for key, value in errors.items()) + raise ValueError( + "The model rejected every supported way of passing a batch of " + f"univariate contexts:\n{details}\nPass an explicit `predict_fn` to " + "control how the model is called." + ) + + def fit(self, X, y=None): + """Do nothing, as a zero-shot model needs no training. + + Parameters + ---------- + X : array-like of shape=(n_ts, sz, d) + Time series dataset. + y : Ignored + + Returns + ------- + self + The estimator, ready to be used for prediction + + """ + warnings.warn( + "ZeroShotForecaster.fit does not train anything: the wrapped " + "model is used as-is for zero-shot forecasting.", + UserWarning, + stacklevel=2, + ) + return self + + def predict(self, X, n=1): + """Forecast ``n`` timestamps ahead of the given series. + + Parameters + ---------- + X : array-like of shape=(n_ts, sz, d) + Time series dataset to forecast. + n : int (default: 1) + The number of timestamps to forecast, a.k.a. the horizon. + + Returns + ------- + array of shape=(n_ts, n, d) + Array of forecasted timestamps + + """ + self._check_layout() + if X is None: + raise ValueError( + "A zero-shot model keeps no state about the data seen at fit " + "time, so `X` is required at predict time." + ) + if n < 1: + raise ValueError(f"`n` must be a positive integer, got {n}.") + X = self._check_input(X) + if self.context_length is not None: + X = X[:, -self.context_length :] + + n_ts, sz, d = X.shape + context = _layout_to_model_input(X, self.input_layout) + call = self._resolve_predict_fn() + + if self.input_layout == "univariate": + forecast = _unwrap_forecast(self._call_univariate(call, context, n)) + forecast = _reduce_to_point_forecast( + forecast, n_ts * d, n, self.quantile, self.horizon_axis + ) + # (n_ts * d, n) -> (n_ts, n, d) + return np.swapaxes(forecast.reshape(n_ts, d, n), 1, 2) + + forecast = np.asarray( + _unwrap_forecast(call(context, n)), dtype=np.float64 + ) + if forecast.shape[:1] != (n_ts,): + raise ValueError( + f"The model returned forecasts for {forecast.shape[0]} series " + f"while {n_ts} were provided." + ) + if self.input_layout == "channels_first": + forecast = np.swapaxes(forecast, 1, 2) + if forecast.ndim != 3 or forecast.shape[2] != d: + raise ValueError( + f"Expected forecasts of shape (n_ts, n, d) = ({n_ts}, {n}, {d}) " + f"with input_layout='{self.input_layout}', got shape " + f"{forecast.shape}. Pass an explicit `predict_fn` to control " + "how the model is called." + ) + return forecast[:, :n] + + +class LinearProbeForecaster(_BaseFoundationForecaster): + """Forecast by fitting a head on a frozen pre-trained model. + + Linear probing keeps the pre-trained model entirely frozen and only fits a + linear map from its representations to the quantity of interest. Compared + to zero-shot use, it adapts the model to the data at hand at a very small + computational cost, and compared to fine-tuning, it leaves the pre-trained + weights untouched, which makes it much cheaper. + + Training pairs are cut out of the provided series with a sliding window: + the ``context_length`` timestamps preceding a cut point are embedded, and + the ``horizon`` timestamps following it are used as the target. + + Parameters + ---------- + model : torch.nn.Module + A pre-trained model used as a frozen feature extractor. + probe : sklearn estimator or None (default: None) + The head fitted on top of the frozen representations. When None, a + :class:`sklearn.linear_model.RidgeCV` is used, which selects its + regularization strength by cross-validation and admits a closed-form + solution. Any multi-output capable regressor can be passed instead. + Frozen representations sometimes have a very small + variance, in which case prepending a + :class:`sklearn.preprocessing.StandardScaler` to the head might help. + context_length : int (default: 128) + Number of timestamps fed to the pre-trained model. + horizon : int (default: 1) + Number of timestamps forecasted at once. This cannot be changed after + fit time, as one linear head is fitted for the whole horizon. + stride : int (default: 1) + Step between two consecutive training windows. Larger values yield + fewer, less redundant training pairs and a faster fit. + layer : int or None (default: None) + Layer to probe, see + :class:`~tslearn.foundation.TimeSeriesFoundationEmbedder`. + layers_path : str or None (default: None) + Dotted path to the stack of layers, see + :class:`~tslearn.foundation.TimeSeriesFoundationEmbedder`. + pooling : {"mean", "max", "token", "last", "flatten"} (default: "mean") + How token representations are aggregated, see + :class:`~tslearn.foundation.TimeSeriesFoundationEmbedder` and + note that ``pooling=None`` is not accepted here. + tokens : slice, (start, stop) pair or None (default: None) + Which tokens to keep before pooling, see + :class:`~tslearn.foundation.TimeSeriesFoundationEmbedder`. Restricting + the average to the tokens that actually represent the context, e.g. + ``tokens=(0, -2)`` for Chronos-2, avoids diluting it with class or + forecast tokens. + token_index : int (default: 0) + Index of the token selected when ``pooling="token"``, see + :class:`~tslearn.foundation.TimeSeriesFoundationEmbedder`. + input_layout : {"univariate", "channels_last", "channels_first"} (default: "univariate") + How the context is laid out when handed over to the model. + ``"univariate"`` forecasts every channel of every series + independently. The other two layouts feed + a ``(n_ts, sz, d)`` or ``(n_ts, d, sz)`` array respectively, for + models that are natively multivariate. + input_name : str or None (default: None) + Name of the ``forward`` argument receiving the context values. + model_kwargs : dict or None (default: None) + Extra keyword arguments passed to every ``forward`` call. + batch_size : int (default: 32) + Number of series embedded at once, see + :class:`~tslearn.foundation.TimeSeriesFoundationEmbedder`. + device : str or None (default: None) + Device on which inference is run, see + :class:`~tslearn.foundation.TimeSeriesFoundationEmbedder`. + verbose : int (default: 0) + When positive, prints progress information. + + Attributes + ---------- + embedder_ : TimeSeriesFoundationEmbedder + The frozen feature extractor. + probe_ : sklearn estimator + The fitted head. + n_features_in_ : int + Number of features (channels) of the series seen during fit. + n_windows_ : int + Number of training windows cut out of the fit data. + + See Also + -------- + ZeroShotForecaster: Use a pre-trained model without any training. + TimeSeriesFoundationEmbedder: The underlying feature extractor, which + can be composed with a classifier through a + :class:`sklearn.pipeline.Pipeline` to build a linear probe model for classification, if needed. + + Examples + -------- + >>> from chronos import Chronos2Pipeline # doctest: +SKIP + >>> pipeline = Chronos2Pipeline.from_pretrained("amazon/chronos-2") # doctest: +SKIP + >>> model = LinearProbeForecaster(pipeline.model, horizon=24) # doctest: +SKIP + >>> model.fit(X_train) # doctest: +SKIP + ... + >>> model.predict(X_test).shape # doctest: +SKIP + (10, 24, 1) + >>> model.predict(X_test, n=12).shape # doctest: +SKIP + (10, 12, 1) + >>> model.predict(X_test, n=36) # doctest: +SKIP + Traceback (most recent call last): + ... + ValueError: This estimator was fitted for a horizon of 24, so `n` must lie in [1, 24], got 36. + Refit with a larger `horizon` to forecast further ahead. + + References + ---------- + .. [1] G. Alain and Y. Bengio. Understanding intermediate layers using + linear classifier probes. ICLR Workshop, 2017. + + """ + + def __init__( + self, + model, + probe=None, + context_length=128, + horizon=1, + stride=1, + layer=None, + layers_path=None, + pooling="mean", + tokens=None, + token_index=0, + input_layout="univariate", + input_name=None, + model_kwargs=None, + batch_size=32, + device=None, + verbose=0, + ): + if torch is None: + raise ValueError( + "Could not use LinearProbeForecaster since torch is not installed" + ) + + self.model = model + self.probe = probe + self.context_length = context_length + self.horizon = horizon + self.stride = stride + self.layer = layer + self.layers_path = layers_path + self.pooling = pooling + self.tokens = tokens + self.token_index = token_index + self.input_layout = input_layout + self.input_name = input_name + self.model_kwargs = model_kwargs + self.batch_size = batch_size + self.device = device + self.verbose = verbose + + @property + def _dtype(self): + try: + return next(self.model.parameters()).dtype + except StopIteration: + return torch.get_default_dtype() + + def _make_embedder(self): + return TimeSeriesFoundationEmbedder( + model=self.model, + layer=self.layer, + layers_path=self.layers_path, + pooling=self.pooling, + tokens=self.tokens, + token_index=self.token_index, + input_layout=self.input_layout, + input_name=self.input_name, + model_kwargs=self.model_kwargs, + batch_size=self.batch_size, + device=self.device, + verbose=self.verbose, + ) + + def _make_probe(self): + if self.probe is None: + return RidgeCV(alphas=np.logspace(-3, 3, 13)) + return clone(self.probe) + + def _make_windows(self, X): + """Cut (context, target) pairs out of a (n_ts, sz, d) dataset.""" + n_ts, sz, d = X.shape + span = self.context_length + self.horizon + if sz < span: + raise ValueError( + f"Series of length at least context_length + horizon = {span} " + f"are required to build training windows, got {sz}." + ) + starts = np.arange(0, sz - span + 1, self.stride) + contexts = np.stack( + [X[:, start : start + self.context_length] for start in starts], axis=1 + ) + targets = np.stack( + [ + X[:, start + self.context_length : start + span] + for start in starts + ], + axis=1, + ) + # (n_ts, n_windows, ...) -> (n_ts * n_windows, ...) + contexts = contexts.reshape(n_ts * len(starts), self.context_length, d) + targets = targets.reshape(n_ts * len(starts), self.horizon * d) + return contexts, targets + + def _validate_params_(self): + self._check_layout() + _check_probe_pooling(self.pooling) + for name in ("context_length", "horizon", "stride"): + value = getattr(self, name) + if not isinstance(value, (int, np.integer)) or value < 1: + raise ValueError( + f"`{name}` must be a positive integer, got {value}." + ) + + def fit(self, X, y=None): + """Fit the probe on top of the frozen pre-trained model. + + Parameters + ---------- + X : array-like of shape=(n_ts, sz, d) + Time series dataset, with ``sz >= context_length + horizon``. + y : Ignored + + Returns + ------- + self + The fitted estimator + + """ + self._validate_params_() + X = self._check_input(X) + contexts, targets = self._make_windows(X) + + self.embedder_ = self._make_embedder() + embeddings = self.embedder_.fit_transform(contexts) + + self.probe_ = self._make_probe() + self.probe_.fit(embeddings, targets) + + self.n_features_in_ = X.shape[2] + self.n_windows_ = contexts.shape[0] + return self + + def predict(self, X, n=None): + """Forecast ``horizon`` (or ``n``, if set) timestamps ahead of the + given series. + + Parameters + ---------- + X : array-like of shape=(n_ts, sz, d) + Time series dataset to forecast, with + ``sz >= context_length``. + n : int or None (default: None) + The number of timestamps to forecast. A single head is fitted for + the whole horizon, so ``n`` may not exceed the ``horizon`` value + used at fit time. When None, ``horizon`` is used. + + Returns + ------- + array of shape=(n_ts, n, d) + Array of forecasted timestamps + + """ + check_is_fitted(self, "probe_") + n = self.horizon if n is None else n + if not 1 <= n <= self.horizon: + raise ValueError( + f"This estimator was fitted for a horizon of {self.horizon}, so " + f"`n` must lie in [1, {self.horizon}], got {n}. Refit with a " + "larger `horizon` to forecast further ahead." + ) + X = self._check_input(X) + if X.shape[2] != self.n_features_in_: + raise ValueError( + f"Series with {self.n_features_in_} features were expected, got " + f"{X.shape[2]}." + ) + if X.shape[1] < self.context_length: + raise ValueError( + f"Series of length at least context_length=" + f"{self.context_length} are required, got {X.shape[1]}." + ) + contexts = X[:, -self.context_length :] + predictions = self.probe_.predict(self.embedder_.transform(contexts)) + predictions = np.asarray(predictions, dtype=np.float64) + predictions = predictions.reshape(X.shape[0], self.horizon, self.n_features_in_) + return predictions[:, :n] + + def score(self, X, y=None): + """Return the coefficient of determination of the forecast. + + Parameters + ---------- + X : array-like of shape=(n_ts, sz, d) + Time series dataset, with ``sz >= context_length + horizon``. + y : Ignored + + Returns + ------- + float + :math:`R^2` of the forecast over all windows cut out of ``X``. + + """ + check_is_fitted(self, "probe_") + X = self._check_input(X) + contexts, targets = self._make_windows(X) + return self.probe_.score(self.embedder_.transform(contexts), targets)