Faster sparse matrix Cholesky

enterprise/signals/signal_base.py

@@ -221,9 +221,17 @@ def __call__(self, xs, phiinv_method="cliques"):
 
             if self.cholesky_sparse:
                 try:
-                    cf = cholesky(TNT + sps.csc_matrix(phiinv))  # cf(Sigma)
-                    expval = cf(TNr)
-                    logdet_sigma = cf.logdet()
+                    Sigma_sp = TNT + sps.csc_matrix(phiinv)
+
+                    if hasattr(self, "cf_sp"):
+                        # Have analytical decomposition already. Just do update
+                        self.cf_sp.cholesky_inplace(Sigma_sp)
+                    else:
+                        # Do analytical and numerical Sparse Cholesky
+                        self.cf_sp = cholesky(Sigma_sp)
+
+                    expval = self.cf_sp(TNr)
+                    logdet_sigma = self.cf_sp.logdet()
                 except CholmodError:  # pragma: no cover
                     return -np.inf
             else:

tests/test_likelihood.py

@@ -326,6 +326,43 @@ def test_compare_ecorr_likelihood(self):
         msg = "Likelihood mismatch between ECORR methods"
         assert np.allclose(l1, l2), msg
 
+    def test_like_sparse_cache(self):
+        """Test likelihood with sparse Cholesky caching"""
+
+        # find the maximum time span to set GW frequency sampling
+        tmin = [p.toas.min() for p in self.psrs]
+        tmax = [p.toas.max() for p in self.psrs]
+        Tspan = np.max(tmax) - np.min(tmin)
+
+        # setup basic model
+        efac = parameter.Constant(1.0)
+        log10_A = parameter.Constant(-15.0)
+        gamma = parameter.Constant(4.33)
+
+        ef = white_signals.MeasurementNoise(efac)
+        pl = utils.powerlaw(log10_A=log10_A, gamma=gamma)
+
+        orf = utils.hd_orf()
+        crn = gp_signals.FourierBasisCommonGP(pl, orf, components=20, name="GW", Tspan=Tspan)
+
+        tm = gp_signals.TimingModel()
+        m = tm + ef + crn
+
+        # Two identical arrays that we'll compare with two sets of parameters
+        pta1 = signal_base.PTA([m(p) for p in self.psrs])
+        pta2 = signal_base.PTA([m(p) for p in self.psrs])
+
+        params_init = parameter.sample(pta1.params)
+        params_check = parameter.sample(pta1.params)
+
+        # First call for pta1 only initializes the sparse decomposition. Second one uses it
+        _ = pta1.get_lnlikelihood(params_init)
+        l1 = pta1.get_lnlikelihood(params_check)
+        l2 = pta2.get_lnlikelihood(params_check)
+
+        msg = "Likelihood mismatch between sparse Cholesky full & inplace"
+        assert np.allclose(l1, l2), msg
+
 
 class TestLikelihoodPint(TestLikelihood):
     @classmethod