Merge pull request #1 from Exabyte-io/improvement/atoms-clone-remove-…

…sites-in-place

improvement: add Atoms.clone() and remove sites by indices, including in-place

.gitignore

@@ -103,4 +103,8 @@ ENV/
 # mypy
 .mypy_cache/
 
-.pytest_cache/
+.pytest_cache/
+
+# pycharm
+.idea
+

jarvis/core/atoms.py

@@ -698,14 +698,23 @@ def from_dict(self, d={}):
 
     def remove_site_by_index(self, site=0):
         """Remove an atom by its index number."""
+        return self.remove_sites_by_indices(indices=[site])
+
+    def remove_sites_by_indices(self, indices=[0], in_place=False):
+        """Remove multiple atoms by their corresponding indices number."""
         new_els = []
         new_coords = []
         new_props = []
         for ii, i in enumerate(self.frac_coords):
-            if ii != site:
+            if ii in indices:
                 new_els.append(self.elements[ii])
                 new_coords.append(self.frac_coords[ii])
                 new_props.append(self.props[ii])
+        if in_place:
+            self.elements = new_els
+            self.coords = new_coords
+            self.props = new_props
+            return self
         return Atoms(
             lattice_mat=self.lattice_mat,
             elements=new_els,
@@ -1431,6 +1440,18 @@ def get_string(self, cart=True, sort_order="X"):
         result = header + middle + rest
         return result
 
+    def clone(self):
+        """Clones the class instance."""
+        return Atoms(
+            lattice_mat=self.lattice_mat,
+            elements=self.elements,
+            coords=self.frac_coords,
+            props=self.props,
+            cartesian=self.cartesian,
+            show_props=self.show_props,
+        )
+
+
 
 class VacuumPadding(object):
     """Adds vaccum padding to make 2D structure or making molecules."""

jarvis/tests/testfiles/core/test_atoms.py

@@ -13,6 +13,12 @@
 import tarfile
 import tempfile
 
+FIXTURES = {
+    "lattice_mat": [[2.715, 2.715, 0], [0, 2.715, 2.715], [2.715, 0, 2.715]],
+    "coords": [[0, 0, 0], [0.25, 0.2, 0.25]],
+    "elements": ["Si", "Si"],
+}
+
 new_file, filename = tempfile.mkstemp()
 
 
@@ -76,10 +82,7 @@ def test_from_cif():
 
 def test_basic_atoms():
 
-    box = [[2.715, 2.715, 0], [0, 2.715, 2.715], [2.715, 0, 2.715]]
-    coords = [[0, 0, 0], [0.25, 0.2, 0.25]]
-    elements = ["Si", "Si"]
-    Si = Atoms(lattice_mat=box, coords=coords, elements=elements)
+    Si = Atoms(lattice_mat=FIXTURES["lattice_mat"], coords=FIXTURES["coords"], elements=FIXTURES["elements"])
     dim = get_supercell_dims(Si)
     build_xanes_poscar(atoms=Si, filename_with_prefix=True)
     assert dim == [3, 3, 3]
@@ -208,6 +211,14 @@ def test_basic_atoms():
     cmd = "rm atoms.xyz POSCAR atoms.cif"
     os.system(cmd)
 
-
-# test_basic_atoms()
-# def test_basic_atoms():
+def test_clone():
+    Si = Atoms(lattice_mat=FIXTURES["lattice_mat"], coords=FIXTURES["coords"], elements=FIXTURES["elements"])
+    Si2 = Si.clone()
+    assert (Si2.lattice_mat == Si.lattice_mat and Si2.coords == Si.coords and Si2.elements == Si.elements
+            and Si2.props == Si.props and Si2.cartesian == Si.cartesian and Si2.show_props == Si.show_props)
+
+def test_remove_sites_by_indices():
+    Si = Atoms(lattice_mat=FIXTURES["lattice_mat"], coords=FIXTURES["coords"], elements=FIXTURES["elements"])
+    Si_supercell = Si.make_supercell([2, 2, 2])
+    Si2_supercell_without_two_atoms = Si_supercell.remove_sites_by_indices(indices=[0, 1])
+    assert (Si2_supercell_without_two_atoms.num_atoms == 14)