From e8313dc3843c3c6ad139a2c1e13bc6edcd4d11cc Mon Sep 17 00:00:00 2001 From: Jakob Gamper <97gamjak@gmail.com> Date: Wed, 1 May 2024 23:48:16 +0200 Subject: [PATCH] added feature select from restart file --- PQAnalysis/topology/api.py | 35 +++++++++++++++++++++++++++++++++-- 1 file changed, 33 insertions(+), 2 deletions(-) diff --git a/PQAnalysis/topology/api.py b/PQAnalysis/topology/api.py index 6d19679e..63afdf61 100644 --- a/PQAnalysis/topology/api.py +++ b/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)