Merge pull request DifferentiableUniverseInitiative#58 from dkirkby/s…

…parse Implement optional sparse Gaussian covariance
LuzGarciaP · Jul 26, 2020 · 060ef0e · 060ef0e
2 parents 99edb4f + 5fc49b6
commit 060ef0e
Show file tree

Hide file tree

Showing 9 changed files with 1,741 additions and 1,161 deletions.
diff --git a/design.md b/design.md
@@ -150,17 +150,18 @@ with code styling again. Here are the steps to follow:
 - Install `black` and `pre-commit`:
 ```bash
 $ pip install --user black pre-commit reorder_python_imports
+$ pre-commit install
 ```
-`pre-commit` will be tasked with automatically running `black` formatting
-whenever you commit some code.
+`pre-commit` will be tasked with automatically running `black` and `reorder_python_imports` formatting
+whenever you commit some code. The import guidelines are documented [here](https://github.com/asottile/reorder_python_imports#what-does-it-do).
 
 - Manually running black formatting:
 ```bash
 $ black .
 ```
 from the root directory.
 
-- Automatically running black at each commit: You actually have nothing
+- Automatically running `black` and `reorder_python_imports` at each commit: You actually have nothing
 else to do. If pre-commit is installed it will happen automatically for
 you.
 

diff --git a/docs/notebooks/jax-cosmo-intro.ipynb b/docs/notebooks/jax-cosmo-intro.ipynb
diff --git a/jax_cosmo/__init__.py b/jax_cosmo/__init__.py
@@ -22,3 +22,4 @@
 import jax_cosmo.transfer as transfer
 from jax_cosmo.core import *
 from jax_cosmo.parameters import *
+import jax_cosmo.sparse as sparse
diff --git a/jax_cosmo/angular_cl.py b/jax_cosmo/angular_cl.py
@@ -122,14 +122,19 @@ def get_noise_cl(inds):
     return lax.map(get_noise_cl, cl_index)
 
 
-def gaussian_cl_covariance(ell, probes, cl_signal, cl_noise, f_sky=0.25):
+def gaussian_cl_covariance(ell, probes, cl_signal, cl_noise, f_sky=0.25, sparse=True):
     """
     Computes a Gaussian covariance for the angular cls of the provided probes
 
+    Set sparse True to return a sparse matrix representation that uses a factor
+    of n_ell less memory and is compatible with the linear algebra operations
+    in :mod:`jax_cosmo.sparse`.
+
     return_cls: (returns covariance)
     """
     ell = np.atleast_1d(ell)
     n_ell = len(ell)
+    one = 1.0 if sparse else np.eye(n_ell)
 
     # Adding noise to auto-spectra
     cl_obs = cl_signal + cl_noise
@@ -144,15 +149,22 @@ def gaussian_cl_covariance(ell, probes, cl_signal, cl_noise, f_sky=0.25):
     def get_cov_block(inds):
         a, b, c, d = inds
         cov = (cl_obs[a] * cl_obs[b] + cl_obs[c] * cl_obs[d]) / norm
-        return cov * np.eye(n_ell)
+        return cov * one
 
+    # Return a sparse representation of the matrix containing only the diagonals
+    # for each of the n_cls x n_cls blocks of size n_ell x n_ell.
+    # We could compress this further using the symmetry of the blocks, but
+    # it is easier to invert this matrix with this redundancy included.
     cov_mat = lax.map(get_cov_block, cov_blocks)
 
     # Reshape covariance matrix into proper matrix
-    cov_mat = cov_mat.reshape((n_cls, n_cls, n_ell, n_ell))
-    cov_mat = cov_mat.transpose(axes=(0, 2, 1, 3)).reshape(
-        (n_ell * n_cls, n_ell * n_cls)
-    )
+    if sparse:
+        cov_mat = cov_mat.reshape((n_cls, n_cls, n_ell))
+    else:
+        cov_mat = cov_mat.reshape((n_cls, n_cls, n_ell, n_ell))
+        cov_mat = cov_mat.transpose(axes=(0, 2, 1, 3)).reshape(
+            (n_ell * n_cls, n_ell * n_cls)
+        )
     return cov_mat
 
 
@@ -163,10 +175,15 @@ def gaussian_cl_covariance_and_mean(
     transfer_fn=tklib.Eisenstein_Hu,
     nonlinear_fn=power.halofit,
     f_sky=0.25,
+    sparse=False,
 ):
     """
     Computes a Gaussian covariance for the angular cls of the provided probes
 
+    Set sparse True to return a sparse matrix representation that uses a factor
+    of n_ell less memory and is compatible with the linear algebra operations
+    in :mod:`jax_cosmo.sparse`.
+
     return_cls: (returns signal + noise cl, covariance)
     """
     ell = np.atleast_1d(ell)
@@ -179,6 +196,6 @@ def gaussian_cl_covariance_and_mean(
     cl_noise = noise_cl(ell, probes)
 
     # retrieve the covariance
-    cov_mat = gaussian_cl_covariance(ell, probes, cl_signal, cl_noise, f_sky)
+    cov_mat = gaussian_cl_covariance(ell, probes, cl_signal, cl_noise, f_sky, sparse)
 
     return cl_signal.flatten(), cov_mat
diff --git a/jax_cosmo/likelihood.py b/jax_cosmo/likelihood.py
@@ -6,27 +6,60 @@
 import jax.numpy as np
 import jax.scipy as sp
 
+import jax_cosmo.sparse as sparse
 from jax_cosmo.angular_cl import gaussian_cl_covariance
 
 
 def gaussian_log_likelihood(data, mu, C, constant_cov=True, inverse_method="inverse"):
     """
-    Computes the likelihood for some cl
+    Computes the log likelihood for a given data vector under a multivariate
+    Gaussian distribution.
+
+    If the covariance C is sparse (according to :meth:`jax_cosmo.sparse.is_sparse`)
+    use sparse inverse and determinant algorithms (and ignore ``inverse_method``).
+
+    Parameters
+    ----------
+    data: array_like
+        Data vector, with shape [N].
+
+    mu: array_like, 1d
+        Mean of the Gaussian likelihood, with shape [N].
+
+    C: array_like or sparse matrix
+        Covariance of Gaussian likelihood with shape [N,N]
+
+    constant_cov: boolean
+        Whether to include the log determinant of the covariance matrix in the
+        likelihood. If `constant_cov` is true, the log determinant is ignored
+        (default: True)
+
+    inverse_method: string
+        Methods for computing the precision matrix. Either "inverse", "cholesky".
+        Note that this option is ignored when the covariance is sparse. (default: "inverse")
     """
     # Computes residuals
     r = mu - data
 
-    # TODO: check what is the fastest and works the best between cholesky+solve
-    # and just inversion
-    if inverse_method == "inverse":
-        y = np.dot(np.linalg.inv(C), r)
-    elif inverse_method == "cholesky":
-        y = sp.linalg.cho_solve(sp.linalg.cho_factor(C, lower=True), r)
+    if sparse.is_sparse(C):
+        r = r.reshape(-1, 1)
+        rT_Cinv_r = sparse.dot(r.T, sparse.inv(C), r)[0, 0]
     else:
-        raise NotImplementedError
+        # TODO: check what is the fastest and works the best between cholesky+solve
+        # and just inversion
+        if inverse_method == "inverse":
+            y = np.dot(np.linalg.inv(C), r)
+        elif inverse_method == "cholesky":
+            y = sp.linalg.cho_solve(sp.linalg.cho_factor(C, lower=True), r)
+        else:
+            raise NotImplementedError
+        rT_Cinv_r = r.dot(y)
 
     if constant_cov:
-        return -0.5 * r.dot(y)
+        return -0.5 * rT_Cinv_r
     else:
-        _, logdet = np.linalg.slogdet(C)
-        return -0.5 * r.dot(y) - 0.5 * logdet
+        if sparse.is_sparse(C):
+            _, logdet = sparse.slogdet(C)
+        else:
+            _, logdet = np.linalg.slogdet(C)
+        return -0.5 * (rT_Cinv_r - logdet)