just some cleanup

lbl-camera · Feb 15, 2024 · bad9304 · bad9304
1 parent 88293ea
commit bad9304
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 4 deletions.
diff --git a/fvgp/gp.py b/fvgp/gp.py
@@ -1207,7 +1207,6 @@ def _update_GPpriorV(self, x_data, x_new, y_data, y_new, hyperparameters, calc_i
             if self.info: print("Transferring the covariance matrix to host done after ", time.time() - st,
                                 " seconds. sparsity = ", Ksparsity, flush=True)
         else:
-            #K = self._compute_K(x_data, x_new, hyperparameters) ####here we could use block matrices to make this faster TODO
             off_diag = self._compute_K(x_data, x_new, hyperparameters)
             K = np.block([
                          [self.K,          off_diag],
@@ -1227,7 +1226,7 @@ def _update_GPpriorV(self, x_data, x_new, y_data, y_new, hyperparameters, calc_i
         return K, KV, KVinvY, KVlogdet, factorization_obj, KVinv, prior_mean_vec, V
 
     ##################################################################################
-    def _compute_gp_linalg(self, vec, KV, calc_inv=False ,try_sparse_LU = False):
+    def _compute_gp_linalg(self, vec, KV, calc_inv=False, try_sparse_LU=False):
         st = time.time()
         if self.gp2Scale:
             from imate import logdet

diff --git a/fvgp/gp2Scale.py b/fvgp/gp2Scale.py
@@ -118,7 +118,6 @@ def update_covariance(self, x_new, hyperparameters, client, cov):
         ranges_ij = list(itertools.product(ranges_data, ranges_input))
         ranges_ij2 = list(itertools.product(ranges_input, ranges_input))
 
-
         kernel_caller = kernel_function_update
 
         #K = np.block([[A, B],
@@ -138,7 +137,8 @@ def update_covariance(self, x_new, hyperparameters, client, cov):
         B = sparse.coo_matrix((data, (i_s, j_s)))
 
         # mirror across diagonal
-
+        ranges_ij2 = [range_ij2 for range_ij2 in ranges_ij2 if range_ij2[0][0] <= range_ij2[1][0]]  # filter lower diagonal
+        kernel_caller = kernel_function
         results = list(map(self.harvest_result,
                           distributed.as_completed(client.map(
                               partial(kernel_caller,
@@ -149,6 +149,10 @@ def update_covariance(self, x_new, hyperparameters, client, cov):
                               [self.x_new_scatter_future] * len(ranges_ij2)),
                               with_results=True)))
         data, i_s, j_s = map(np.hstack, zip(*results))
+        diagonal_mask = i_s != j_s
+        data, i_s, j_s = np.hstack([data, data[diagonal_mask]]), \
+            np.hstack([i_s, j_s[diagonal_mask]]), \
+            np.hstack([j_s, i_s[diagonal_mask]])
         D = sparse.coo_matrix((data, (i_s, j_s)))
 
         res = block_array([[cov,           B],