Serialization and representation GP + Small fixes

flare/gp.py

@@ -8,6 +8,8 @@
 from flare.gp_algebra import get_ky_mat, get_ky_and_hyp, \
     get_like_from_ky_mat, get_like_grad_from_mats, get_neg_likelihood, \
     get_neg_like_grad, get_ky_and_hyp_par
+from flare.kernels import str_to_kernel
+from flare.mc_simple import str_to_mc_kernel
 
 
 class GaussianProcess:
@@ -36,17 +38,23 @@ def __init__(self, kernel: Callable,
         self.hyp_labels = hyp_labels
         self.cutoffs = cutoffs
         self.algo = opt_algorithm
-        self.l_mat = None
-        self.alpha = None
+
         self.training_data = []
         self.training_labels = []
         self.training_labels_np = np.empty(0, )
         self.maxiter = maxiter
-        self.likelihood = None
-        self.likelihood_gradient = None
         self.par = par
         self.output = output
 
+        # Parameters set during training
+        self.ky_mat = None
+        self.l_mat = None
+        self.alpha = None
+        self.ky_mat_inv = None
+        self.l_mat_inv = None
+        self.likelihood = None
+        self.likelihood_gradient = None
+
     # TODO unit test custom range
     def update_db(self, struc: Structure, forces: list,
                   custom_range: List[int] = ()):
@@ -128,6 +136,7 @@ def train(self, output=None, custom_bounds=None,
         args = (self.training_data, self.training_labels_np,
                 self.kernel_grad, self.cutoffs, output,
                 self.par)
+        res = None
 
         if self.algo == 'L-BFGS-B':
 
@@ -166,7 +175,8 @@ def train(self, output=None, custom_bounds=None,
                            options={'disp': False,
                                     'maxiter': self.maxiter,
                                     'xtol': x_tol})
-
+        if res is None:
+            raise RuntimeError("Optimization failed for some reason.")
         self.hyps = res.x
         self.set_L_alpha()
         self.likelihood = -res.fun
@@ -185,12 +195,6 @@ def predict(self, x_t: AtomicEnvironment, d: int) -> [float, float]:
         pred_var = self_kern - \
                    np.matmul(np.matmul(k_v, self.ky_mat_inv), k_v)
 
-        # # get predictive variance (possibly slow)
-        # v_vec = solve_triangular(self.l_mat, k_v, lower=True)
-        # self_kern = self.kernel(x_t, x_t, self.bodies, d, d, self.hyps,
-        #                         self.cutoffs)
-        # pred_var = self_kern - np.matmul(v_vec, v_vec)
-
         return pred_mean, pred_var
 
     def predict_local_energy(self, x_t: AtomicEnvironment) -> float:
@@ -292,8 +296,8 @@ def set_L_alpha(self):
         self.ky_mat_inv = ky_mat_inv
         self.l_mat_inv = l_mat_inv
 
-        self.like = like
-        self.like_grad = like_grad
+        self.likelihood = like
+        self.likelihood_gradient = like_grad
 
     def update_L_alpha(self):
         """
@@ -331,3 +335,82 @@ def update_L_alpha(self):
         self.alpha = alpha
         self.ky_mat_inv = ky_mat_inv
         self.l_mat_inv = l_mat_inv
+
+    def __str__(self):
+
+        thestr = "GaussianProcess Object\n"
+        thestr += 'Kernel: {}\n'.format(self.kernel_name)
+        thestr += "Training points: {}\n".format(len(self.training_data))
+        thestr += 'Cutoffs: {}\n'.format(self.cutoffs)
+        thestr += 'Model Likelihood: {}\n'.format(self.likelihood)
+
+        thestr += 'Hyperparameters: \n'
+        if self.hyp_labels is None:
+            # Put unlabeled hyperparameters on one line
+            thestr = thestr[:-1]
+            thestr += str(self.hyps) + '\n'
+        else:
+            for hyp, label in zip(self.hyps, self.hyp_labels):
+                thestr += "{}: {}\n".format(label, hyp)
+
+        return thestr
+
+    def as_dict(self):
+
+        out_dict = dict(vars(self))
+
+        out_dict['training_data'] = [env.as_dict() for env in
+                                     self.training_data]
+        # Remove the callables
+        del out_dict['kernel']
+        del out_dict['kernel_grad']
+
+        return out_dict
+
+    @staticmethod
+    def from_dict(dictionary):
+
+        if 'mc' in dictionary['kernel_name']:
+            force_kernel, grad = str_to_mc_kernel(dictionary['kernel_name'],
+                                               include_grad=True)
+        else:
+            force_kernel, grad = str_to_kernel(dictionary['kernel_name'],
+                                               include_grad=True)
+
+        if dictionary['energy_kernel'] is not None:
+            energy_kernel = str_to_kernel(dictionary['energy_kernel'])
+        else:
+            energy_kernel = None
+
+        if dictionary['energy_force_kernel'] is not None:
+            energy_force_kernel = str_to_kernel(dictionary[
+                                                    'energy_force_kernel'])
+        else:
+            energy_force_kernel = None
+
+        new_gp = GaussianProcess(kernel=force_kernel,
+                                 kernel_grad=grad,
+                                 energy_kernel=energy_kernel,
+                                 energy_force_kernel=energy_force_kernel,
+                                 cutoffs=np.array(dictionary['cutoffs']),
+                                 hyps=np.array(dictionary['hyps']),
+                                 hyp_labels=dictionary['hyp_labels'],
+                                 par=dictionary['par'],
+                                 maxiter=dictionary['maxiter'],
+                                 opt_algorithm=dictionary['algo'])
+
+        # Save time by attempting to load in computed attributes
+        new_gp.l_mat = np.array(dictionary.get('l_mat', None))
+        new_gp.l_mat_inv = np.array(dictionary.get('l_mat_inv', None))
+        new_gp.alpha = np.array(dictionary.get('alpha', None))
+        new_gp.ky_mat = np.array(dictionary.get('ky_mat', None))
+        new_gp.ky_mat_inv = np.array(dictionary.get('ky_mat_inv', None))
+
+        new_gp.training_data = [AtomicEnvironment.from_dict(env) for env in
+                                dictionary['training_data']]
+        new_gp.training_labels = dictionary['training_labels']
+
+        new_gp.likelihood = dictionary['likelihood']
+        new_gp.likelihood_gradient = dictionary['likelihood_gradient']
+
+        return new_gp

flare/gp_from_aimd.py

@@ -245,7 +245,7 @@ def train_gp(self):
 
         self.output.write_hyps(self.gp.hyp_labels, self.gp.hyps,
                                self.start_time,
-                               self.gp.like, self.gp.like_grad)
+                               self.gp.likelihood, self.gp.likelihood_gradient)
         self.train_count += 1
 
     def is_std_in_bound(self, frame: Structure)->(bool, List[int]):

flare/kernels.py

@@ -869,3 +869,34 @@ def triplet_force_en_kernel(ci1, ci2, ri1, ri2, ri3, rj1, rj2, rj3,
     print(kern_finite_diff)
     print(kern_analytical)
     assert(np.isclose(kern_finite_diff, kern_analytical, atol=tol))
+
+_str_to_kernel = {'two_body': two_body,
+                  'two_body_en': two_body_en,
+                  'two_body_force_en': two_body_force_en,
+                  'three_body': three_body,
+                  'three_body_en': three_body_en,
+                  'three_body_force_en': three_body_force_en,
+                  'two_plus_three_body': two_plus_three_body,
+                  'two_plus_three_en': two_plus_three_en,
+                  'two_plus_three_force_en': two_plus_three_force_en
+                  }
+
+
+def str_to_kernel(string: str, include_grad: bool=False):
+
+    if string not in _str_to_kernel.keys():
+        raise ValueError("Kernel {} not found in list of available "
+                         "kernels{}:".format(string,_str_to_kernel.keys()))
+
+    if not include_grad:
+        return _str_to_kernel[string]
+    else:
+        if 'two' in string and 'three' in string:
+            return _str_to_kernel[string], two_plus_three_body_grad
+        elif 'two' in string and 'three' not in string:
+            return _str_to_kernel[string], two_body_grad
+        elif 'two' not in string and 'three' in string:
+            return _str_to_kernel[string], three_body_grad
+        else:
+            raise ValueError("Gradient callable for {} not found".format(
+                string))

flare/mc_simple.py

@@ -822,3 +822,35 @@ def two_body_mc_en_jit(bond_array_1, c1, etypes1,
                 kern += fi * fj * sig2 * exp(-r11 * r11 * ls1)
 
     return kern
+
+
+_str_to_kernel = {'two_body_mc': two_body_mc,
+                  'two_body_en_mc': two_body_mc_en,
+                  'two_body_mc_force_en': two_body_mc_force_en,
+                  'three_body_mc': three_body_mc,
+                  'three_body_mc_en': three_body_mc_en,
+                  'three_body_mc_force_en': three_body_mc_force_en,
+                  'two_plus_three_body_mc': two_plus_three_body_mc,
+                  'two_plus_three_mc_en': two_plus_three_mc_en,
+                  'two_plus_three_mc_force_en': two_plus_three_mc_force_en
+                  }
+
+
+def str_to_mc_kernel(string: str, include_grad: bool=False):
+
+    if string not in _str_to_kernel.keys():
+        raise ValueError("Kernel {} not found in list of available "
+                         "kernels{}:".format(string, _str_to_kernel.keys()))
+
+    if not include_grad:
+        return _str_to_kernel[string]
+    else:
+        if 'two' in string and 'three' in string:
+            return _str_to_kernel[string], two_plus_three_body_mc_grad
+        elif 'two' in string and 'three' not in string:
+            return _str_to_kernel[string], two_body_mc_grad
+        elif 'two' not in string and 'three' in string:
+            return _str_to_kernel[string], three_body_mc_grad
+        else:
+            raise ValueError("Gradient callable for {} not found".format(
+                string))

flare/otf.py

@@ -259,7 +259,7 @@ def train_gp(self):
         self.gp.train(self.output)
         self.output.write_hyps(self.gp.hyp_labels, self.gp.hyps,
                                self.start_time,
-                               self.gp.like, self.gp.like_grad)
+                               self.gp.likelihood, self.gp.likelihood_gradient)
 
     def is_std_in_bound(self):
         # set uncertainty threshold

flare/util.py

@@ -128,7 +128,8 @@
 
 _Z_to_element = {z: elt for elt, z in _element_to_Z.items()}
 
-def element_to_Z(element:str)->int:
+
+def element_to_Z(element: str)->int:
     """
     Returns the atomic number Z associated with an elements 1-2 letter name.
     Returns the same integer if an integer is passed in.