Merge pull request #672 from ISISNeutronMuon/maciej/simplify-atom-sel…

…ection

New atom selection mechanism

MDANSE/Src/MDANSE/Chemistry/ChemicalSystem.py

@@ -15,7 +15,7 @@
 #
 
 from __future__ import annotations
-from typing import List, Tuple, Dict, Any
+from typing import List, Tuple, Dict, Any, Set
 import copy
 from functools import reduce
 
@@ -229,6 +229,11 @@ def number_of_atoms(self) -> int:
         """The number of non-ghost atoms in the ChemicalSystem."""
         return self._total_number_of_atoms
 
+    @property
+    def all_indices(self) -> Set[int]:
+        """The number of non-ghost atoms in the ChemicalSystem."""
+        return set(self._atom_indices)
+
     @property
     def total_number_of_atoms(self) -> int:
         """The number of all atoms in the ChemicalSystem, including ghost ones."""

MDANSE/Src/MDANSE/Framework/AtomSelector/__init__.py

@@ -13,4 +13,3 @@
 #    You should have received a copy of the GNU General Public License
 #    along with this program.  If not, see <https://www.gnu.org/licenses/>.
 #
-from .selector import Selector

MDANSE/Src/MDANSE/Framework/AtomSelector/atom_selection.py

@@ -0,0 +1,79 @@
+#    This file is part of MDANSE.
+#
+#    MDANSE is free software: you can redistribute it and/or modify
+#    it under the terms of the GNU General Public License as published by
+#    the Free Software Foundation, either version 3 of the License, or
+#    (at your option) any later version.
+#
+#    This program is distributed in the hope that it will be useful,
+#    but WITHOUT ANY WARRANTY; without even the implied warranty of
+#    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+#    GNU General Public License for more details.
+#
+#    You should have received a copy of the GNU General Public License
+#    along with this program.  If not, see <https://www.gnu.org/licenses/>.
+#
+
+from collections.abc import Sequence
+from typing import Optional
+
+from MDANSE.MolecularDynamics.Trajectory import Trajectory
+
+
+def select_atoms(
+    trajectory: Trajectory,
+    *,
+    index_list: Optional[Sequence[int]] = None,
+    index_range: Optional[Sequence[int]] = None,
+    index_slice: Optional[Sequence[int]] = None,
+    atom_types: Sequence[str] = (),
+    atom_names: Sequence[str] = (),
+    **_kwargs: str,
+) -> set[int]:
+    """Select specific atoms in the trajectory.
+
+    Atoms can be selected based
+    on indices, atom type or trajectory-specific atom name.
+    The atom type is normally the chemical element, while
+    the atom name can be more specific and depend on the
+    force field used.
+
+    Parameters
+    ----------
+    trajectory : Trajectory
+        A trajectory instance to which the selection is applied
+    index_list : Sequence[int]
+        a list of indices to be selected
+    index_range : Sequence[int]
+        a pair of (first, last+1) indices defining a range
+    index_slice : Sequence[int]
+        a sequence of (first, last+1, step) indices defining a slice
+    atom_types : Sequence[str]
+        a list of atom types (i.e. chemical elements) to be selected, given as string
+    atom_names : Sequence[str]
+        a list of atom names (as used by the MD engine, force field, etc.) to be selected
+
+    Returns
+    -------
+    set[int]
+        A set of indices which have been selected
+
+    """
+    selection = set()
+    system = trajectory.chemical_system
+    element_list = system.atom_list
+    name_list = system.name_list
+    indices = system.all_indices
+    if index_list is not None:
+        selection |= indices & set(index_list)
+    if index_range is not None:
+        selection |= indices & set(range(*index_range))
+    if index_slice is not None:
+        selection |= indices & set(range(*index_slice))
+    if atom_types:
+        new_indices = {index for index in indices if element_list[index] in atom_types}
+        selection |= new_indices
+    if atom_names:
+        new_indices = {index for index in indices if name_list[index] in atom_names}
+        selection |= new_indices
+    return selection