954 modified Atoms.index

changes/954.change.rst

@@ -0,0 +1,5 @@
+Improved Atoms.index usability
+
+Now one gets an empty list if no atoms are found.
+One can also get the combined indices of several
+atom types, e.g., `atoms.index("C", "H")`.

src/sisl/_core/atom.py

@@ -759,8 +759,40 @@
         return uZ[self.species]
 
     def index(self, atom):
-        """Return the indices of the atom object"""
-        return (self._species == self.species_index(atom)).nonzero()[0]
+        """Return indices of atoms matching the specified atom identifier(s).
+
+        Parameters
+        ----------
+        atom : str, Atom, int, or lists hereof
+            Atom identifier (element string, Atom instance, or species_index).
+            It can also be a list of identifiers for matching against any of the entries.
+
+        Returns
+        -------
+        numpy.ndarray
+            Indices of all matching atoms. Returns an empty array if no match is found.
+
+        Notes
+        -----
+        Unlike Python's ``list.index``, this method does **not** raise a ``ValueError``
+        when no match is found. Instead, it returns an empty array.
+        To check for existence, use e.g. ``if idx.size > 0: ...``.
+        """
+
+        if not isinstance(atom, (list, tuple, np.ndarray)):
+            atom = [atom]
+
+        idx = []
+        for a in atom:
+            try:
+                idx.append((self._species == self.species_index(a)).nonzero()[0])
+            except KeyError:
+                pass
+
+        if len(idx) > 0:
+            idx = np.unique(idx)
+
+        return idx
 
     def species_index(self, atom):
         """Return the species index of the atom object"""

src/sisl/_core/tests/test_atoms.py

@@ -297,5 +297,10 @@ def test_charge_diff():
 
 def test_index1():
     atom = Atoms(["C", "Au"])
-    with pytest.raises(KeyError):
-        atom.index(Atom("B"))
+    assert atom.index("B") == []
+    assert atom.index("C") == [0]
+    assert atom.index(Atom(79)) == [1]
+    assert atom.index(1) == [1]
+    assert atom.index(["C", "C"]) == [0]
+    assert atom.index(["B", "C", "C"]) == [0]
+    assert np.allclose(atom.index(["Au", "C"]), [0, 1])