refactor(timesfm3): simplify detrending column handling

Author: guerrantif

sdm/models/timesfm3/model.py

@@ -42,7 +42,7 @@
 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 TIME_COLUMN, default_recipe
+from sdm.models.timesfm3.recipe import default_recipe
 from sdm.models.timesfm3.util import (
     gather_future_patches,
     get_running_stats,
@@ -52,7 +52,7 @@
 from sdm.tensor.table import TableSchema
 
 
-def _with_time_column(table: TableTensor, start: int) -> TableTensor:
+def _with_id_column(table: TableTensor, start: int) -> TableTensor:
     time_steps = torch.arange(
         start, start + table.size(-2), device=table.device
     ).expand(*table.size()[:-2], -1)
@@ -62,7 +62,7 @@ def _with_time_column(table: TableTensor, start: int) -> TableTensor:
             (
                 table,
                 TableTensor(
-                    columns={Stype.id: [TIME_COLUMN]},
+                    columns={Stype.id: ["__timesfm3_id__"]},
                     id=ColumnarTensor((time_steps,)),
                 ),
             ),
@@ -71,9 +71,9 @@ def _with_time_column(table: TableTensor, start: int) -> TableTensor:
     )
 
 
-def _drop_time_column(table: TableTensor) -> TableTensor:
-    if TIME_COLUMN in table.columns[Stype.id]:
-        return table.drop_columns(TIME_COLUMN)
+def _drop_id_column(table: TableTensor) -> TableTensor:
+    if "__timesfm3_id__" in table.columns[Stype.id]:
+        return table.drop_columns("__timesfm3_id__")
     return table
 
 
@@ -204,13 +204,13 @@ def forward(
         x_query = expand_query(x_context.schema, x_query)
 
         out = super().forward(
-            _with_time_column(x_context, 0),
-            _with_time_column(y_context, 0),
-            _with_time_column(x_query, context_length),
+            _with_id_column(x_context, 0),
+            _with_id_column(y_context, 0),
+            _with_id_column(x_query, context_length),
             *args,
             **kwargs,
         )
-        return _drop_time_column(out)
+        return _drop_id_column(out)
 
     def fit(
         self,
@@ -228,7 +228,7 @@ def fit(
         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
+            _with_id_column(x, 0), _with_id_column(y, 0), *args, **kwargs
         )
 
     def predict(
@@ -249,18 +249,18 @@ def predict(
         cached_kwargs = cast(dict[str, Any], self._cache["kwargs"])
         x = expand_query(cached_kwargs["_x_context_schema"], x)
         out = super().predict(
-            _with_time_column(x, cached_kwargs["_context_length"]),
+            _with_id_column(x, cached_kwargs["_context_length"]),
             *args,
             **kwargs,
         )
-        return _drop_time_column(out)
+        return _drop_id_column(out)
 
     def _prepare_context(
         self,
         context: MemberContext,
         callbacks: Sequence[Callback],
     ) -> MemberContext:
-        context = context._replace(y=_drop_time_column(context.y))
+        context = context._replace(y=_drop_id_column(context.y))
         return super()._prepare_context(context, callbacks)
 
     def _forward(
@@ -307,7 +307,7 @@ def _forward(
             if x_context is not None
             else kwargs["_context_length"]
         )
-        return _with_time_column(out, context_length)
+        return _with_id_column(out, context_length)
 
 
 class _TimesFM3(torch.nn.Module):

sdm/models/timesfm3/recipe.py

@@ -3,12 +3,10 @@
 
 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),
+        features=sp.LinearDetrend("__timesfm3_id__", threshold=0.5),
+        target=sp.LinearDetrend("__timesfm3_id__", threshold=0.5),
         output=sp.AverageEstimators(),
     )

sdm/processing/numerical/linear_detrend.py

@@ -2,7 +2,6 @@
 # SPDX-License-Identifier: Apache-2.0
 
 import torch
-from torch import Tensor
 
 from sdm import Stype, TableTensor
 from sdm.processing import InvertibleMixin, Processor
@@ -11,13 +10,13 @@
 class LinearDetrend(Processor, InvertibleMixin):
     """Remove a linear trend fitted to observed numerical values.
 
-    The ``time_column`` contains numerical step coordinates stored as an ID
+    The ``id_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.
+        id_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.
@@ -28,20 +27,17 @@ class LinearDetrend(Processor, InvertibleMixin):
 
     def __init__(
         self,
-        time_column: str,
+        id_column: str,
         *,
         threshold: float | None = None,
     ) -> None:
         super().__init__()
-        self.time_column = time_column
+        self.id_column = id_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,
@@ -51,7 +47,7 @@ def _fit(
         values = table.numerical.to(
             dtype=torch.promote_types(table.numerical.dtype, torch.float32)
         )
-        time = self._time(table)
+        time = table[self.id_column].id[..., 0]
         self.origin = time[..., -1:].clone()
         time = (time - self.origin).to(dtype=values.dtype).unsqueeze(-1)
         valid = values.isfinite()
@@ -80,24 +76,23 @@ def _fit(
             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)
+        time = (table[self.id_column].id[..., 0] - 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]
+        time = (
+            (table[self.id_column].id[..., 0] - self.origin)
+            .to(dtype=self.slope.dtype)
+            .unsqueeze(-1)
+        )
+        trend = self.slope * time + self.intercept  # [..., F, T]
         values = table.numerical.unflatten(-1, (trend.size(-1), -1))
         numerical = (values + trend.unsqueeze(-1)).flatten(-2)
         return table.replace_blocks(
@@ -107,5 +102,5 @@ def _inverse_transform(self, table: TableTensor) -> TableTensor:
     def __repr__(self, *, indent: int = 0) -> str:
         return (
             f"{' ' * indent}{self.__class__.__name__}("
-            f"time_column={self.time_column!r}, threshold={self.threshold})"
+            f"id_column={self.id_column!r}, threshold={self.threshold})"
         )

test/models/timesfm3/test_detrending.py

@@ -5,7 +5,6 @@
 
 from sdm import ColumnarTensor, TableTensor
 from sdm.models.timesfm3 import TimesFM3
-from sdm.models.timesfm3.recipe import TIME_COLUMN
 from sdm.processing.execution import RecipeExecution
 from sdm.testing import withCUDA
 
@@ -15,7 +14,7 @@ def test_restores_each_target_trend_across_quantiles(
     device: torch.device,
 ) -> None:
     context = TableTensor(
-        columns={"numerical": ["a", "b"], "id": [TIME_COLUMN]},
+        columns={"numerical": ["a", "b"], "id": ["__timesfm3_id__"]},
         numerical=torch.tensor(
             [[1.0, 10.0], [5.0, 20.0], [9.0, 30.0]], device=device
         ),
@@ -36,7 +35,7 @@ def test_restores_each_target_trend_across_quantiles(
                 for name in ("a", "b")
                 for q in quantiles
             ],
-            "id": [TIME_COLUMN],
+            "id": ["__timesfm3_id__"],
         },
         numerical=residual.flatten(-2),
         id=ColumnarTensor((torch.arange(3, 5, device=device),)),