src/taurenmd/cli_trajedit.py
"""
# Edit a trajectory.
*In short*, takes a trajectory and apply a modification:
1. frame slicing
2. atom selection
3. unwrap
4. align
5. file format change
Saves the result to a new trajectory.
All operations are performed with `MDAnalysis <https://www.mdanalysis.org>`_.
## Examples
Changes trajectory file format
taurenmd trajedit top.pdb traj.xtd -d traj.dcd
Extracts a part of the trajectory atoms, in this example ``segid A``,
the option ``-o`` saves the first frame of the new trajectory to a topology
file:
taurenmd trajedit top.pdb traj.xtc -d tsegidA.dcd -o -l "segid A"
You can slice the trajectory by appending the following ``-s``, ``-e`` or
``-p`` options, this saves only every 100 frames:
[...] -p 100
You can align the trajectory to a part of the system, for example,
align the whole system to one of its subunits:
taurenmd trajedit top.pdb traj.dcd -d alignedA.dcd -a "segid A and name CA"
further restrain the output to a specific subselection with ``-l``:
[...] -l "segid A or segid B"
``trajedit`` also implements the ``unwrap`` method from which is an
alternative approach to the ``imagemol`` client, that implements from
``MDTraj``. See references section.
taurenmd trajedit top.pdb traj.dcd -d unwrapped.dcd -w -o unwrapped_frame0.pdb
## References
""" # noqa: E501
import argparse
import functools
import MDAnalysis as mda
from taurenmd import _BANNER, Path, _controlled_exit
from taurenmd import core as tcore
from taurenmd import log
from taurenmd.libs import libcli, libio, libmda
from taurenmd.logger import S, T
__author__ = 'Joao M.C. Teixeira'
__email__ = 'joaomcteixeira@gmail.com'
__maintainer__ = 'Joao M.C. Teixeira'
__credits__ = ['Joao M.C. Teixeira']
__status__ = 'Production'
__doc__ += (
f'{tcore.ref_mda}'
f'{tcore.ref_mda_selection}'
f'{tcore.ref_mda_unwrap}'
f'{tcore.ref_mda_alignto}'
)
_help = 'Edits trajectory in many different ways.'
_name = 'trajedit'
ap = libcli.CustomParser(
description=_BANNER + __doc__,
formatter_class=argparse.RawDescriptionHelpFormatter,
)
libcli.add_version_arg(ap)
libcli.add_topology_arg(ap)
libcli.add_trajectories_arg(ap)
libcli.add_insort_arg(ap)
libcli.add_atom_selection_arg(ap)
libcli.add_slice_arg(ap)
libcli.add_traj_output_arg(ap)
libcli.add_top_output_arg(ap)
ap.add_argument(
'-a',
'--align',
help=(
'Align system to a atom group. '
'Alignmed RMSD is compared to the Topology coordinates. '
'Uses MDAnalysis.analysis.align.alignto. '
'If given without argument defaults to \'all\'. '
'Defaults to ``False``.'
),
default=False,
const='all',
nargs='?',
)
ap.add_argument(
'-w',
'--unwrap',
help=(
'Unwraps selection according to: '
'https://www.mdanalysis.org/docs/documentation_pages/core/groups.html#MDAnalysis.core.groups.AtomGroup.unwrap' # noqa: E501
),
action='store_true',
)
ap.add_argument(
'--unwrap-reference',
help=(
'The unwrap method reference parameter. '
'Has effect only if \'-w\' is given. '
'Defaults to `None`.'
),
default=None,
)
ap.add_argument(
'--unwrap-compound',
help=(
'The unwrap method compound parameter. '
'Has effect only if \'-w\' is given. '
'Defaults to `fragments`.'
),
default='fragments',
type=str,
)
def _ap():
return ap
def main(
topology,
trajectories,
insort=False,
start=None,
stop=None,
step=None,
selection='all',
traj_output='traj_out.dcd',
top_output=None,
unwrap=False,
unwrap_reference='com',
unwrap_compound='fragments',
align=False,
**kwargs,
):
"""Execute main client logic."""
log.info(T('editing trajectory'))
topology = Path(topology)
trajectories = [Path(t) for t in trajectories]
u = libmda.load_universe(
topology,
*trajectories,
insort=insort,
guess_bonds=unwrap,
)
if unwrap:
log.info(T('unwrapping'))
log.info(S('set to: {}', unwrap))
log.info(S('reference: {}', unwrap_reference))
log.info(S('compound: {}', unwrap_compound))
if align:
log.info(T('Alignment'))
log.info(S('trajectory selection will be aligned to subselection:'))
log.info(S('- {}', align, indent=2))
align_reference = mda.Universe(Path(topology).str())
frame_slice = libmda.get_frame_slices(u, start, stop, step)
log.info(S('selecting: {}', selection))
atom_selection = u.select_atoms(selection)
log.info(S('with {} atoms', atom_selection.n_atoms, indent=2))
log.info(T('saving trajectory'))
traj_output = Path(traj_output)
log.info(S('destination: {}', traj_output.resolve().str()))
total_frames = len(u.trajectory[frame_slice])
with mda.Writer(traj_output.str(), atom_selection.n_atoms) as W, \
libcli.ProgressBar(total_frames, suffix='frames') as pb:
for _ts in u.trajectory[frame_slice]:
if unwrap:
log.debug(S('unwrapping', indent=2))
atom_selection.unwrap(
reference=unwrap_reference,
compound=unwrap_compound,
)
if align:
try:
libmda.mdaalignto(u, align_reference, selection=align)
except (ValueError, ZeroDivisionError):
_controlled_exit()
W.write(atom_selection)
pb.increment()
log.info(S('trajectory saved'))
if top_output:
log.info(T('saving topology'))
fout = libio.parse_top_output(top_output, traj_output)
log.info(S('saving frame 0 to: {}', fout.resolve()))
with mda.Writer(fout.str(), atom_selection.n_atoms) as W:
for _ts in u.trajectory[frame_slice][0:1]:
if unwrap:
log.debug(S('unwrapping for topology', indent=2))
atom_selection.unwrap(
reference=unwrap_reference,
compound=unwrap_compound,
)
W.write(atom_selection)
log.info(S('Done'))
return
maincli = functools.partial(libcli.maincli, ap, main)
if __name__ == '__main__':
maincli()