Skip to content

Commit 97c150d

Browse files
BlockCovariance: compute precision_ block-wise (O(p*b^2)), not a dense O(p^3) inverse (#54)
The inherited base precision_ densified the full p x p covariance and inverted it globally, defeating the estimator's memory/compute thesis (and contradicting its own docstring). Override precision_ to invert each small block and assemble block-diagonally -- positive-definite by construction, never forming or inverting the dense matrix. Factor the shared per-block PD construction into _pd_block. Add a test: precision_ is block-diagonal, PD, and a genuine inverse of the covariance (P @ C = I). Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
1 parent 7a651bd commit 97c150d

2 files changed

Lines changed: 31 additions & 1 deletion

File tree

precise/block_covariance.py

Lines changed: 17 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -60,11 +60,27 @@ def _update_state(self, s: dict, x: np.ndarray) -> dict:
6060
subs.append(ewa_update(sub, x[a:b]))
6161
return {**s, "n_samples": s["n_samples"] + 1, "subs": subs}
6262

63+
@staticmethod
64+
def _pd_block(sub: dict) -> np.ndarray:
65+
return make_pos_def(to_symmetric(np.asarray(sub["cov"], dtype=float)))
66+
6367
def _state_to_cov(self, state: dict) -> np.ndarray:
6468
p = state["n_dim"]
6569
out = np.zeros((p, p))
6670
for sub, (a, b) in zip(state["subs"], state["slices"]):
67-
out[a:b, a:b] = make_pos_def(to_symmetric(np.asarray(sub["cov"], dtype=float)))
71+
out[a:b, a:b] = self._pd_block(sub)
72+
return out
73+
74+
@property
75+
def precision_(self) -> np.ndarray:
76+
"""Block-wise inverse: invert each (small) block and assemble block-diagonally, so the
77+
precision is obtained in ``O(p*b^2)`` without ever forming or inverting the dense ``p x p``
78+
matrix (the base class would densify and invert globally, ``O(p^3)``)."""
79+
state = self._fitted_state()
80+
p = state["n_dim"]
81+
out = np.zeros((p, p))
82+
for sub, (a, b) in zip(state["subs"], state["slices"]):
83+
out[a:b, a:b] = np.linalg.inv(self._pd_block(sub))
6884
return out
6985

7086
def _state_to_mean(self, state: dict) -> np.ndarray:

tests/test_block_covariance.py

Lines changed: 14 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -45,6 +45,20 @@ def test_state_is_subquadratic():
4545
assert "block_covs" in state and len(state["block_covs"]) == nb
4646

4747

48+
def test_precision_is_block_diagonal_inverse():
49+
# precision_ must be the block-wise inverse: block-diagonal, PD, and P@C=I (not a dense inverse)
50+
p, nb = 24, 4
51+
e = BlockCovariance(n_blocks=nb, r=0.05).fit(_data(p=p))
52+
P, C = e.precision_, e.covariance_
53+
block_id = np.empty(p, dtype=int)
54+
for bi, idx in enumerate(np.array_split(np.arange(p), nb)):
55+
block_id[idx] = bi
56+
off_block = block_id[:, None] != block_id[None, :]
57+
assert np.all(P[off_block] == 0.0) # precision is block-diagonal
58+
assert np.all(np.linalg.eigvalsh(P) > 0) # positive-definite
59+
assert np.allclose(P @ C, np.eye(p), atol=1e-8) # genuine inverse of the covariance
60+
61+
4862
def test_state_is_json_serializable_and_roundtrips():
4963
e = BlockCovariance(n_blocks=4).fit(_data())
5064
state = e.get_state()

0 commit comments

Comments
 (0)