Merge pull request #31 from MolarVerse/feature/select_from_restart_file

added feature select from restart file

PQAnalysis/topology/api.py

@@ -3,8 +3,9 @@
 """
 
 from PQAnalysis.traj import Trajectory
-from PQAnalysis.io import FileWritingMode
-from .selection import SelectionCompatible
+from PQAnalysis.io import FileWritingMode, read_restart_file
+from PQAnalysis.types import Np1DIntArray
+from .selection import SelectionCompatible, Selection
 from .shake_topology import ShakeTopologyGenerator
 
 
@@ -40,3 +41,33 @@ def generate_shake_topology_file(trajectory: Trajectory,
     generator = ShakeTopologyGenerator(selection, use_full_atom_info)
     generator.generate_topology(trajectory)
     generator.write_topology(output, mode)
+
+
+def select_from_restart_file(selection: SelectionCompatible,
+                             restart_file: str,
+                             moldescriptor_file: str | None = None,
+                             use_full_atom_info: bool = False,
+                             ) -> Np1DIntArray:
+    """
+    Selects atoms from a restart file and writes them to a new file.
+
+    Parameters
+    ----------
+    selection : SelectionCompatible
+        Selection is either a selection object or any object that can be 
+        initialized via 'Selection(selection)'.
+    restart_file : str
+        The restart file to read the atoms from.
+    moldescriptor_file : str | None, optional
+        The moldescriptor file to read the atom types from, by default None
+    use_full_atom_info : bool, optional
+        If True, the full atom information (name, index, mass) is used
+        for the selection, by default False
+        Is always ignored if atoms is not a list of atom objects.
+    """
+
+    system = read_restart_file(restart_file, moldescriptor_file)
+
+    selection = Selection(selection)
+
+    return selection.select(system.topology, use_full_atom_info=use_full_atom_info)