Add train_hyps to the on-the-fly settings

A variable controlling when the training is performed is added.
This can give improvements if the hyperparameter training gets
stucked into local minima due to few environments in the dataset.

Author: bastonero

Modules/Ensemble.py

@@ -194,6 +194,7 @@ def __init__(self, dyn0, T0, supercell = None, **kwargs):
         self.checkpt_files = None
         self.write_model = None
         self.init_atoms = None
+        self.train_hyps = None
 
         # The frequencies and polarizations of the ensemble
         # In q space
@@ -4239,6 +4240,7 @@ def set_otf(
         update_threshold: float | None = None,
         # other args
         build_mode="bayesian",
+        train_hyps: tuple = (100,120), 
     ):
         """Set on-the-fly training.
         
@@ -4299,6 +4301,12 @@ def set_otf(
         ]
 
         self.write_model = write_model
+        
+        if train_hyps[0] == "inf":
+            train_hyps[0] = np.inf
+        if train_hyps[1] == "inf":
+            train_hyps[1] = np.inf
+        self.train_hyps = train_hyps
 
 
 #-------------------------------------------------------------------------------

Modules/aiida_ensemble.py

@@ -171,7 +171,8 @@ def compute_ensemble( # pylint: disable=arguments-renamed
                 # ================ TRAIN SECTION ================ #
                 if self.gp_model is not None:
                     if dft_counts > 0:
-                        self._train_gp()
+                        if self.train_hyps[0] <= len(self.gp_model.training_data) <= self.train_hyps[1]:
+                            self._train_gp()
                         self._write_model()
 
                 sys.stdout.flush()

tests/aiida_ensemble/test_otf_flare.py

@@ -74,6 +74,7 @@ def test_no_otf(generate_ensemble):
     assert ensemble.checkpt_files is None
     assert ensemble.write_model is None 
     assert ensemble.init_atoms is None 
+    assert ensemble.train_hyps is None 
 
 
 def test_set_otf(generate_ensemble):
@@ -83,6 +84,7 @@ def test_set_otf(generate_ensemble):
     ensemble.set_otf(flare_calc, max_atoms_added=-1)
 
     assert ensemble.gp_model is not None
+    assert ensemble.train_hyps == (100,120)
 
 
 def test_compute_properties(generate_ensemble):