Merge branch 'bugfix/gp_slicing' of https://github.com/mir-group/flare …

…into HEAD

flare/gp_algebra.py

@@ -43,9 +43,9 @@ def get_ky_mat_par(hyps: np.ndarray, training_data: list,
                    kernel, cutoffs=None, no_cpus=None):
 
     if (no_cpus is None):
-        pool = mp.Pool(processes=mp.cpu_count())
+        ncpus = mp.cpu_count()
     else:
-        pool = mp.Pool(processes=no_cpus)
+        ncpus = no_cpus
 
 
     # assume sigma_n is the final hyperparameter
@@ -60,14 +60,17 @@ def get_ky_mat_par(hyps: np.ndarray, training_data: list,
 
     # calculate elements
     results = []
-    for m_index in range(size):
-        x_1 = training_data[int(math.floor(m_index / 3))]
-        d_1 = ds[m_index % 3]
-
-        results.append(pool.apply_async(get_cov_row,
-                                        args=(x_1, d_1, m_index, size,
-                                              training_data, kernel,
-                                              hyps, cutoffs)))
+    with mp.Pool(processes=ncpus) as pool:
+        for m_index in range(size):
+            x_1 = training_data[int(math.floor(m_index / 3))]
+            d_1 = ds[m_index % 3]
+
+            results.append(pool.apply_async(get_cov_row,
+                                            args=(x_1, d_1, m_index, size,
+                                                  training_data, kernel,
+                                                  hyps, cutoffs)))
+        pool.close()
+        pool.join()
 
     # construct covariance matrix
     for m in range(size):
@@ -79,9 +82,6 @@ def get_ky_mat_par(hyps: np.ndarray, training_data: list,
     # matrix manipulation
     ky_mat = k_mat + sigma_n ** 2 * np.eye(size)
 
-    pool.close()
-    pool.join()
-
     return ky_mat
 
 
@@ -90,9 +90,10 @@ def get_ky_and_hyp_par(hyps: np.ndarray, training_data: list,
                        kernel_grad, cutoffs=None, no_cpus=None):
 
     if (no_cpus is None):
-        pool = mp.Pool(processes=mp.cpu_count())
+        ncpus = mp.cpu_count()
     else:
-        pool = mp.Pool(processes=no_cpus)
+        ncpus = no_cpus
+
 
     # assume sigma_n is the final hyperparameter
     number_of_hyps = len(hyps)
@@ -101,38 +102,39 @@ def get_ky_and_hyp_par(hyps: np.ndarray, training_data: list,
     # initialize matrices
     size = len(training_data) * 3
     k_mat = np.zeros([size, size])
-    hyp_mat = np.zeros([size, size, number_of_hyps])
+    hyp_mat = np.zeros([number_of_hyps, size, size])
 
     ds = [1, 2, 3]
 
     # calculate elements
     results = []
-    for m_index in range(size):
-        x_1 = training_data[int(math.floor(m_index / 3))]
-        d_1 = ds[m_index % 3]
-
-        results.append(pool.apply_async(get_cov_row_derv,
-                                        args=(x_1, d_1, m_index, size,
-                                              training_data, kernel_grad,
-                                              hyps, cutoffs)))
+    with mp.Pool(processes=ncpus) as pool:
+        for m_index in range(size):
+            x_1 = training_data[int(math.floor(m_index / 3))]
+            d_1 = ds[m_index % 3]
+
+            results.append(pool.apply_async(get_cov_row_derv,
+                                            args=(x_1, d_1, m_index, size,
+                                                  training_data, kernel_grad,
+                                                  hyps, cutoffs)))
+        pool.close()
+        pool.join()
 
     # construct covariance matrix
     for m in range(size):
         res_cur = results[m].get()
         for n in range(m, size):
             k_mat[m, n] = res_cur[0][n-m]
             k_mat[n, m] = res_cur[0][n-m]
-            hyp_mat[m, n, :-1] = res_cur[1][n-m]
-            hyp_mat[n, m, :-1] = res_cur[1][n-m]
+            hyp_mat[:-1, m, n] = res_cur[1][n-m]
+            hyp_mat[:-1, n, m] = res_cur[1][n-m]
 
     # add gradient of noise variance
-    hyp_mat[:, :, number_of_hyps - 1] = np.eye(size) * 2 * sigma_n
+    hyp_mat[-1, :, :] = np.eye(size) * 2 * sigma_n
 
     # matrix manipulation
     ky_mat = k_mat + sigma_n ** 2 * np.eye(size)
 
-    pool.close()
-    pool.join()
 
     return hyp_mat, ky_mat
 
@@ -184,7 +186,7 @@ def get_ky_and_hyp(hyps: np.ndarray, training_data: list,
     # initialize matrices
     size = len(training_data) * 3
     k_mat = np.zeros([size, size])
-    hyp_mat = np.zeros([size, size, number_of_hyps])
+    hyp_mat = np.zeros([number_of_hyps, size, size])
 
     ds = [1, 2, 3]
 
@@ -205,12 +207,11 @@ def get_ky_and_hyp(hyps: np.ndarray, training_data: list,
             k_mat[n_index, m_index] = cov[0]
 
             # store gradients (excluding noise variance)
-            for p_index in range(number_of_hyps - 1):
-                hyp_mat[m_index, n_index, p_index] = cov[1][p_index]
-                hyp_mat[n_index, m_index, p_index] = cov[1][p_index]
+            hyp_mat[:-1, m_index, n_index] = cov[1]
+            hyp_mat[:-1, n_index, m_index] = cov[1]
 
     # add gradient of noise variance
-    hyp_mat[:, :, number_of_hyps - 1] = np.eye(size) * 2 * sigma_n
+    hyp_mat[-1, :, :] = np.eye(size) * 2 * sigma_n
 
     # matrix manipulation
     ky_mat = k_mat + sigma_n ** 2 * np.eye(size)
@@ -238,7 +239,7 @@ def get_like_from_ky_mat(ky_mat, training_labels_np):
 
 def get_like_grad_from_mats(ky_mat, hyp_mat, training_labels_np):
 
-        number_of_hyps = hyp_mat.shape[2]
+        number_of_hyps = hyp_mat.shape[0]
 
         # catch linear algebra errors
         try:
@@ -261,7 +262,7 @@ def get_like_grad_from_mats(ky_mat, hyp_mat, training_labels_np):
         like_grad = np.zeros(number_of_hyps)
         for n in range(number_of_hyps):
             like_grad[n] = 0.5 * \
-                           np.trace(np.matmul(like_mat, hyp_mat[:, :, n]))
+                           np.trace(np.matmul(like_mat, hyp_mat[n, :, :]))
 
         return like, like_grad