diff --git a/flare/gp_algebra.py b/flare/gp_algebra.py index 4f43bf809..e1ef2d49e 100644 --- a/flare/gp_algebra.py +++ b/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]