Source code for prody.apps.prody_apps.prody_catdcd

"""Concatenate, slice, and/or select DCD files."""

from ..apptools import *

__all__ = ['prody_catdcd']

def prody_catdcd(*dcd, **kwargs):
    """Concatenate *dcd* files.

    :arg select: atom selection

    :arg align: atom selection for aligning frames

    :arg pdb: PDB file used in atom selections and as reference for alignment

    :arg psf: PSF file used in atom selections

    :arg output: output filename

    :arg first: index of the first output frame

    :arg last: index of the last output frame

    :arg stride: number of steps between output frames"""

    import prody
    LOGGER = prody.LOGGER
    if kwargs.get('numframes', False):
        from sys import stdout
        for fn in dcd:
            stdout.write(str(prody.DCDFile(fn).numFrames()) + '\n')
        return

    from os.path import splitext

    ag = None
    psf = kwargs.get('psf', None)
    if psf:
        ag = prody.parsePSF(psf)

    pdb = kwargs.get('pdb', None)
    if pdb:
        if ag:
            ag = prody.parsePDB(pdb, ag=ag)
        else:
            ag = prody.parsePDB(pdb)

    align = kwargs.get('align', None)
    select = kwargs.get('select', None)
    if select == 'all':
        select = None
    if ag is None and (align or select):
        raise ValueError('one of PSF or PDB files must be provided for '
                         'align and select options to work')
    dcd = list(dcd)
    traj = prody.Trajectory(dcd.pop(0))
    while dcd:
        traj.addFile(dcd.pop(0))

    if ag:
        traj.link(ag)
        if ag.numCoordsets():
            traj.setCoords(ag.getCoords())
        if select:
            select = ag.select(select)
            if select is None:
                raise ValueError('{0} did not match any atoms'
                                 .format(repr(kwargs.get('select'))))
        else:
            select = ag
        LOGGER.info('{0} atoms are selected for writing output.'
                    .format(len(select)))
        if align:
            align = ag.select(align)
            traj.setAtoms(align)
            LOGGER.info('{0} atoms are selected for aligning frames.'
                        .format(len(align)))

    output = kwargs.get('output', 'trajectory.dcd')
    out = prody.DCDFile(output, 'w')
    count = 0
    stride = kwargs.get('stride', 1)
    goto = stride != 1
    slc = slice(kwargs.get('first', 0), kwargs.get('last', -1),
                stride).indices(len(traj)+1)
    for i in range(*slc):
        if goto:
            traj.goto(i)
        frame = next(traj)
        if align:
            frame.superpose()
        if select:
            out.write(select._getCoords(), frame.getUnitcell())
        else:
            out.write(frame._getCoords(), frame.getUnitcell())
        count += 1
    traj.close()
    out.close()
    LOGGER.info("{0} frames are written into {1}."
                .format(count, output))

def addCommand(commands):

    subparser = commands.add_parser('catdcd',
        help='concatenate dcd files')

    subparser.add_argument('--quiet', help="suppress info messages to stderr",
        action=Quiet, nargs=0)

    subparser.add_argument('--examples', action=UsageExample, nargs=0,
        help='show usage examples and exit')

    subparser.set_defaults(usage_example=
    """Concatenate two DCD files and output all atmos:

  $ prody catdcd mdm2.dcd mdm2sim2.dcd

Concatenate two DCD files and output backbone atoms:

  $ prody catdcd mdm2.dcd mdm2sim2.dcd --pdb mdm2.pdb -s bb""")

    subparser.add_argument('-s', '--select', default='all', type=str,
        dest='select', metavar='SEL',
        help='atom selection (default: %(default)s)')

    subparser.add_argument('-o', '--output', type=str, metavar='FILE',
        default='trajectory.dcd',
        help='output filename (default: %(default)s)')

    subparser.add_argument('-n', '--num', default=False, action='store_true',
        dest='numframes',
        help='print the number of frames in each file and exit')

    subparser.add_argument('--psf',
        help='PSF filename (must have same number of atoms as DCDs)')
    subparser.add_argument('--pdb',
        help='PDB filename (must have same number of atoms as DCDs)')

    subparser.add_argument('--first', metavar='INT', type=int, default=0,
        help='index of the first output frame, default: %(default)s')
    subparser.add_argument('--last', metavar='INT', type=int, default=-1,
        help='index of the last output frame, default: %(default)s')
    subparser.add_argument('--stride', metavar='INT', type=int, default=1,
        help='number of steps between output frames, default: %(default)s')

    subparser.add_argument('--align', metavar='SEL', type=str,
        help='atom selection for aligning frames, a PSF or PDB file must be '
             'provided, if a PDB is provided frames will be superposed onto '
             'PDB coordinates')

    subparser.add_argument('dcd', nargs='+',
        help='DCD filename(s) (all must have same number of atoms)')

    subparser.set_defaults(subparser=subparser)
    subparser.set_defaults(func=lambda ns: prody_catdcd(
                                    *ns.__dict__.pop('dcd'), **ns.__dict__))