diff --git a/PQAnalysis/topology/api.py b/PQAnalysis/topology/api.py index 63afdf61..f5f292a8 100644 --- a/PQAnalysis/topology/api.py +++ b/PQAnalysis/topology/api.py @@ -2,6 +2,8 @@ A module including the API for the topology subpackage. """ +from beartype.typing import List + from PQAnalysis.traj import Trajectory from PQAnalysis.io import FileWritingMode, read_restart_file from PQAnalysis.types import Np1DIntArray @@ -23,7 +25,7 @@ def generate_shake_topology_file(trajectory: Trajectory, trajectory : Trajectory The trajectory to generate the shake topology for. selection : SelectionCompatible, optional - Selection is either a selection object or any object that can be + Selection is either a selection object or any object that can be initialized via 'Selection(selection)'. default None (all atoms) output : str | None, optional The output file. If not specified, the output is printed to stdout. @@ -54,7 +56,7 @@ def select_from_restart_file(selection: SelectionCompatible, Parameters ---------- selection : SelectionCompatible - Selection is either a selection object or any object that can be + 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. @@ -71,3 +73,36 @@ def select_from_restart_file(selection: SelectionCompatible, selection = Selection(selection) return selection.select(system.topology, use_full_atom_info=use_full_atom_info) + + +def selection_from_restart_file_as_list(selection: SelectionCompatible, + restart_file: str, + moldescriptor_file: str | None = None, + use_full_atom_info: bool = False, + ) -> List[int]: + """ + 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. + """ + + return list( + select_from_restart_file( + selection, + restart_file, + moldescriptor_file, + use_full_atom_info + ) + )