# -*- Mode: python; tab-width: 4; indent-tabs-mode:nil; coding:utf-8 -*-
# vim: tabstop=4 expandtab shiftwidth=4 softtabstop=4 fileencoding=utf-8
#
# MDAnalysis --- http://mdanalysis.googlecode.com
# Copyright (c) 2006-2015 Naveen Michaud-Agrawal, Elizabeth J. Denning, Oliver Beckstein
# and contributors (see AUTHORS for the full list)
#
# Released under the GNU Public Licence, v2 or any higher version
#
# Please cite your use of MDAnalysis in published work:
#
# N. Michaud-Agrawal, E. J. Denning, T. B. Woolf, and O. Beckstein.
# MDAnalysis: A Toolkit for the Analysis of Molecular Dynamics Simulations.
# J. Comput. Chem. 32 (2011), 2319--2327, doi:10.1002/jcc.21787
#
"""
Base classes for the selection writers
======================================
Specialized SelectionWriters are derived from
:class:`SelectionWriter`. Override the :meth:`~SelectionWriter._write_head`,
:meth:`~SelectionWriter._translate`, and :meth:`~SelectionWriter._write_tail`
methods.
.. autoclass:: SelectionWriter
:members: __init__, write, _translate, _write_head, _write_tail, comment
.. autofunction:: join
.. autofunction:: get_writer
"""
import os.path
import MDAnalysis.core.util as util
import MDAnalysis.selections
[docs]def join(seq, string="", func=None):
"""Create a list from sequence.
*string* is appended to each element but the last.
*func* is applied to every element before appending *string*.
"""
if func is None:
func = lambda x: x
return [func(x) + string for x in seq[:-1]] + [func(seq[-1])]
[docs]def get_writer(filename, defaultformat):
"""Return a :class:`SelectionWriter` for *filename* or a *defaultformat*."""
if filename:
format = os.path.splitext(filename)[1][1:] # strip initial dot!
format = format or defaultformat # use default if no fmt from fn
format = format.strip().lower() # canonical for lookup
try:
return MDAnalysis.selections._selection_writers[format]
except KeyError:
raise NotImplementedError("Writing as %r is not implemented; only %r will work."
% (format, MDAnalysis.selections._selection_writers.keys()))
[docs]class SelectionWriter(object):
"""Export a selection in MDAnalysis to a format usable in an external package.
The :class:`SelectionWriter` writes a selection string to a file
that can be used in another package such as `VMD`_, `PyMOL`_,
`Gromacs`_ or `CHARMM`_. In this way, analysis and visualization
can be done with the best or most convenient tools at hand.
:class:`SelectionWriter` is a base class and child classes are
derived with the appropriate customizations for the package file
format.
.. _VMD: http://www.ks.uiuc.edu/Research/vmd/
.. _PyMol: http://www.pymol.org/
.. _CHARMM: http://www.charmm.org/
.. _Gromacs: http://www.gromacs.org/
"""
#: Name of the format.
format = None
#: Extension of output files.
ext = None
#: Special character to continue a line across a newline.
continuation = ''
#: Comment format string; should contain '%s' or ``None`` for no comments.
commentfmt = None
default_numterms = 8
[docs] def __init__(self, filename, mode="wa", numterms=None, preamble=None, **kwargs):
"""Set up for writing to *filename*.
:Arguments:
*filename*
output file
*mode*
overwrite ("w") for every write, append ("a") to existing
file, or overwrite an existing file and then append ("wa") ["wa"]
*numterms*
number of individual index numbers per line for output
formats that write multiple entries in one line. If set
to 0 or ``False`` then no special formatting is done [8]
*preamble*
string that is written as a comment at the top of the file []
*kwargs*
use as defaults for :meth:`write`
"""
self.filename = util.filename(filename, ext=self.ext)
if not mode in ('a', 'w', 'wa'):
raise ValueError("mode must be one of 'w', 'a', 'wa', not %r" % mode)
self.mode = mode
self._current_mode = mode[0]
if numterms is None or numterms < 0:
self.numterms = self.default_numterms
elif numterms is False:
self.numterms = 0
else:
self.numterms = numterms
self.preamble = preamble
self.otherargs = kwargs # hack
self.number = 0
self.write_preamble()
def write_preamble(self):
"""Write a header, depending on the file format."""
if self.preamble is None:
return
with open(self.filename, self._current_mode) as out:
out.write(self.comment(self.preamble))
self._current_mode = 'a'
[docs] def write(self, selection, number=None, name=None, frame=None, mode=None):
"""Write selection to the output file.
:Arguments:
*selection*
a :class:`MDAnalysis.core.AtomGroup.AtomGroup`
*number*
selection will be named "mdanalysis<number>"
(``None`` auto increments between writes; useful
when appending) [``None``]
*name*
selection will be named *name* (instead of numbered)
[``None``]
*frame*
write selection of this frame (or the current one if
``None`` [``None``]
"""
u = selection.universe
if frame is not None:
u.trajectory[frame] # advance to frame
else:
try:
frame = u.trajectory.ts.frame
except AttributeError:
frame = 1 # should catch cases when we are analyzing a single PDB (?)
name = name or self.otherargs.get('name', None)
if name is None:
if number is None:
self.number += 1
number = self.number
name = "mdanalysis%(number)03d" % vars()
# build whole selection in one go (cleaner way to deal with
# to deal with line breaks after self.numterms entries)
# selection_list must contain entries to be joined with spaces or linebreaks
selection_terms = self._translate(selection.atoms)
step = self.numterms or len(selection.atoms)
with open(self.filename, self._current_mode) as out:
self._write_head(out, name=name)
for iatom in xrange(0, len(selection.atoms), step):
line = selection_terms[iatom:iatom + step]
out.write(" ".join(line))
if len(line) == step and not iatom + step == len(selection.atoms):
out.write(' ' + self.continuation + '\n')
out.write(' ') # safe so that we don't have to put a space at the start of tail
self._write_tail(out)
out.write('\n') # always terminate with newline
if self.mode == 'wa':
self._current_mode = 'a' # switch after first write
elif self.mode == 'w':
self._current_mode = 'w' # switch back after eg preamble
[docs] def _translate(self, atoms, **kwargs):
"""Translate atoms into a list of native selection terms.
- build list of ALL selection terms as if this was a single line, e.g.
``['index 12 |', 'index 22 |', 'index 33']``
- only one term per atom!!
- terms *must* be strings
- something like::
" ".join(terms)
must work
"""
raise NotImplementedError
[docs] def _write_head(self, out, **kwargs):
"""Initial output to open file object *out*."""
pass
[docs] def _write_tail(self, out, **kwargs):
"""Last output to open file object *out*."""
pass