Performance optimizations

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/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)

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,

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,

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)

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)