go back to precomputing measure-scaling

Author: jwdink

torchcast/exp_smooth/exp_smooth.py

@@ -22,24 +22,23 @@ def __init__(self,
                  measure_covariance: Optional[Covariance] = None,
                  smoothing_matrix: Optional[SmoothingMatrix] = None,
                  measure_funs: Optional[dict[str, str]] = None,
-                 adaptive_measure_var: bool = False):
+                 adaptive_scaling: bool = False):
 
         super().__init__(
             processes=processes,
             measures=measures,
             measure_covariance=measure_covariance,
             measure_funs=measure_funs,
-            adaptive_measure_var=adaptive_measure_var,
+            adaptive_scaling=adaptive_scaling,
         )
         if smoothing_matrix is None:
             smoothing_matrix = SmoothingMatrix.from_measures_and_processes(measures=measures, processes=processes)
         self.smoothing_matrix = smoothing_matrix.set_id('smoothing_matrix')
 
     def initial_covariance(self, inputs: dict, num_groups: int, num_times: int, _ignore_input: bool = False) -> Tensor:
-        # initial covariance is always zero. this will be replaced by the 1-step-ahead covariance in the first call to
-        # predict
-        ms = self._get_measure_scaling()
-        return torch.zeros((num_groups, num_times, self.state_rank, self.state_rank), dtype=ms.dtype, device=ms.device)
+        # initial covariance is always zero. this will be replaced by the 1-step covariance in the first call to predict
+        m = list(self.processes.values())[0].initial_mean  # get a parameter, any parameter, to get device
+        return torch.zeros((num_groups, num_times, self.state_rank, self.state_rank), dtype=m.dtype, device=m.device)
 
     def _mask_mats(self,
                    groups: torch.Tensor,
@@ -105,10 +104,14 @@ def _predict_cov(self,
                      cov: torch.Tensor,
                      transition_mat: torch.Tensor,
                      cov1step: torch.Tensor,
+                     scaling: Optional[torch.Tensor] = None,
                      mask: Optional[torch.Tensor] = None) -> torch.Tensor:
         # new_cov will at least be cov1step (see note above in _update_step)
         new_cov = cov1step
 
+        if scaling is not None:
+            raise NotImplementedError
+
         # fastpath: if the call to update returned the zero-dim tensor (see _update above) then we are done
         if len(cov.shape):
             if mask is None or mask.all():

torchcast/kalman_filter/binomial_filter.py

@@ -22,7 +22,7 @@ def __init__(self,
                  measure_covariance: Optional[Covariance] = None,
                  process_covariance: Optional[Covariance] = None,
                  initial_covariance: Optional[Covariance] = None,
-                 adaptive_measure_var: bool = False):
+                 adaptive_scaling: bool = False):
 
         if binary_measures is None:
             binary_measures = list(measures)
@@ -41,7 +41,7 @@ def __init__(self,
             process_covariance=process_covariance,
             measure_covariance=measure_covariance,
             initial_covariance=initial_covariance,
-            adaptive_measure_var=adaptive_measure_var,
+            adaptive_scaling=adaptive_scaling,
             measure_funs={m: 'ilogit' for m in binary_measures},
         )
 

torchcast/kalman_filter/kalman_filter.py

@@ -24,7 +24,7 @@ def __init__(self,
                  process_covariance: Optional[Covariance] = None,
                  initial_covariance: Optional[Covariance] = None,
                  measure_funs: Optional[dict[str, str]] = None,
-                 adaptive_measure_var: bool = False):
+                 adaptive_scaling: bool = False):
 
         if initial_covariance is None:
             initial_covariance = Covariance.from_processes(processes, cov_type='initial')
@@ -37,7 +37,7 @@ def __init__(self,
             measures=measures,
             measure_covariance=measure_covariance,
             measure_funs=measure_funs,
-            adaptive_measure_var=adaptive_measure_var,
+            adaptive_scaling=adaptive_scaling,
         )
         self.process_covariance = process_covariance.set_id('process_covariance')
         self.initial_covariance = initial_covariance.set_id('initial_covariance')
@@ -46,6 +46,7 @@ def _predict_cov(self,
                      cov: torch.Tensor,
                      transition_mat: torch.Tensor,
                      Q: torch.Tensor,
+                     scaling: Optional[torch.Tensor] = None,
                      mask: Optional[torch.Tensor] = None) -> torch.Tensor:
         if mask is None or mask.all():
             mask = slice(None)
@@ -100,20 +101,24 @@ def _parse_kwargs(self,
         )
 
         # process-variance:
-        measure_scaling = torch.diag_embed(self._get_measure_scaling().unsqueeze(0))
         pcov_kwargs = {}
         if self.process_covariance.expected_kwargs:
             pcov_kwargs = {k: kwargs[k] for k in self.process_covariance.expected_kwargs}
         used_keys |= set(pcov_kwargs)
+
+        mcov = self.measure_covariance({}, num_groups=1, num_times=1, _ignore_input=True)[0, 0]
+        measure_std = mcov.diagonal(dim1=-2, dim2=-1).sqrt()
+        for idx in self.measure_covariance.empty_idx:
+            measure_std[idx] = torch.ones_like(measure_std[idx])  # empty measures have no variance, so set to 1
+
         if pcov_kwargs:
-            measure_scaling = measure_scaling.unsqueeze(0)
             pcov_raw = self.process_covariance(pcov_kwargs, num_groups=num_groups, num_times=num_timesteps)
-            Qs = measure_scaling @ pcov_raw @ measure_scaling
+            Qs = self._apply_cov_scaling(pcov_raw, scaling=measure_std, is_process_cov=True)
             predict_kwargs['Q'] = Qs.unbind(1)
         else:
             # faster if not time-varying
-            pcov_raw = self.process_covariance(pcov_kwargs, num_groups=num_groups, num_times=1)
-            Qs = measure_scaling @ pcov_raw.squeeze(1) @ measure_scaling
+            pcov_raw = self.process_covariance(pcov_kwargs, num_groups=num_groups, num_times=1).squeeze(1)
+            Qs = self._apply_cov_scaling(pcov_raw, scaling=measure_std, is_process_cov=True)
             predict_kwargs['Q'] = [Qs] * num_timesteps
 
         return predict_kwargs, update_kwargs, used_keys

…cast/state_space/adaptive_measure_var.py torchcast/state_space/adaptive_scaling.pytorchcast/state_space/adaptive_measure_var.py renamed to torchcast/state_space/adaptive_scaling.py torchcast/state_space/adaptive_measure_var.py renamed to torchcast/state_space/adaptive_scaling.py

@@ -12,17 +12,17 @@
 from torchcast.process.utils import Bounded
 
 
-class AdaptiveMeasureVar(nn.Module):
+class AdaptiveScaler(nn.Module):
     def reset(self):
         raise NotImplementedError
 
     def forward(self, residuals: torch.Tensor, skip_mask: torch.Tensor) -> torch.Tensor:
         raise NotImplementedError
 
 
-class EWMAdaptiveMeasureVar(AdaptiveMeasureVar):
+class EWMAdaptiveScaler(AdaptiveScaler):
     """
-    Exponentially Weighted Moving Average (EWM) based adaptive measure variance.
+    Exponentially Weighted Moving Average (EWM) based adaptive scaling.
     """
 
     def __init__(self,