WIP: Docstrings and typehints

flare/gp_algebra.py

@@ -3,9 +3,29 @@
 import multiprocessing as mp
 import time
 
+from typing import List
+
+#######################################
+##### COVARIANCE MATRIX FUNCTIONS
+#######################################
+
+def get_cov_row(x_1, d_1, m_index: int, size: int, training_data: list,
+                kernel: callable, kern_hyps: np.array,
+                cutoffs: np.array)-> list:
+    """
+    Compute an individual row of a covariance matrix.
+
+    :param x_1: Atomic environment to compare to At. Envs. in training data
+    :param d_1: Index of force component (x,y,z), indexed as (1,2,3)
+    :param m_index:
+    :param size:
+    :param training_data: Set of atomic environments to compare against
+    :param kernel: Kernel function to compare x_1 against training data with
+    :param kern_hyps: Hyperparameters which parameterize kernel function
+    :param cutoffs: The cutoff values used for the atomic environments
+    :return: covs, list of covariance matrix row elements
+    """
 
-def get_cov_row(x_1, d_1, m_index, size, training_data, kernel,
-                kern_hyps, cutoffs):
     covs = []
     ds = [1, 2, 3]
     for n_index in range(m_index, size):
@@ -20,8 +40,23 @@ def get_cov_row(x_1, d_1, m_index, size, training_data, kernel,
     return covs
 
 
-def get_cov_row_derv(x_1, d_1, m_index, size, training_data, kernel_grad,
-                     kern_hyps, cutoffs):
+def get_cov_row_derv(x_1, d_1: int, m_index: int, size: int,
+                     training_data: list,
+                     kernel_grad: callable,
+                     kern_hyps: np.array, cutoffs: np.array)-> (list, list):
+    """
+
+    :param x_1: Atomic Environment to compare to At. Envs. in training data
+    :param d_1: Index of force
+    :param m_index:
+    :param size:
+    :param training_data: Set of atomic environments to compare against
+    :param kernel: Kernel callable to compare x_1 against training data with
+    :param kernel_grad: Callable for gradient of kernel
+    :param kern_hyps: Hyperparameters which parameterize kernel function
+    :param cutoffs: The cutoff values used for the atomic environments
+    :return: covs, hyps
+    """
     covs = []
     hyps = []
     ds = [1, 2, 3]
@@ -37,10 +72,109 @@ def get_cov_row_derv(x_1, d_1, m_index, size, training_data, kernel_grad,
 
     return covs, hyps
 
+#######################################
+##### KY MATRIX FUNCTIONS
+#######################################
 
-def get_ky_mat_par(hyps: np.ndarray, training_data: list,
+
+def get_ky_mat(hyps: np.array, training_data: list,
+               kernel: callable, cutoffs=None):
+    """
+    Compute covariance matrix K by comparing training data with itself
+
+    :param hyps:
+    :param training_data:
+    :param training_labels_np:
+    :param kernel:
+    :param cutoffs:
+    :return:
+    """
+
+    # assume sigma_n is the final hyperparameter
+    sigma_n = hyps[-1]
+
+    # initialize matrices
+    size = len(training_data) * 3
+    k_mat = np.zeros([size, size])
+
+    ds = [1, 2, 3]
+
+    # calculate elements
+    for m_index in range(size):
+        x_1 = training_data[int(math.floor(m_index / 3))]
+        d_1 = ds[m_index % 3]
+
+        for n_index in range(m_index, size):
+            x_2 = training_data[int(math.floor(n_index / 3))]
+            d_2 = ds[n_index % 3]
+
+            # calculate kernel and gradient
+            kern_curr = kernel(x_1, x_2, d_1, d_2, hyps,
+                               cutoffs)
+
+            # store kernel value
+            k_mat[m_index, n_index] = kern_curr
+            k_mat[n_index, m_index] = kern_curr
+
+    # matrix manipulation
+    ky_mat = k_mat + sigma_n ** 2 * np.eye(size)
+
+    return ky_mat
+
+
+def get_ky_and_hyp(hyps: np.ndarray, training_data: list,
                    training_labels_np: np.ndarray,
+                   kernel_grad, cutoffs=None):
+    # assume sigma_n is the final hyperparameter
+    number_of_hyps = len(hyps)
+    sigma_n = hyps[number_of_hyps - 1]
+
+    # initialize matrices
+    size = len(training_data) * 3
+    k_mat = np.zeros([size, size])
+    hyp_mat = np.zeros([number_of_hyps, size, size])
+
+    ds = [1, 2, 3]
+
+    # calculate elements
+    for m_index in range(size):
+        x_1 = training_data[int(math.floor(m_index / 3))]
+        d_1 = ds[m_index % 3]
+
+        for n_index in range(m_index, size):
+            x_2 = training_data[int(math.floor(n_index / 3))]
+            d_2 = ds[n_index % 3]
+
+            # calculate kernel and gradient
+            cov = kernel_grad(x_1, x_2, d_1, d_2, hyps, cutoffs)
+
+            # store kernel value
+            k_mat[m_index, n_index] = cov[0]
+            k_mat[n_index, m_index] = cov[0]
+
+            # store gradients (excluding noise variance)
+            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[-1, :, :] = np.eye(size) * 2 * sigma_n
+
+    # matrix manipulation
+    ky_mat = k_mat + sigma_n ** 2 * np.eye(size)
+
+    return hyp_mat, ky_mat
+
+def get_ky_mat_par(hyps: np.ndarray, training_data: list,
                    kernel, cutoffs=None, no_cpus=None):
+    """
+    Parallelized version of function which computes ky matrix
+    :param hyps:
+    :param training_data:
+    :param kernel:
+    :param cutoffs:
+    :param no_cpus:
+    :return:
+    """
 
     if (no_cpus is None):
         ncpus = mp.cpu_count()
@@ -139,84 +273,6 @@ def get_ky_and_hyp_par(hyps: np.ndarray, training_data: list,
     return hyp_mat, ky_mat
 
 
-def get_ky_mat(hyps: np.ndarray, training_data: list,
-               training_labels_np: np.ndarray,
-               kernel, cutoffs=None):
-
-    # assume sigma_n is the final hyperparameter
-    number_of_hyps = len(hyps)
-    sigma_n = hyps[number_of_hyps - 1]
-
-    # initialize matrices
-    size = len(training_data) * 3
-    k_mat = np.zeros([size, size])
-
-    ds = [1, 2, 3]
-
-    # calculate elements
-    for m_index in range(size):
-        x_1 = training_data[int(math.floor(m_index / 3))]
-        d_1 = ds[m_index % 3]
-
-        for n_index in range(m_index, size):
-            x_2 = training_data[int(math.floor(n_index / 3))]
-            d_2 = ds[n_index % 3]
-
-            # calculate kernel and gradient
-            kern_curr = kernel(x_1, x_2, d_1, d_2, hyps,
-                               cutoffs)
-
-            # store kernel value
-            k_mat[m_index, n_index] = kern_curr
-            k_mat[n_index, m_index] = kern_curr
-
-    # matrix manipulation
-    ky_mat = k_mat + sigma_n ** 2 * np.eye(size)
-
-    return ky_mat
-
-
-def get_ky_and_hyp(hyps: np.ndarray, training_data: list,
-                   training_labels_np: np.ndarray,
-                   kernel_grad, cutoffs=None):
-    # assume sigma_n is the final hyperparameter
-    number_of_hyps = len(hyps)
-    sigma_n = hyps[number_of_hyps - 1]
-
-    # initialize matrices
-    size = len(training_data) * 3
-    k_mat = np.zeros([size, size])
-    hyp_mat = np.zeros([number_of_hyps, size, size])
-
-    ds = [1, 2, 3]
-
-    # calculate elements
-    for m_index in range(size):
-        x_1 = training_data[int(math.floor(m_index / 3))]
-        d_1 = ds[m_index % 3]
-
-        for n_index in range(m_index, size):
-            x_2 = training_data[int(math.floor(n_index / 3))]
-            d_2 = ds[n_index % 3]
-
-            # calculate kernel and gradient
-            cov = kernel_grad(x_1, x_2, d_1, d_2, hyps, cutoffs)
-
-            # store kernel value
-            k_mat[m_index, n_index] = cov[0]
-            k_mat[n_index, m_index] = cov[0]
-
-            # store gradients (excluding noise variance)
-            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[-1, :, :] = np.eye(size) * 2 * sigma_n
-
-    # matrix manipulation
-    ky_mat = k_mat + sigma_n ** 2 * np.eye(size)
-
-    return hyp_mat, ky_mat
 
 def get_ky_mat_update_row(params):
     '''
@@ -263,6 +319,9 @@ def get_ky_mat_update(ky_mat_old, training_data, get_kernel_vector, hyps, par):
     ky_mat[n:, n:] += sigma_n ** 2 * np.eye(3 * m)
     return ky_mat
 
+#######################################
+##### LIKELIHOOD + LIKELIHOOD GRADIENT
+#######################################
 
 def get_like_from_ky_mat(ky_mat, training_labels_np):
         # catch linear algebra errors
@@ -282,6 +341,7 @@ def get_like_from_ky_mat(ky_mat, training_labels_np):
         return like
 
 
+
 def get_like_grad_from_mats(ky_mat, hyp_mat, training_labels_np):
 
         number_of_hyps = hyp_mat.shape[0]