Test test_fit_precision_cholesky_approximate

Author: dafeda

tests/test_precision_estimation.py

@@ -1,3 +1,4 @@
+import networkx as nx
 import numpy as np
 import pytest
 import scipy as sp
@@ -55,6 +56,92 @@ def test_closed_form_matches_iterative_solver():
     np.testing.assert_allclose(diag_cf, diag_iter, rtol=1e-4)
 
 
+def test_fit_precision_cholesky_approximate_two_hop_recovers_fill_in():
+    rng = np.random.default_rng(42)
+    n = 50000
+
+    # The true precision used below is sparse on the 4-cycle graph: it only
+    # couples immediate graph neighbors (0-1, 1-2, 2-3, 3-0), while the
+    # opposite corners (0, 2) and (1, 3) are exactly zero.
+    #
+    # The method under test does not fit the precision entries directly. It
+    # fits a lower-triangular factor C such that Prec ~= C.T @ C, and the
+    # allowed nonzeros in each row of C come from the graph neighborhood.
+    #
+    # For this matrix, in the natural ordering, the exact factor used by this
+    # code path is approximately
+    #
+    #   [[ 1.3615,  0.    ,  0.    ,  0.    ],
+    #    [-0.3112,  1.3648,  0.    ,  0.    ],
+    #    [-0.0667, -0.3706,  1.3491,  0.    ],
+    #    [-0.2121,  0.    , -0.4243,  1.4142]]
+    #
+    # so C[2, 0] is a nonzero fill-in term even though Prec[0, 2] = 0.
+    # Intuitively, row 3 already creates an indirect 0-2 coupling inside
+    # C.T @ C, so row 2 needs a compensating C[2, 0] to cancel it. That
+    # nonzero C[2, 0] is exactly the fill-in that the 2-hop expansion must
+    # allow. Nodes 0 and 2 are not adjacent in the original cycle, but they
+    # are two hops apart (0 -> 1 -> 2, or 0 -> 3 -> 2).
+    #
+    # Algebraically,
+    #
+    #   (C.T @ C)[0, 2] = sum_r C[r, 0] * C[r, 2].
+    #
+    # Because C is lower triangular, only rows 2 and 3 can contribute, so
+    #
+    #   (C.T @ C)[0, 2] = C[2, 0] * C[2, 2] + C[3, 0] * C[3, 2].
+    #
+    # Row 3 contributes
+    #
+    #   C[3, 0] * C[3, 2] = (-0.3 / sqrt(2)) * (-0.6 / sqrt(2)) = 0.09.
+    #
+    # But the true precision entry (0, 2) is zero, so row 2 must cancel that:
+    #
+    #   C[2, 0] * C[2, 2] = -0.09.
+    #
+    # The diagonal entry Prec[2, 2] = 2 fixes C[2, 2] through
+    #
+    #   C[2, 2]^2 + C[3, 2]^2 = 2,
+    #
+    # giving C[2, 2] = sqrt(2 - (-0.6 / sqrt(2)) ** 2) = sqrt(1.82) ~= 1.3491,
+    # and therefore C[2, 0] = -0.09 / 1.3491 ~= -0.0667.
+    #
+    # With only the 1-hop graph, row 2 can use column 1 but not column 0, so
+    # the solver cannot represent the true Cholesky sparsity pattern. With a
+    # 2-hop expansion, the missing 0-2 connection becomes available, so the
+    # solver can include that fill-in term and recover the precision much more
+    # accurately.
+    prec_true = np.array(
+        [
+            [2.0, -0.4, 0.0, -0.3],
+            [-0.4, 2.0, -0.5, 0.0],
+            [0.0, -0.5, 2.0, -0.6],
+            [-0.3, 0.0, -0.6, 2.0],
+        ]
+    )
+    cov_true = np.linalg.inv(prec_true)
+    U = rng.multivariate_normal(np.zeros(4), cov_true, size=n)
+    G = nx.cycle_graph(4)
+
+    prec_1_hop = precest.fit_precision_cholesky_approximate(
+        U=U, G=G, neighbourhood_expansion=1, use_tqdm=False
+    ).toarray()
+    prec_2_hop = precest.fit_precision_cholesky_approximate(
+        U=U, G=G, neighbourhood_expansion=2, use_tqdm=False
+    ).toarray()
+
+    # Without the 2-hop expansion, the solver cannot represent the fill-in and
+    # leaks a spurious (0, 2) precision entry.
+    assert abs(prec_1_hop[0, 2]) > 0.07
+    assert abs(prec_2_hop[0, 2]) < 0.02
+
+    err_1_hop = np.max(np.abs(prec_1_hop - prec_true))
+    err_2_hop = np.max(np.abs(prec_2_hop - prec_true))
+    assert err_1_hop > err_2_hop
+    assert err_1_hop > 0.07
+    assert err_2_hop < 0.03
+
+
 if __name__ == "__main__":
     import pytest