Tag atoms and count their exchanges

ase_sim.py

@@ -3,6 +3,7 @@ import numpy as np
 import math
 
 import ase.lattice.cubic
+import ase.data
 import ase.io
 import ase.units
 
@@ -18,6 +19,7 @@ class AseSim:
         self.atoms = None
         self.dyn = None
         self.run = None
+        self.tagged_elements = {}
 
     def create_lattice(self, lattice_size, basis, empty_space, pot_filename, rand_ox=False):
         self.atoms = ase.lattice.cubic.BodyCenteredCubic(basis, size=lattice_size, pbc=(True, True, False))
@@ -41,23 +43,22 @@ class AseSim:
         nbor_list = NeighborList(nbor_cutoffs, skin=0.01, self_interaction=False, bothways=True)
         nbor_list.update(self.atoms)
 
-        pot_path = os.path.join('data', pot_filename)
-
-        self.dyn = SajuhakDynamics(self.atoms, nbor_list, pot_path, rand_ox)
-        self.run = self.dyn.run
-        self.atoms.set_calculator(KokkoPotential(nbor_list, pot_path))
-
         # initialize atom modifications
-        symbols = self.dyn.atoms.get_chemical_symbols()
-        positions = self.dyn.atoms.get_positions()
+        symbols = self.atoms.get_chemical_symbols()
+        positions = self.atoms.get_positions()
         max_miller = lattice_size[2] * (1 - empty_space)
 
         # replace atoms with empty spaces above a certain z position
         for i, position in enumerate(positions):
-            if position[2] > self.dyn.atoms.millerbasis[2][2]*max_miller:
+            if position[2] > self.atoms.millerbasis[2][2]*max_miller:
                 symbols[i] = 'X'
 
-        self.dyn.atoms.set_chemical_symbols(symbols)
+        # store modifications and initialize dyn and calc
+        pot_path = os.path.join('data', pot_filename)
+        self.atoms.set_chemical_symbols(symbols)
+        self.dyn = SajuhakDynamics(self.atoms, nbor_list, pot_path, rand_ox)
+        self.run = self.dyn.run
+        self.atoms.set_calculator(KokkoPotential(nbor_list, pot_path))
 
     # @profile  # add -m memory_profiler to interpreter options
     def dope(self, dopants, basis):
@@ -109,3 +110,26 @@ class AseSim:
                 if self.dyn.atoms.get_atomic_numbers()[nbor_idx] == 0:
                     empty_nbor_list[idx] += 1
         self.dyn.atoms.set_array('empty_nbors', empty_nbor_list)
+
+    def tag_elements(self):
+        """
+        Change the tags of random elements in the lattice so they can be tracked
+
+        :return: None
+        """
+        tag_target = 5  # how many atoms to tag of each element
+        number_list = self.atoms.get_atomic_numbers()
+        element_list = np.unique(number_list)
+        tag_list = self.atoms.get_tags()
+        tagged_counts = {}  # dict of amount of tagged atoms per element
+        tag_index = 1
+        while sum(tagged_counts.values()) < tag_target * len(element_list):
+            rand_idx = np.random.randint(len(self.atoms))
+            rand_number = number_list[rand_idx]
+            tagged_count = tagged_counts.get(rand_number, 0)
+            if not tagged_count >= tag_target:
+                tag_list[rand_idx] = tag_index
+                self.tagged_elements[tag_index] = rand_number
+                tag_index += 1
+                tagged_counts[rand_number] = tagged_count + 1
+        self.atoms.set_tags(tag_list)