LU for gp2Scale < 100 000 data points

lbl-camera · Sep 25, 2023 · 068a7ed · 068a7ed
1 parent 88546ba
commit 068a7ed
Show file tree

Hide file tree

Showing 4 changed files with 570 additions and 32 deletions.
diff --git a/fvgp/gp.py b/fvgp/gp.py
@@ -963,28 +963,30 @@ def _compute_GPpriorV(self, x_data, y_data, hyperparameters, calc_inv = False):
 
         #get Kinv/KVinvY, LU, Chol, logdet(KV)
         if self.gp2Scale:
-            #LU = splu(KV.tocsc())
-            #factorization_obj = ("LU", LU)
-            #KVinvY = LU.solve(y_data - prior_mean_vec)
-            KVinvY, exit_code = cg(KV.tocsc(),y_data - prior_mean_vec)
-            if exit_code != 0:
-                warnings.warn("CG solve not successful in gp2Scale. Trying MINRES")
-                KVinvY, exit_code = minres(KV.tocsc(),y_data - prior_mean_vec)
-            #upper_diag = abs(LU.U.diagonal())
-            #KVlogdet = np.sum(np.log(upper_diag))
-            KVlogdet, info_slq = logdet(KV, method='slq', min_num_samples=50, max_num_samples=200,
+            if len(x_data) < 100000:
+                LU = splu(KV.tocsc())
+                factorization_obj = ("LU", LU)
+                KVinvY = LU.solve(y_data - prior_mean_vec)
+                upper_diag = abs(LU.U.diagonal())
+                KVlogdet = np.sum(np.log(upper_diag))
+            else:
+                KVinvY, exit_code = cg(KV.tocsc(),y_data - prior_mean_vec)
+                factorization_obj = ("gp2Scale",None)
+                if exit_code != 0:
+                    warnings.warn("CG solve not successful in gp2Scale. Trying MINRES")
+                    KVinvY, exit_code = minres(KV.tocsc(),y_data - prior_mean_vec)
+                KVlogdet, info_slq = logdet(KV, method='slq', min_num_samples=50, max_num_samples=200,
                               lanczos_degree=80, error_rtol=0.1,
                               return_info=True, plot=False, verbose=self.info)
-            if self.info: print("Solve and logdet done after ",time.time() - st,"seconds.")
+            if self.info: print("Solve and logdet/LU done after ",time.time() - st,"seconds.")
             KVinv = None
-            factorization_obj = ("gp2Scale",None)
         elif self.sparse_mode and self._is_sparse(KV):
             KV = csc_matrix(KV)
             LU = splu(KV)
             factorization_obj = ("LU", LU)
-            kvinvy = lu.solve(y_data - prior_mean_vec)
+            KVinvY = lu.solve(y_data - prior_mean_vec)
             upper_diag = abs(lu.u.diagonal())
-            kvlogdet = np.sum(np.log(upper_diag))
+            KVlogdet = np.sum(np.log(upper_diag))
             KVinv = None
         else:
             c, l = cho_factor(KV)

diff --git a/fvgp/gp2Scale.py b/fvgp/gp2Scale.py
@@ -116,7 +116,7 @@ def compute_covariance(self, hyperparameters,client):
                 ##make workers available that are not actively computing
                 while not idle_workers:
                     idle_workers, futures, finished_futures = self.free_workers(futures, finished_futures)
-                    time.sleep(0.1)
+                    if not idle_workers: time.sleep(0.1)
 
                 ####collect finished workers but only if actor is not busy, otherwise do it later
                 if len(finished_futures) >= 2000:
@@ -126,7 +126,7 @@ def compute_covariance(self, hyperparameters,client):
 
                 #get idle worker and submit work
                 current_worker = self.get_idle_worker(idle_workers)
-                data = {"scattered_data": self.scatter_future,"hps": hyperparameters, "kernel" :self.kernel,  "range_i": (beg_i,end_i), "range_j": (beg_j,end_j), "mode": "prior","gpu": 0}
+                data = {"scattered_data": self.scatter_future,"hps": hyperparameters, "kernel" :self.kernel,  "range_i": (beg_i,end_i), "range_j": (beg_j,end_j)}
                 futures.append(client.submit(kernel_function, data, workers = current_worker))
                 self.assign_future_2_worker(futures[-1].key,current_worker)
                 if self.info and ( i % 10) == 0.:
@@ -159,6 +159,7 @@ def free_workers(self, futures, finished_futures):
                 free_workers.add(self.future_worker_assignments[future.key])
                 del self.future_worker_assignments[future.key]
             else: remaining_futures.append(future)
+        del futures
         return free_workers, remaining_futures, finished_futures
 
     def assign_future_2_worker(self, future_key, worker_address):
@@ -517,19 +518,13 @@ def neg_log_likelihood(self, y):
 def kernel_function(data):  # pragma: no cover
     st = time.time()
     hps= data["hps"]
-    mode = data["mode"]
     kernel = data["kernel"]
     worker = distributed.get_worker()
-    if mode == "prior":
-        x1 = data["scattered_data"]["x_data"][data["range_i"][0]:data["range_i"][1]]
-        x2 = data["scattered_data"]["x_data"][data["range_j"][0]:data["range_j"][1]]
-        range1 = data["range_i"]
-        range2 = data["range_j"]
-        k = kernel(x1,x2,hps, None)
-    else:
-        x1 = data["x_data"]
-        x2 = data["x_data"]
-        k = kernel(x1,x2,hps, None)
+    x1 = data["scattered_data"]["x_data"][data["range_i"][0]:data["range_i"][1]]
+    x2 = data["scattered_data"]["x_data"][data["range_j"][0]:data["range_j"][1]]
+    range1 = data["range_i"]
+    range2 = data["range_j"]
+    k = kernel(x1,x2,hps, None)
     k_sparse = sparse.coo_matrix(k)
     #print(time.time() - st, flush = True)
     return k_sparse, (data["range_i"][0],data["range_j"][0]), time.time() - st, worker.address