Snapshot test for precision estimation (#160)

* Consistent args in prec-est. Wrote snapshot test
* Use natural instead of metis ordering - for reproducibility

Author: tommyod

graphite_maps/precision_estimation.py

@@ -482,7 +482,7 @@ def fit_precision_cholesky(
 
 def fit_precision_cholesky_approximate(
     U: NDArray[np.floating],
-    G: nx.Graph,
+    Graph_u: nx.Graph,
     neighbourhood_expansion: int = 2,
     use_tqdm: bool = True,
 ) -> csc_array:
@@ -512,7 +512,7 @@ def fit_precision_cholesky_approximate(
         The optimized sparse Cholesky factor of the precision matrix.
     """
     # The k'th power is a "neighborhood expansion", see docs of nx.power
-    G = nx.power(G, k=neighbourhood_expansion)
+    G = nx.power(Graph_u, k=neighbourhood_expansion)
     G.add_edges_from((n, n) for n in G.nodes())  # Add edges (1, 1), (2, 2), ...
 
     C = optimize_sparse_affine_kr_map(

tests/test_enif.py

@@ -116,7 +116,7 @@ def run(U_in):
         # Estimate precision matrix
         Prec_u = fit_precision_cholesky_approximate(
             U=U_in,
-            G=Graph_u,
+            Graph_u=Graph_u,
             use_tqdm=False,
         )
         enif = EnIF(Prec_u=Prec_u, Prec_eps=Prec_eps, H=H)
@@ -167,7 +167,7 @@ def test_snapshot_lowlevel():
     # Estimate precision matrix
     Prec_u = fit_precision_cholesky_approximate(
         U=U,
-        G=Graph_u,
+        Graph_u=Graph_u,
         neighbourhood_expansion=2,
         use_tqdm=True,
     )

tests/test_precision_estimation.py

@@ -1,10 +1,71 @@
+import networkx as nx
 import numpy as np
 import pytest
 import scipy as sp
 from graphite_maps import precision_estimation as precest
 from scipy.optimize import minimize
 
 
+def get_precision_data():
+    rng = np.random.default_rng(8)
+    n = 10  # Size
+    density = 0.4
+
+    # Create G indicating sparsity pattern
+    G = rng.uniform(size=(n, n)) < (density / 2)
+    G = G.T + G
+    np.fill_diagonal(G, G.diagonal() + 1)
+    G_matrix = (G > 0).astype(int)
+    Graph_u = nx.from_scipy_sparse_array(sp.sparse.csc_array(G_matrix))
+
+    # Create data U
+    U = rng.normal(size=(999, n))
+
+    return U, Graph_u, G_matrix
+
+
+def test_snapshot_fit_precision_cholesky():
+
+    U, Graph_u, G_matrix = get_precision_data()
+
+    # Estimate precision with fit_precision_cholesky.
+    # Cannot use METIS (not reproducible across OSes), use 'natural'
+    Prec_est, *_ = precest.fit_precision_cholesky(
+        U=U, Graph_u=Graph_u, ordering_method="natural"
+    )
+    Prec_est = Prec_est.todense()
+
+    entries_at_one = Prec_est[G_matrix > 0]
+    entries_at_zero = Prec_est[G_matrix == 0]
+
+    desired = np.array([1.03393041, -0.05494438, 1.02092228, 1.00923821])
+    np.testing.assert_allclose(entries_at_one[::9], desired, atol=1e-8)
+
+    desired = np.array([0.0, 0.0, 0.0, 0.01763788, -0.03135188, 0.0, 0.0, 0.0])
+    np.testing.assert_allclose(entries_at_zero[::9], desired, atol=1e-8)
+
+
+def test_snapshot_fit_precision_cholesky_approximate():
+
+    U, Graph_u, G_matrix = get_precision_data()
+    # Estimate precision with fit_precision_cholesky
+    Prec_est = precest.fit_precision_cholesky_approximate(
+        U=U, Graph_u=Graph_u, neighbourhood_expansion=2
+    )
+    Prec_est = Prec_est.todense()
+
+    entries_at_one = Prec_est[G_matrix > 0]
+    entries_at_zero = Prec_est[G_matrix == 0]
+
+    desired = np.array([1.03392773, -0.05606645, 1.022626, 1.00883855])
+    np.testing.assert_allclose(entries_at_one[::9], desired, atol=1e-8)
+
+    desired = np.array(
+        [0.0, -0.04588378, -0.00330536, -0.00124644, -0.0301345, -0.02301515, 0.0, 0.0]
+    )
+    np.testing.assert_allclose(entries_at_zero[::9], desired, atol=1e-8)
+
+
 def test_objective_twice():
     # A regression test: ensure that two calls return the same result.
     rng = np.random.default_rng(42)