Sampling arbitrary Precision matrices, not just diagonal (#164)

Author: tommyod

graphite_maps/utils.py

@@ -1,8 +1,10 @@
 import logging
 
 import numpy as np
+import scipy as sp
 from numpy.typing import NDArray
 from scipy.sparse import sparray
+from sksparse.cholmod import cholesky
 
 log = logging.getLogger(__name__)
 log.addHandler(logging.NullHandler())
@@ -25,28 +27,59 @@ def generate_gaussian_noise(
     -------
     np.ndarray
         The Gaussian noise array of shape (n, m), where Prec has shape (m, m).
+
+    Examples
+    --------
+    >>> F = sp.sparse.random_array(shape=(6, 6), density=0.3, random_state=42)
+    >>> Prec = sp.sparse.csc_array(F.T @ F + np.diag([1, 2, 4, 8, 16, 32]))
+    >>> generate_gaussian_noise(n=3, Prec=Prec, seed=1).round(2)
+    array([[ 0.35,  0.19, -0.33, -0.19, -0.32,  0.1 ],
+           [ 0.82,  0.02, -0.07,  0.2 ,  0.23,  0.01],
+           [ 0.33,  0.36, -0.19,  0.11,  0.11, -0.05]])
+
+    >>> Prec = sp.sparse.diags_array(Prec.diagonal())
+    >>> generate_gaussian_noise(n=3, Prec=Prec, seed=2).round(2)
+    array([[ 0.19, -0.34, -0.17, -0.84,  0.45,  0.2 ],
+           [-0.33,  0.51,  0.12, -0.19,  0.24, -0.05],
+           [-0.33, -0.52,  0.19, -0.03,  0.14, -0.1 ]])
     """
     if n < 1:
         raise ValueError(f"`n` should be g.e. 1, got {n}")
 
     m = Prec.shape[0]
     rng = np.random.default_rng(seed)
-    eps = rng.normal(
-        loc=0,
-        scale=np.sqrt(np.linalg.inv(Prec.toarray()).diagonal()),
-        size=(n, m),
-    )
 
-    # standard_normal_samples = rng.normal(size=(n, m))
-    # cholesky_factor = cholesky(Prec)
+    # If precision matrix is diagonal
+    row_idx, col_idx, _ = sp.sparse.find(Prec)
+    if np.all(row_idx == col_idx):
+        # Scale is the standard deviations. diag(Prec) = 1 / variances
+        scale = np.sqrt(1 / Prec.diagonal())
+        return rng.normal(
+            loc=0,
+            scale=scale,
+            size=(n, m),
+        )
+
+    # General case: precision matrix is not diagonal
+    z = rng.normal(size=(m, n))
+
+    # The simplest, but naive, way to sample is to compute:
+    # cov = np.linalg.inv(Prec.todense())
+    # C = np.linalg.cholesky(cov)
+    # assert np.allclose(C @ C.T, cov)
+    # return (np.linalg.cholesky(cov) @ z).T
 
-    # Transform the samples using the inverse Cholesky factor
-    # This transformation results in samples from N(0, Prec^-1)
-    # eps = cholesky_factor.solve_A(standard_normal_samples.T).T
+    # Suppose C @ C.T = Cov, then our samples are eps = C @ z.
+    # If we take cholesky of Prec, we obtain L @ L.T = P @ Prec @ P.T.
+    # Rearrange to (P.T @ L) @ (P.T @ L).T = Prec, then take inverse
+    # to see that C = inv((P.T @ L).T). Thus the equation eps = C @ z
+    # becomes the system (P.T @ L).T @ eps = z, which we solve for eps below
+    factor = cholesky(Prec)
+    v = factor.solve_Lt(z, use_LDLt_decomposition=False)
+    return factor.apply_Pt(v).T
 
-    assert eps.shape == (n, m), "Sampling returns wrong size"
 
-    log.info("Sampling with seed=%s", seed)
-    log.info("Sampled Gaussian noise with shape=%s", eps.shape)
+if __name__ == "__main__":
+    import pytest
 
-    return eps
+    pytest.main(args=[__file__, "--doctest-modules", "-v"])