Merge pull request #16 from ronpandolfi/covariance_return_fix

Fix covariance matrix return

fvgp/gp.py

@@ -876,9 +876,13 @@ def posterior_covariance(self, x_iset, variance_only = False):
 
         k = self.kernel(self.x_data,p,self.hyperparameters,self)
         kk = self.kernel(p, p,self.hyperparameters,self)
-        if self.use_inv is True:
-            if variance_only is True: v = np.diag(kk) - np.einsum('ij,jk,ki->i', k.T, self.K_inv, k); S = False
-            if variance_only is False:  S = kk - (k.T @ self.K_inv @ k); v = np.array(np.diag(S))
+        if self.use_inv:
+            if variance_only:
+                S = False
+                v = np.diag(kk) - np.einsum('ij,jk,ki->i', k.T, self.K_inv, k)
+            else:
+                S = kk - (k.T @ self.K_inv @ k)
+                v = np.array(np.diag(S))
         else:
             k_cov_prod = self.solve(self.prior_covariance,k)
             S = kk - (k_cov_prod.T @ k)
@@ -890,7 +894,8 @@ def posterior_covariance(self, x_iset, variance_only = False):
             or double check the hyperparameter optimization bounds. This will not
             terminate the algorithm, but expect anomalies."""))
             v[v<0.0] = 0.0
-            if S is not False: S = np.fill_diagonal(S,v)
+            if not variance_only:
+                np.fill_diagonal(S, v)
 
         return {"x": p,
                 "v(x)": v,