From 3ed365d3b0cfa1f5f2ef3954546fb424035befe8 Mon Sep 17 00:00:00 2001
From: Pierre-Francois Leget <leget@astro.princeton.edu>
Date: Thu, 16 May 2024 18:11:07 -0400
Subject: [PATCH] saved cholesky output to avoid re-doing same computation

---
 treegp/gp_interp.py | 12 ++++++++----
 1 file changed, 8 insertions(+), 4 deletions(-)

diff --git a/treegp/gp_interp.py b/treegp/gp_interp.py
index 62378fc..2356ac1 100644
--- a/treegp/gp_interp.py
+++ b/treegp/gp_interp.py
@@ -172,10 +172,11 @@ def return_gp_predict(self, y, X1, X2, kernel, y_err, return_cov=False):
         :param y_err:  Error of y. (n_samples)
         """
         HT = kernel.__call__(X2, Y=X1)
-        K = kernel.__call__(X1) + np.eye(len(y)) * y_err**2
-        factor = (cholesky(K, overwrite_a=True, lower=False), False)
-        alpha = cho_solve(factor, y, overwrite_b=False)
-        y_predict = np.dot(HT, alpha.reshape((len(alpha), 1))).T[0]
+        if self._alpha is None:
+            K = kernel.__call__(X1) + np.eye(len(y)) * y_err**2
+            factor = (cholesky(K, overwrite_a=True, lower=False), False)
+            self._alpha = cho_solve(factor, y, overwrite_b=False)
+        y_predict = np.dot(HT, self._alpha.reshape((len(self._alpha), 1))).T[0]
         if return_cov:
             fact = cholesky(
                 K, lower=True
@@ -215,6 +216,9 @@ def initialize(self, X, y, y_err=None):
             self._mean = np.mean(y - self._spatial_average)
         else:
             self._mean = 0.0
+        # Initialize alpha to None so that we know to recompute it if we change the
+        # input data.
+        self._alpha = None
 
     def _build_average_meanify(self, X):
         """Compute spatial average from meanify output for a given coordinate using KN interpolation.