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/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..4baeba8 100644 --- a/torchcast/internals/batch_design/measurement_model.py +++ b/torchcast/internals/batch_design/measurement_model.py @@ -40,7 +40,7 @@ def __init__(self, def is_nonlinear(self) -> bool: return bool(self.nonlinear_processes) or 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 +53,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: @@ -196,8 +197,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 +215,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/kalman_filter/kalman_filter.py b/torchcast/kalman_filter/kalman_filter.py index f092d45..c21ca9f 100644 --- a/torchcast/kalman_filter/kalman_filter.py +++ b/torchcast/kalman_filter/kalman_filter.py @@ -110,6 +110,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/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)