improvement: add Atoms.clone() and remove sites by indices, including in-place. Use develop branch for the PR.

.gitignore

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

jarvis/core/atoms.py

@@ -1,4 +1,5 @@
 """This module provides classes to specify atomic structure."""
+
 import numpy as np
 from jarvis.core.composition import Composition
 from jarvis.core.specie import Specie, atomic_numbers_to_symbols
@@ -698,14 +699,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 not 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,
@@ -1268,49 +1278,50 @@ def make_supercell(self, dim=[2, 2, 2]):
         dim = np.array(dim)
         if dim.shape == (3, 3):
             dim = np.array([int(np.linalg.norm(v)) for v in dim])
-        coords = self.frac_coords
-        all_symbs = self.elements  # [i.symbol for i in s.species]
-        nat = len(coords)
-
-        new_nat = nat * dim[0] * dim[1] * dim[2]
-        new_coords = np.zeros((new_nat, 3))
-        new_symbs = []  # np.chararray((new_nat))
-        props = []  # self.props
-
-        ct = 0
-        for i in range(nat):
-            for j in range(dim[0]):
-                for k in range(dim[1]):
-                    for m in range(dim[2]):
-                        props.append(self.props[i])
-                        new_coords[ct][0] = (coords[i][0] + j) / float(dim[0])
-                        new_coords[ct][1] = (coords[i][1] + k) / float(dim[1])
-                        new_coords[ct][2] = (coords[i][2] + m) / float(dim[2])
-                        new_symbs.append(all_symbs[i])
-                        ct = ct + 1
-
-        nat = new_nat
-
-        nat = len(coords)  # int(s.composition.num_atoms)
-        lat = np.zeros((3, 3))
-        box = self.lattice_mat
-        lat[0][0] = dim[0] * box[0][0]
-        lat[0][1] = dim[0] * box[0][1]
-        lat[0][2] = dim[0] * box[0][2]
-        lat[1][0] = dim[1] * box[1][0]
-        lat[1][1] = dim[1] * box[1][1]
-        lat[1][2] = dim[1] * box[1][2]
-        lat[2][0] = dim[2] * box[2][0]
-        lat[2][1] = dim[2] * box[2][1]
-        lat[2][2] = dim[2] * box[2][2]
-        super_cell = Atoms(
-            lattice_mat=lat,
-            coords=new_coords,
-            elements=new_symbs,
-            props=props,
-            cartesian=False,
-        )
-        return super_cell
+        return self.make_supercell_matrix(dim)
+        # coords = self.frac_coords
+        # all_symbs = self.elements  # [i.symbol for i in s.species]
+        # nat = len(coords)
+
+        # new_nat = nat * dim[0] * dim[1] * dim[2]
+        # new_coords = np.zeros((new_nat, 3))
+        # new_symbs = []  # np.chararray((new_nat))
+        # props = []  # self.props
+
+        # ct = 0
+        # for i in range(nat):
+        #    for j in range(dim[0]):
+        #        for k in range(dim[1]):
+        #            for m in range(dim[2]):
+        #                props.append(self.props[i])
+        #                new_coords[ct][0] = (coords[i][0] + j) / float(dim[0])
+        #                new_coords[ct][1] = (coords[i][1] + k) / float(dim[1])
+        #                new_coords[ct][2] = (coords[i][2] + m) / float(dim[2])
+        #                new_symbs.append(all_symbs[i])
+        #                ct = ct + 1
+
+        # nat = new_nat
+
+        # nat = len(coords)  # int(s.composition.num_atoms)
+        # lat = np.zeros((3, 3))
+        # box = self.lattice_mat
+        # lat[0][0] = dim[0] * box[0][0]
+        # lat[0][1] = dim[0] * box[0][1]
+        # lat[0][2] = dim[0] * box[0][2]
+        # lat[1][0] = dim[1] * box[1][0]
+        # lat[1][1] = dim[1] * box[1][1]
+        # lat[1][2] = dim[1] * box[1][2]
+        # lat[2][0] = dim[2] * box[2][0]
+        # lat[2][1] = dim[2] * box[2][1]
+        # lat[2][2] = dim[2] * box[2][2]
+        # super_cell = Atoms(
+        #    lattice_mat=lat,
+        #    coords=new_coords,
+        #    elements=new_symbs,
+        #    props=props,
+        #    cartesian=False,
+        # )
+        # return super_cell
 
     def get_lll_reduced_structure(self):
         """Get LLL algorithm based reduced structure."""
@@ -1431,6 +1442,17 @@ 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."""
@@ -1860,18 +1882,18 @@ def to_optimade(
         info_at["cartesian_site_positions"] = atoms.cart_coords[order].tolist()
         info_at["nperiodic_dimensions"] = 3
         # info_at["species"] = atoms.elements
-        info_at[
-            "species"
-        ] = self.get_optimade_species()  # dict(atoms.composition.to_dict())
+        info_at["species"] = (
+            self.get_optimade_species()
+        )  # dict(atoms.composition.to_dict())
         info_at["elements_ratios"] = list(
             atoms.composition.atomic_fraction.values()
         )
         info_at["structure_features"] = []
         info_at["last_modified"] = str(now)
         # info_at["more_data_available"] = True
-        info_at[
-            "chemical_formula_descriptive"
-        ] = atoms.composition.reduced_formula
+        info_at["chemical_formula_descriptive"] = (
+            atoms.composition.reduced_formula
+        )
         info_at["dimension_types"] = [1, 1, 1]
         info["attributes"] = info_at
         return info

jarvis/tests/testfiles/core/test_atoms.py

@@ -7,12 +7,18 @@
     OptimadeAdaptor,
 )
 
-
+import numpy as np
 import os
 from jarvis.db.figshare import get_jid_data, data
 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()
 
 
@@ -55,11 +61,26 @@
     "POSCAR",
 )
 
-cif_example = os.path.join(os.path.dirname(__file__), "1000052.cif",)
-cif_example2 = os.path.join(os.path.dirname(__file__), "Bacomp.cif",)
-cif_example3 = os.path.join(os.path.dirname(__file__), "mock.cif",)
-cif_example4 = os.path.join(os.path.dirname(__file__), "exp_000034.cif",)
-cif_example5 = os.path.join(os.path.dirname(__file__), "1000000.cif",)
+cif_example = os.path.join(
+    os.path.dirname(__file__),
+    "1000052.cif",
+)
+cif_example2 = os.path.join(
+    os.path.dirname(__file__),
+    "Bacomp.cif",
+)
+cif_example3 = os.path.join(
+    os.path.dirname(__file__),
+    "mock.cif",
+)
+cif_example4 = os.path.join(
+    os.path.dirname(__file__),
+    "exp_000034.cif",
+)
+cif_example5 = os.path.join(
+    os.path.dirname(__file__),
+    "1000000.cif",
+)
 
 
 def test_from_cif():
@@ -76,10 +97,11 @@ 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]
@@ -109,8 +131,8 @@ def test_basic_atoms():
     prim = Si.get_primitive_atoms
     print(prim.cart_coords)
     conv = Si.get_conventional_atoms
-    spgn=Si.get_spacegroup
-    comp=compare_atoms(atoms1=prim,atoms2=conv)
+    spgn = Si.get_spacegroup
+    comp = compare_atoms(atoms1=prim, atoms2=conv)
     assert round(prim.cart_coords[0][0], 2) == round(4.37815150, 2)
     # print ('raw_distance_matrix', prim.raw_distance_matrix)
     # print ('raw_distance_matrix', Si.raw_distance_matrix)
@@ -209,5 +231,37 @@ def test_basic_atoms():
     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()
+    np.testing.assert_array_equal(Si2.lattice_mat, Si.lattice_mat)
+    np.testing.assert_array_equal(Si2.coords, Si.coords)
+    assert Si2.props == Si.props
+    assert Si2.elements == Si.elements
+    assert Si2.cartesian == Si.cartesian
+    assert 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])
+    print("Si_supercell", Si_supercell)
+    Si2_supercell_without_two_atoms = Si_supercell.remove_sites_by_indices(
+        indices=[0, 1]
+    )
+    print(
+        "Si2_supercell_without_two_atoms.num_atoms",
+        Si2_supercell_without_two_atoms.num_atoms,
+    )
+    assert Si2_supercell_without_two_atoms.num_atoms == 14
+
+
+# test_remove_sites_by_indices()