gp: use Cholesky factor rather than explicit inverse

Author: asher-m

UncertainSCI/gp/__init__.py

@@ -7,6 +7,7 @@
 from . import kernel
 
 import numpy as np
+from scipy import linalg
 
 NUGGET = 1e-6
 
@@ -17,7 +18,7 @@ class ScalarGaussianProcess():
     x_obs: np.ndarray | None = None
     y_obs: np.ndarray | None = None
     sigma: np.ndarray | None = None
-    inv: np.ndarray | None = None
+    cho_factor: tuple[np.ndarray, bool] | None = None
 
     def __init__(self,
                  mu: wrapper.ScalarFunction,
@@ -46,24 +47,26 @@ def condition(self, x: np.ndarray,
         self.x_obs = x
         self.y_obs = y
         self.sigma = sigma
-        self.inv = np.linalg.inv(self.k(x) + sigma * np.eye(len(x)))
+        self.cho_factor = linalg.cho_factor(self.k(x) + sigma * np.eye(len(x)))
+
+    def tune(self):
+        pass
 
     def mu_posterior(self, x: np.ndarray):
-        if (self.x_obs is None) or (self.y_obs is None) or (self.inv is None):
+        if (self.x_obs is None) or (self.y_obs is None) or (self.cho_factor is None):
             raise ValueError('GP not conditioned: call ScalarGaussianProcess.'
                              'condition with observations before computing '
                              'posterior!')
-        return self.mu(x) + self.k(x, self.x_obs) @ self.inv @ \
-            (self.y_obs - self.mu(self.x_obs))
+        return self.mu(x) + self.k(x, self.x_obs) @ \
+            linalg.cho_solve(self.cho_factor, self.y_obs - self.mu(self.x_obs))
 
     def k_posterior(self, x1: np.ndarray, x2: np.ndarray | None = None):
-        if (self.x_obs is None) or (self.y_obs is None) or (self.inv is None):
+        if (self.x_obs is None) or (self.y_obs is None) or (self.cho_factor is None):
             raise ValueError('GP not conditioned: call ScalarGaussianProcess.'
                              'condition with observations before computing '
                              'posterior!')
         return (self.k(x1) if x2 is None else self.k(x1, x2)) - \
-            self.k(x1, self.x_obs) @ self.inv @ \
-            (self.k(self.x_obs, x1) if x2 is None else self.k(self.x_obs, x2))
+            self.k(x1, self.x_obs) @ linalg.cho_solve(self.cho_factor, self.k(self.x_obs, x1) if x2 is None else self.k(self.x_obs, x2))
 
     def sample_prior(self, x: np.ndarray, n: int = 1) -> np.ndarray:
         ell = np.linalg.cholesky(self.k(x) + NUGGET * np.eye(len(x)))
@@ -84,7 +87,7 @@ class VectorGaussianProcess():
     x_obs: np.ndarray | None = None
     y_obs: np.ndarray | None = None
     sigma: np.ndarray | None = None
-    inv: np.ndarray | None = None
+    cho_factor: tuple[np.ndarray, bool] | None = None
 
     def __init__(self,
                  mu: wrapper.VectorFunction,
@@ -121,23 +124,26 @@ def condition(self, x: np.ndarray,
         self.x_obs = x
         self.y_obs = y
         self.sigma = sigma
-        self.inv = np.linalg.inv(self.k(x) + sigma.flatten() * np.eye(len(x) * self.cdim))
+        self.cho_factor = linalg.cho_factor(self.k(x) + sigma.flatten() * np.eye(len(x) * self.cdim))
+
+    def tune(self):
+        pass
 
     def mu_posterior(self, x: np.ndarray):
-        if (self.x_obs is None) or (self.y_obs is None) or (self.inv is None):
+        if (self.x_obs is None) or (self.y_obs is None) or (self.cho_factor is None):
             raise ValueError('GP not conditioned: call VectorGaussianProcess.'
                              'condition with observations before computing '
                              'posterior!')
         return self.mu(x) + \
-            (self.k(x, self.x_obs) @ self.inv @ (self.y_obs - self.mu(self.x_obs)).flatten()).reshape((len(x), self.cdim))
+            (self.k(x, self.x_obs) @ linalg.cho_solve(self.cho_factor, (self.y_obs - self.mu(self.x_obs)).flatten())).reshape((len(x), self.cdim))
 
     def k_posterior(self, x1: np.ndarray, x2: np.ndarray | None = None):
-        if (self.x_obs is None) or (self.y_obs is None) or (self.inv is None):
+        if (self.x_obs is None) or (self.y_obs is None) or (self.cho_factor is None):
             raise ValueError('GP not conditioned: call VectorGaussianProcess.'
                              'condition with observations before computing '
                              'posterior!')
         return (self.k(x1) if x2 is None else self.k(x1, x2)) - \
-            self.k(x1, self.x_obs) @ self.inv @ (self.k(self.x_obs, x1) if x2 is None else self.k(self.x_obs, x2))
+            self.k(x1, self.x_obs) @ linalg.cho_solve(self.cho_factor, self.k(self.x_obs, x1) if x2 is None else self.k(self.x_obs, x2))
 
     def sample_prior(self, x: np.ndarray, n: int = 1) -> np.ndarray:
         if n > 1:
@@ -147,6 +153,7 @@ def sample_prior(self, x: np.ndarray, n: int = 1) -> np.ndarray:
             sn = (len(x) * self.cdim)
             sr = (len(x), self.cdim)
 
+        # TODO: make this faster (use precomputed cho_factor and Schur complement)
         ell = np.linalg.cholesky(self.k(x) + NUGGET * np.eye(self.cdim * len(x)))
         y = (self.mu(x)[..., None] if n > 1 else self.mu(x)) + \
             (ell @ np.random.normal(0, 1, sn)).reshape(sr)
@@ -160,6 +167,7 @@ def sample_posterior(self, x: np.ndarray, n: int = 1) -> tuple[np.ndarray, np.nd
             sn = (len(x) * self.cdim)
             sr = (len(x), self.cdim)
 
+        # TODO: make this faster (use precomputed cho_factor and Schur complement)
         ell = np.linalg.cholesky(self.k_posterior(x) + NUGGET * np.eye(self.cdim * len(x)))
         y = (self.mu_posterior(x)[..., None] if n > 1 else self.mu_posterior(x)) + \
             (ell @ np.random.normal(0, 1, (sn))).reshape(sr)