# -*- 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 Lesser GNU Public Licence, v2.1 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 ` 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 --------------- Definition 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 ------- """ import itertools import os import errno import numpy as np import warnings import logging logger = logging.getLogger("MDAnalysis.coordinates.XYZ") from . import base from .timestep import Timestep from ..lib import util from ..lib.util import cached, store_init_arguments from ..exceptions import NoDataError from ..version import __version__ 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`_. Notes ----- By default, the XYZ writer will attempt to use the input :class:`~MDAnalysis.core.groups.AtomGroup` or :class:`~MDAnalysis.core.universe.Universe` ``elements`` record to assign atom names in the XYZ file. If the ``elements`` record is missing, then the ``name`` record will be used. In the event that neither of these are available, the atoms will all be named ``X``. Please see, the `User Guide`_ for more information on how to add topology attributes if you wish to add your own elements / atom names to a :class:`~MDAnalysis.core.universe.Universe`. .. Links .. _xyzplugin: http://www.ks.uiuc.edu/Research/vmd/plugins/molfile/xyzplugin.html .. _User Guide: https://userguide.mdanalysis.org/examples/constructing_universe.html#Adding-topology-attributes .. versionchanged:: 1.0.0 Use elements attribute instead of names attribute, if present. .. versionchanged:: 2.0.0 Support for passing timestep to the writer was deprecated in 1.0 and has now been removed. As a consequence, custom names can no longer be passed to the writer, these should be added to the :class:`~MDAnalysis.core.universe.Universe`, or :class:`~MDAnalysis.core.groups.AtomGroup` before invoking the writer. """ format = "XYZ" multiframe = True # these are assumed! units = {"time": "ps", "length": "Angstrom"} def __init__( self, filename, n_atoms=None, convert_units=True, remark=None, precision=5, **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. 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 and frame precision: int (optional) set precision of saved trjactory to this number of decimal places. .. versionadded:: 2.9.0 .. versionchanged:: 1.0.0 Removed :code:`default_remark` variable (Issue #2692). .. versionchanged:: 2.0.0 Due to the removal of timestep as an input for writing, the atoms parameter is no longer relevant and has been removed. If passing an empty universe, please use ``add_TopologyAttr`` to add in the required elements or names. """ self.filename = filename self.remark = remark self.n_atoms = n_atoms self.convert_units = convert_units self.precision = precision # 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""" try: return atoms.atoms.elements except (AttributeError, NoDataError): try: return atoms.atoms.names except (AttributeError, NoDataError): wmsg = ( "Input AtomGroup or Universe does not have atom " "elements or names attributes, writer will default " "atom names to 'X'" ) warnings.warn(wmsg) return itertools.cycle(("X",)) 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 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. .. versionchanged:: 2.0.0 Deprecated support for Timestep argument has now been removed. Use AtomGroup or Universe as an input instead. """ # 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: errmsg = "Input obj is neither an AtomGroup or Universe" raise TypeError(errmsg) from 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_frame(ts) def _write_next_frame(self, ts=None): """ Write coordinate information in *ts* to the trajectory .. versionchanged:: 1.0.0 Print out :code:`remark` if present, otherwise use generic one (Issue #2692). Renamed from `write_next_timestep` to `_write_next_frame`. """ 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 unknown atoms. " 'Expected {} atoms, got {}. Try using "write" if you are ' 'using "_write_next_frame" 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 # Write number of atoms self._xyz.write("{0:d}\n".format(ts.n_atoms)) # Write remark if self.remark is None: remark = ( "frame {} | Written by MDAnalysis {} (release {})\n".format( ts.frame, self.__class__.__name__, __version__ ) ) self._xyz.write(remark) else: self._xyz.write(self.remark.strip() + "\n") # Write content for atom, (x, y, z) in zip(self.atomnames, coordinates): self._xyz.write( "{0!s:>8} {1:10.{p}f} {2:10.{p}f} {3:10.{p}f}\n" "".format(atom, x, y, z, p=self.precision) ) 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 `. .. 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 = Timestep @store_init_arguments 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: raise EOFError(err) from 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 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) def close(self): """Close xyz trajectory file if it was open.""" if self.xyzfile is None: return self.xyzfile.close() self.xyzfile = None