works when tested on C6H6

gofee.py

@@ -66,6 +66,10 @@ class GOFEE():
 
     kappa: 
     
+    N_relax_final_pop: int or None.
+    If not None the best 'N_relax_final_pop' structures in the
+    population is relaxed after the specified number of search
+    iterations is reached.
     """
     def __init__(self, structures=None,
                  calc=None,
@@ -74,14 +78,15 @@ class GOFEE():
                  candidate_generator=None,
                  trajectory=None,
                  kappa=2,
-                 Niter=200,
+                 max_steps=200,
                  Ninit=10,
                  dmax_cov=3.5,
                  Ncandidates=30,
                  population_size=5,
                  dualpoint=True,
                  min_certainty=0.7,
-                 restart=None):
+                 restart=None,
+                 N_relax_final_pop=None):
         
         if structures is None:
             assert startgenerator is not None
@@ -124,7 +129,7 @@ class GOFEE():
                     indices_fixed = constraint.get_indices()
                     self.n_to_optimize -= len(indices_fixed)
                     break
-            
+        
         if candidate_generator is not None:
             self.candidate_generator = candidate_generator
         else:
@@ -141,13 +146,21 @@ class GOFEE():
         self.master = self.comm.rank == 0
 
         self.kappa = kappa
-        self.Niter = Niter
+        self.max_steps = max_steps
         self.Ninit = Ninit
         self.dmax_cov = dmax_cov
         self.Ncandidates = Ncandidates
         self.dualpoint = dualpoint
         self.min_certainty = min_certainty
         self.restart = restart
+        if N_relax_final_pop is None:
+            self.N_relax_final_pop = 0
+        else:
+            try:
+                self.N_relax_final_pop = int(N_relax_final_pop)
+            except TypeError:
+                raise
+                
         
         if isinstance(trajectory, str):
             self.trajectory = Trajectory(filename=trajectory, mode='a', master=self.master)
@@ -177,6 +190,7 @@ class GOFEE():
         generated using self.startgenerator (if supplied), otherwise
         they are generated by heavily rattling the supplied structures.
         """
+        
         # Collect potentially supplied structures and evaluate
         # energies and forces if not present.
         structures_init = []
@@ -191,9 +205,9 @@ class GOFEE():
                         continue
                 a = self.evaluate(a)
                 structures_init.append(a)
-
+        
         Nremaining = self.Ninit - len(structures_init)
-    
+        
         if Nremaining > 0 and self.startgenerator is None:
             # Initialize rattle-mutation for all atoms.
             rattle = RattleMutation(self.n_to_optimize,
@@ -221,7 +235,7 @@ class GOFEE():
     def run(self):
         self.get_initial_structures()
 
-        while self.steps < self.Niter:
+        while self.steps < self.max_steps:
             self.print_master('steps:', self.steps)
             self.train_surrogate()
             self.update_population()
@@ -259,6 +273,8 @@ class GOFEE():
                 print('anew pred:', anew.info['key_value_pairs']['Epred'], anew.info['key_value_pairs']['Epred_std'])
                 print('E_true:', [a.get_potential_energy() for a in a_add])
                 print('pop:', [a.get_potential_energy() for a in self.population.pop])
+
+        self.relax_final_population()
             
     def get_dualpoint(self, a, lmax=0.10, Fmax_flat=5):
         """Returns dual-point structure, i.e. the original structure
@@ -401,6 +417,14 @@ class GOFEE():
 
         return a
 
+    def relax_final_population(self):
+        if self.N_relax_final_pop > 0:
+            relaxed_population_trajectory = Trajectory(filename='relaxed_final_population.traj',
+                                                       mode='a', master=self.master)
+            for a in self.population.pop[:self.N_relax_final_pop]:
+                a = relax(a, self.calc, Fmax=0.05, steps_max=20)
+                relaxed_population_trajectory.write(a)
+
     def write(self, a):
         if self.trajectory is not None:
             self.trajectory.write(a)