Source code for MDAnalysis.coordinates.XYZ

# -*- Mode: python; tab-width: 4; indent-tabs-mode:nil; coding:utf-8 -*-
# vim: tabstop=4 expandtab shiftwidth=4 softtabstop=4
#
# MDAnalysis --- https://www.mdanalysis.org
# Copyright (c) 2006-2017 The MDAnalysis Development Team and contributors
# (see the file AUTHORS for the full list of names)
#
# Released under the GNU Public Licence, v2 or any higher version
#
# Please cite your use of MDAnalysis in published work:
#
# R. J. Gowers, M. Linke, J. Barnoud, T. J. E. Reddy, M. N. Melo, S. L. Seyler,
# D. L. Dotson, J. Domanski, S. Buchoux, I. M. Kenney, and O. Beckstein.
# MDAnalysis: A Python package for the rapid analysis of molecular dynamics
# simulations. In S. Benthall and S. Rostrup editors, Proceedings of the 15th
# Python in Science Conference, pages 102-109, Austin, TX, 2016. SciPy.
# doi: 10.25080/majora-629e541a-00e
#
# 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
#


"""XYZ trajectory reader --- :mod:`MDAnalysis.coordinates.XYZ`
==============================================================

The :ref:`XYZ format <xyz-format>` is a loosely defined, simple
coordinate trajectory format. The implemented format definition was
taken from the `VMD xyzplugin`_ and is therefore compatible with VMD.

Note the following:

* Comments are not allowed in the XYZ file (we neither read nor write
  them to remain compatible with VMD).
* The atom name (first column) is ignored during reading.
* The coordinates are assumed to be space-delimited rather than fixed
  width (this may cause issues - see below).
* All fields to the right of the z-coordinate are ignored.
* The unitcell information is all zeros since this is not recorded in
  the XYZ format.

.. rubric:: Units

* Coordinates are in Angstroms.
* The length of a timestep can be set by passing the *dt* argument,
  it's assumed to be in ps (default: 1 ps).

There appears to be no rigid format definition so it is likely users
will need to tweak this class.

.. _xyz-format:

XYZ File format
---------------

Definiton used by the :class:`XYZReader` and :class:`XYZWriter` (and
the `VMD xyzplugin`_ from whence the definition was taken)::

    [ comment line            ] !! NOT IMPLEMENTED !! DO NOT INCLUDE
    [ N                       ] # of atoms, required by this xyz reader plugin  line 1
    [ molecule name           ] name of molecule (can be blank)                 line 2
    atom1 x y z [optional data] atom name followed by xyz coords                line 3
    atom2 x y z [ ...         ] and (optionally) other data.
    ...
    atomN x y z [ ...         ]                                                 line N+2


Note
----
* comment lines not implemented (do not include them)
* molecule name: the line is required but the content is ignored
  at the moment
* optional data (after the coordinates) are presently ignored


.. Links
.. _`VMD xyzplugin`: http://www.ks.uiuc.edu/Research/vmd/plugins/molfile/xyzplugin.html

Classes
-------

"""
from __future__ import division, absolute_import
import six
from six.moves import range, zip

import itertools
import os
import errno
import numpy as np
import logging
logger = logging.getLogger('MDAnalysis.coordinates.XYZ')

from . import base
from ..lib import util
from ..lib.util import cached
from ..exceptions import NoDataError
from ..version import __version__


[docs]class XYZWriter(base.WriterBase): """Writes an XYZ file The XYZ file format is not formally defined. This writer follows the VMD implementation for the molfile `xyzplugin`_. .. _xyzplugin: http://www.ks.uiuc.edu/Research/vmd/plugins/molfile/xyzplugin.html .. versionchanged: 0.21.0 Use elements attribute instead of names attribute, if present """ format = 'XYZ' multiframe = True # these are assumed! units = {'time': 'ps', 'length': 'Angstrom'} def __init__(self, filename, n_atoms=None, atoms=None, convert_units=True, remark=None, **kwargs): """Initialize the XYZ trajectory writer Parameters ---------- filename: str filename of trajectory file. If it ends with "gz" then the file will be gzip-compressed; if it ends with "bz2" it will be bzip2 compressed. n_atoms: int (optional) Number of atoms in trajectory. By default assume that this is None and that this file is used to store several different models instead of a single trajectory. If a number is provided each written TimeStep has to contain the same number of atoms. atoms: str | list (optional) Provide atom names: This can be a list of names or an :class:`AtomGroup`. If none is provided, atoms will be called 'X' in the output. These atom names will be used when a trajectory is written from raw :class:`Timestep` objects which do not contain atom information. If you write a :class:`AtomGroup` with :meth:`XYZWriter.write` then atom information is taken at each step and *atoms* is ignored. convert_units : bool (optional) convert quantities to default MDAnalysis units of Angstrom upon writing [``True``] remark: str (optional) single line of text ("molecule name"). By default writes MDAnalysis version """ self.filename = filename self.n_atoms = n_atoms self.convert_units = convert_units self.atomnames = self._get_atoms_elements_or_names(atoms) default_remark = "Written by {0} (release {1})".format( self.__class__.__name__, __version__) self.remark = default_remark if remark is None else remark # can also be gz, bz2 self._xyz = util.anyopen(self.filename, 'wt') def _get_atoms_elements_or_names(self, atoms): """Return a list of atom elements (if present) or fallback to atom names""" # Default case if atoms is None: return itertools.cycle(('X',)) # Single atom name provided elif isinstance(atoms, six.string_types): return itertools.cycle((atoms,)) # List of atom names providded elif isinstance(atoms, list): return atoms # AtomGroup or Universe, grab the names else default # (AtomGroup.atoms just returns AtomGroup) try: return atoms.atoms.elements except (AttributeError, NoDataError): try: return atoms.atoms.names except (AttributeError, NoDataError): return itertools.cycle(('X',))
[docs] def close(self): """Close the trajectory file and finalize the writing""" if self._xyz is not None: self._xyz.write("\n") self._xyz.close() self._xyz = None
[docs] def write(self, obj): """Write object `obj` at current trajectory frame to file. Atom elements (or names) in the output are taken from the `obj` or default to the value of the `atoms` keyword supplied to the :class:`XYZWriter` constructor. Parameters ---------- obj : Universe or AtomGroup The :class:`~MDAnalysis.core.groups.AtomGroup` or :class:`~MDAnalysis.core.universe.Universe` to write. """ # prepare the Timestep and extract atom names if possible # (The way it is written it should be possible to write # trajectories with frames that differ in atom numbers # but this is not tested.) try: atoms = obj.atoms except AttributeError: if isinstance(obj, base.Timestep): ts = obj else: six.raise_from(TypeError("No Timestep found in obj argument"), None) else: if hasattr(obj, 'universe'): # For AtomGroup and children (Residue, ResidueGroup, Segment) ts_full = obj.universe.trajectory.ts if ts_full.n_atoms == atoms.n_atoms: ts = ts_full else: # Only populate a time step with the selected atoms. ts = ts_full.copy_slice(atoms.indices) elif hasattr(obj, 'trajectory'): # For Universe only --- get everything ts = obj.trajectory.ts # update atom names self.atomnames = self._get_atoms_elements_or_names(atoms) self.write_next_timestep(ts)
[docs] def write_next_timestep(self, ts=None): """Write coordinate information in *ts* to the trajectory""" if ts is None: if not hasattr(self, 'ts'): raise NoDataError('XYZWriter: no coordinate data to write to ' 'trajectory file') else: ts = self.ts if self.n_atoms is not None: if self.n_atoms != ts.n_atoms: raise ValueError('n_atoms keyword was specified indicating ' 'that this should be a trajectory of the ' 'same model. But the provided TimeStep has a ' 'different number ({}) then expected ({})' ''.format(ts.n_atoms, self.n_atoms)) else: if (not isinstance(self.atomnames, itertools.cycle) and len(self.atomnames) != ts.n_atoms): logger.info('Trying to write a TimeStep with unkown atoms. ' 'Expected {}, got {}. Try using "write" if you are ' 'using "write_next_timestep" directly'.format( len(self.atomnames), ts.n_atoms)) self.atomnames = np.array([self.atomnames[0]] * ts.n_atoms) if self.convert_units: coordinates = self.convert_pos_to_native( ts.positions, inplace=False) else: coordinates = ts.positions self._xyz.write("{0:d}\n".format(ts.n_atoms)) self._xyz.write("frame {0}\n".format(ts.frame)) for atom, (x, y, z) in zip(self.atomnames, coordinates): self._xyz.write("{0!s:>8} {1:10.5f} {2:10.5f} {3:10.5f}\n" "".format(atom, x, y, z))
[docs]class XYZReader(base.ReaderBase): """Reads from an XYZ file :Data: :attr:`ts` Timestep object containing coordinates of current frame :Methods: ``len(xyz)`` return number of frames in xyz ``for ts in xyz:`` iterate through trajectory .. Note: this can read both compressed (foo.xyz) and compressed (foo.xyz.bz2 or foo.xyz.gz) files; uncompression is handled on the fly and also reads streams via :class:`~MDAnalysis.lib.util.NamedStream`. The XYZ file format follows VMD's xyzplugin_ and is also described under :ref:`XYZ format <xyz-format>`. .. versionchanged:: 0.11.0 Frames now 0-based instead of 1-based. Added *dt* and *time_offset* keywords (passed to :class:`Timestep`) """ # Phil Fowler: # Validation: the geometric centre of 1284 atoms was calculated over # 500 frames using both MDAnalysis and a VMD Tcl script. There was # exact agreement (measured to three decimal places). bzipped and # gzipped versions of the XYZ file were also tested format = "XYZ" # these are assumed! units = {'time': 'ps', 'length': 'Angstrom'} _Timestep = base.Timestep def __init__(self, filename, **kwargs): super(XYZReader, self).__init__(filename, **kwargs) # the filename has been parsed to be either be foo.xyz or foo.xyz.bz2 by # coordinates::core.py so the last file extension will tell us if it is # bzipped or not root, ext = os.path.splitext(self.filename) self.xyzfile = util.anyopen(self.filename) self.compression = ext[1:] if ext[1:] != "xyz" else None self._cache = dict() self.ts = self._Timestep(self.n_atoms, **self._ts_kwargs) # Haven't quite figured out where to start with all the self._reopen() # etc. # (Also cannot just use seek() or reset() because that would break # with urllib2.urlopen() streams) self._read_next_timestep() @property @cached('n_atoms') def n_atoms(self): """number of atoms in a frame""" with util.anyopen(self.filename) as f: n = f.readline() # need to check type of n return int(n) @property @cached('n_frames') def n_frames(self): try: return self._read_xyz_n_frames() except IOError: return 0 def _read_xyz_n_frames(self): # the number of lines in the XYZ file will be 2 greater than the # number of atoms linesPerFrame = self.n_atoms + 2 counter = 0 offsets = [] with util.anyopen(self.filename) as f: line = True while line: if not counter % linesPerFrame: offsets.append(f.tell()) line = f.readline() counter += 1 # need to check this is an integer! n_frames = int(counter / linesPerFrame) self._offsets = offsets return n_frames def _read_frame(self, frame): self.xyzfile.seek(self._offsets[frame]) self.ts.frame = frame - 1 # gets +1'd in next return self._read_next_timestep() def _read_next_timestep(self, ts=None): # check that the timestep object exists if ts is None: ts = self.ts f = self.xyzfile try: # we assume that there are only two header lines per frame f.readline() f.readline() # convert all entries at the end once for optimal speed tmp_buf = [] for i in range(self.n_atoms): tmp_buf.append(f.readline().split()[1:4]) ts.positions = tmp_buf ts.frame += 1 return ts except (ValueError, IndexError) as err: six.raise_from(EOFError(err), None) def _reopen(self): self.close() self.open_trajectory() def open_trajectory(self): if self.xyzfile is not None: raise IOError( errno.EALREADY, 'XYZ file already opened', self.filename) self.xyzfile = util.anyopen(self.filename) # reset ts ts = self.ts ts.frame = -1 return self.xyzfile
[docs] def Writer(self, filename, n_atoms=None, **kwargs): """Returns a XYZWriter for *filename* with the same parameters as this XYZ. Parameters ---------- filename: str filename of the output trajectory n_atoms: int (optional) number of atoms. If none is given use the same number of atoms from the reader instance is used **kwargs: See :class:`XYZWriter` for additional kwargs Returns ------- :class:`XYZWriter` (see there for more details) See Also -------- :class:`XYZWriter` """ if n_atoms is None: n_atoms = self.n_atoms return XYZWriter(filename, n_atoms=n_atoms, **kwargs)
[docs] def close(self): """Close xyz trajectory file if it was open.""" if self.xyzfile is None: return self.xyzfile.close() self.xyzfile = None