solving some bugs with the pbc feature and creating utilities like resetting and finding NNs

new/adAtom.py

@@ -16,12 +16,12 @@ class AdAtom:
         self.bcs.append(bc)
 
 
-    def add_position_condition(self, position_min=0.0, position_max=1.0):
+    def add_position_condition(self, position_min=-np.inf, position_max=np.inf):
         pc = PositionCondition(self, position_min, position_max)
         self.add_boundary_condition(pc)
 
 
-    def add_distance_condition(self, neighbour, distance_min=0.0, distance_max=np.infty):
+    def add_distance_condition(self, neighbour, distance_min=0.0, distance_max=np.inf):
         dc = DistanceCondition(self, neighbour, distance_min, distance_max)
         self.add_boundary_condition(dc)
 

new/boundaryCondition.py

@@ -31,7 +31,7 @@ class PositionCondition(BoundaryCondition):
     """
     Derived class for position boundary conditions.
     """
-    def __init__(self, atomA='X', minPos=0.0, maxPos=np.infty):
+    def __init__(self, atomA='X', minPos=-np.inf, maxPos=np.inf):
         self.bc_type = 'pos'
         self.atoms = [atomA]
         self.minValue = minPos
@@ -50,7 +50,7 @@ class DistanceCondition(BoundaryCondition):
     """
     Derived class for distance boundary conditions.
     """
-    def __init__(self, atomA='X', atomB='Y', minDist=0.0, maxDist=np.infty):
+    def __init__(self, atomA='X', atomB='Y', minDist=0.0, maxDist=np.inf):
         self.bc_type = 'dist'
         self.atoms = [atomA, atomB]
         self.minValue = minDist

new/codered.py

@@ -4,6 +4,7 @@ from .adAtom import AdAtom
 from .sideChainGen import SideChain
 from .boundaryCondition import *
 from .reduction import *
+from .utils import get_periodic_images,find_NNs
 from .output import printProgressBar
 
 
@@ -13,6 +14,7 @@ class CodeRed:
         self.n_adAtoms = 0
         self.current_index = 0
         self.combinations = []
+        self.structures = []
         self.surface = Atoms()
         self.surfaceRed = Atoms()
         self.sideChains = []
@@ -40,9 +42,36 @@ class CodeRed:
         return adAtom
 
 
-    def add_sidechain(self, sidechain, connector=0, connect_to=-1, distance=2.0, mode='auto'):
+    def get_adAtom(self, key):
+        if isinstance(key, AdAtom):
+            if key in self.adAtoms:
+                return key
+            else:
+                return None
+        elif isinstance(key, int):
+            try:
+                adAtom = self.adAtoms[key]
+            except IndexError:
+                adAtom = None
+                pass
+            return adAtom
+        elif isinstance(key, str):
+            found = []
+            for adAtom in self.adAtoms:
+                if adAtom.label==key or adAtom.species==key:
+                    found.append(adAtom)
+                    
+            if len(found)==0:
+                return None
+            elif len(found)==1:
+                return found[0]
+            else:
+                return found
+
+
+    def add_sidechain(self, sidechain, base=np.zeros(3), connect_to=-1, distance=2.0, mode='auto'):
         # create sidechain object and add it to the list of side chains
-        new_sc = SideChain(sidechain, connector, connect_to, distance, mode)
+        new_sc = SideChain(sidechain, base, connect_to, distance, mode)
         self.sideChains.append(new_sc)
         return new_sc
 
@@ -107,19 +136,8 @@ class CodeRed:
 
             # add periodic images of the positions if not the current atom
             if atom!=self.adAtoms[self.current_index] and atom!=self.current_index:
-                # create array with all possible lattice vector combinations
-                pzm = np.array([-1,0,1])
-                pbc_combinations = np.array(np.meshgrid(pzm,pzm,pzm,indexing='ij')).reshape(3,27).T
-                # extract combinations that are compatible with pbc settings
-                pbc_combinations = np.unique(pbc_combinations*self.surfaceRed.pbc, axis=0)
-                # convert to actual position shifts
-                pbc_shifts = np.dot(pbc_combinations, self.surfaceRed.cell)
-            
-                # add positions accordingly
-                pbc_pos = []
-                for shift in pbc_shifts:
-                    pbc_pos.append(atom_positions+shift)
-            
+                # calculate periodic images
+                pbc_pos = get_periodic_images(atom_positions,surf.get_cell(),surf.get_pbc())
                 # add positions to the list
                 positions.append(np.hstack(pbc_pos))
 
@@ -183,7 +201,7 @@ class CodeRed:
                 # make sure that position conditions are met
                 # (in case generate_mesh was not used for mesh generation)
                 for pc in adatom.get_position_conditions():
-                    adatom.mesh = adatom.mesh[pc.check_condition(mesh)]
+                    adatom.mesh = adatom.mesh[pc.check_condition(adatom.mesh)]
                     # stop if there are no positions left
                     if len(adatom.mesh)==0:
                         return
@@ -249,6 +267,9 @@ class CodeRed:
             # finally increase current index by 1 (only happens if everything worked out)
             self.current_index += 1
 
+        # save created structures
+        self.structures = self.get_structures()
+
         # generate informative output
         print("Done! Generated {:d} structures in total.".format(len(self.combinations)))
 
@@ -290,3 +311,8 @@ class CodeRed:
         print("Done! Reduced to {:d} structures in total.".format(len(self.combinations)))
 
         return
+
+
+    def reset(self):
+        self.current_index = 0
+        self.combinations = []

new/sideChainGen.py

@@ -179,9 +179,11 @@ def findDIAdirs(basis, NNpos):
 
 
 class SideChain:
-    def __init__(self, atoms=Atoms(), connector=0, connect_to=-1, distance=0., mode='auto'):
-        self.atoms = atoms.copy()
-        self.connector = connector
+    def __init__(self, atoms=Atoms(), base=np.zeros(3), connect_to=-1, distance=0., mode='auto'):
+        sc_atoms = atoms.copy()
+        sc_atoms.set_positions(sc_atoms.get_positions()-base)
+        self.atoms = sc_atoms.copy()
+        self.base = base
         self.connect_to = connect_to
         self.distance = distance
         self.mode = mode

new/utils.py

+import numpy as np
+from ase.cell import Cell
+
+
+def get_periodic_images(positions, cell=Cell(np.eye(3)), pbc=np.array([True,True,True])):
+    # create array with all possible lattice vector combinations
+    pzm = np.array([-1,0,1])
+    pbc_combinations = np.array(np.meshgrid(pzm,pzm,pzm,indexing='ij')).reshape(3,27).T
+    
+    # extract combinations that are compatible with pbc settings
+    pbc_combinations = np.unique(pbc_combinations*pbc, axis=0)
+    
+    # convert to actual position shifts
+    pbc_shifts = np.dot(pbc_combinations, cell)
+    
+    # add positions accordingly
+    pbc_pos = []
+    for shift in pbc_shifts:
+        pbc_pos.append(positions+shift)
+    
+    # return all positions (periodic images and original ones in one array)
+    return pbc_pos
+
+
+def find_NNs(center, candidates, cutoff=np.infty, n_max=np.infty, cell=Cell(np.eye(3)), pbc=np.array([False,False,False])):
+    # add periodic images of the candidates
+    pbc_candidates = np.reshape(np.array(get_periodic_images(candidates, cell, pbc)), (-1,3))
+    
+    # get nearest neighbours by distance
+    distances = np.linalg.norm(pbc_candidates-center, axis=-1)
+    NNs = pbc_candidates[distances<cutoff]
+
+    # reduce to the n_max nearest neighbours
+    if n_max<np.infty:
+        NNs = NNs[np.argsort(distances)[:n_max]]
+    
+    return NNs