Develop

CHANGELOG.md

@@ -1,5 +1,25 @@
 # CHANGELOG
 
+## v1.2.0 (2026-06-08)
+
+### New Features
+
+**EKF**
+
+The KalmanFilter's support for extended-kalman-filtering, introduced silently in [v1.1.0](https://github.com/onesixsolutions/torchcast/releases/tag/v1.1.0), is now public, focusing on two newly documented features:
+
+- `SaturatedLinearModel` process is now documented and exported from the `torchcast.process` top-level namespace.
+- `BinomialFilter` is now documented and exported from ``torchcast.kalman_filter``.
+
+**`TimeSeriesDataset and DataLoader`**
+
+- Adds `standardize()` method to `TimeSeriesDataset` that centers and scales one or more tensors in a dataset. 
+- `TimeSeriesDataLoader` has been refactored to improve support for subclasses, which can override ``_collate_fn`` to determine how each group's DataFrame gets collated and then transformed into a `TimeSeriesDataset`.
+
+### Deprecations / Breaking Changes
+
+- **Covariance module restructured**: The `torchcast/covariance` module directory has been consolidated into a single `torchcast/covariance.py` module. Existing imports from `torchcast.covariance.base` (typically in pickled models) will continue to work via a backwards-compatibility shim but will emit a `DeprecationWarning`.
+
 ## v1.1.1 (2026-04-17)
 
 ### Bug fix: LBFGS default optimizer regression with PyTorch >= 2.10

docs/api/api.rst

@@ -6,6 +6,7 @@ API
 
    state_space
    kalman_filter
+   binomial_filter
    exp_smooth
    processes
    covariance

docs/api/binomial_filter.rst

@@ -0,0 +1,8 @@
+Binomial Filter
+=============
+
+.. automodule:: torchcast.kalman_filter.binomial_filter
+   :members: BinomialFilter
+   :show-inheritance:
+
+.. include:: ../macros.hrst

docs/api/kalman_filter.rst

@@ -3,7 +3,6 @@ Kalman Filter
 
 .. automodule:: torchcast.kalman_filter.kalman_filter
    :members: KalmanFilter
-   :exclude-members: ss_step_cls
    :show-inheritance:
 
 .. include:: ../macros.hrst

docs/api/processes.rst

@@ -4,5 +4,5 @@ Processes
 .. include:: ../macros.hrst
 
 .. automodule:: torchcast.process
-   :members: LocalLevel, LocalTrend, Season, LinearModel
+   :members: LocalLevel, LocalTrend, Season, LinearModel, SaturatedLinearModel
    :exclude-members: forward

docs/conf.py

@@ -68,6 +68,8 @@
 ]
 
 #
+autodoc_inherit_docstrings = False
+
 intersphinx_mapping = {
     'python': ('https://docs.python.org/3/', None),
     'PyTorch': ('https://pytorch.org/docs/stable/', None),

tests/test_covariance.py

@@ -4,7 +4,7 @@
 
 @torch.no_grad()
 def test_from_log_cholesky():
-    module = Covariance(id='test', rank=3)
+    module = Covariance(id='test', rank=3, init_diag_multi=0.1)
 
     module.state_dict()['cholesky_log_diag'][:] = torch.arange(1., 3.1)
     module.state_dict()['cholesky_off_diag'][:] = torch.arange(1., 3.1)
@@ -18,7 +18,7 @@ def test_from_log_cholesky():
 
 @torch.no_grad()
 def test_empty_idx():
-    module = Covariance(id='test', rank=3, empty_idx=[0])
+    module = Covariance(id='test', rank=3, empty_idx=[0], init_diag_multi=0.1)
     cov = module({}, num_groups=1, num_times=1)
     cov = cov.squeeze()
     assert (cov[0, :] == 0).all()

tests/test_data_utils.py

@@ -1,9 +1,8 @@
-from warnings import warn
-
 import numpy as np
 import torch
 
 from torchcast.utils.data import TimeSeriesDataset
+import pandas as pd
 
 
 def test_time_series_dataset():
@@ -16,11 +15,7 @@ def test_time_series_dataset():
         measures=[['y1', 'y2']],
         dt_unit=None
     )
-    try:
-        import pandas as pd
-    except ImportError:
-        warn("Not testing TimeSeriesDataset.to_dataframe, pandas not installed.")
-        return
+
     df1 = batch.to_dataframe()
 
     df2 = pd.concat([
@@ -31,8 +26,7 @@ def test_time_series_dataset():
 
 
 def test_pad_x(num_times: int = 10):
-    from pandas import DataFrame
-    df = DataFrame({'x1': np.random.randn(num_times), 'x2': np.random.randn(num_times)})
+    df = pd.DataFrame({'x1': np.random.randn(num_times), 'x2': np.random.randn(num_times)})
     df['y'] = 1.5 * df['x1'] + -.5 * df['x2'] + .1 * np.random.randn(num_times)
     df['time'] = df.index.values
     df['group'] = '1'
@@ -56,3 +50,53 @@ def test_pad_x(num_times: int = 10):
     assert not torch.isnan(dataset1.tensors[1]).any()
     assert not torch.isnan(dataset2.tensors[1]).any()
     assert (dataset1.tensors[1] == dataset2.tensors[1]).all()
+
+
+def test_standardize():
+    y = torch.randn((3, 20, 1))
+    X = torch.randn((3, 20, 2)) * 5 + 10
+    ds = TimeSeriesDataset(
+        y, X,
+        group_names=['a', 'b', 'c'],
+        start_times=[0, 0, 0],
+        measures=[['y'], ['x1', 'x2']],
+        dt_unit=None
+    )
+
+    # standardizing self: X tensor should have ~0 mean and ~1 std; y tensor unchanged
+    ds_std = ds.standardize(which=(1,))
+    assert ds_std.tensors[1].mean().abs() < 1e-5  # mean g2g
+    assert abs(ds_std.tensors[1].std().item() - 1.0) < 0.05  # std-dev g2g
+    assert torch.allclose(ds_std.tensors[0], y)  # first tensor unaffect ('which' arg)
+
+    # standardizing a separate dataset:
+    Xtrain = torch.as_tensor([[-1, 0, 1]], dtype=torch.float)
+    Xtrain = torch.stack([Xtrain, Xtrain + 1], -1)
+    ds_train = TimeSeriesDataset(
+        Xtrain,
+        group_names=['a'],
+        start_times=[0],
+        measures=[['x1', 'x2']],
+        dt_unit=None
+    )
+    ds_val = TimeSeriesDataset(
+        Xtrain * 2 + 1,
+        group_names=['a'],
+        start_times=[0],
+        measures=[['x1', 'x2']],
+        dt_unit=None
+    )
+    ds_val_std = ds_train.standardize(ds_val, which=(0,))
+    Xval_std = ds_val_std.tensors[0]
+    assert torch.allclose(Xval_std[:, :, 0].mean(), torch.as_tensor(1.))
+    assert torch.allclose(Xval_std[:, :, 0].std(), torch.as_tensor(2.))
+    assert torch.allclose(Xval_std[:, :, 1].mean(), torch.as_tensor(2.))
+    assert torch.allclose(Xval_std[:, :, 1].std(), torch.as_tensor(2.))
+
+# def test_different_behavior():
+#     Xtrain = torch.as_tensor([[-1, 0, 1]], dtype=torch.float)
+#     Xtrain = torch.stack([Xtrain, Xtrain + 1], -1)
+#     torch_result = Xtrain.std(dim=(0,1))
+#     np_result = Xtrain.numpy().std(axis=(0, 1), ddof=1)
+#     print(torch_result)
+#     print(np_result)

tests/test_process.py

@@ -3,6 +3,7 @@
 import numpy as np
 from torchcast.kalman_filter import KalmanFilter
 from torchcast.process.season import Season
+import pytest
 
 
 @torch.no_grad()
@@ -19,3 +20,68 @@ def test_fourier_season():
     pred = kf(data, start_offsets=start_datetimes)
     for g in range(6):
         assert torch.abs(pred.means[g] - data[g]).mean() < .01, f"Group {g} failed"
+
+
+@pytest.mark.parametrize("period,K,fixed,decay", [
+    (12.0, 2, True, None),    # no decay (fixed season)
+    (12.0, 2, False, None),   # no decay (evolving season)
+    (12.0, 2, False, True),   # learnable decay
+    (24.0, 4, True, None),    # larger K, no decay
+])
+@torch.no_grad()
+def test_season_get_initial_mean_integer_offsets(period, K, fixed, decay):
+    """get_initial_mean([i]) == F^i @ initial_mean for integer-unit offsets."""
+    torch.manual_seed(0)
+    season = Season(id='test', period=period, dt_unit=None, K=K, fixed=fixed, decay=decay)
+    season.initial_mean.data = torch.randn(season.rank)
+
+    F = season.get_transition_matrix()
+    offsets = np.array([0., 1., 4., 7., int(period) - 1])
+
+    result = season.get_initial_mean(offsets)
+
+    for g, offset in enumerate(offsets):
+        expected = torch.matrix_power(F, int(offset)) @ season.initial_mean
+        assert torch.allclose(result[g], expected, atol=1e-5), (
+            f"Mismatch at offset={offset}, period={period}, K={K}, fixed={fixed}, decay={decay}"
+        )
+
+
+@torch.no_grad()
+def test_season_get_initial_mean_datetime_offsets():
+    """get_initial_mean works correctly when start_offsets are np.datetime64 values."""
+    torch.manual_seed(0)
+    season = Season(id='test', period='7D', dt_unit='D', K=3, fixed=True)
+    season.initial_mean.data = torch.randn(season.rank)
+
+    F = season.get_transition_matrix()
+
+    # Days 0–6 relative to the Unix epoch; day 7 wraps back to 0.
+    epoch = np.datetime64('1970-01-01')
+    offsets = np.array([epoch + np.timedelta64(i, 'D') for i in range(7)])
+    result = season.get_initial_mean(offsets)
+
+    for day in range(7):
+        expected = torch.matrix_power(F, day) @ season.initial_mean
+        assert torch.allclose(result[day], expected, atol=1e-5), (
+            f"Mismatch at day={day}"
+        )
+
+
+@torch.no_grad()
+def test_season_get_initial_mean_multiple_groups():
+    """get_initial_mean handles a batch of groups, each potentially at a different offset."""
+    torch.manual_seed(0)
+    season = Season(id='test', period=24.0, dt_unit=None, K=3, fixed=True)
+    season.initial_mean.data = torch.randn(season.rank)
+
+    F = season.get_transition_matrix()
+    offsets = np.array([0., 3., 7., 12., 18., 23.])
+    result = season.get_initial_mean(offsets)
+
+    assert result.shape == (len(offsets), season.rank)
+    for g, offset in enumerate(offsets):
+        expected = torch.matrix_power(F, int(offset)) @ season.initial_mean
+        assert torch.allclose(result[g], expected, atol=1e-5), (
+            f"Mismatch for group {g} at offset={offset}"
+        )

torchcast/__init__.py

@@ -1 +1 @@
-__version__ = '1.1.1'
+__version__ = '1.1.2'

torchcast/covariance/base.py → torchcast/covariance.py

@@ -1,17 +1,21 @@
 import math
+import sys
+import types
 
-from typing import List, Optional, Sequence, Dict, Union
+from typing import List, Optional, Sequence, Dict, Union, Collection
 from warnings import warn
 
 import torch
-
 from torch import Tensor, nn, jit
 
 from torchcast.process.utils import Identity
-from torchcast.covariance.util import num_off_diag, mini_cov_mask
 from torchcast.internals.utils import is_near_zero, validate_gt_shape
 from torchcast.process.process import Process
 
+DEFAULT_MCOV_MULTI = 1.0
+DEFAULT_PCOV_MULTI = 0.1  # less than measure-cov by default
+DEFAULT_ICOV_MULTI = 0.5  # somewhere in between
+
 
 class Covariance(nn.Module):
     """
@@ -73,11 +77,10 @@ def from_processes(cls,
                     no_cov_idx.append(state_rank + i)
             state_rank += len(p.state_elements)
 
-        if cov_type == 'process':
-            # by default, assume process cov is less than measure cov:
-            if 'init_diag_multi' not in kwargs:
-                kwargs['init_diag_multi'] = .05
-        elif cov_type != 'initial':
+        if 'init_diag_multi' not in kwargs:
+            kwargs['init_diag_multi'] = DEFAULT_PCOV_MULTI if cov_type == 'process' else DEFAULT_ICOV_MULTI
+
+        if cov_type not in {'initial', 'process'}:
             raise ValueError(f"Unrecognized cov_type {cov_type}, expected 'initial' or 'process'.")
 
         if predict_variance is True:
@@ -118,7 +121,7 @@ def from_measures(cls,
         if 'method' not in kwargs and len(measures) > 5:
             kwargs['method'] = 'low_rank'
         if 'init_diag_multi' not in kwargs:
-            kwargs['init_diag_multi'] = 1.0
+            kwargs['init_diag_multi'] = DEFAULT_MCOV_MULTI
 
         if predict_variance is True:
             predict_variance = Identity()
@@ -129,17 +132,19 @@ def from_measures(cls,
 
     def __init__(self,
                  rank: int,
+                 init_diag_multi: float,
                  method: str = 'log_cholesky',
                  empty_idx: List[int] = (),
                  predict_variance: Optional[nn.Module] = None,
                  expected_kwargs: Optional[Sequence[str]] = None,
-                 id: Optional[str] = None,
-                 init_diag_multi: float = 0.1):
+                 id: Optional[str] = None):
         """
         You should rarely call this directly. Instead, call :func:`Covariance.from_measures` and
         :func:`Covariance.from_processes`.
 
         :param rank: The number of elements along the diagonal.
+        :param init_diag_multi: A float that will be applied as a multiplier to the initial values along the diagonal.
+         This can be useful to provide intelligent starting-values to speed up optimization.
         :param method: The parameterization for the covariance. The default, "log_cholesky", parameterizes the
          covariance using the cholesky factorization (which is itself split into two tensors: the log-transformed
          diagonal elements and the off-diagonal). The other currently supported option is "low_rank", which
@@ -152,8 +157,6 @@ def __init__(self,
          at ``forward()``.
         :param id: Identifier for this covariance. Typically left ``None`` and set when passed to the
          :class:`.StateSpaceModel`.
-        :param init_diag_multi: A float that will be applied as a multiplier to the initial values along the diagonal.
-         This can be useful to provide intelligent starting-values to speed up optimization.
         """
 
         super().__init__()
@@ -268,3 +271,41 @@ def forward(self,
         mask = self.mask.unsqueeze(0).unsqueeze(0)
 
         return mask @ mini_cov @ mask.transpose(-1, -2)
+
+
+def num_off_diag(rank: int) -> int:
+    return int(rank * (rank - 1) / 2)
+
+
+def cov2corr(cov: Tensor) -> Tensor:
+    std_ = torch.sqrt(torch.diagonal(cov, dim1=-2, dim2=-1))
+    # TODO: cov / std_.unsqueeze(-1) / std_.unsqueeze(-2)
+    return cov / (std_.unsqueeze(-1) @ std_.unsqueeze(-2))
+
+
+def mini_cov_mask(rank: int, empty_idx: Collection[int], **kwargs) -> Tensor:
+    param_rank = rank - len(empty_idx)
+    mask = torch.zeros((rank, param_rank), **kwargs)
+    c = 0
+    for r in range(rank):
+        if r not in empty_idx:
+            mask[r, c] = 1.
+            c += 1
+    return mask
+
+
+# backwards compat shim
+class _DeprecatedBase(types.ModuleType):
+    def __getattr__(self, name):
+        if name != '__file__':
+            warn(
+                f"`torchcast.covariance.base.{name}` is deprecated, instead just import from "
+                f"`torchcast.covariance.{name}`.",
+                DeprecationWarning,
+                stacklevel=2
+            )
+        return globals()[name]
+
+
+_base = _DeprecatedBase('covariance.base')
+sys.modules['torchcast.covariance.base'] = _base