feat(timesfm3): add conditional detrending

Author: guerrantif

sdm/models/timesfm3/model.py

@@ -26,22 +26,57 @@
 import torch
 from torch import Tensor
 
-from sdm import Recipe, RelatedTables, Stype, StypeLike, TableTensor, Task
+from sdm import (
+    ColumnarTensor,
+    Recipe,
+    RelatedTables,
+    Stype,
+    StypeLike,
+    TableTensor,
+    Task,
+)
 from sdm.cache import Cache
+from sdm.models import ICLModel
 from sdm.models._huggingface import download_checkpoint
-from sdm.models.base import ICLModel
+from sdm.models.callback import Callback
 from sdm.models.timesfm3.block import ResidualBlock
 from sdm.models.timesfm3.ckpt import remap_ckpt
 from sdm.models.timesfm3.icl import ICLBlock
-from sdm.models.timesfm3.recipe import default_recipe
+from sdm.models.timesfm3.recipe import TIME_COLUMN, default_recipe
 from sdm.models.timesfm3.util import (
     gather_future_patches,
     get_running_stats,
     revin,
 )
+from sdm.processing.execution import MemberContext
 from sdm.tensor.table import TableSchema
 
 
+def _with_time_column(table: TableTensor, start: int) -> TableTensor:
+    time_steps = torch.arange(
+        start, start + table.size(-2), device=table.device
+    ).expand(*table.size()[:-2], -1)
+    return cast(
+        TableTensor,
+        torch.cat(
+            (
+                table,
+                TableTensor(
+                    columns={Stype.id: [TIME_COLUMN]},
+                    id=ColumnarTensor((time_steps,)),
+                ),
+            ),
+            dim=-1,
+        ),
+    )
+
+
+def _drop_time_column(table: TableTensor) -> TableTensor:
+    if TIME_COLUMN in table.columns[Stype.id]:
+        return table.drop_columns(TIME_COLUMN)
+    return table
+
+
 class TimesFM3(ICLModel):
     r"""The multivariate forecasting foundation model from `"TimesFM-3: A
     Zero-shot Foundation Model for Multivariate Forecasting"
@@ -149,58 +184,84 @@ def _load_from_pretrained(
 
         return self
 
-    def forward(self, *args: Any, **kwargs: Any) -> TableTensor:
+    def forward(
+        self,
+        x_context: Tensor | TableTensor,
+        y_context: Tensor | TableTensor,
+        x_query: Tensor | TableTensor,
+        *args: Any,
+        **kwargs: Any,
+    ) -> TableTensor:
         r""":meta private:"""  # noqa: D415
-        x_context = kwargs["x_context"] if "x_context" in kwargs else args[0]
         if not isinstance(x_context, TableTensor):
             x_context = TableTensor.from_tensor(x_context)
-        kwargs["_x_context_schema"] = x_context.schema
-
-        x_query = kwargs["x_query"] if "x_query" in kwargs else args[2]
+        if not isinstance(y_context, TableTensor):
+            y_context = TableTensor.from_tensor(y_context)
         if not isinstance(x_query, TableTensor):
             x_query = TableTensor.from_tensor(x_query)
-        kwargs["_x_query_schema"] = x_query.schema
 
+        context_length = x_context.size(-2)
         x_query = expand_query(x_context.schema, x_query)
 
-        if "x_query" in kwargs:
-            kwargs["x_query"] = x_query
-        else:
-            args = (*args[:2], x_query, *args[3:])
-
-        return super().forward(*args, **kwargs)  # type: ignore
+        out = super().forward(
+            _with_time_column(x_context, 0),
+            _with_time_column(y_context, 0),
+            _with_time_column(x_query, context_length),
+            *args,
+            **kwargs,
+        )
+        return _drop_time_column(out)
 
-    def fit(self, *args: Any, **kwargs: Any) -> None:
+    def fit(
+        self,
+        x: Tensor | TableTensor,
+        y: Tensor | TableTensor,
+        *args: Any,
+        **kwargs: Any,
+    ) -> None:
         r""":meta private:"""  # noqa: D415
-        x = kwargs["x"] if "x" in kwargs else args[0]
         if not isinstance(x, TableTensor):
             x = TableTensor.from_tensor(x)
-        kwargs["_x_context_schema"] = x.schema
+        if not isinstance(y, TableTensor):
+            y = TableTensor.from_tensor(y)
 
-        super().fit(*args, **kwargs)
+        kwargs["_x_context_schema"] = x.schema
+        kwargs["_context_length"] = x.size(-2)
+        super().fit(
+            _with_time_column(x, 0), _with_time_column(y, 0), *args, **kwargs
+        )
 
-    def predict(self, *args: Any, **kwargs: Any) -> TableTensor:
+    def predict(
+        self,
+        x: Tensor | TableTensor,
+        *args: Any,
+        **kwargs: Any,
+    ) -> TableTensor:
         r""":meta private:"""  # noqa: D415
         if self._cache is None:
             raise RuntimeError(
                 f"{self.__class__.__name__!r} not yet fitted. Make sure to "
                 f"call '{self.__class__.__name__}.fit()' before."
             )
-
-        x = kwargs["x"] if "x" in kwargs else args[0]
         if not isinstance(x, TableTensor):
             x = TableTensor.from_tensor(x)
 
         cached_kwargs = cast(dict[str, Any], self._cache["kwargs"])
-        cached_kwargs["_x_query_schema"] = x.schema
         x = expand_query(cached_kwargs["_x_context_schema"], x)
+        out = super().predict(
+            _with_time_column(x, cached_kwargs["_context_length"]),
+            *args,
+            **kwargs,
+        )
+        return _drop_time_column(out)
 
-        if "x" in kwargs:
-            kwargs["x"] = x
-        else:
-            args = (x, *args[1:])
-
-        return super().predict(*args, **kwargs)
+    def _prepare_context(
+        self,
+        context: MemberContext,
+        callbacks: Sequence[Callback],
+    ) -> MemberContext:
+        context = context._replace(y=_drop_time_column(context.y))
+        return super()._prepare_context(context, callbacks)
 
     def _forward(
         self,
@@ -212,7 +273,7 @@ def _forward(
         cache: Cache | None,
         generator: torch.Generator | None,
         **kwargs: Any,
-    ) -> TableTensor:  # [..., R_query, Y * 9]
+    ) -> TableTensor:  # [..., R_query, T * Q]
 
         if y_context is not None:
             columns = y_context.columns[Stype.numerical]
@@ -227,18 +288,26 @@ def _forward(
             assert x_context is not None
             size = (*x_context.size()[:-2], 0)
 
-        return TableTensor(
+        out = TableTensor(
             columns={
                 Stype.numerical: [
-                    f"{name}__q{i}"
-                    for name, i in product(columns, range(10, 100, 10))
+                    f"{name}__q{int(100 * q)}"
+                    for name, q in product(columns, self.model.quantiles)
                 ]
             },
             numerical=torch.zeros(
-                (*size, len(columns) * 9),
+                (*size, len(columns) * len(self.model.quantiles)),
                 device=next(self.parameters()).device,
             ),
         )
+        if x_query is None:
+            return out
+        context_length = (
+            x_context.size(-2)
+            if x_context is not None
+            else kwargs["_context_length"]
+        )
+        return _with_time_column(out, context_length)
 
 
 class _TimesFM3(torch.nn.Module):
@@ -310,7 +379,9 @@ def _preprocess(
         count, mean, std = get_running_stats(x, mask)
 
         # Gather the prediction patches and hide unavailable values
-        horizon_x, past_end = gather_future_patches(x, self.num_horizon_patches)
+        horizon_x, past_end = gather_future_patches(
+            x, self.num_horizon_patches
+        )
         horizon_mask, _ = gather_future_patches(mask, self.num_horizon_patches)
         horizon_mask.logical_or_(context_only[..., None])
         horizon_mask.logical_or_(past_end)

sdm/models/timesfm3/recipe.py

@@ -3,8 +3,12 @@
 
 import sdm.processing as sp
 
+TIME_COLUMN = "__timesfm3_time__"
+
 
 def default_recipe() -> sp.Recipe:  # noqa: D103
     return sp.Recipe(
+        features=sp.LinearDetrend(TIME_COLUMN, threshold=0.5),
+        target=sp.LinearDetrend(TIME_COLUMN, threshold=0.5),
         output=sp.AverageEstimators(),
     )

sdm/processing/__init__.py

@@ -34,6 +34,7 @@
     QuantileTransform,
     RankGaussian,
     Standardize,
+    LinearDetrend,
     RobustScale,
     FlipSign,
     DropConstantColumns,
@@ -79,6 +80,7 @@
     "QuantileTransform",
     "RankGaussian",
     "Standardize",
+    "LinearDetrend",
     "RobustScale",
     "FlipSign",
     "DropConstantColumns",

sdm/processing/numerical/__init__.py

@@ -13,6 +13,7 @@
 from sdm.processing.numerical.quantile import QuantileTransform
 from sdm.processing.numerical.rank_gaussian import RankGaussian
 from sdm.processing.numerical.standardize import Standardize
+from sdm.processing.numerical.linear_detrend import LinearDetrend
 from sdm.processing.numerical.robust_scale import RobustScale
 from sdm.processing.numerical.flip_sign import FlipSign
 from sdm.processing.numerical.constant import DropConstantColumns
@@ -30,6 +31,7 @@
     "QuantileTransform",
     "RankGaussian",
     "Standardize",
+    "LinearDetrend",
     "RobustScale",
     "FlipSign",
     "DropConstantColumns",

sdm/processing/numerical/linear_detrend.py

@@ -0,0 +1,111 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+import torch
+from torch import Tensor
+
+from sdm import Stype, TableTensor
+from sdm.processing import InvertibleMixin, Processor
+
+
+class LinearDetrend(Processor, InvertibleMixin):
+    """Remove a linear trend fitted to observed numerical values.
+
+    The ``time_column`` contains numerical step coordinates stored as an ID
+    column. Non-finite values are ignored when fitting and preserved during
+    transformation. For predictions with multiple outputs per fitted column,
+    the outputs for each column must be adjacent.
+
+    Args:
+        time_column: Column containing the step coordinate for each row.
+        threshold: Detrend only when the residual standard deviation is less
+            than this fraction of the original standard deviation. ``None``
+            always applies the fitted trend.
+    """
+
+    handles_stypes = frozenset({Stype.numerical})
+    requires_fit = True
+
+    def __init__(
+        self,
+        time_column: str,
+        *,
+        threshold: float | None = None,
+    ) -> None:
+        super().__init__()
+        self.time_column = time_column
+        self.threshold = threshold
+        self.register_buffer("origin", torch.empty(0, dtype=torch.long))
+        self.register_buffer("slope", torch.empty(0))
+        self.register_buffer("intercept", torch.empty(0))
+
+    def _time(self, table: TableTensor) -> Tensor:
+        return table[self.time_column].id[..., 0]
+
+    def _fit(
+        self,
+        table: TableTensor,
+        *,
+        generator: torch.Generator | None = None,
+    ) -> None:
+        values = table.numerical.to(
+            dtype=torch.promote_types(table.numerical.dtype, torch.float32)
+        )
+        time = self._time(table)
+        self.origin = time[..., -1:].clone()
+        time = (time - self.origin).to(dtype=values.dtype).unsqueeze(-1)
+        valid = values.isfinite()
+        count = valid.sum(dim=-2, keepdim=True, dtype=torch.int32)
+        count.clamp_min_(1)
+        sum_t = torch.where(valid, time, 0.0).sum(dim=-2, keepdim=True)
+        sum_t2 = torch.where(valid, time.square(), 0.0).sum(
+            dim=-2, keepdim=True
+        )
+        centered = values.masked_fill(~valid, 0.0)
+        mean = centered.sum(dim=-2, keepdim=True) / count
+        centered.sub_(mean).masked_fill_(~valid, 0.0)
+        time_ss = (sum_t2 - sum_t.square() / count).clamp_min_(0.0)
+        cross = time.transpose(-2, -1) @ centered
+
+        constant_time = time_ss == 0
+        self.slope = cross / torch.where(constant_time, 1.0, time_ss)
+        self.slope.masked_fill_(constant_time, 0.0)
+        self.intercept = mean - self.slope * (sum_t / count)
+
+        if self.threshold is not None:
+            value_ss = centered.square_().sum(dim=-2, keepdim=True)
+            # OLS residual sum of squares: value_ss - slope * cross.
+            residual_ss = (value_ss - self.slope * cross).clamp_min_(0.0)
+            apply = residual_ss.sqrt() < self.threshold * value_ss.sqrt()
+            self.slope.masked_fill_(~apply, 0.0)
+            self.intercept.masked_fill_(~apply, 0.0)
+
+    def _trend(self, table: TableTensor) -> Tensor:
+        time = (
+            (self._time(table) - self.origin)
+            .to(dtype=self.slope.dtype)
+            .unsqueeze(-1)
+        )
+        return self.slope * time + self.intercept
+
+    def _transform(self, table: TableTensor) -> TableTensor:
+        time = (self._time(table) - self.origin).to(dtype=self.slope.dtype)
+        numerical = table.numerical - self.intercept
+        numerical.addcmul_(self.slope, time.unsqueeze(-1), value=-1)
+        return table.replace_blocks(
+            numerical=numerical.to(table.numerical.dtype)
+        )
+
+    def _inverse_transform(self, table: TableTensor) -> TableTensor:
+        trend = self._trend(table)  # [..., F, T]
+        values = table.numerical.unflatten(-1, (trend.size(-1), -1))
+        numerical = (values + trend.unsqueeze(-1)).flatten(-2)
+        return table.replace_blocks(
+            numerical=numerical.to(table.numerical.dtype)
+        )
+
+    def __repr__(self, *, indent: int = 0) -> str:
+        return (
+            f"{' ' * indent}{self.__class__.__name__}("
+            f"time_column={self.time_column!r}, threshold={self.threshold})"
+        )