Source code for moldesign.molecules.trajectory

# Copyright 2016 Autodesk Inc.
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# See the License for the specific language governing permissions and
# limitations under the License.
import os
import time
from cStringIO import StringIO

import numpy as np

import moldesign as mdt
from moldesign import helpers, utils
from moldesign import units as u

from .molecule import MolecularProperties
from . import toplevel

[docs]class SubSelection(object): """ Descriptors to get bits of the trajectory trajectory.atoms[3].position -> array of positions trajectory.atoms[5].distance( mol.chains['B'].residues[5] ) -> array of distances trajectory.chains['A'].residue[23].com -> array of COMs NOT IMPLEMENTED YET """
[docs]class Frame(utils.DotDict): """ A snapshot of a molecule during its motion. This is really just a dictionary of properties. These properties are those accessed as ````, and can vary substantially depending on the origin of the trajectory. They also include relevant dynamical data and metadata (such as ``time``, ``momenta``, ``minimization_step``, etc.) Properties can be accessed either as attributes (``frame.property_name``) or as keys (``frame['property_name']``) Some properties with specific meaning are described below: Attributes: annotation (str): text describing this frame (will be displayed automatically when visualized) minimization_step (int): for minimization trajectories time (u.Scalar[time]): time during a dynamical trajectory Example: >>> mol = mdt.from_name('benzene') >>> mol.set_potential_model(moldesign.methods.models.RHF(basis='3-21g')) >>> traj = mol.minimize() >>> starting_frame = traj.frames[0] >>> assert starting_frame.potential_energy >= traj.frames[-1].potential_energy >>> assert starting_frame.minimization_step == 0 """ pass
class _TrajAtom(object): """ A helper class for querying individual atoms' dynamics """ ATOMIC_ARRAYS = {'position': 'positions', 'momentum': 'momenta', 'force': 'forces'} def __init__(self, traj, index): self.traj = traj self.index = index self.real_atom = self.traj.mol.atoms[self.index] def __getattr__(self, item): if item in ('traj', 'index', 'real_atom'): raise AttributeError('_TrajAtom.%s not assigned (pickle issue?)' % item) if item in self.ATOMIC_ARRAYS: is_array = True item = self.ATOMIC_ARRAYS[item] else: is_array = False try: # try to get a time-dependent version trajslice = getattr(self.traj, item) except AttributeError: # not time-dependent - look for an atomic property return getattr(self.real_atom, item) if is_array: return trajslice[:, self.index, :] else: if trajslice[0]['type'] != 'atomic': raise ValueError('%s is not an atomic quantity' % item) else: return u.array([f[self.real_atom] for f in trajslice]) def __dir__(self): attrs = [self.ATOMIC_ARRAYS.get(x, x) for x in dir(self.traj)] attrs.extend(dir(self.atom)) def distance(self, other): return mdt.geom.distance(self, other) @toplevel
[docs]class Trajectory(object): """ A ``Trajectory`` stores information about a molecule's motion and how its properties change as it moves. A trajectory object contains 1. a reference to the :class:`moldesign.Molecule` it describes, and 2. a list of :class:`Frame` objects, each one containing a snapshot of the molecule at a particular point in its motion. Args: mol (moldesign.Molecule): the trajectory will describe the motion of this molecule unit_system (u.UnitSystem): convert all attributes to this unit system (default: ``moldesign.units.default``) first_frame(bool): Create the trajectory's first :class:`Frame` from the molecule's current position Attributes: mol (moldesign.Molecule): the molecule object that this trajectory comes from frames (List[Frame]): a list of the trajectory frames in the order they were created info (str): text describing this trajectory unit_system (u.UnitSystem): convert all attributes to this unit system """ def __init__(self, mol, unit_system=None, first_frame=False): self._init = True = "Trajectory" self.frames = [] self.mol = mol self.unit_system = utils.if_not_none(unit_system, mdt.units.default) self._property_keys = None self._tempmol = mdt.Molecule(self.mol.atoms, copy_atoms=True) self._tempmol.dynamic_dof = self.mol.dynamic_dof self._viz = None self.atoms = [_TrajAtom(self, i) for i in xrange(self.mol.num_atoms)] if first_frame: self.new_frame() # TODO: the current implementation does not support cases where the molecular topology changes # TODO: allow caching to disk for very large trajectories MOL_ATTRIBUTES = ['positions', 'momenta', 'time'] """List[str]: Always store these molecular attributes""" def __getstate__(self): return self.__dict__.copy() def __setstate__(self, d): self.__dict__.update(d) @property def num_frames(self): """int: number of frames in this trajectory""" return len(self) def __len__(self): """overrides len(trajectory) to return number of frames""" return len(self.frames) @utils.kwargs_from(mdt.widgets.trajectory.TrajectoryViewer)
[docs] def draw3d(self, **kwargs): """TrajectoryViewer: create a trajectory visualization Args: **kwargs (dict): kwargs for :class:`moldesign.widgets.trajectory.TrajectoryViewer` """ self._viz = mdt.widgets.trajectory.TrajectoryViewer(self, **kwargs) return self._viz
draw = draw3d # synonym for backwards compatibility
[docs] def draw_orbitals(self, align=True): """TrajectoryOrbViewer: create a trajectory visualization""" from moldesign import widgets for frame in self: if 'wfn' not in frame: raise ValueError("Can't draw orbitals - orbital information missing in at least " "one frame. It must be calculated with a QM method.") if align: self.align_orbital_phases() self._viz = widgets.trajectory.TrajectoryOrbViewer(self) return self._viz
def __str__(self): return 'Trajectory for molecule "%s" (%d frames)' % (self.mol, self.num_frames) def __repr__(self): try: return '<%s>' % str(self) except Exception: return '<Trajectory object @ %s (exception in repr)>' % hex(id(self)) def __add__(self, other): newtraj = Trajectory(self.mol, unit_system=self.unit_system) newtraj.frames = self.frames + other.frames return newtraj
[docs] def new_frame(self, properties=None, **additional_data): """ Create a new frame, EITHER from the parent molecule or from a list of properties Args: properties (dict): dictionary of properties (i.e. {'positions':[...], 'potential_energy':...}) **additional_data (dict): any additional data to be added to the trajectory frame Returns: int: frame number (0-based) """ # TODO: callbacks to update a status display - allows monitoring a running simulation if properties is None: new_frame = Frame() for attr in self.MOL_ATTRIBUTES: val = getattr(self.mol, attr) try: # take numpy arrays' values, not a reference val = val.copy() except AttributeError: pass if val is not None: new_frame[attr] = val new_frame.update( else: new_frame = Frame(properties) for key, val in additional_data.iteritems(): assert key not in new_frame, ( "Can't overwrite molecule's properties with additional_data key %s" % key) new_frame[key] = val self._update_property_keys(new_frame) self.frames.append(new_frame)
def _get_traj_atom(self, a): if a is None: return None elif isinstance(a, mdt.Atom): return self.atoms[a.index] else: assert isinstance(a, _TrajAtom)
[docs] def distance(self, a1, a2): a1, a2 = map(self._get_traj_atom, (a1, a2)) return mdt.distance(a1, a2)
[docs] def angle(self, a1, a2, a3): a1, a2, a3 = map(self._get_traj_atom, (a1, a2, a3)) return mdt.angle(a1, a2, a3)
[docs] def dihedral(self, a1, a2, a3=None, a4=None): a1, a2, a3, a4 = map(self._get_traj_atom, (a1, a2, a3, a4)) return mdt.dihedral(a1, a2, a3, a4)
def _update_property_keys(self, new_frame=None): """ Update the internal list of molecular properties that can be sliced. If a frame is passed, update from that frame's properties. Otherwise, update from the entire list of stored frames """ if self._property_keys is None: self._property_keys = set() if new_frame is None: for frame in self: self._property_keys.update(frame.keys()) else: self._property_keys.update(new_frame) def __dir__(self): return list(self._property_keys.union(dir(self.__class__)).union(self.__dict__)) def __getitem__(self, item): return self.frames[item] def __getattr__(self, item): # TODO: move slicing to frames, so that this will work with __getitem__ as well # TODO: prevent identical recursion (so __getattr__(foo) can't call __getattr__(foo)) if self._property_keys is None: self._update_property_keys() if item in self._property_keys: # return a list of properties for each frame return self.slice_frames(item) else: raise AttributeError('Frame %s has no attribute %s' % (self, item))
[docs] def rmsd(self, atoms=None, reference=None): r""" Calculate root-mean-square displacement for each frame in the trajectory. The RMSD between times :math:`t` and :math:`t0` is given by :math:`\text{RMSD}(t;t_0) =\sqrt{\sum_{i \in \text{atoms}} \left( \mathbf{R}_i(t) - \mathbf{R}_i(t_0) \right)^2}`, where :math:`\mathbf{R}_i(t)` is the position of atom *i* at time *t*. Args: atoms (list[mdt.Atom]): list of atoms to calculate the RMSD for (all atoms in the ``Molecule``) reference (u.Vector[length]): Reference positions for RMSD. (default: ``traj.frames[0].positions``) Returns: u.Vector[length]: list of RMSD displacements for each frame in the trajectory """ if reference is None: refpos = self.frames[0].positions else: refpos = reference.positions atoms = mdt.utils.if_not_none(atoms, self.mol.atoms) indices = np.array([atom.index for atom in atoms]) rmsds = [] for f in self.frames: diff = (refpos[indices] - f.positions[indices]) rmsds.append(np.sqrt((diff*diff).sum()/len(atoms))) return u.array(rmsds).defunits()
[docs] def slice_frames(self, key, missing=None): """ Return an array of giving the value of ``key`` at each frame. Args: key (str): name of the property, e.g., time, potential_energy, annotation, etc missing: value to return if a given frame does not have this property Returns: moldesign.units.Vector: vector containing the value at each frame, or the value given in the ``missing`` keyword) (len= `len(self)` ) """ has_units = True result = [] for f in self.frames: val = f.get(key, None) if not issubclass(type(val), u.MdtQuantity): has_units = False result.append(val) if has_units: result = u.array([frame.get(key, None) for frame in self.frames]) return u.default.convert(result) else: return np.array(result)
@property def kinetic_energy(self): convert_units = True energies = [] for frame in self.frames: if 'momenta' in frame: energies.append( helpers.kinetic_energy(frame.momenta, self.mol.dim_masses)) else: convert_units = False energies.append(None) if convert_units: arr = u.array(energies) return u.default.convert(arr) else: return energies @property def kinetic_temperature(self): convert_units = True temps = [] energies = self.kinetic_energy dof = self.mol.dynamic_dof for energy, frame in zip(energies, self.frames): if energy is not None: temps.append(helpers.kinetic_temperature(energy, dof)) else: convert_units = False temps.append(None) if convert_units: arr = u.array(temps) return u.default.convert(arr) else: return temps DONOTAPPLY = set(['kinetic_energy'])
[docs] def apply_frame(self, frame): """ Reconstruct the underlying molecule with the given frame. Right now, any data not passed is ignored, which may result in properties that aren't synced up with each other ... """ # TODO: need to prevent multiple things using the _tempmol from conflicting with each other m = self._tempmol for prop in self.MOL_ATTRIBUTES: if prop in self.DONOTAPPLY: continue if prop in frame: setattr(m, prop, frame[prop]) = MolecularProperties(m) for attr in frame:[attr] = frame[attr]
# TODO: need to fix this - it creates import problems #@utils.args_from(mdt.converters.write_trajectory, # allexcept=['traj'], # append_docstring_description=True)
[docs] def write(self, *args, **kwargs): return mdt.fileio.write_trajectory(self, *args, **kwargs)
[docs] def plot(self, x, y, **kwargs): """ Create a matplotlib plot of property x against property y Args: x,y (str): names of the properties **kwargs (dict): kwargs for :meth:`matplotlib.pylab.plot` Returns: List[matplotlib.lines.Lines2D]: the lines that were plotted """ from matplotlib import pylab xl = yl = None if type(x) is str: strx = x x = getattr(self, x) xl = '%s / %s' % (strx, x.units) if type(y) is str: stry = y y = getattr(self, y) yl = '%s / %s' % (stry, y.units) plt = pylab.plot(x, y, **kwargs) pylab.xlabel(xl); pylab.ylabel(yl); pylab.grid() return plt
[docs] def align_orbital_phases(self, reference_frame=None): """ Try to remove orbital sign flips between frames. If `reference_frame` is not passed, we'll start with frame 0 and align successive pairs of orbitals. If `reference_frame` is an int, we'll move forwards and backwards from that frame number. Otherwise, we'll try to align every orbital frame to those in reference_frame Args: reference_frame (int or Frame): ``Frame`` containing the orbitals to align with (default: align each frame with the previous one) """ if reference_frame is None: iframe = 0 relative_alignment = True elif type(reference_frame) == int: iframe = reference_frame relative_alignment = True else: relative_alignment = False if relative_alignment: for i in xrange(iframe+1, self.num_frames): self.frames[i].wfn.align_orbital_phases( self.frames[i-1].wfn) for i in xrange(iframe-1, -1, -1): self.frames[i].wfn.align_orbital_phases( self.frames[i+1].wfn) else: for frame in self.frames: frame.wfn.align_orbital_phases(reference_frame.wfn)