Source code for prody.trajectory.dcdfile

# -*- coding: utf-8 -*-
"""This module defines classes for handling trajectory files in `DCD format`_.

.. _DCD format: http://www.ks.uiuc.edu/Research/namd/2.6/ug/node13.html"""

import os
from time import time
from struct import calcsize, unpack, pack
from os.path import getsize
import datetime

import numpy as np
from numpy import float32, fromstring

from prody.atomic import Atomic
from prody.ensemble import Ensemble
from prody.utilities import checkCoords
from prody import LOGGER, PY2K

from .frame import Frame
from .trajbase import TrajBase
from .trajfile import TrajFile
from .trajectory import Trajectory

if PY2K:
    range = xrange

__all__ = ['DCDFile', 'parseDCD', 'writeDCD']

now = datetime.datetime.now

DEBUG = False

PISQUARE = np.pi ** 2
RECSCALE32BIT = 1
RECSCALE64BIT = 2

[docs]class DCDFile(TrajFile): """A class for reading and writing DCD files. DCD header and first frame is parsed at instantiation. Coordinates from the first frame is set as the reference coordinate set. This class has been tested for 32-bit DCD files. 32-bit floating-point coordinate array can be casted automatically to a specified type, such as 64-bit float, using *astype* keyword argument, i.e. ``astype=float``, using :meth:`ndarray.astype` method.""" def __init__(self, filename, mode='rb', **kwargs): TrajFile.__init__(self, filename, mode) self._astype = kwargs.get('astype', None) if not self._mode.startswith('w'): self._parseHeader() __init__.__doc__ = TrajFile.__init__.__doc__ def _parseHeader(self): """Read the header information from a dcd file. Input: fd - a file struct opened for binary reading. Output: 0 on success, negative error code on failure. Side effects: *natoms set to number of atoms per frame *nsets set to number of frames in dcd file *istart set to starting timestep of dcd file *nsavc set to timesteps between dcd saves *delta set to value of trajectory timestep *nfixed set to number of fixed atoms *freeind may be set to heap-allocated space *reverse set to one if reverse-endian, zero if not. *charmm set to internal code for handling charmm data. """ dcd = self._file endian = b'' #'=' # native endian rec_scale = RECSCALE32BIT charmm = None dcdcordmagic = unpack(endian + b'i', b'CORD')[0] # Check magic number in file header and determine byte order bits = dcd.read(calcsize('ii')) temp = unpack(endian + b'ii', bits) if temp[0] + temp[1] == 84: LOGGER.info('Detected CHARMM -i8 64-bit DCD file of native ' 'endianness.') rec_scale = RECSCALE64BIT elif temp[0] == 84 and temp[1] == dcdcordmagic: pass #LOGGER.info('Detected standard 32-bit DCD file of native ' # 'endianness.') else: if unpack(b'>ii', bits) == temp: endian = '>' else: endian = '<' temp = unpack(endian + b'ii', bits) if temp[0] + temp[1] == 84: rec_scale = RECSCALE64BIT LOGGER.info('Detected CHARMM -i8 64-bit DCD file of opposite ' 'endianness.') else: endian = '' temp = unpack(endian + b'ii', bits) if temp[0] == 84 and temp[1] == dcdcordmagic: LOGGER.info('Detected standard 32-bit DCD file of ' 'opposite endianness.') else: raise IOError('Unrecognized DCD header or unsupported ' 'DCD format.') # check for magic string, in case of long record markers if rec_scale == RECSCALE64BIT: raise IOError('CHARMM 64-bit DCD files are not yet supported.'); temp = unpack(b'I', dcd.read(calcsize('I'))) if temp[0] != dcdcordmagic: raise IOError('Failed to find CORD magic in CHARMM -i8 64-bit ' 'DCD file.'); # Buffer the entire header for random access bits = dcd.read(80) # CHARMm-genereate DCD files set the last integer in the # header, which is unused by X-PLOR, to its version number. # Checking if this is nonzero tells us this is a CHARMm file # and to look for other CHARMm flags. temp = unpack(endian + b'i'*20 , bits) if temp[-1] != 0: charmm = True if charmm: #LOGGER.info('CHARMM format DCD file (also NAMD 2.1 and later).') temp = unpack(endian + b'i'*9 + b'f' + b'i'*10 , bits) else: LOGGER.info('X-PLOR format DCD file (also NAMD 2.0 and earlier) ' 'is not supported.') return None # Store the number of sets of coordinates (NSET) self._n_csets = temp[0] # Store ISTART, the starting timestep self._first_ts = temp[1] # Store NSAVC, the number of timesteps between dcd saves self._framefreq = temp[2] # Store NAMNF, the number of fixed atoms self._n_fixed = temp[8] if self._n_fixed > 0: raise IOError('DCD files with fixed atoms is not yet supported.') # Read in the timestep, DELTA # Note: DELTA is stored as double with X-PLOR but as float with CHARMm self._timestep = temp[9] self._unitcell = temp[10] == 1 # Get the end size of the first block if unpack(endian + b'i', dcd.read(rec_scale * calcsize('i')))[0] != 84: raise IOError('Unrecognized DCD format.') # Read in the size of the next block temp = unpack(endian + b'i', dcd.read(rec_scale * calcsize('i'))) if (temp[0] - 4) % 80 != 0: raise IOError('Unrecognized DCD format.') noremarks = temp[0] == 84 # Read NTITLE, the number of 80 character title strings there are temp = unpack(endian + b'i', dcd.read(rec_scale * calcsize('i'))) self._dcdtitle = dcd.read(80) if not noremarks: self._remarks = dcd.read(80) # Get the ending size for this block temp = unpack(endian + b'i', dcd.read(rec_scale * calcsize('i'))) if (temp[0] - 4) % 80 != 0: raise IOError('Unrecognized DCD format.') # Read in an integer '4' if unpack(endian + b'i', dcd.read(rec_scale * calcsize('i')))[0] != 4: raise IOError('Unrecognized DCD format.') # Read in the number of atoms self._n_atoms = unpack(endian + b'i', dcd.read(rec_scale*calcsize('i')))[0] # Read in an integer '4' if unpack(endian + b'i', dcd.read(rec_scale * calcsize('i')))[0] != 4: raise IOError('Bad DCD format.') self._is64bit = rec_scale == RECSCALE64BIT self._endian = endian self._n_floats = (self._n_atoms + 2) * 3 if self._is64bit: if self._unitcell: self._bytes_per_frame = 56 + self._n_floats * 8 else: self._bytes_per_frame = self._n_floats * 8 LOGGER.warning('Reading of 64 bit DCD files has not been tested. ' 'Please report any problems that you may find.') self._dtype = np.float64 self._itemsize = 8 else: if self._unitcell: self._bytes_per_frame = 56 + self._n_floats * 4 else: self._bytes_per_frame = self._n_floats * 4 self._dtype = np.float32 self._itemsize = 4 self._first_byte = self._file.tell() n_csets = (getsize(self._filename) - self._first_byte ) // self._bytes_per_frame if n_csets != self._n_csets: LOGGER.warning('DCD header claims {0} frames, file size ' 'indicates there are actually {1} frames.' .format(self._n_csets, n_csets)) self._n_csets = n_csets self._coords = self.nextCoordset() self._file.seek(self._first_byte) self._nfi = 0
[docs] def hasUnitcell(self): return self._unitcell
hasUnitcell.__doc__ = TrajBase.hasUnitcell.__doc__
[docs] def getRemarks(self): """Returns remarks parsed from DCD file.""" return self._remarks
def __next__(self): if self._closed: raise ValueError('I/O operation on closed file') nfi = self._nfi if nfi < self._n_csets: unitcell = self._nextUnitcell() coords = self._nextCoordset() if self._ag is None: frame = Frame(self, nfi, coords, unitcell) else: frame = self._frame Frame.__init__(frame, self, nfi, None, unitcell) return frame __next__.__doc__ = TrajBase.__next__.__doc__ next = __next__
[docs] def nextCoordset(self): """Returns next coordinate set.""" if self._closed: raise ValueError('I/O operation on closed file') if self._nfi < self._n_csets: #Skip extended system coordinates (unit cell data) if self._unitcell: self._file.seek(56, 1) if self._indices is None: return self._nextCoordset() else: return self._nextCoordset()[self._indices]
def _nextCoordset(self): n_floats = self._n_floats n_atoms = self._n_atoms xyz = fromstring(self._file.read(self._itemsize * n_floats), self._dtype) if len(xyz) != n_floats: return None xyz = xyz.reshape((3, n_atoms+2)).T[1:-1,:] xyz = xyz.reshape((n_atoms, 3)) if self._ag is not None: self._ag._setCoords(xyz, self._title + ' frame ' + str(self._nfi), overwrite=True) self._nfi += 1 if self._astype is not None and self._astype != xyz.dtype: xyz= xyz.astype(self._astype) return xyz nextCoordset.__doc__ = TrajBase.nextCoordset.__doc__ def _nextUnitcell(self): if self._unitcell: self._file.read(4) unitcell = fromstring(self._file.read(48), dtype=np.float64) unitcell = unitcell[[0,2,5,1,3,4]] if np.all(abs(unitcell[3:]) <= 1): # This file was generated by CHARMM, or by NAMD > 2.5, with the angle */ # cosines of the periodic cell angles written to the DCD file. */ # This formulation improves rounding behavior for orthogonal cells */ # so that the angles end up at precisely 90 degrees, unlike acos(). */ unitcell[3:] = 90. - np.arcsin(unitcell[3:]) * 90 / PISQUARE self._file.read(4) return unitcell
[docs] def getCoordsets(self, indices=None): """Returnss coordinate sets at given *indices*. *indices* may be an integer, a list of integers or **None**. **None** returns all coordinate sets.""" if self._closed: raise ValueError('I/O operation on closed file') if (self._indices is None and (indices is None or indices == slice(None))): nfi = self._nfi self.reset() n_floats = self._n_floats + self._unitcell * 14 n_atoms = self._n_atoms n_csets = self._n_csets data = self._file.read(self._itemsize * n_floats * n_csets) data = fromstring(data, self._dtype) if len(data) > n_floats * n_csets: n_csets = len(data)/n_floats data = data[:n_csets] LOGGER.warning('DCD is corrupt, {0} out of {1} frames ' 'were parsed.'.format(n_csets, self._n_csets)) data = data.reshape((n_csets, n_floats)) if self._unitcell: data = data[:, 14:] data = data.reshape((n_csets, 3, n_atoms+2)) data = data[:, :, 1:-1] data = data.transpose(0, 2, 1) self.goto(nfi) if self._astype is not None and self._astype != data.dtype: data = data.astype(self._astype) return data else: return TrajFile.getCoordsets(self, indices)
getCoordsets.__doc__ = TrajBase.getCoordsets.__doc__
[docs] def write(self, coords, unitcell=None, **kwargs): """Write *coords* to a file open in 'a' or 'w' mode. *coords* may be a NUmpy array or a ProDy object that stores or points to coordinate data. Note that all coordinate sets of ProDy object will be written. Number of atoms will be determined from the file or based on the size of the first coordinate set written. If *unitcell* is provided for the first coordinate set, it will be expected for the following coordinate sets as well. If *coords* is an :class:`~.Atomic` or :class:`~.Ensemble` all coordinate sets will be written. Following keywords are used when writing the first coordinate set: :arg timestep: timestep used for integration, default is 1 :arg firsttimestep: number of the first timestep, default is 0 :arg framefreq: number of timesteps between frames, default is 1""" if self._closed: raise ValueError('I/O operation on closed file') if self._mode == 'r': raise IOError('File not open for writing') try: coords = coords._getCoordsets() except AttributeError: try: coords = coords._getCoords() except AttributeError: checkCoords(coords, csets=True, dtype=None) else: if unitcell is None: try: coords = coords.getUnitcell() except AttributeError: pass if coords.dtype != float32: coords = coords.astype(float32) n_atoms = coords.shape[-2] if self._n_atoms == 0: self._n_atoms = n_atoms elif self._n_atoms != n_atoms: raise ValueError('coords does not have correct number of atoms') if coords.ndim == 2: coords = [coords] dcd = self._file pack_i_4N = pack('i', self._n_atoms * 4) if self._n_csets == 0: if unitcell is None: self._unitcell = False else: self._unitcell = True timestep = float(kwargs.get('timestep', 1.0)) first_ts = int(kwargs.get('firsttimestep', 0)) framefreq = int(kwargs.get('framefreq', 1)) n_fixed = 0 pack_i_0 = pack(b'i', 0) pack_ix4_0x4 = pack(b'i'*4, 0, 0, 0, 0) pack_i_1 = pack(b'i', 1) pack_i_2 = pack(b'i', 2) pack_i_4 = pack(b'i', 4) pack_i_84 = pack(b'i', 84) pack_i_164 = pack(b'i', 164) dcd.write(pack_i_84) dcd.write(b'CORD') dcd.write(pack_i_0) # 0 Number of frames in file, none written yet dcd.write(pack(b'i', first_ts)) # 1 Starting timestep dcd.write(pack(b'i', framefreq)) # 2 Timesteps between frames dcd.write(pack_i_0) # 3 Number of timesteps in simulation dcd.write(pack_i_0) # 4 NAMD writes NSTEP or ISTART - NSAVC here? dcd.write(pack_ix4_0x4) # 5, 6, 7, 8 dcd.write(pack('f', timestep)) # 9 timestep dcd.write(pack('i', int(self._unitcell))) # 10 with unitcell dcd.write(pack_ix4_0x4) # 11, 12, 13, 14 dcd.write(pack_ix4_0x4) # 15, 16, 17, 18 dcd.write(pack('i', 24)) # 19 Pretend to be CHARMM version 24 dcd.write(pack_i_84) dcd.write(pack_i_164) dcd.write(pack_i_2) dcd.write(b'Created by ProDy'.ljust(80)) temp = now().strftime('%d %B, %Y at %H:%M') try: temp = bytes(temp, encoding='utf-8') except TypeError: pass dcd.write((b'REMARKS Created ' + temp).ljust(80)) dcd.write(pack_i_164) dcd.write(pack_i_4) dcd.write(pack(b'i', n_atoms)) dcd.write(pack_i_4) self._first_byte = dcd.tell() if self._unitcell: if unitcell is None: raise TypeError('unitcell data is expected') else: uc = unitcell uc[3:] = np.sin((PISQUARE/90) * (90-uc[3:])) uc = uc[[0,3,1,4,5,2]] pack_i_48 = pack('i', 48) dcd.seek(0, 2) for xyz in coords: if self._unitcell: dcd.write(pack_i_48) uc.tofile(dcd) dcd.write(pack_i_48) xyz = xyz.T dcd.write(pack_i_4N) xyz[0].tofile(dcd) dcd.write(pack_i_4N) dcd.write(pack_i_4N) xyz[1].tofile(dcd) dcd.write(pack_i_4N) dcd.write(pack_i_4N) xyz[2].tofile(dcd) dcd.write(pack_i_4N) self._n_csets += 1 dcd.seek(8, 0) dcd.write(pack('i', self._n_csets)) dcd.seek(0, 2) self._nfi = self._n_csets
[docs] def flush(self): """Flush the internal output buffer.""" if self._mode != 'r': self._file.flush() os.fsync(self._file.fileno())
[docs]def parseDCD(filename, start=None, stop=None, step=None, astype=None): """Parse CHARMM format DCD files (also NAMD 2.1 and later). Returns an :class:`Ensemble` instance. Conformations in the ensemble will be ordered as they appear in the trajectory file. Use :class:`DCDFile` class for parsing coordinates of a subset of atoms. :arg filename: DCD filename :type filename: str :arg start: index of first frame to read :type start: int :arg stop: index of the frame that stops reading :type stop: int :arg step: steps between reading frames, default is 1 meaning every frame :type step: int :arg astype: cast coordinate array to specified type :type astype: type""" dcd = DCDFile(filename, astype=astype) time_ = time() n_frames = dcd.numFrames() LOGGER.info('DCD file contains {0} coordinate sets for {1} atoms.' .format(n_frames, dcd.numAtoms())) ensemble = dcd[slice(start,stop,step)] dcd.close() time_ = time() - time_ or 0.01 dcd_size = 1.0 * dcd.numFrames() * dcd._bytes_per_frame / (1024*1024) LOGGER.info('DCD file was parsed in {0:.2f} seconds.'.format(time_)) LOGGER.info('{0:.2f} MB parsed at input rate {1:.2f} MB/s.' .format(dcd_size, dcd_size/time_)) LOGGER.info('{0} coordinate sets parsed at input rate {1} frame/s.' .format(n_frames, int(n_frames/time_))) return ensemble
[docs]def writeDCD(filename, trajectory, start=None, stop=None, step=None, align=False): """Write 32-bit CHARMM format DCD file (also NAMD 2.1 and later). *trajectory can be an :class:`Trajectory`, :class:`DCDFile`, or :class:`Ensemble` instance. *filename* is returned upon successful output of file.""" if not isinstance(trajectory, (TrajBase, Ensemble, Atomic)): raise TypeError('{0} is not a valid type for trajectory' .format(type(trajectory))) irange = list(range(*slice(start, stop,step) .indices(trajectory.numCoordsets()))) n_csets = len(irange) if n_csets == 0: raise ValueError('trajectory does not have any coordinate sets, or ' 'no coordinate sets are selected') if isinstance(trajectory, Atomic): isEnsemble = False isAtomic = True n_atoms = trajectory.numAtoms() else: isEnsemble = True isAtomic = False n_atoms = trajectory.numSelected() if n_atoms == 0: raise ValueError('no atoms are selected in the trajectory') if isinstance(trajectory, TrajBase): isTrajectory = True unitcell = trajectory.hasUnitcell() nfi = trajectory.nextIndex() trajectory.reset() pack_i_48 = pack('i', 48) if isinstance(trajectory, Trajectory): timestep = trajectory.getTimestep()[0] first_ts = trajectory.getFirstTimestep()[0] framefreq = trajectory.getFrameFreq()[0] n_fixed = trajectory.numFixed()[0] else: timestep = trajectory.getTimestep() first_ts = trajectory.getFirstTimestep() framefreq = trajectory.getFrameFreq() n_fixed = trajectory.numFixed() else: isTrajectory = False unitcell = False if isinstance(trajectory, Ensemble): frame = trajectory[0] else: frame = trajectory acsi = trajectory.getACSIndex() timestep = 1 first_ts = 0 framefreq = 1 n_fixed = 0 dcd = DCDFile(filename, mode='w') LOGGER.progress('Writing DCD', len(irange), '_prody_writeDCD') prev = -1 uc = None time_ = time() for j, i in enumerate(irange): diff = i - prev if diff > 1: trajectory.skip(diff-1) prev = i if isTrajectory: frame = next(trajectory) if frame is None: break if unitcell: uc = frame._getUnitcell() uc[3:] = np.sin((PISQUARE/90) * (90-uc[3:])) uc = uc[[0,3,1,4,5,2]] elif isEnsemble: frame._index = i else: frame.setACSIndex(i) if align: frame.superpose() if j == 0: dcd.write(frame._getCoords(), uc, timestep=timestep, firsttimestep=first_ts, framefreq=framefreq) else: dcd.write(frame._getCoords(), uc) LOGGER.update(i, label='_prody_writeDCD') if isAtomic: trajectory.setACSIndex(acsi) j += 1 LOGGER.finish() dcd.close() time_ = time() - time_ or 0.01 dcd_size = 1.0 * (56 + (n_atoms * 3 + 6) * 4 ) * n_csets / (1024*1024) LOGGER.info('DCD file was written in {0:.2f} seconds.'.format(time_)) LOGGER.info('{0:.2f} MB written at input rate {1:.2f} MB/s.' .format(dcd_size, dcd_size/time_)) LOGGER.info('{0} coordinate sets written at output rate {1} frame/s.' .format(n_csets, int(n_csets/time_))) if j != n_csets: LOGGER.warn('Warning: {0} frames expected, {1} written.' .format(n_csets, j)) if isTrajectory: trajectory.goto(nfi) return filename