BlockCovariance: block-wise precision_ (fix O(p^3) densify)

precise/block_covariance.py

@@ -60,11 +60,27 @@ def _update_state(self, s: dict, x: np.ndarray) -> dict:
             subs.append(ewa_update(sub, x[a:b]))
         return {**s, "n_samples": s["n_samples"] + 1, "subs": subs}
 
+    @staticmethod
+    def _pd_block(sub: dict) -> np.ndarray:
+        return make_pos_def(to_symmetric(np.asarray(sub["cov"], dtype=float)))
+
     def _state_to_cov(self, state: dict) -> np.ndarray:
         p = state["n_dim"]
         out = np.zeros((p, p))
         for sub, (a, b) in zip(state["subs"], state["slices"]):
-            out[a:b, a:b] = make_pos_def(to_symmetric(np.asarray(sub["cov"], dtype=float)))
+            out[a:b, a:b] = self._pd_block(sub)
+        return out
+
+    @property
+    def precision_(self) -> np.ndarray:
+        """Block-wise inverse: invert each (small) block and assemble block-diagonally, so the
+        precision is obtained in ``O(p*b^2)`` without ever forming or inverting the dense ``p x p``
+        matrix (the base class would densify and invert globally, ``O(p^3)``)."""
+        state = self._fitted_state()
+        p = state["n_dim"]
+        out = np.zeros((p, p))
+        for sub, (a, b) in zip(state["subs"], state["slices"]):
+            out[a:b, a:b] = np.linalg.inv(self._pd_block(sub))
         return out
 
     def _state_to_mean(self, state: dict) -> np.ndarray:

tests/test_block_covariance.py

@@ -45,6 +45,20 @@ def test_state_is_subquadratic():
     assert "block_covs" in state and len(state["block_covs"]) == nb
 
 
+def test_precision_is_block_diagonal_inverse():
+    # precision_ must be the block-wise inverse: block-diagonal, PD, and P@C=I (not a dense inverse)
+    p, nb = 24, 4
+    e = BlockCovariance(n_blocks=nb, r=0.05).fit(_data(p=p))
+    P, C = e.precision_, e.covariance_
+    block_id = np.empty(p, dtype=int)
+    for bi, idx in enumerate(np.array_split(np.arange(p), nb)):
+        block_id[idx] = bi
+    off_block = block_id[:, None] != block_id[None, :]
+    assert np.all(P[off_block] == 0.0)                       # precision is block-diagonal
+    assert np.all(np.linalg.eigvalsh(P) > 0)                 # positive-definite
+    assert np.allclose(P @ C, np.eye(p), atol=1e-8)          # genuine inverse of the covariance
+
+
 def test_state_is_json_serializable_and_roundtrips():
     e = BlockCovariance(n_blocks=4).fit(_data())
     state = e.get_state()