Feat/patchtst fm

mkdocs.yml

@@ -30,6 +30,7 @@ nav:
           - examples/chronos-family.ipynb
           - examples/cryptocurrency-quickstart.ipynb
           - examples/sktime.ipynb
+          - examples/patchtst-fm.ipynb
   - Experiments:
       - experiments/gift-eval.md
       - experiments/fev.md

pyproject.toml

@@ -88,7 +88,7 @@ dependencies = [
   "tabpfn-time-series==1.0.3 ; python_full_version < '3.13'",
   "tensorboard>=2.20.0",
   "timecopilot-chronos-forecasting>=0.2.0",
-  "timecopilot-granite-tsfm>=0.1.1",
+  "timecopilot-granite-tsfm>=0.1.2",
   "timecopilot-timesfm>=0.2.1",
   "timecopilot-tirex>=0.1.0 ; python_full_version >= '3.11'",
   "timecopilot-toto>=0.1.5",

tests/models/conftest.py

@@ -6,6 +6,7 @@
 from timecopilot.models.foundation.chronos import Chronos
 from timecopilot.models.foundation.flowstate import FlowState
 from timecopilot.models.foundation.moirai import Moirai
+from timecopilot.models.foundation.patchtst_fm import PatchTSTFM
 from timecopilot.models.foundation.timesfm import TimesFM
 from timecopilot.models.foundation.toto import Toto
 from timecopilot.models.ml import AutoLGBM
@@ -69,6 +70,9 @@ def disable_mps_session(monkeypatch):
         repo_id="ibm-granite/granite-timeseries-flowstate-r1",
         alias="FlowState-Granite",
     ),
+    PatchTSTFM(
+        context_length=2_048,
+    ),
     Toto(context_length=256, batch_size=2),
     Moirai(
         context_length=256,

tests/models/test_models.py

@@ -194,7 +194,7 @@ def test_using_quantiles(model):
     qs = [round(i * 0.1, 1) for i in range(1, 10)]
     df = generate_series(n_series=3, freq="D")
     if model.alias in ["AutoLGBM", "AutoNHITS", "AutoTFT"]:
-        # AutoLGBM does not support quantiles yet
+        # These models do not support quantiles yet
         with pytest.raises(ValueError) as excinfo:
             model.forecast(
                 df=df,
@@ -239,8 +239,13 @@ def test_using_quantiles(model):
 def test_using_level(model):
     level = [0, 20, 40, 60, 80]  # corresponds to qs [0.1, 0.2, ..., 0.9]
     df = generate_series(n_series=2, freq="D")
-    if model.alias in ["AutoLGBM", "AutoNHITS", "AutoTFT"]:
-        # AutoLGBM does not support quantiles yet
+    if model.alias in [
+        "AutoLGBM",
+        "AutoNHITS",
+        "AutoTFT",
+        "PatchTST-FM",
+    ]:
+        # These models do not support levels yet
         with pytest.raises(ValueError) as excinfo:
             model.forecast(
                 df=df,

timecopilot/models/foundation/flowstate.py

@@ -3,16 +3,15 @@
 import numpy as np
 import pandas as pd
 import torch
-from gluonts.transform import LastValueImputation
 from tqdm import tqdm
 from tsfm_public import FlowStateForPrediction
 from tsfm_public.models.flowstate.utils.utils import get_fixed_factor
 
-from ..utils.forecaster import Forecaster, QuantileConverter
+from ..utils.forecaster import Forecaster, QuantileConverter, _DataProcessor
 from .utils import TimeSeriesDataset
 
 
-class FlowState(Forecaster):
+class FlowState(Forecaster, _DataProcessor):
     """
     FlowState is the first time-scale adjustable Time Series Foundation Model (TSFM),
     open-sourced by IBM Research. Combining a State Space Model (SSM) Encoder with a
@@ -99,49 +98,6 @@ def _get_model(self) -> FlowStateForPrediction:
             del model
             torch.cuda.empty_cache()
 
-    def _left_pad_and_stack_1D(self, tensors: list[torch.Tensor]) -> torch.Tensor:
-        max_len = max(len(c) for c in tensors)
-        padded = []
-        for c in tensors:
-            assert isinstance(c, torch.Tensor)
-            assert c.ndim == 1
-            padding = torch.full(
-                size=(max_len - len(c),),
-                fill_value=torch.nan,
-                device=c.device,
-                dtype=c.dtype,
-            )
-            padded.append(torch.concat((padding, c), dim=-1))
-        return torch.stack(padded)
-
-    def _prepare_and_validate_context(
-        self,
-        context: list[torch.Tensor] | torch.Tensor,
-    ) -> torch.Tensor:
-        if isinstance(context, list):
-            context = self._left_pad_and_stack_1D(context)
-        assert isinstance(context, torch.Tensor)
-        if context.ndim == 1:
-            context = context.unsqueeze(0)
-        assert context.ndim == 2
-        return context
-
-    def _maybe_impute_missing(self, batch: torch.Tensor) -> torch.Tensor:
-        if torch.isnan(batch).any():
-            batch = batch.float().numpy()
-            imputed_rows = []
-            for i in range(batch.shape[0]):
-                row = batch[i]
-                imputed_row = LastValueImputation()(row)
-                imputed_rows.append(imputed_row)
-            batch = np.vstack(imputed_rows)
-            batch = torch.tensor(
-                batch,
-                dtype=self.dtype,
-                device=self.device,
-            )
-        return batch
-
     def _predict_batch(
         self,
         model: FlowStateForPrediction,