position constraint that both constraints mutations and surrogate relaxations

bfgslinesearch_constrained.py

@@ -13,8 +13,9 @@ import sys
 class BFGSLineSearch_constrained(BFGSLineSearch):
     def __init__(self, atoms, pos_init=None, restart=None, logfile='-', maxstep=.2,
                  trajectory=None, c1=0.23, c2=0.46, alpha=10.0, stpmax=50.0,
-                 master=None, force_consistent=None, zlim=None,
-                 blmin=None, blmax=None, max_relax_dist=4.0, rk=None):
+                 master=None, force_consistent=None,
+                 blmin=None, blmax=None, max_relax_dist=4.0, 
+                 position_constraint=None, rk=None):
         """
         add maximum displacement of single atoms to BFGSLineSearch:
 
@@ -24,9 +25,9 @@ class BFGSLineSearch_constrained(BFGSLineSearch):
 
         self.rk = rk  # for testing
         
-        self.zlim = zlim
         self.blmin = blmin
         self.blmax = blmax
+        self.position_constraint=position_constraint
         
         self.max_relax_dist = max_relax_dist
         if pos_init is not None:
@@ -59,12 +60,12 @@ class BFGSLineSearch_constrained(BFGSLineSearch):
     def check_constraints(self):
         # Check if stop due to large displacement
         valid_displace = self.check_displacement()
-        # Check if stop due to zlim
-        valid_z = self.check_zcoord()
+        # Check if stop due to position-constraint
+        valid_pos = self.check_positions()
         # Check if stop due to invalid bondlengths
         valid_bondlengths = self.check_bondlengths()
         
-        if not valid_displace or not valid_z or not valid_bondlengths:
+        if not valid_displace or not valid_pos or not valid_bondlengths:
             #self.r0 = self.atoms_prior.get_positions()
             return True
     
@@ -77,19 +78,18 @@ class BFGSLineSearch_constrained(BFGSLineSearch):
                 valid_displace = False
         return valid_displace
 
-    def check_zcoord(self):
-        valid_z = True
-        if self.zlim is not None:
+    def check_positions(self):
+        valid_pos = True
+        if self.position_constraint is not None:
             # get indices of non-fixed atoms
             indices = np.arange(self.atoms.get_number_of_atoms())
             for constraint in self.atoms.constraints:
                 if isinstance(constraint, FixAtoms):
                     indices_fixed = constraint.get_indices()
-                    indices = np.delete(np.arange(self.atoms.get_number_of_atoms()), indices_fixed)
-            pos_z = self.atoms.positions[indices,2]
-            if np.any(pos_z < self.zlim[0]) or np.any(pos_z > self.zlim[1]):
-                valid_z = False
-        return valid_z
+                    indices_not_fixed = np.delete(np.arange(self.atoms.get_number_of_atoms()), indices_fixed)
+            pos_not_fixed = self.atoms.positions[indices_not_fixed]
+            valid_pos = self.position_constraint.check_if_valid(pos_not_fixed)
+        return valid_pos
 
     def check_bondlengths(self):
         valid_bondlengths = True
@@ -100,7 +100,7 @@ class BFGSLineSearch_constrained(BFGSLineSearch):
         return valid_bondlengths
 
 
-def relax(structure, calc, Fmax=0.05, steps_max=200, max_relax_dist=None):
+def relax(structure, calc, Fmax=0.05, steps_max=200, max_relax_dist=None, position_constraint=None):
     a = structure.copy()
     # Set calculator 
     a.set_calculator(calc)
@@ -111,7 +111,8 @@ def relax(structure, calc, Fmax=0.05, steps_max=200, max_relax_dist=None):
         dyn = BFGSLineSearch_constrained(a,
                                          logfile=None,
                                          pos_init=pos_init,
-                                         max_relax_dist=max_relax_dist)
+                                         max_relax_dist=max_relax_dist,
+                                         position_constraint=position_constraint)
         dyn.run(fmax = Fmax, steps = steps_max)
     except Exception as err:
         print('Error in surrogate-relaxation:', err, flush=True)

candidate_operations/basic_mutations.py

@@ -95,7 +95,7 @@ class RattleMutation(OffspringOperation):
         # Perform rattle operations in sequence.
         for i in indices_to_rattle:
             posi_0 = np.copy(a.positions[i])
-            for _ in range(100):
+            for _ in range(200):
                 # Perform rattle
                 pos_add = pos_add_sphere(self.rattle_range)
                 a.positions[i] += pos_add

candidate_operations/candidate_generation.py

@@ -118,6 +118,9 @@ class OffspringOperation(ABC):
                                Atom {index_shortest_bond} has bond with d={d_shortest_bond}d_covalent"""
                     warnings.warn(text)
 
+    def set_constraints(self, constraints):
+        self.constraints = constraints
+
     def check_constraints(self, indices=None):
         if self.constraints is not None:
             valid = self.constraints.check_if_valid(indices)
@@ -180,12 +183,18 @@ class CandidateGenerator():
                 anew = op_list[to_use].finalize(anew, successfull=False)
         return anew
 
+    def set_constraints(self, constraints):
+        for op_list in self.operations:
+            for op in op_list:
+                op.set_constraints(constraints)
+        self.constraints = constraints
+
     def train(self, data):
         """ Method to train all trainable operations in
         self.operations.
         """
-        for oplist in self.operations:
-            for operation in oplist:
+        for op_list in self.operations:
+            for operation in op_list:
                 operation.train(data)
 
 
@@ -231,9 +240,14 @@ class OperationConstraint():
             pos = positions.reshape(-1,3)
         else:
             pos = positions
-        
+
         if self.box is not None:
-            pass
+            p0, V = self.box
+            p_rel = pos - p0  # positions relative to box anchor.
+            V_inv = np.linalg.inv(V)
+            p_box = p_rel @ V_inv  # positions in box-vector basis.
+            if (np.any(p_box < 0) or np.any(p_box > 1)):
+                return False
         if self.xlim is not None:
             if (np.any(pos[:,0] < self.xlim[0]) or 
                     np.any(pos[:,0] > self.xlim[1])):

gofee.py

@@ -100,6 +100,11 @@ class GOFEE():
         considdered for evaluation. (in units of the maximum possible
         uncertainty.)
 
+    position_constraint: OperationConstraint object
+        Enforces constraints on the positions of the "free" atoms
+        in the search. The constraint is enforces both during
+        mutation/crossover operations and during surrogate-relaxation.
+
     restart: str
         Filename for restart file.
     """
@@ -108,8 +113,6 @@ class GOFEE():
                  gpr=None,
                  startgenerator=None,
                  candidate_generator=None,
-                 trajectory='structures.traj',
-                 logfile='search.log',
                  kappa=2,
                  max_steps=200,
                  Ninit=10,
@@ -118,6 +121,9 @@ class GOFEE():
                  population_size=10,
                  dualpoint=True,
                  min_certainty=0.7,
+                 position_constraint=None,
+                 trajectory='structures.traj',
+                 logfile='search.log',
                  restart='restart.pickl'):
         
         if structures is None:
@@ -146,6 +152,7 @@ class GOFEE():
         else:
             self.startgenerator = startgenerator
 
+        # Determine atoms to optimize
         if startgenerator is not None:
             self.n_to_optimize = len(self.startgenerator.stoichiometry)
         else:
@@ -156,6 +163,7 @@ class GOFEE():
                     self.n_to_optimize -= len(indices_fixed)
                     break
         
+        # Set up candidate-generator if not supplied
         if candidate_generator is not None:
             self.candidate_generator = candidate_generator
         else:
@@ -175,8 +183,12 @@ class GOFEE():
         self.Ncandidates = Ncandidates
         self.dualpoint = dualpoint
         self.min_certainty = min_certainty
+        self.position_constraint = position_constraint
         self.restart = restart
         
+        # Add position-constraint to candidate-generator
+        self.candidate_generator.set_constraints(position_constraint)
+
         if isinstance(trajectory, str):
             self.trajectory = Trajectory(filename=trajectory, mode='a', master=self.master)
             if self.restart:
@@ -192,7 +204,6 @@ class GOFEE():
         self.logfile = logfile
         self.log_msg = ''
 
-
         if restart is None or not isfile(restart):
             self.initialize()
 
@@ -207,6 +218,7 @@ class GOFEE():
             self.read()
             self.comm.barrier()
 
+
     def initialize(self):
         self.get_initial_structures()
         self.steps = 0
@@ -373,7 +385,9 @@ class GOFEE():
         surrogate potential.
         """
         calc = self.gpr.get_calculator(kappa)
-        a_relaxed = relax(a, calc, max_relax_dist=self.max_relax_dist, Fmax=Fmax, steps_max=steps)
+        a_relaxed = relax(a, calc, Fmax=Fmax, steps_max=steps,
+                          max_relax_dist=self.max_relax_dist,
+                          position_constraint=self.position_constraint)
         # Evaluate uncertainty
         E, Estd = self.gpr.predict_energy(a_relaxed, eval_std=True)