# -*- coding: utf-8 -*-
"""This module defines base class for trajectory handling."""
from numbers import Integral
from numpy import ndarray, unique
from prody.ensemble import Ensemble
from prody.utilities import checkCoords, checkWeights
from .frame import Frame
__all__ = ['TrajBase']
[docs]class TrajBase(object):
"""Base class for :class:`.Trajectory` and :class:`.TrajFile`. Derived
classes must implement functions described in this class."""
def __init__(self, title='Unknown'):
self._title = str(title).strip()
self._coords = None # reference
self._n_atoms = 0
self._n_csets = 0 # number of conformations/frames/coordinate sets
self._weights = None
self._ag = None
self._atoms = None
self._indices = None # indices of selected atoms
self._frame = None # if atoms are set, always return the same frame
self._nfi = 0
self._closed = False
def __iter__(self):
if self._closed:
raise ValueError('I/O operation on closed file')
while self._nfi < self._n_csets:
yield next(self)
def __str__(self):
return '{0} {1}'.format(self.__class__.__name__, self._title)
def __getitem__(self, index):
if self._closed:
raise ValueError('I/O operation on closed file')
if isinstance(index, Integral):
return self.getFrame(index)
elif isinstance(index, (slice, list, ndarray)):
if isinstance(index, slice):
ens = Ensemble('{0} ({1[0]}:{1[1]}:{1[2]})'.format(
self._title, index.indices(len(self))))
else:
ens = Ensemble('{0} slice'.format(self._title))
ens.setCoords(self.getCoords())
if self._weights is not None:
ens.setWeights(self._weights.copy())
ens.addCoordset(self.getCoordsets(index))
ens.setAtoms(self._atoms)
return ens
else:
raise IndexError('invalid index')
def __len__(self):
return self._n_csets
def __enter__(self):
return self
def __exit__(self, type, value, tb):
self.close()
[docs] def getTitle(self):
"""Returns title of the ensemble."""
return self._title
[docs] def setTitle(self, title):
"""Set title of the ensemble."""
self._title = str(title)
[docs] def numAtoms(self):
"""Returns number of atoms."""
return self._n_atoms
[docs] def numFrames(self):
"""Returns number of frames."""
return self._n_csets
numCoordsets = numFrames
[docs] def numSelected(self):
"""Returns number of selected atoms. A subset of atoms can be selected
by passing a selection to :meth:`setAtoms`."""
return self._n_atoms if self._indices is None else len(self._indices)
[docs] def getAtoms(self):
"""Returns associated/selected atoms."""
return self._atoms
[docs] def setAtoms(self, atoms):
"""Set *atoms* or specify a selection of atoms to be considered in
calculations and coordinate requests. When a selection is set,
corresponding subset of coordinates will be considered in, for
example, alignments and RMSD calculations. Setting atoms also
allows some functions to access atomic data when needed. For
example, :class:`.Trajectory` and :class:`.Frame` instances become
suitable arguments for :func:`.writePDB`. Passing **None** as *atoms*
argument will deselect atoms. Note that setting atoms does not change
the reference coordinates of the trajectory. To change the reference,
use :meth:`.setCoords` method."""
if atoms is None:
self._atoms = self._indices = None
return
try:
atoms.getACSIndex()
except AttributeError:
raise TypeError('atoms must be an Atomic instance')
n_atoms = self._n_atoms
if n_atoms:
if atoms.numAtoms() > n_atoms:
raise ValueError('atoms must be same size or smaller than '
'the trajectory')
try:
dummies = atoms.numDummies()
except AttributeError:
pass
else:
if dummies:
raise ValueError('atoms must not have any dummies')
else:
indices = atoms._getIndices()
if indices != unique(indices):
raise ValueError('atoms must be ordered by indices')
if atoms.numAtoms() == n_atoms:
self._atoms = atoms
self._indices = None
else:
try:
ag = atoms.getAtomGroup()
except AttributeError:
raise ValueError('atoms must indicate a subset or must '
'match the trajectory size')
else:
if ag.numAtoms() != n_atoms:
raise ValueError('atoms must point to an AtomGroup '
'of the same size as the trajectory')
self._atoms = atoms
self._indices = atoms.getIndices()
else:
self._n_atoms = atoms.numAtoms()
self._atoms = atoms
[docs] def link(self, *ag):
"""Link, return, or unlink an :class:`.AtomGroup` instance. When a
link to *ag* is established, coordinates of new frames parsed from the
trajectory file will be set as the coordinates of *ag* and this will
update coordinates of all selections and atom subsets pointing to it.
At link time, if *ag* does not have any coordinate sets and reference
coordinates of the trajectory is set, reference coordinates of the
trajectory will be passed to *ag*. To break an established link, pass
**None** argument, or to return the linked atom group instance, call
with no arguments.
.. warning::
Every time a frame is parsed from the trajectory, all coordinate
sets present in the linked :class:`.AtomGroup` will be overwritten.
"""
if not ag:
return self._ag
else:
if len(ag) > 1:
raise TypeError('link() takes at most 1 argument ({0} given)'
.format(len(ag)))
ag = ag[0]
try:
ag.getACSIndex()
except AttributeError:
raise TypeError('ag must be an AtomGroup instance')
try:
ag.getAtomGroup()
except AttributeError:
pass
else:
raise TypeError('ag must be an AtomGroup instance')
if self.numAtoms() != ag.numAtoms():
raise ValueError('atom group and trajectory must '
'have same number of atoms')
self._ag = ag
if ag._getCoords() is None and self._coords is not None:
self._ag._setCoords(self._coords)
self._frame = Frame(None, None, None)
self._atoms = self._atoms or ag
[docs] def getLinked(self):
"""Returns linked :class:`.AtomGroup` instance, or **None** if a link
is not established."""
return self._ag
[docs] def isLinked(self):
"""Returns **True** if trajectory is linked to an :class:`.AtomGroup`
instance."""
return self._ag is not None
[docs] def getCoords(self):
"""Returns a copy of reference coordinates for (selected) atoms."""
if self._coords is None:
return None
if self._indices is None:
return self._coords.copy()
return self._coords[self._indices]
def _getCoords(self):
"""Returns a view of reference coordinates for (selected) atoms."""
if self._coords is None:
return None
if self._indices is None:
return self._coords
return self._coords[self._indices]
[docs] def setCoords(self, coords):
"""Set *coords* as the trajectory reference coordinate set. *coords*
must be an object with :meth:`getCoords` method, or a Numpy array with
suitable data type, shape, and dimensionality."""
atoms = coords
try:
coords = atoms.getCoords()
except AttributeError:
pass
else:
if coords is None:
raise ValueError('coordinates of {0} are not set'
.format(str(atoms)))
try:
checkCoords(coords, natoms=self._n_atoms)
except TypeError:
raise TypeError('coords must be a numpy array or an object '
'with `getCoords` method')
self._coords = coords
[docs] def setWeights(self, weights):
"""Set atomic weights."""
if self._n_atoms == 0:
raise AttributeError('coordinates must be set first')
self._weights = checkWeights(weights, self._n_atoms, None)
[docs] def getWeights(self):
"""Returns a copy of weights of (selected) atoms."""
if self._weights is not None:
if self._indices is None:
return self._weights.copy()
else:
return self._weights[self._indices]
def _getWeights(self):
if self._weights is not None:
if self._indices is None:
return self._weights
else:
return self._weights[self._indices]
[docs] def nextIndex(self):
"""Returns the index of the next frame."""
return self._nfi
[docs] def iterCoordsets(self):
"""Yield coordinate sets for (selected) atoms. Reference coordinates
are not included. Iteration starts from the next frame in line."""
if self._closed:
raise ValueError('I/O operation on closed file')
while self._nfi < self._n_csets:
yield self.nextCoordset()
[docs] def getCoordsets(self, indices=None):
"""Returns coordinate sets at given *indices*. *indices* may be an
integer, a list of ordered integers or **None**. **None** returns all
coordinate sets. If a list of indices is given, unique numbers will
be selected and sorted. That is, this method will always return unique
coordinate sets in the order they appear in the trajectory file.
Shape of the coordinate set array is (n_sets, n_atoms, 3)."""
pass
[docs] def getFrame(self, index):
"""Returns frame at given *index*."""
pass
[docs] def nextCoordset(self):
"""Returns next coordinate set."""
pass
def __next__(self):
"""Returns next coordinate set in a :class:`.Frame` instance. Note that
when atoms are set for the trajectory, this method will return the same
frame instance after updating its coordinates."""
pass
next = __next__
[docs] def goto(self, n):
"""Go to the frame at index *n*. ``n=0`` will rewind the trajectory
to the beginning, same as calling :meth:`reset` method. ``n=-1``
will go to the last frame. Frame *n* will not be parsed until one
of :meth:`next` or :meth:`nextCoordset` methods is called."""
pass
[docs] def skip(self, n):
"""Skip *n* frames. *n* must be a positive integer. Skipping some
frames will only change the next frame index (:meth:`nextIndex`)
Next frame will not be parsed until one of :meth:`next` or
:meth:`nextCoordset` methods is called."""
pass
[docs] def reset(self):
"""Go to first frame at index 0. First frame will not be parsed until
one of :meth:`next` or :meth:`nextCoordset` methods is called."""
pass
[docs] def close(self):
"""Close trajectory file."""
pass
[docs] def hasUnitcell(self):
"""Returns **True** if trajectory has unitcell data."""
pass
