Repository: mir-group/flare

Issue:
Finer control over GP optimization parameters by instance
The TrajectoryTrainer and OTF classes contain calls to the `gp.train` method but do not pass in arguments that the `train` method can accept. In order to expose more control to the user, being able to set optimization parameters (such as `grad_tol` and `x_tol`) on an individual instance would be helpful, since it makes sense for those to be parameters associated with a given `GaussianProcess` instance. This could involve, for instance, re-factoring the train method to draw from an internal `opt_parameter_kwarg` dictionary which could be passed when instantiating the `GaussianProcess`.

Comments:

Patch:
diff --git a/flare/gp.py b/flare/gp.py
index 07a2eb7a..876db51d 100644
--- a/flare/gp.py
+++ b/flare/gp.py
@@ -25,7 +25,8 @@ class GaussianProcess:
                  energy_force_kernel: Callable = None,
                  energy_kernel: Callable = None,
                  opt_algorithm: str = 'L-BFGS-B',
-                 maxiter=10, par=False,
+                 opt_params: dict = None,
+                 maxiter=10, par: bool=False,
                  output=None):
         """Initialize GP parameters and training data."""
 
@@ -38,6 +39,7 @@ class GaussianProcess:
         self.hyp_labels = hyp_labels
         self.cutoffs = cutoffs
         self.algo = opt_algorithm
+        self.opt_params = opt_params if opt_params is not None else {}
 
         self.training_data = []
         self.training_labels = []
@@ -119,14 +121,17 @@ class GaussianProcess:
 
         return forces_np
 
-    def train(self, output=None, custom_bounds=None,
-              grad_tol: float = 1e-4,
-              x_tol: float = 1e-5,
-              line_steps: int = 20):
-        """Train Gaussian Process model on training data. Tunes the \
+    def train(self, output=None, opt_param_override: dict = None):
+        """
+        Train Gaussian Process model on training data. Tunes the \
 hyperparameters to maximize the likelihood, then computes L and alpha \
-(related to the covariance matrix of the training set)."""
+(related to the covariance matrix of the training set).
 
+        :param output: Output to write to
+        :param opt_param_override: Dictionary of parameters to override
+        instace's optimzation parameters.
+        :return:
+        """
         x_0 = self.hyps
 
         args = (self.training_data, self.training_labels_np,
@@ -134,7 +139,21 @@ hyperparameters to maximize the likelihood, then computes L and alpha \
                 self.par)
         res = None
 
-        if self.algo == 'L-BFGS-B':
+        # Make local copy of opt params so as to not overwrite with override
+        opt_param_temp = dict(self.opt_params)
+        if opt_param_override is not None:
+            for key, value in opt_param_override.items():
+                opt_param_temp[key] = value
+
+        grad_tol = opt_param_temp.get('grad_tol', 1e-4)
+        x_tol = opt_param_temp.get('x_tol', 1e-5)
+        line_steps = opt_param_temp.get('max_ls', 20)
+        max_iter = opt_param_temp.get('max_iter', self.maxiter)
+        algo = opt_param_temp.get('algorithm', self.algo)
+        disp = opt_param_temp.get('disp', False)
+
+
+        if algo == 'L-BFGS-B':
 
             # bound signal noise below to avoid overfitting
             bounds = np.array([(1e-6, np.inf)] * len(x_0))
@@ -144,35 +163,36 @@ hyperparameters to maximize the likelihood, then computes L and alpha \
             try:
                 res = minimize(get_neg_like_grad, x_0, args,
                                method='L-BFGS-B', jac=True, bounds=bounds,
-                               options={'disp': False, 'gtol': grad_tol,
+                               options={'disp': disp, 'gtol': grad_tol,
                                         'maxls': line_steps,
-                                        'maxiter': self.maxiter})
+                                        'maxiter': max_iter})
             except:
                 print("Warning! Algorithm for L-BFGS-B failed. Changing to "
                       "BFGS for remainder of run.")
-                self.algo = 'BFGS'
+                self.opt_params['algorithm'] = 'BFGS'
 
-        if custom_bounds is not None:
+        if opt_param_temp.get('custom_bounds', None) is not None:
             res = minimize(get_neg_like_grad, x_0, args,
-                           method='L-BFGS-B', jac=True, bounds=custom_bounds,
-                           options={'disp': False, 'gtol': grad_tol,
+                           method='L-BFGS-B', jac=True,
+                           bounds=opt_param_temp.get('custom_bonuds'),
+                           options={'disp': disp, 'gtol': grad_tol,
                                     'maxls': line_steps,
-                                    'maxiter': self.maxiter})
+                                    'maxiter': max_iter})
 
-        elif self.algo == 'BFGS':
+        elif algo == 'BFGS':
             res = minimize(get_neg_like_grad, x_0, args,
                            method='BFGS', jac=True,
-                           options={'disp': False, 'gtol': grad_tol,
-                                    'maxiter': self.maxiter})
+                           options={'disp': disp, 'gtol': grad_tol,
+                                    'maxiter': max_iter})
 
-        elif self.algo == 'nelder-mead':
+        elif algo == 'nelder-mead':
             res = minimize(get_neg_likelihood, x_0, args,
                            method='nelder-mead',
-                           options={'disp': False,
-                                    'maxiter': self.maxiter,
+                           options={'disp': disp,
+                                    'maxiter': max_iter,
                                     'xtol': x_tol})
         if res is None:
-            raise RuntimeError("Optimization failed for some reason.")
+            raise RuntimeError("Optimization failed for unknown reason.")
         self.hyps = res.x
         self.set_L_alpha()
         self.likelihood = -res.fun
@@ -410,6 +430,7 @@ environment and the environments in the training set."""
                                  hyp_labels=dictionary['hyp_labels'],
                                  par=dictionary['par'],
                                  maxiter=dictionary['maxiter'],
+                                 opt_params=dictionary['opt_params'],
                                  opt_algorithm=dictionary['algo'])
 
         # Save time by attempting to load in computed attributes


Test patch:
diff --git a/tests/test_gp.py b/tests/test_gp.py
index b014a74f..85270d2e 100644
--- a/tests/test_gp.py
+++ b/tests/test_gp.py
@@ -242,6 +242,16 @@ def test_serialization_method(two_body_gp, test_point):
 
         if isinstance(x, np.ndarray):
             assert np.equal(x, y).all()
+
+        elif isinstance(x, dict):
+            xkeys = set(x.keys())
+            ykeys = set(y.keys())
+            assert xkeys == ykeys
+
+            # Once keys are same determine if all values are equal
+            for xk in sorted(list(xkeys)):
+                assert x[xk] == y[xk]
+
         elif hasattr(x, '__len__'):
 
             if isinstance(x[0], np.ndarray):
diff --git a/tests/test_gp_from_aimd.py b/tests/test_gp_from_aimd.py
index 4ea6986a..d6203546 100644
--- a/tests/test_gp_from_aimd.py
+++ b/tests/test_gp_from_aimd.py
@@ -21,8 +21,7 @@ def methanol_gp():
                                             1.70172923e-03]),
                              cutoffs=np.array([7, 7]),
                              hyp_labels=['l2', 's2', 'l3', 's3', 'n0'],
-                             maxiter=1,
-                             opt_algorithm='L-BFGS-B')
+                             maxiter=1)
     with open('./test_files/methanol_envs.json') as f:
         dicts = [loads(s) for s in f.readlines()]