Merge pull request #30 from strongio/feature/bern-filter

binomial filter

torchcast/__init__.py

@@ -1 +1 @@
-__version__ = '0.4.3'
+__version__ = '0.5.0'

torchcast/covariance/base.py

@@ -8,7 +8,7 @@
 from torch import Tensor, nn, jit
 
 from torchcast.process.utils import Identity
-from torchcast.covariance.util import num_off_diag
+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.base import Process
 
@@ -80,7 +80,7 @@ def from_processes(cls,
             if 'init_diag_multi' not in kwargs:
                 kwargs['init_diag_multi'] = .01
             if 'method' in kwargs and kwargs['method'] == 'low_rank':
-                warn("``method='low_rank'`` not recommended for processes, consider 'low_rank+block_diag'")
+                warn("``method='low_rank'`` not recommended for processes")
         elif cov_type == 'initial':
             if (state_rank - len(no_cov_idx)) >= 10:
                 # by default, use low-rank parameterization for initial cov:
@@ -173,12 +173,7 @@ def __init__(self,
         empty_idx = set(empty_idx)
         assert all(isinstance(x, int) for x in empty_idx)
         self.param_rank = self.rank - len(empty_idx)
-        mask = torch.zeros((self.rank, self.param_rank))
-        c = 0
-        for r in range(self.rank):
-            if r not in empty_idx:
-                mask[r, c] = 1.
-                c += 1
+        mask = mini_cov_mask(rank=self.rank, empty_idx=empty_idx)
         self.register_buffer('mask', mask)
 
         self._set_params(method, init_diag_multi)

torchcast/covariance/util.py

@@ -1,3 +1,5 @@
+from typing import Collection
+
 import torch
 from torch import Tensor
 
@@ -9,3 +11,14 @@ def num_off_diag(rank: int) -> int:
 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