diff --git a/CHANGELOG.md b/CHANGELOG.md index 4c05213..806693c 100644 --- a/CHANGELOG.md +++ b/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 diff --git a/docs/api/api.rst b/docs/api/api.rst index a070690..87eee54 100644 --- a/docs/api/api.rst +++ b/docs/api/api.rst @@ -6,6 +6,7 @@ API state_space kalman_filter + binomial_filter exp_smooth processes covariance diff --git a/docs/api/binomial_filter.rst b/docs/api/binomial_filter.rst new file mode 100644 index 0000000..80642d8 --- /dev/null +++ b/docs/api/binomial_filter.rst @@ -0,0 +1,8 @@ +Binomial Filter +============= + +.. automodule:: torchcast.kalman_filter.binomial_filter + :members: BinomialFilter + :show-inheritance: + +.. include:: ../macros.hrst \ No newline at end of file diff --git a/docs/api/kalman_filter.rst b/docs/api/kalman_filter.rst index 8209830..f924995 100644 --- a/docs/api/kalman_filter.rst +++ b/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 \ No newline at end of file diff --git a/docs/api/processes.rst b/docs/api/processes.rst index 1793261..4229caa 100644 --- a/docs/api/processes.rst +++ b/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 diff --git a/docs/conf.py b/docs/conf.py index cb56a86..e262ffb 100644 --- a/docs/conf.py +++ b/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), diff --git a/tests/test_covariance.py b/tests/test_covariance.py index b7e8cb0..ca9ad4e 100644 --- a/tests/test_covariance.py +++ b/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() diff --git a/tests/test_data_utils.py b/tests/test_data_utils.py index e49a59a..15aeeca 100644 --- a/tests/test_data_utils.py +++ b/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) diff --git a/tests/test_process.py b/tests/test_process.py index a225e82..3bd83b6 100644 --- a/tests/test_process.py +++ b/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}" + ) diff --git a/torchcast/__init__.py b/torchcast/__init__.py index b3ddbc4..7b344ec 100644 --- a/torchcast/__init__.py +++ b/torchcast/__init__.py @@ -1 +1 @@ -__version__ = '1.1.1' +__version__ = '1.1.2' diff --git a/torchcast/covariance/base.py b/torchcast/covariance.py similarity index 88% rename from torchcast/covariance/base.py rename to torchcast/covariance.py index 0595212..738f793 100644 --- a/torchcast/covariance/base.py +++ b/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 diff --git a/torchcast/covariance/__init__.py b/torchcast/covariance/__init__.py deleted file mode 100644 index 0c8a18e..0000000 --- a/torchcast/covariance/__init__.py +++ /dev/null @@ -1 +0,0 @@ -from .base import Covariance diff --git a/torchcast/covariance/util.py b/torchcast/covariance/util.py deleted file mode 100644 index 92f7549..0000000 --- a/torchcast/covariance/util.py +++ /dev/null @@ -1,24 +0,0 @@ -from typing import Collection - -import torch -from torch import Tensor - - -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)) - 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 diff --git a/torchcast/exp_smooth/exp_smooth.py b/torchcast/exp_smooth/exp_smooth.py index 207c87f..d2c2c60 100644 --- a/torchcast/exp_smooth/exp_smooth.py +++ b/torchcast/exp_smooth/exp_smooth.py @@ -69,6 +69,7 @@ def _parse_kwargs(self, if self.smoothing_matrix.expected_kwargs: smat_kwargs = {k: kwargs[k] for k in self.smoothing_matrix.expected_kwargs} used_keys |= set(smat_kwargs) + # todo: instead of branching here, clean up Covariance.forward(): if smat_kwargs: Ks = self.smoothing_matrix(smat_kwargs, num_groups=num_groups, num_times=num_timesteps) update_kwargs['K'] = Ks.unbind(1) diff --git a/torchcast/internals/batch_design/measurement_model.py b/torchcast/internals/batch_design/measurement_model.py index b192ad0..b663257 100644 --- a/torchcast/internals/batch_design/measurement_model.py +++ b/torchcast/internals/batch_design/measurement_model.py @@ -9,6 +9,7 @@ from torchcast.internals.utils import normalize_index, compute_index_result_shape + if TYPE_CHECKING: from torchcast.process import Process @@ -38,9 +39,9 @@ def __init__(self, @property def is_nonlinear(self) -> bool: - return bool(self.nonlinear_processes) or self.measure_funs + return bool(self.nonlinear_processes) or bool(self.measure_funs) - @property + @cached_property def nonlinear_processes(self) -> list['Process']: return [p for p in self.processes.values() if not p.linear_measurement] @@ -53,10 +54,11 @@ def __call__(self, nl_procs_and_means = list(self._get_nonlinear_processes_and_means(mean)) - measured_mean_adj = self.adjust_measured_mean(measured_mean, nl_procs_and_means, time) - measure_mat_adj = self._adjust_measure_mat(measure_mat, nl_procs_and_means, measured_mean, time) + if self.is_nonlinear: + measured_mean = self.adjust_measured_mean(measured_mean, nl_procs_and_means, time) + measure_mat = self._adjust_measure_mat(measure_mat, nl_procs_and_means, measured_mean, time) - return measured_mean_adj, measure_mat_adj + return measured_mean, measure_mat @cached_property def extended_measure_mat(self) -> torch.Tensor: @@ -65,15 +67,17 @@ def extended_measure_mat(self) -> torch.Tensor: linear_mmat = self._get_linear_measure_mat(0) - nonlinear_rank = sum(p.rank for p in self.nonlinear_processes) - extension = torch.zeros( - (self.num_groups, nonlinear_rank, self.state_rank), - dtype=self.dtype, - device=self.device - ) - if nonlinear_rank: - extension[:, :, -nonlinear_rank:] = torch.eye(nonlinear_rank).unsqueeze(0) - return torch.cat([linear_mmat, extension], dim=1) + extension = [] + for proc in self.nonlinear_processes: + slc = self.process2slice[proc.id] + for i in range(slc.start, slc.stop): + x = torch.zeros(self.state_rank, device=linear_mmat.device, dtype=linear_mmat.dtype) + x[i] = 1.0 + extension.append(x) + if extension: + extension = torch.stack(extension).unsqueeze(0).expand(linear_mmat.shape[0], -1, -1) + return torch.cat([linear_mmat, extension], dim=1) + return linear_mmat @property def extended_mmat_slices(self) -> dict[str, slice]: @@ -196,8 +200,10 @@ def _cache_per_process(self) -> dict[str, dict]: @cached_property def _measure_mats(self) -> Sequence[torch.Tensor]: + is_time_varying = any(p.measurement_kwargs for p in self.processes.values()) + n_times = self.num_timesteps if is_time_varying else 1 H = torch.zeros( - (self.num_groups, self.num_timesteps, len(self.measures), self.state_rank), + (self.num_groups, n_times, len(self.measures), self.state_rank), device=self.device, dtype=self.dtype ) @@ -212,11 +218,16 @@ def _measure_mats(self) -> Sequence[torch.Tensor]: if len(value.shape) == 1: value = value.unsqueeze(0).unsqueeze(0) elif len(value.shape) != 3: + assert not is_time_varying raise ValueError(f"for process {pid}, measurement matrix expected to be a vector or have shape" f"(num_groups, num_times, rank). Instead got {value.shape}. ") - # todo is all this masking inefficient? H[:, :, midx, pidx] = value + if not is_time_varying: + # much faster for backward-step + H0 = H.squeeze(1) + return [H0] * self.num_timesteps + return H.unbind(1) def _copy(self, diff --git a/torchcast/internals/batch_design/transition_model.py b/torchcast/internals/batch_design/transition_model.py index 593029c..dfd0f23 100644 --- a/torchcast/internals/batch_design/transition_model.py +++ b/torchcast/internals/batch_design/transition_model.py @@ -20,8 +20,11 @@ def __init__(self, ) self.measures = measures + is_time_varying = False # could be supported in the future + n_times = self.num_timesteps if is_time_varying else 1 + F = torch.zeros( - (self.num_groups, self.num_timesteps, self.state_rank, self.state_rank), + (self.num_groups, n_times, self.state_rank, self.state_rank), device=self.device, dtype=self.dtype ) @@ -31,11 +34,13 @@ def __init__(self, F[:, :, pidx, pidx] = process.get_transition_matrix() else: raise NotImplementedError - self._transition_mats = F - @cached_property - def transition_mats(self) -> Sequence[torch.Tensor]: - return self._transition_mats.to(device=self.device, dtype=self.dtype).unbind(1) + if is_time_varying: + self.transition_mats = F.unbind(1) + else: + # much faster for backward-step: + F0 = F.squeeze(1) + self.transition_mats = [F0] * self.num_timesteps def __call__(self, mean: torch.Tensor, diff --git a/torchcast/internals/monte_carlo.py b/torchcast/internals/monte_carlo.py index 863e41c..fc29ca7 100644 --- a/torchcast/internals/monte_carlo.py +++ b/torchcast/internals/monte_carlo.py @@ -15,6 +15,9 @@ def __init__(self, self._num_samples = num_samples self.reset(random_state) + def __repr__(self) -> str: + return "FixedWhiteNoise(num_samples={})".format(self.num_samples) + @property def num_samples(self) -> int: return self._num_samples diff --git a/torchcast/internals/utils.py b/torchcast/internals/utils.py index 80cbae1..cb700fc 100644 --- a/torchcast/internals/utils.py +++ b/torchcast/internals/utils.py @@ -1,5 +1,5 @@ import functools -from typing import Union, Any, Tuple, Sequence, List, Optional, Iterable, Collection, Type +from typing import Union, Any, Tuple, Sequence, List, Optional, Iterable, Type import torch diff --git a/torchcast/kalman_filter/__init__.py b/torchcast/kalman_filter/__init__.py index 2eae6f3..f4919bb 100644 --- a/torchcast/kalman_filter/__init__.py +++ b/torchcast/kalman_filter/__init__.py @@ -1 +1,2 @@ from .kalman_filter import KalmanFilter +from .binomial_filter import BinomialFilter diff --git a/torchcast/kalman_filter/binomial_filter.py b/torchcast/kalman_filter/binomial_filter.py index 0a2449b..68734d2 100644 --- a/torchcast/kalman_filter/binomial_filter.py +++ b/torchcast/kalman_filter/binomial_filter.py @@ -1,10 +1,15 @@ +""" +The :class:`.BinomialFilter` is a :class:`.KalmanFilter` with (1) one or more sigmoid measurement-functions +and (2) with a log_prob that use binomial likelihood (using monte-carlo approximation). +""" + from math import log import torch from torch.distributions import Binomial from typing import Sequence, TYPE_CHECKING, Optional, Union -from torchcast.covariance import Covariance +from torchcast.covariance import Covariance, DEFAULT_MCOV_MULTI from torchcast.kalman_filter import KalmanFilter from torchcast.state_space import Predictions from torchcast.internals.batch_design import MeasurementModel, Sigmoid @@ -14,12 +19,28 @@ class BinomialFilter(KalmanFilter): + """ + :param processes: A list of :class:`.Process` modules. + :param measures: A list of strings specifying the names of the dimensions of the time-series being measured. + :param binary_measures: A subset of ``measures`` with binary (binomial) outcomes. + :param observed_counts: If True, then ``y`` is interpreted as counts; if False then as probabilities 0-1. + :param do_post_hoc_correction: Default True. Apply correction to the binary residuals during the update step, to + compensate for the linearization error introduced by the EKF approximation. The exact correction is learned: a + learned baseline correction level is modulated by a learned increase/decrease according to `num_obs`. + :param measure_covariance: A module created with ``Covariance.from_measures(measures)``. + :param process_covariance: A module created with ``Covariance.from_processes(processes, type='process')``. + :param initial_covariance: A module created with ``Covariance.from_processes(measures, type='initial')``. + :param adaptive_scaling: Experimental feature to adaptively scale the covariance as a function of residuals. This + is useful if different groups have very different magnitudes. + """ + def __init__(self, processes: Sequence['Process'], measures: Optional[Sequence[str]], binary_measures: Optional[Sequence[str]] = None, observed_counts: Optional[bool] = None, - measure_covariance: Optional[Covariance] = None, + do_post_hoc_correction: bool = True, + measure_covariance: Optional[Union[Covariance, dict]] = None, process_covariance: Optional[Covariance] = None, initial_covariance: Optional[Covariance] = None, adaptive_scaling: bool = False): @@ -45,6 +66,14 @@ def __init__(self, measure_funs={m: 'ilogit' for m in binary_measures}, ) + if do_post_hoc_correction: + self.post_correction_module = torch.nn.Linear(1, 1, bias=True) + with torch.no_grad(): + self.post_correction_module.bias.normal_(mean=-.5, std=.1) + self.post_correction_module.weight.normal_(std=.1) + else: + self.post_correction_module = None + @classmethod def _validate_measure_cov(cls, measures: Sequence[str], @@ -61,7 +90,7 @@ def _validate_measure_cov(cls, mcov_empty_idx = [i for i, m in enumerate(measures) if m in binary_measures] if measure_covariance is None: - measure_covariance = {} + measure_covariance = {'init_diag_multi': DEFAULT_MCOV_MULTI} if isinstance(measure_covariance, dict): measure_covariance['id'] = 'measure_covariance' measure_covariance['rank'] = len(measures) @@ -195,7 +224,9 @@ def _update_step(self, # validate num_obs, use to normalize input if observed_counts=True: if self.observed_counts is None: - if num_obs is not None and (num_obs != 1).any(): + if num_obs is None: + num_obs = 1 + elif (num_obs != 1).any(): raise ValueError( "If `num_obs` is supplied, must specify whether observed values are counts (observed_counts=True) " "or proportions (observed_counts=False)." @@ -218,6 +249,22 @@ def _update_step(self, ) measure_cov = measure_cov + bin_measure_cov + if self.do_post_hoc_correction: + # super takes input and mean, not resid. + # we want to multiply the resid, so we'll just do that then apply that adjustment to the measured-mean: + raw_resid = input - measured_mean + + resid = torch.zeros_like(input) + resid[..., binary_idx] = self._binomial_post_hoc_correction( + raw_resid[..., binary_idx], + binary_measured_mean, + num_obs + ) + other_idx = [x for x in range(measured_mean.shape[-1]) if x not in binary_idx] + if other_idx: + resid[..., other_idx] = raw_resid[..., other_idx] + measured_mean = input - resid + return super()._update_step( input=input, mean=mean, @@ -228,6 +275,36 @@ def _update_step(self, **kwargs ) + @property + def do_post_hoc_correction(self) -> bool: + return getattr(self, 'post_correction_module', None) is not None + + def _binomial_post_hoc_correction(self, + resid: torch.Tensor, + measured_mean: torch.Tensor, + num_obs: torch.Tensor) -> torch.Tensor: + measured_mean = torch.clamp(measured_mean, min=1e-6, max=1 - 1e-6) + Hx = torch.log(measured_mean / (1 - measured_mean)) + + # only applies when residual is heading us towards extremes (so neg resid if logit < 0, pos resid if > 0) + consistent = torch.sign(resid) == torch.sign(Hx) + + # sharpness is learned empirically: + sharpness = torch.exp(self.post_correction_module(torch.log(num_obs[consistent]).unsqueeze(-1))).squeeze(-1) + + # multi is 1 if ~consistent, or its set by normalized derivative of sigmoid + multi = torch.ones_like(resid) + deriv_at_Hx = self._sigmoid_deriv(Hx[consistent], s=sharpness) + deriv_at_zero = self._sigmoid_deriv(torch.as_tensor(0.), s=sharpness) + multi[consistent] = (deriv_at_zero / deriv_at_Hx) + + return resid * multi + + @staticmethod + def _sigmoid_deriv(x: torch.Tensor, s: torch.Tensor) -> torch.Tensor: + exp_neg = torch.exp(-s * x) + return s * exp_neg / (exp_neg + 1) ** 2 + @torch.jit.ignore() def forward(self, y: Optional[torch.Tensor] = None, @@ -358,7 +435,7 @@ def _log_prob(self, return gaussian_lp + binary_lp -def main(num_groups: int = 50, num_timesteps: int = 365, bias: float = -1, prop_common: float = 1.): +def main(num_groups: int = 50, num_timesteps: int = 100, bias: float = -2, prop_common: float = 0.5): from torchcast.process import LocalLevel, Season from torchcast.utils import TimeSeriesDataset from scipy.special import expit @@ -402,17 +479,21 @@ def main(num_groups: int = 50, num_timesteps: int = 365, bias: float = -1, prop_ bf = BinomialFilter( processes=[LocalLevel(id=f'level_{m}', measure=m) for m in measures] - + [Season(id=f'season_{m}', measure=m, dt_unit='D', period=7, K=2) for m in measures], + # + [Season(id=f'season_{m}', measure=m, dt_unit='D', period=7, K=2) for m in measures] + , measures=measures, binary_measures=binary_measures, - observed_counts=False + observed_counts=False, + do_post_hoc_correction=False ) y = dataset.tensors[0] - bf.fit(y, start_offsets=dataset.start_offsets) + bf.fit(y, start_offsets=dataset.start_offsets, + stopping={'monitor_params': True}, + ) _kwargs = {} - # if TOTAL_COUNT != 1: - # _kwargs['num_obs'] = TOTAL_COUNT + if TOTAL_COUNT != 1: + _kwargs['num_obs'] = TOTAL_COUNT preds = bf( dataset.tensors[0], start_offsets=dataset.start_offsets, @@ -437,6 +518,7 @@ def main(num_groups: int = 50, num_timesteps: int = 365, bias: float = -1, prop_ ).show() # preds._white_noise = torch.zeros((1, len(binary_measures))) # print(preds.log_prob(y).mean()) + # with correction tensor(-1.3281, grad_fn=) if __name__ == '__main__': diff --git a/torchcast/kalman_filter/kalman_filter.py b/torchcast/kalman_filter/kalman_filter.py index f092d45..b00fcea 100644 --- a/torchcast/kalman_filter/kalman_filter.py +++ b/torchcast/kalman_filter/kalman_filter.py @@ -1,6 +1,6 @@ """ The :class:`.KalmanFilter` is a :class:`torch.nn.Module` which generates forecasts using the full kalman-filtering -algorithm. +algorithm (or optionally extended-kalman filtering, if any measure-funs or nonlinear processes are used). This class inherits most of its methods from :class:`torchcast.state_space.StateSpaceModel`. """ @@ -17,6 +17,16 @@ class KalmanFilter(StateSpaceModel): + """ + :param processes: A list of :class:`.Process` modules. + :param measures: A list of strings specifying the names of the dimensions of the time-series being measured. + :param measure_covariance: A module created with ``Covariance.from_measures(measures)``. + :param process_covariance: A module created with ``Covariance.from_processes(processes, type='process')``. + :param initial_covariance: A module created with ``Covariance.from_processes(measures, type='initial')``. + :param measure_funs: A dictionary mapping measure-names to measurement-functions. Currently only supports 'sigmoid'. + :param adaptive_scaling: Experimental feature to adaptively scale the covariance as a function of residuals. This + is useful if different groups have very different magnitudes. + """ def __init__(self, processes: Sequence['Process'], measures: Sequence[str], @@ -110,6 +120,7 @@ def _parse_kwargs(self, measure_scaling = self._get_measure_scaling() + # todo: instead of branching here, clean up Covariance.forward(): if pcov_kwargs: pcov_raw = self.process_covariance(pcov_kwargs, num_groups=num_groups, num_times=num_timesteps) Qs = self._apply_cov_scaling(pcov_raw, scaling=measure_scaling, is_process_cov=True) diff --git a/torchcast/process/__init__.py b/torchcast/process/__init__.py index 759bfbb..8fc229a 100644 --- a/torchcast/process/__init__.py +++ b/torchcast/process/__init__.py @@ -6,11 +6,12 @@ * :class:`.Season` - a process with seasonal structure, implementing the fourier-series based model from `De Livera, A.M., Hyndman, R.J., & Snyder, R. D. (2011)`. * :class:`.LinearModel` - a linear-model allowing for external predictors. +* :class:`.SaturatedLinearModel` - a linear model that allows for saturation effects (via EKF). ---------- """ from .process import Process -from .regression import LinearModel +from .regression import LinearModel, SaturatedLinearModel from .local import LocalLevel, LocalTrend from .season import Season diff --git a/torchcast/process/regression.py b/torchcast/process/regression.py index 8c9c67d..01f6e2b 100644 --- a/torchcast/process/regression.py +++ b/torchcast/process/regression.py @@ -4,6 +4,8 @@ from typing import Sequence, Optional, Union, Collection +from torch.nn.functional import softplus + from torchcast.process import Process from torchcast.process.utils import ProcessKwarg, StateElement, standardize_decay @@ -32,6 +34,7 @@ def __init__(self, fixed: Union[bool, Collection[str]] = True, decay: Optional[tuple[float, float]] = None, model_mat_kwarg_name: str = 'X'): + if isinstance(fixed, str): raise ValueError(f"`fixed` should be a collection of strings not a single string.") elif hasattr(fixed, '__contains__'): @@ -41,6 +44,8 @@ def __init__(self, else: fixed = list(predictors) if fixed else [] + self.predictors = predictors + super().__init__( id=id, state_elements=self._init_state_elements(predictors, fixed), @@ -84,7 +89,37 @@ def get_measurement_matrix(self, **kwargs) -> torch.Tensor: class SaturatedLinearModel(LinearModel): + """ + Similar to :class:`.LinearModel`, except an additional ceiling state-element allows for saturating effects. That + is, if ``yhat = X @ state`` and in a normal linear model ``measured_mean = y_hat``, the saturated linear model + still has ``measured_mean = y_hat`` when far from the ceiling, but has ``measured_mean = ceiling`` when close. + + The measurement-function this process uses is: + + .. code-block:: text + + measured_mean = yhat - (1. / s) * softplus(s * (yhat - ceiling)) + + With yhat defined above and ``s`` being a sharpness parameter which is scaled to the inverse of the ceiling height, + so that, as the ceiling lowers, sharpness increases. This allows the ``yhat -> measured_mean`` relationship to be + consistent when yhat is far from the ceiling (i.e., the ceiling won't impact where yhat crosses the origin). + + :param id: Unique identifier for the process + :param predictors: A sequence of strings with predictor-names. + :param measure: The name of the measure for this process. + :param fixed: By default, the regression-coefficients are assumed to be fixed: we are initially + uncertain about their value at the start of each series, but we gradually grow more confident. See LinearModel. + :param fix_ceiling: Like ``fixed``, but for the ceiling state-element. + :param decay: See :class:`.LinearModel` + :param model_mat_kwarg_name: See :class:`.LinearModel` + :param ceiling_init_value: The initial value for the ceiling prior. Defaults to 1 +/- jitter. If your measure is + very much not centered and/or scaled, optimization could be improved by putting an informative guess here. + :param anchor: If we start with yhat near the ceiling and reduce it, ``anchor`` is the yhat value at which yhat + converges to ``measured_mean``. Typically, you want to leave this at zero, but that implicitly assumes your + predictors are centered. + """ linear_measurement = False + base_sharpness = 6.0 def __init__(self, id: str, @@ -93,7 +128,9 @@ def __init__(self, fixed: Union[bool, Collection[str]] = True, fix_ceiling: bool = True, decay: Optional[tuple[float, float]] = None, - model_mat_kwarg_name: str = 'X'): + model_mat_kwarg_name: str = 'X', + ceiling_init_value: Optional[float] = None, + anchor: float = 0.0): self.fix_ceiling = fix_ceiling super().__init__( id=id, @@ -104,6 +141,16 @@ def __init__(self, model_mat_kwarg_name=model_mat_kwarg_name ) + # ceiling initial value: + if ceiling_init_value is None: + ceiling_init_value = 1.0 + torch.randn(1).item() / 10 + with torch.no_grad(): + assert list(self.state_elements)[-1] == '_ceiling' + self.initial_mean[-1] = ceiling_init_value + + # yhat value below which y~=yhat (regardless of ceiling value) + self.anchor = anchor + def _init_state_elements(self, predictors: Sequence[str], fixed: Sequence[str]) -> Sequence[StateElement]: @@ -119,9 +166,14 @@ def get_measurement_matrix(self, X: torch.Tensor) -> torch.Tensor: @property def num_predictors(self) -> int: - return self.rank - 1 + return len(self.predictors) - def prepare_measurement_cache(self, X: torch.Tensor) -> dict: + def prepare_measurement_cache(self, **kwargs) -> dict: + X = kwargs.pop(self.model_mat_kwarg_name, None) + if X is None: + raise TypeError(f"{self.id}.prepare_measurement_cache() missing `{self.model_mat_kwarg_name}` argument") + if kwargs: + raise ValueError(f"{self.id}.prepare_measurement_cache() received unexpected kwargs: {set(kwargs)}") assert not torch.isnan(X).any() assert not torch.isinf(X).any() if X.shape[-1] != self.num_predictors: @@ -129,21 +181,88 @@ def prepare_measurement_cache(self, X: torch.Tensor) -> dict: f"process '{self.id}' received X that has shape {X.shape}, but expected last dim to " f"match len(predictors) {self.num_predictors}." ) - return {'X': X.unbind(1)} + return { + 'X': X.unbind(1), + 's': [None] * X.shape[1], + 'yhat': [None] * X.shape[1] + } def get_measured_mean(self, mean: torch.Tensor, time: int, cache: dict) -> torch.Tensor: - # TODO: reparameterize X = cache['X'][time] coefs = mean[:, :self.num_predictors] ceiling = mean[:, self.num_predictors] - cache['yhat'] = (X * coefs).sum(-1) - return cache['yhat'] - torch.nn.functional.softplus(cache['yhat'] - ceiling) + yhat = cache['yhat'][time] = (X * coefs).sum(-1) + return _sat_measured_mean( + yhat, + ceiling=ceiling, + sharpness=self.base_sharpness, + anchor=self.anchor, + ) def get_measurement_jacobian(self, mean: torch.Tensor, time: int, cache: dict) -> torch.Tensor: - # TODO: reparameterize - X = cache['X'][time] - ceiling = mean[:, self.num_predictors] - ceil_derivs = torch.sigmoid((cache['yhat'] - ceiling).clamp(min=-10, max=10)) - coef_derivs = X * (1 - ceil_derivs.unsqueeze(-1)) - jacobian = torch.cat([coef_derivs, ceil_derivs.unsqueeze(-1)], dim=-1) - return jacobian + return _sat_jacobian( + X=cache['X'][time], + yhat=cache['yhat'][time], + ceiling=mean[:, self.num_predictors], + sharpness=self.base_sharpness, + anchor=self.anchor, + ) + + +def _sat_measured_mean(yhat: torch.Tensor, + ceiling: torch.Tensor, + sharpness: float, + anchor: float) -> torch.Tensor: + d = (ceiling - anchor).clamp(min=1e-6) + s = sharpness / d + nl_mask = yhat > ceiling / sharpness + adjustment = torch.zeros_like(yhat) + adjustment[nl_mask] = softplus(s[nl_mask] * (yhat[nl_mask] - ceiling[nl_mask])) / s[nl_mask] + # quick note on `nl_mask`: + # the idea with this saturation approach is that when yhat << ceiling, the problem should reduce to linear + # (`adjustment` asymptotes to 0), and when yhat >> ceiling, then the output is just the ceiling (i.e. the + # ``adjustment = yhat + ceiling``, so ``yhat - adjustment`` just equals ceiling). + # the problem is that the adjustment doesn't quite asymptote towards zero when yhat << ceiling. specifically if + # yhat is reasonable but ceiling is growing, the adjustment instead asymptotes towards + # ``(ceiling / sharpness) * exp(-sharpness)``. this means, counter-intuitively, that in this regime, larger ceiling + # values actually push *down* our output relative to yhat (instead of bringing it upwards). this is as confusing + # for us as it is for the optimizer. + # to help with this, we want to ignore the adjustment around where this counter-intuitive behavior starts. where is + # that? for a pair of ceilings, it starts at the yhat where their two ``adjustment`` curves cross, (i.e. where + # ceil_larger switches from having higher outputs to having lower outputs than ceil_smaller's outputs). This is + # ``ceil_larger*ceil_smaller * log(ceil_smaller/ceil_larger) / (sharpness * (ceil_smaller - ceil_larger))``. If we + # then use taylor expansions to understand the limiting behavior as the two ceilings get closer together, we + # get a fairly principled cutoff for when we should start ignoring the adjustment, which is how ``nl_mask`` is + # defined above. + return yhat - adjustment + + +def _sat_jacobian(X: torch.Tensor, + yhat: torch.Tensor, + ceiling: torch.Tensor, + sharpness: float, + anchor: float) -> torch.Tensor: + d = ceiling - anchor + ac_mask = d > 1e-6 # mask indicating ceiling is above anchor (below that `s` doesnt work) + d = d.clamp(min=1e-6) + s = sharpness / d + + u = yhat - ceiling + su = (s * u).clamp(min=-20, max=20) + sigma = torch.sigmoid(su) + + nl_mask = yhat > ceiling / sharpness # mask indicating we're in nonlinear regime (see _sat_measured_mean) + + # ceil derivs: + _mask = nl_mask & ac_mask + path2 = torch.zeros_like(sigma) + path2[_mask] = -softplus(su[_mask]) / (s[_mask] * d[_mask]) + (u[_mask] / d[_mask]) * sigma[_mask] + ceil_derivs = torch.zeros_like(sigma) + ceil_derivs[nl_mask] = (sigma[nl_mask] + path2[nl_mask].clamp(min=-1., max=1.)).clamp(min=0.) + + # coef derivs: + coef_deriv_multi = torch.ones_like(sigma) + coef_deriv_multi[nl_mask] = (1. - sigma[nl_mask]) + coef_derivs = X * coef_deriv_multi.unsqueeze(-1) + + return torch.cat([coef_derivs, ceil_derivs.unsqueeze(-1)], dim=-1) diff --git a/torchcast/process/season.py b/torchcast/process/season.py index f5e1149..824a35d 100644 --- a/torchcast/process/season.py +++ b/torchcast/process/season.py @@ -6,7 +6,6 @@ import torch -from torchcast.internals.utils import update_tensor from torchcast.process.process import Process from torchcast.process.utils import Multi, standardize_decay, StateElement, NoInputSequential @@ -145,15 +144,38 @@ def get_initial_mean(self, start_offsets: np.ndarray) -> torch.Tensor: start_offsets = start_offsets.round() num_groups = len(start_offsets) - if self.linear_transition: - out = [] - zeros = torch.zeros((num_groups, self.rank), device=self.initial_mean.device) - F = self.get_transition_matrix().expand(len(start_offsets), -1, -1) - mean = self.initial_mean.expand(len(start_offsets), -1) - for i in range(int(self.period) + 1): - maski = (start_offsets == i) - out.append(update_tensor(zeros, new=mean[maski], mask=maski)) - mean = (F @ mean.unsqueeze(-1)).squeeze(-1) - return torch.stack(out, 0).sum(0) - else: + if not self.linear_transition: raise NotImplementedError + + # Each 2x2 block of F for harmonic j is: decay * [[cos(lam_j), -sin(lam_j)], + # [sin(lam_j), cos(lam_j)]] + # so F^i has blocks: decay^i * R(i * lam_j). + # This lets us compute F^i @ initial_mean directly via trig instead of stepping + # through i matrix multiplications (the old approach was O(period) per forward pass). + F = self.get_transition_matrix() # (rank, rank) + K = self.rank // 2 + even = torch.arange(0, self.rank, 2, device=F.device) # indices of s_j components + F_cos = F[even, even] # decay * cos(lam_j), shape (K,) + F_sin = F[even + 1, even] # decay * sin(lam_j), shape (K,) + lam = torch.atan2(F_sin, F_cos) # rotation angle per harmonic, (K,) + decay = torch.sqrt(F_cos ** 2 + F_sin ** 2) # magnitude per harmonic, (K,) + + offsets = torch.as_tensor( + start_offsets, dtype=self.initial_mean.dtype, device=self.initial_mean.device + ) # (num_groups,) + + angles = offsets[:, None] * lam[None, :] # (num_groups, K) + decay_factors = decay[None, :] ** offsets[:, None] # (num_groups, K) + + mean0 = self.initial_mean.view(K, 2) # (K, 2): [[s1, s*1], [s2, s*2], ...] + s0 = mean0[:, 0] # (K,) + s_star0 = mean0[:, 1] # (K,) + + cos_a = torch.cos(angles) # (num_groups, K) + sin_a = torch.sin(angles) # (num_groups, K) + + new_s = (s0 * cos_a - s_star0 * sin_a) * decay_factors # (num_groups, K) + new_s_star = (s0 * sin_a + s_star0 * cos_a) * decay_factors # (num_groups, K) + + # Interleave back to [s1, s*1, s2, s*2, ...] layout + return torch.stack([new_s, new_s_star], dim=-1).reshape(num_groups, self.rank) diff --git a/torchcast/state_space/state_space.py b/torchcast/state_space/state_space.py index 99cd709..8179fbd 100644 --- a/torchcast/state_space/state_space.py +++ b/torchcast/state_space/state_space.py @@ -133,15 +133,6 @@ def forward(self, :func:`Predictions.to_dataframe()` methods. """ - if y is None: - if out_timesteps is None: - raise RuntimeError("If no y is passed, must specify `out_timesteps`") - else: - if not torch.is_floating_point(y): - raise ValueError(f"Expected y to be a float tensor, got {y.dtype}") - if torch.isinf(y).any(): - raise ValueError("y contains infinite values.") - initial_state = self._prepare_initial_state( initial_state, start_offsets=start_offsets, @@ -157,15 +148,7 @@ def forward(self, for k, v in kwargs.items() } - if isinstance(n_step, float): - if not n_step.is_integer(): - raise ValueError("`n_step` must be an int.") - n_step = int(n_step) - if isinstance(out_timesteps, float): - if not out_timesteps.is_integer(): - raise ValueError("`out_timesteps` must be an int.") - out_timesteps = int(out_timesteps) - + n_step = int(n_step) assert n_step > 0 meanu, covu, inputs, num_groups, out_timesteps = self._standardize_input( @@ -559,6 +542,10 @@ def _standardize_input(self, if covu.shape[0] == 1: covu = repeat(covu, times=num_groups, dim=0) else: + if not torch.is_floating_point(input): + raise ValueError(f"Expected input to be a float tensor, got {input.dtype}") + if torch.isinf(input).any(): + raise ValueError("input contains infinite values.") if len(input.shape) != 3: raise ValueError(f"Expected len(input.shape) == 3 (group,time,measure)") if input.shape[-1] != len(self.measures): @@ -792,8 +779,8 @@ def get_laplace_mvnorm(self, except (RuntimeError, ValueError) as e: warn( f"Unable to get valid covariance from optimized parameters (see error below)." - f"If you haven't already, fit the model with ``monitor_params=True`` (see the ``stopping`` argument" - f" of ``fit()``)." + f"If you haven't already tried, scale your data, and fit the model with ``monitor_params=True`` " + f"(see the ``stopping`` argument of ``fit()``)." f"\n{str(e)}" ) fake_cov = torch.diag(torch.diag(hess).pow(-1).clip(min=1E-5)) @@ -870,12 +857,29 @@ def __init__(self, self.kwargs = kwargs self.callable_kwargs = callable_kwargs self.get_loss = get_loss + self._bad_count = 0 + self._max_bad_count = self.optimizer.param_groups[0].get('max_eval', 10) def __call__(self): self.optimizer.zero_grad() self.kwargs.update({k: v() for k, v in self.callable_kwargs.items()}) - pred = self.ss_model(self.y, **self.kwargs) - loss = self.get_loss(pred, self.y) + + try: + pred = self.ss_model(self.y, **self.kwargs) + loss = self.get_loss(pred, self.y) + except torch.linalg.LinAlgError: + # linalgerror means bad covs. most common case is LBFGS line-search, which will respond to infinite loss + # by back-tracking and trying a different (hopefully more stable) parameter proposal. + # for simpler optimizers, will stall out and falsely converge, but that would have happened anyways. + self._bad_count += 1 + if self._bad_count > self._max_bad_count: + raise RuntimeError( + "Optimizer cannot find a region of param-space where all covs are valid. " + "Try again, potentially with a lower learning-rate." + ) + return torch.tensor(float('inf')) + self._bad_count = 0 + loss.backward() self.prog.update() self.prog.set_description( diff --git a/torchcast/utils/__init__.py b/torchcast/utils/__init__.py index b096fd2..f6012d4 100644 --- a/torchcast/utils/__init__.py +++ b/torchcast/utils/__init__.py @@ -9,11 +9,10 @@ * Simple trainer classes for PyTorch models, with specialized subclasses for torchcast's model-classes, as well as a special class for training neural networks to embed complex seasonal patterns into lower dimensional embeddings. * A 'Stopping' class for controlling convergence/stopping for the `fit()` method in state-space models. ---- """ from .features import add_season_features from .data import TimeSeriesDataset, TimeSeriesDataLoader, complete_times from .baseline import make_baseline -from .training import SimpleTrainer, StateSpaceTrainer, SeasonalEmbeddingsTrainer +from .training import SimpleTrainer, StateSpaceTrainer, ModelMatEmbeddingsTrainer, SeasonalEmbeddingsTrainer from .stopping import Stopping diff --git a/torchcast/utils/data.py b/torchcast/utils/data.py index a2fc203..7f7933c 100644 --- a/torchcast/utils/data.py +++ b/torchcast/utils/data.py @@ -94,6 +94,48 @@ def __repr__(self) -> str: kwargs.append("{}={!r}".format(k, v)) return "{}({})".format(type(self).__name__, ", ".join(kwargs)) + @torch.no_grad() + def standardize(self, + dataset: Optional['TimeSeriesDataset'] = None, + which: Sequence[int] = (1,)) -> 'TimeSeriesDataset': + """ + Take a TimeSeriesDataset and standardize its tensors. If no dataset is passed, a copy of this one is returned; + if a dataset *is* passed, the mean/std from this dataset is used (and again, a copy of the input is returned). + + :param dataset: A TimeSeriesDataset whose tensors will be standardized (centered and scaled by mean). + :param which: Which tensors to standardize (ints with 0-indexing). Defaults to (1,) (i.e. just the 2nd tensor). + :return: If no dataset was passed a copy of this dataset with standardized tensors. If a dataset was passed a + copy of that, standardized using the mean/std of *this* dataset. + + >>> ds_train, ds_val = ds.train_val_split() + >>> ds_train_std = ds_train.standardize() + >>> ds_val_std = ds_train.standardize(dataset=ds_val) # use train mean/std to avoid leakage + """ + + if isinstance(which, int): + which = (which,) + if not which: + which = range(len(self.tensors)) + which = list(which) + + if dataset is None: + dataset = self + else: + for i in which: + assert dataset.measures[i] == self.measures[i], "dataset-measures must match" + + tensors = [x.clone() for x in dataset.tensors] + for i in which: + # get mean/std from self: + m = self.tensors[i].nanmean([0, 1], keepdims=True) + s = np.nanstd(self.tensors[i].numpy(), axis=(0, 1), keepdims=True, ddof=1) + + # standardize cloned: + tensors[i] -= m + tensors[i] /= s + + return dataset.with_new_tensors(*tensors) + @property def sizes(self) -> Sequence: return [t.size() for t in self.tensors] @@ -529,19 +571,58 @@ class TimeSeriesDataLoader(DataLoader): def __init__(self, dataset: '_DataFrameGroupByDataset', batch_size: Optional[int], + dt_unit: Optional[str], + pad_X: Union[float, str, None] = 'ffill', + dtype: torch.dtype = torch.float32, + device: Optional[torch.device] = None, **kwargs): + + self.X_colnames = dataset.X_colnames + self.y_colnames = dataset.y_colnames + self.group_colname = dataset.group_colname + self.time_colname = dataset.time_colname + self.dt_unit = dt_unit + self.pad_X = pad_X + self.dtype = dtype + self.device = device + + if kwargs.pop('collate_fn', None) is not None: + raise TypeError( + f"{type(self).__name__} does not support custom `collate_fn`, please subclass and " + f"override ``_collate_fn()``" + ) + super().__init__( dataset=dataset, batch_size=batch_size, - collate_fn=self._collate, + collate_fn=self._collate_fn, **kwargs ) - @staticmethod - def _collate(batch: Sequence['TimeSeriesDataset']) -> 'TimeSeriesDataset': - if len(batch) == 1: - return batch[0] - raise NotImplementedError + def _collate_fn(self, batch: Sequence[pd.DataFrame]) -> 'TimeSeriesDataset': + df_group = pd.concat(batch).reset_index(drop=True) + X_colnames = dfX = None + if callable(self.X_colnames): + dfX = self.X_colnames(df_group) + if not dfX.index.equals(df_group.index): + raise ValueError("`X_colnames` function must return a dataframe with the same index as the input.") + elif self.X_colnames: + dfX = df_group[self.X_colnames] + df = df_group[[self.group_colname, self.time_colname] + list(self.y_colnames)] + if dfX is not None: + X_colnames = dfX.columns.tolist() + df = pd.concat([df, dfX], axis=1) + return TimeSeriesDataset.from_dataframe( + df, + group_colname=self.group_colname, + time_colname=self.time_colname, + dt_unit=self.dt_unit, + y_colnames=self.y_colnames, + X_colnames=X_colnames, + pad_X=self.pad_X, + dtype=self.dtype, + device=self.device + ) @classmethod def from_dataframe(cls, @@ -583,41 +664,32 @@ def from_dataframe(cls, df=dataframe, group_colname=group_colname, time_colname=time_colname, - dt_unit=dt_unit, y_colnames=measure_colnames or y_colnames, - X_colnames=X_colnames, - pad_X=pad_X, - device=device, - dtype=dtype + X_colnames=X_colnames ), + dt_unit=dt_unit, + pad_X=pad_X, + dtype=dtype, + device=device, **kwargs ) class _DataFrameGroupByDataset(Dataset): """ - Util class for ``TimeSeriesDataLoader``, to allow lazy (and so memory-efficient) transformations to the data before - converting to a ``TimeSeriesDataset``. + Util class for ``TimeSeriesDataLoader``. """ def __init__(self, df: pd.DataFrame, group_colname: str, time_colname: str, - dt_unit: Optional[str], y_colnames: Sequence[str], - X_colnames: Optional[Union[Sequence[str], Callable]], - pad_X: Union[float, str, None] = 'ffill', - dtype: torch.dtype = torch.float32, - device: Optional[torch.device] = None): + X_colnames: Optional[Union[Sequence[str], Callable]]): self.group_colname = group_colname self.time_colname = time_colname - self.dt_unit = dt_unit self.y_colnames = y_colnames self.X_colnames = X_colnames - self.pad_X = pad_X - self.dtype = dtype - self.device = device self.group_dfs = {g: dfg for g, dfg in df.groupby(self.group_colname, sort=False)} self.groups = df[group_colname].dropna().unique() @@ -627,36 +699,9 @@ def __len__(self) -> int: def __getitem__(self, idx) -> TimeSeriesDataset: raise NotImplementedError - def __getitems__(self, indices: list) -> Sequence[TimeSeriesDataset]: + def __getitems__(self, indices: list) -> Sequence[pd.DataFrame]: group_names = [self.groups[idx] for idx in indices] - df_group = pd.concat([self.group_dfs[group_name] for group_name in group_names]).reset_index(drop=True) - - # extract/create model-matrix: - X_colnames = dfX = None - if callable(self.X_colnames): - dfX = self.X_colnames(df_group) - if not dfX.index.equals(df_group.index): - raise ValueError("`X_colnames` function must return a dataframe with the same index as the input.") - elif self.X_colnames: - dfX = df_group[self.X_colnames] - - df = df_group[[self.group_colname, self.time_colname] + list(self.y_colnames)] - if dfX is not None: - X_colnames = dfX.columns.tolist() - df = pd.concat([df, dfX], axis=1) - - # return a batch of size 1 since this is what collate_fn expects, but basically it's already pre-collated - return [TimeSeriesDataset.from_dataframe( - df, - group_colname=self.group_colname, - time_colname=self.time_colname, - dt_unit=self.dt_unit, - y_colnames=self.y_colnames, - X_colnames=X_colnames, - pad_X=self.pad_X, - dtype=self.dtype, - device=self.device - )] + return [self.group_dfs[group_name] for group_name in group_names] def complete_times(data: 'DataFrame', diff --git a/torchcast/utils/outliers.py b/torchcast/utils/outliers.py index 8081118..d7e73e6 100644 --- a/torchcast/utils/outliers.py +++ b/torchcast/utils/outliers.py @@ -1,5 +1,4 @@ import torch -from torch.linalg import LinAlgError def get_outlier_multi(resid: torch.Tensor, diff --git a/torchcast/utils/stopping.py b/torchcast/utils/stopping.py index 3009040..c608716 100644 --- a/torchcast/utils/stopping.py +++ b/torchcast/utils/stopping.py @@ -29,11 +29,19 @@ def __init__(self, raise ValueError("At least one of `monitor_loss` or `monitor_params` must be True") self.monitor_loss = monitor_loss self.monitor_params = monitor_params - self.module = module self.abstol = abstol - self.values = [] self.patience = patience self.max_iter = max_iter + + self.module = None + self.values = None + self._patience_counter = None + self.last_change = None + self.reset(module) + + def reset(self, module: Optional[torch.nn.Module] = None): + self.module = module + self.values = [] self._patience_counter = 0 self.last_change = float('nan') diff --git a/torchcast/utils/training.py b/torchcast/utils/training.py index cd586ce..2a5d1df 100644 --- a/torchcast/utils/training.py +++ b/torchcast/utils/training.py @@ -14,7 +14,7 @@ from torch.utils.data import DataLoader, Dataset from torch.optim import Optimizer import torch.nn as nn -from typing import Generator, Union, Type, Sequence, Tuple, Dict, Optional, Callable +from typing import Generator, Union, Type, Sequence, Tuple, Dict, Optional, Callable, Any from tqdm.auto import tqdm @@ -44,10 +44,10 @@ def to(self, device) -> 'BaseTrainer': self._device = device return self - def _get_closure(self, batch: any, forward_kwargs: dict) -> callable: + def _get_closure(self, batch: Any, forward_kwargs: dict) -> Callable: raise NotImplementedError - def _get_batch_numel(self, batch: any) -> int: + def _get_batch_numel(self, batch: Any) -> int: raise NotImplementedError def __call__(self, @@ -106,11 +106,11 @@ class SimpleTrainer(BaseTrainer): def __init__(self, module: nn.Module, optimizer: Union[Optimizer, Type[Optimizer]] = torch.optim.Adam, - loss_fn: callable = torch.nn.MSELoss()): + loss_fn: Callable = torch.nn.MSELoss()): self.loss_fn = loss_fn super().__init__(module=module, optimizer=optimizer) - def _get_closure(self, batch: Dataset, forward_kwargs: dict) -> callable: + def _get_closure(self, batch: Dataset, forward_kwargs: dict) -> Callable: inputs, targets, *_other = batch if len(_other) and not self._warned: warnings.warn("Ignoring additional tensors in batch.") @@ -134,8 +134,8 @@ def _get_batch_numel(self, batch: Dataset) -> int: class StateSpaceTrainer(BaseTrainer): """ - A trainer for a :``StateSpaceModel``. This is for contexts in which the data are too large for - ``StateSpaceModel.fit()`` to be practical. Rather than the base DataLoader, this class takes a + A trainer for a :class:`torchcast.state_space.StateSpaceModel`. This is for contexts in which the data are too + large for ``StateSpaceModel.fit()`` to be practical. Rather than the base DataLoader, this class takes a :class:`torchcast.utils.TimeSeriesDataLoader`. Usage: @@ -196,7 +196,7 @@ def _batch_to_args(self, batch: TimeSeriesDataset) -> Tuple[torch.Tensor, dict]: kwargs[k] = t return y, kwargs - def _get_closure(self, batch: TimeSeriesDataset, forward_kwargs: dict) -> callable: + def _get_closure(self, batch: TimeSeriesDataset, forward_kwargs: dict) -> Callable: def closure(): # we call _batch_to_args from inside the closure in case `dataset_to_kwargs` is callable & involves grad. @@ -212,7 +212,7 @@ def closure(): return closure - def _get_batch_numel(self, batch: any) -> int: + def _get_batch_numel(self, batch: Any) -> int: return batch.tensors[0].numel() @@ -242,7 +242,7 @@ class ModelMatEmbeddingsTrainer(BaseTrainer): def __init__(self, module: nn.Module, - loss_fn: callable = torch.nn.MSELoss(), + loss_fn: Callable = torch.nn.MSELoss(), getX: Optional[Callable[[TimeSeriesDataset], torch.Tensor]] = None, **kwargs): @@ -322,7 +322,7 @@ def _solve_and_predict(cls, coefs = cls._l2_solve(y=y, X=X, prior_precision=prior_precision) return (coefs.transpose(-1, -2) * X).sum(-1).unsqueeze(-1) - def _get_closure(self, batch: TimeSeriesDataset, forward_kwargs: dict) -> callable: + def _get_closure(self, batch: TimeSeriesDataset, forward_kwargs: dict) -> Callable: X, y = self._getXy(batch) def closure(): @@ -335,7 +335,7 @@ def closure(): return closure - def _get_batch_numel(self, batch: any) -> int: + def _get_batch_numel(self, batch: Any) -> int: return batch.tensors[0].numel() def predict(self, batch: TimeSeriesDataset) -> torch.Tensor: @@ -360,7 +360,7 @@ def __init__(self, weekly: int, daily: int, other: Sequence[Tuple[np.timedelta64, int]] = (), - loss_fn: callable = torch.nn.MSELoss(), + loss_fn: Callable = torch.nn.MSELoss(), **kwargs): super().__init__(module=module, loss_fn=loss_fn, **kwargs)