Skip to content

Commit

Permalink
added feature select from restart file
Browse files Browse the repository at this point in the history
  • Loading branch information
97gamjak committed May 1, 2024
1 parent 81c41bd commit e8313dc
Showing 1 changed file with 33 additions and 2 deletions.
35 changes: 33 additions & 2 deletions PQAnalysis/topology/api.py
Original file line number Diff line number Diff line change
Expand Up @@ -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


Expand Down Expand Up @@ -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)

Check warning on line 69 in PQAnalysis/topology/api.py

View check run for this annotation

Codecov / codecov/patch

PQAnalysis/topology/api.py#L69

Added line #L69 was not covered by tests

selection = Selection(selection)

Check warning on line 71 in PQAnalysis/topology/api.py

View check run for this annotation

Codecov / codecov/patch

PQAnalysis/topology/api.py#L71

Added line #L71 was not covered by tests

return selection.select(system.topology, use_full_atom_info=use_full_atom_info)

0 comments on commit e8313dc

Please sign in to comment.