LOOS  v2.3.2
trajectories.py
1 """
2 Python-based trajectory classes that wrap loos.Trajectory objects
3 
4 """
5 import loos
6 import copy
7 
8 ## Python-based wrapper for LOOS Trajectories
9 # This class turns a loos Trajectory into something more
10 # python-like. Behind the scenes, it wraps a loos::AtomicGroup and
11 # a loos::Trajectory. The behavior of the trajectory can be controlled
12 # through passed keywords,
13 #
14 # Keyword | Description
15 # -----------|------------------------------------------------------------------------------
16 # skip=n | Skip the first n-frames of the wrapped trajectory
17 # stride=n | Step through the wrapped trajectory n-frames at a time
18 # iterator=i | Use the python iterator object i to select frames from the wrapped trajectory
19 # subset=s | Use 's' to select a subset of the model to use for each frame
20 #
21 # Remember that all atoms are shared. If you want to decouple the
22 # trajectory from other groups, pass it a copy of the model.
23 #
24 # examples:
25 # \code
26 # model = loos.createSystem('foo.pdb')
27 # traj = loos.pyloos.Trajectory('foo.dcd', model)
28 # calphas = loos.selectAtoms(model, 'name == "CA"')
29 # for frame in traj:
30 # print calphas.centroid()
31 # \endcode
32 #
33 # The same thing but skipping the first 50 frames and taking every other frame
34 # \code
35 # traj = loos.pyloos.Trajectory('foo.dcd', model, skip=50, stride=2)
36 # \endcode
37 #
38 # Only use frames 19-39 (i.e. the 20th through 40th frames)
39 # \code
40 # traj = loos.pyloos.Trajectory('foo.dcd', model, iterator=range(19,40))
41 # \endcode
42 #
43 # An alternative way of only iterating over a subset...
44 # \code
45 # model = loos.createSystem('foo.pdb')
46 # traj = loos.pyloos.Trajectory('foo.dcd', model, subset='name == "CA"')
47 # \endcode
48 #
49 # Decouple the model stored in the trajectory,
50 # \code
51 # traj = loos.pyloos.Trajectory('foo.dcd', model.copy())
52 # \endcode
53 #
54 
55 class Trajectory(object):
56  """
57  Python-based wrapper for LOOS Trajectories
58  >>> model = loos.createSystem('foo.pdb')
59  >>> traj = loos.pyloos.Trajectory('foo.dcd', model)
60 
61  keyword args:
62  skip = # of frames to skip from start
63  stride = # of frames to step through
64  iterator = Python iterator used to pick frame (overrides skip and stride)
65  subset = Selection used to pick subset for each frame
66 
67  See the Doxygen documentation for more details.
68  """
69 
70  ## Instantiate a Trajectory object.
71  def __init__(self, fname, model, **kwargs):
72 
73  self._skip = 0
74  self._stride = 1
75  self._iterator = None
76 
77  if 'skip' in kwargs:
78  self._skip = kwargs['skip']
79  if 'stride' in kwargs:
80  self._stride = kwargs['stride']
81  if 'iterator' in kwargs:
82  self._iterator = kwargs['iterator']
83  if 'subset' in kwargs:
84  self._subset = loos.selectAtoms(model, kwargs['subset'])
85  else:
86  self._subset = model
87 
88  self._model = model
89  self._fname = fname
90  self._traj = loos.createTrajectory(fname, model)
91 
92  self._stale = 1
93  self._initFrameList()
94 
95 
96  def _initFrameList(self):
97  self._framelist = []
98  if self._iterator is None:
99  it = range(self._skip, self._traj.nframes(), self._stride)
100  else:
101  it = iter(self._iterator)
102 
103  for i in it:
104  self._framelist.append(i)
105 
106  self._index = 0
107  self._stale = 0
108 
109  def stride(self, n):
110  """
111  Step through the trajectory by this number of frames
112  """
113  self._stride = n
114 
115  def skip(self, n):
116  """
117  Skip this number of frames at the start of the trajectory
118  """
119  self._skip = n
120 
121  def fileName(self):
122  """
123  Return the filename that this Trajectory represents
124  """
125  return(self._fname)
126 
127  def setSubset(self, selection):
128  """
129  Set the subset used when iterating over a trajectory.
130  The selection is a LOOS selection string.
131  """
132  self._subset = loos.selectAtoms(self._model, selection)
133 
134 
135  def __iter__(self):
136  if self._stale:
137  self._initFrameList()
138  self._index = 0
139  return(self)
140 
141  def __len__(self):
142  """
143  Number of frames in the trajectory
144  """
145  if self._stale:
146  self._initFrameList()
147  return(len(self._framelist))
148 
149 
150  def reset(self):
151  """Reset the iterator"""
152  self._index = 0
153 
154  def next(self):
155  if (self._index >= len(self._framelist)):
156  raise StopIteration
157  frame = self.__getitem__(self._index)
158 
159  self._index += 1
160  return(frame)
161 
162 
163  def trajectory(self):
164  """Access the wrapped loos.Trajectory"""
165  return(self._traj)
166 
167  def model(self):
168  """Return the current model"""
169  return(self._model)
170 
171  def readFrame(self, i):
172  """Read a frame and update the model"""
173  if self._stale:
174  self._initFrameList()
175  if (i < 0 or i >= len(self._framelist)):
176  raise IndexError
177  self._traj.readFrame(self._framelist[i])
178  self._traj.updateGroupCoords(self._model)
179  return(self._subset)
180 
181  def frame(self):
182  """Return the current frame (subset)"""
183  return(self._subset)
184 
185  def realIndex(self):
186  """The 'real' frame in the trajectory for this index"""
187  if self._stale:
188  self._initFrameList()
189  return(self._framelist[self._index-1])
190 
191  def index(self):
192  """The state of the iterator"""
193  return(self._index-1)
194 
195 
196  def frameNumber(self, i):
197  """
198  Returns the real frame numbers corresponding to the passed indices. Can accept
199  either an integer or a list of integers.
200 
201  For example:
202  >>> t = loos.pyloos.Trajectory('foo.dcd', model, skip=50)
203  >>> t.frameNumber(0)
204  50
205  >>> t.frameNumber(range(0,2))
206  [50, 51]
207  """
208 
209  if self._stale:
210  self._initFrameList()
211  if type(i) is int:
212  if (i < 0):
213  i += len(self._framelist)
214  return(self._framelist[i])
215 
216  indices = [x if x >=0 else len(self._framelist)+x for x in i]
217  framenos = [self._framelist[x] for x in indices]
218  return(framenos)
219 
220 
221  def _getSlice(self, s):
222  if self._stale:
223  self._initFrameList()
224  indices = list(range(*s.indices(self.__len__())))
225  ensemble = []
226  for i in indices:
227  self._traj.readFrame(self._framelist[i])
228  self._traj.updateGroupCoords(self._model)
229  dup = self._subset.copy()
230  ensemble.append(dup)
231  return(ensemble)
232 
233 
234  def __getitem__(self, i):
235  """Handle array indexing and slicing. Negative indices are
236  relative to the end of the trajectory"""
237  if self._stale:
238  self._initFrameList()
239  if isinstance(i, slice):
240  return(self._getSlice(i))
241 
242  if (i < 0):
243  i += len(self._framelist)
244  if (i >= len(self._framelist) or i < 0):
245  raise IndexError
246  self._traj.readFrame(self._framelist[i])
247  self._traj.updateGroupCoords(self._model)
248  return(self._subset)
249 
250 
251 
252 ## Virtual trajectory composed of multiple Trajectory objects
253 # This class can combine multiple loos.pyloos.Trajectory objects
254 # into one big "virtual" trajectory. Any skips or strides set in
255 # the contained trajectories will be honored. In addition, a skip
256 # and a stride for the whole meta-trajectory are available. These
257 # can be set via keyword arguments when creating a VirtualTrajectory,
258 #
259 # Keyword | Description
260 # -----------|------------------------------------------------------------------------------
261 # skip=n | Skip the first n-frames of the virtual trajectory
262 # stride=n | Step through the virtual trajectory n frames at a time
263 # iterator=i | Use the python iterator object i to select frames from the virtual trajectory
264 #
265 # There is no requirement that the subsets used for all trajectories
266 # must be the same. Ideally, the frame (subset) that is returned
267 # should be compatible (e.g. same atoms in the same order), but the
268 # models used for each trajectory (and the corresponding subset
269 # selection) can vary. Why would you want to do this? Imagine
270 # combining three different GPCRs where the subsets are the common
271 # trans-membrane C-alphas. This makes processing all of the
272 # ensembles together easier.
273 #
274 # <h2>WARNING</h2>
275 # Since each contained trajectory can have a different set of shared
276 # atoms it updates, care must be taken when pre-selecting atoms.
277 #
278 #
279 # Examples:
280 # \code
281 # model = loos.createSystem('foo.pdb')
282 # Takes a filename (\a fname) and an AtomicGroup model (\a model).
283 # Takes a filename (\a fname) and an AtomicGroup model (\a model).
284 # Takes a filename (\a fname) and an AtomicGroup model (\a model).
285 # traj1 = loos.pyloos.Trajectory('foo-1.dcd', model)
286 # traj2 = loos.pyloos.Trajectory('foo-2.dcd', model)
287 # vtraj = loos.pyloos.VirtualTrajectory(traj1, traj2)
288 #
289 # for frame in vtraj:
290 # print frame.centroid()
291 # \endcode
292 #
293 # Adding another trajectory in with its own stride and skip
294 # \code
295 # traj3 = loos.pyloos.Trajectory('foo-3.dcd', skip=50, stride=2)
296 # vtraj.append(traj3)
297 # \endcode
298 #
299 # Same as above but stride through the combined trajectory
300 # \code
301 # vtraj10 = loos.pyloos.VirtualTrajectory(traj1, traj2, stride=10)
302 # \endcode
303 
304 class VirtualTrajectory(object):
305  """
306  Combines multiple loos.pyloos.Trajectory objects into one big virtual trajectory
307  >>> vtraj = loos.pyloos.VirtualTrajectory(traj1, traj2, traj3)
308  Keyword arguments:
309  skip = # of frames to skip at start of composite traj
310  stride = # of frames to step through in the composite traj
311  iterator = Python iterator used to pick frames from the composite traj
312 
313  See the Doxygen documentation for more details.
314  """
315 
316 
317  def __init__(self, *trajs, **kwargs):
318  self._skip = 0
319  self._stride = 1
320  self._nframes = 0
321  self._iterator = None
322  self._trajectories = list(trajs)
323 
324  self._index = 0
325  self._framelist = []
326  self._trajlist = []
327  self._stale = 1
328 
329  if 'skip' in kwargs:
330  self._skip = kwargs['skip']
331  if 'stride' in kwargs:
332  self._stride = kwargs['stride']
333  if 'iterator' in kwargs:
334  self._iterator = kwargs['iterator']
335  if 'subset' in kwargs:
336  self.setSubset(kwargs['subset'])
337 
338 
339  def append(self, *traj):
340  """
341  Add a trajectory to the end of the virtual trajectory. Resets
342  the iterator state
343  """
344  self._trajectories.extend(traj)
345  self._stale = 1
346 
347  def stride(self, n):
348  """
349  Set the stride of the combined trajectory
350  """
351  self._stride = n
352  self._stale = 1
353 
354  def skip(self, n):
355  """
356  Set the skip of the combined trajectory
357  """
358  self._skip = n
359  self._stale = 1
360 
361  def allStride(self, n):
362  """
363  Sets the stride of all contained trajectories
364  """
365  self._stale = 1
366  for t in self._trajectories:
367  t.stride(n)
368 
369  def allSkip(self, n):
370  """
371  Sets the skip of all contained trajectories
372  """
373  self._stale = 1
374  for t in self._trajectories:
375  t.skip(n)
376 
377  def setSubset(self, selection):
378  """
379  Set the subset selection for all managed trajectories
380  """
381  for t in self._trajectories:
382  t.setSubset(selection)
383 
384  def frame(self):
385  """
386  Return the current frame/model. If the iterator is past the
387  end of the trajectory list, return the last valid frame.
388  """
389  if self._stale:
390  self._initFrameList()
391 
392  if self._index >= len(self._framelist):
393  i = len(self._framelist) - 1
394  else:
395  i = self._index
396 
397  return(self._trajectories[self._trajlist[i]].frame())
398 
399  def index(self):
400  """
401  Return index into composite trajectory for current frame
402  """
403  return(self._index-1)
404 
405 
406 
407  ## Returns information about the ith frame in the VirtualTrajectory
408  # The tuple returned has the following format:
409  # \code
410  # (frame-index, traj-index, trajectory, real-frame-within-trajectory)
411  # \endcode
412  #
413  # Consider the following,
414  # \code
415  # t1 = loos.pyloos.Trajectory('foo.dcd', model) # 50 frames
416  # t2 = loos.pyloos.Trajectory('bar.dcd', model) # 25 frames
417  # vt = loos.pyloos.VirtualTrajectory(t1, t2)
418  # \endcode
419  #
420  # * \c frame-index is the index into the corresponding trajectory object. For example,
421  # frameLocation(50) would have a frame-index of 0 because vt[50] would return the first
422  # frame from t2.
423  #
424  # * \c traj-index is the index into the list of managed trajectories for the frame.
425  # In the above example, the traj-index will be 1.
426  # * \c trajectory is the actual loos.pyloos.Trajectory object that contains the frame.
427  # * \c real-frame-within-trajectory is the same as calling trajectory.frameNumber(frame-index).
428  #
429  # Instead of the \c t1 above, imagine it was setup this way,
430  # \code
431  # t1 = loos.pyloos.Trajectory('foo.dcd', model, skip=25)
432  # \endcode
433  # Now, <tt>vt.frameLocation(0)</tt> will return <tt>(0, 0, t1, 25)</tt>,
434  # and <tt>vt.frameLocation(25)</tt> will return <tt>(25, 1, t2, 0)</tt>
435  #
436  # Python documentation:
437 
438  def frameLocation(self, i):
439  """
440  Return info about where a frame comes from.
441  >>> (frame-index, traj-index, trajectory, real-frame-within-trajectory) = vtraj.frameLocation(i)
442  """
443  if (self._stale):
444  self._initFrameList()
445 
446  if (i < 0):
447  i += len(self._framelist)
448 
449  t = self._trajectories[self._trajlist[i]]
450  return( self._framelist[i], self._trajlist[i], t, t.frameNumber(self._framelist[i]))
451 
452  def _initFrameList(self):
453  frames = []
454  trajs = []
455  for j in range(len(self._trajectories)):
456  t = self._trajectories[j]
457  for i in range(len(t)):
458  frames.append(i)
459  trajs.append(j)
460 
461  self._framelist = []
462  self._trajlist = []
463  n = len(frames)
464  if (self._iterator is None):
465  it = iter(range(self._skip, n, self._stride))
466  else:
467  it = iter(iterator)
468  for i in it:
469  self._framelist.append(frames[i])
470  self._trajlist.append(trajs[i])
471 
472  self._index = 0
473  self._stale = 0
474 
475  def __len__(self):
476  """
477  Total number of frames
478  """
479  if self._stale:
480  self._initFrameList()
481  return(len(self._framelist))
482 
483 
484  def __getitem__(self, i):
485  """
486  Return the ith frame in the composite trajectory. Supports
487  Python slicing. Negative indices are relative to the end of
488  the composite trajectory.
489  """
490  if self._stale:
491  self._initFrameList()
492 
493  if isinstance(i, slice):
494  return(self._getSlice(i))
495 
496  if (i < 0):
497  i += len(self)
498  if (i >= len(self)):
499  raise IndexError
500 
501  return(self._trajectories[self._trajlist[i]][self._framelist[i]])
502 
503 
504  def __iter__(self):
505  if self._stale:
506  self._initFrameList()
507  self._index = 0
508  return(self)
509 
510  def reset(self):
511  self._index = 0
512 
513  def next(self):
514  if self._stale:
515  self._initFrameList()
516  if (self._index >= len(self._framelist)):
517  raise StopIteration
518  frame = self.__getitem__(self._index)
519  self._index += 1
520  return(frame)
521 
522  def _getSlice(self, s):
523  indices = list(range(*s.indices(self.__len__())))
524  ensemble = []
525  for i in indices:
526  frame = self._trajectories[self._trajlist[i]][self._framelist[i]].copy()
527  ensemble.append(frame)
528  return(ensemble)
529 
530 
531 
532 ## A virtual trajectory that supports iterative alignment.
533 # Only the
534 # transformation needed to align each frame is stored. When a frame
535 # is accessed, it is automatically transformed into the aligned
536 # orientation. All keywords from VirtualTrajectory are supported,
537 # along with the following new ones,
538 #
539 # Keyword | Description
540 # ------------|------------------------------------------------------------------------------
541 # alignwith=s | Use 's' to select what part of the model is used for aligning.
542 # reference=g | Use the AtomicGroup g as a reference structure. All frames will be aligned to it.
543 #
544 # There are two ways that a trajectory can be aligned. The first
545 # uses in iterative alignment method (the same used in LOOS). This
546 # is the default method. In order to do the alignment, the
547 # alignwith subset must be read into memory and temporarily stored.
548 # This can potentially use a lot of memory and create delays in
549 # execution. Once the alignment is complete, however, those cached
550 # frames are released and subsequent frame accesses will be quick.
551 #
552 # The second method is to align each frame to a reference
553 # structure. This method is selected when a reference structure is
554 # passed to the constructor (with the 'reference' keyword), or when
555 # setReference() is called. Note that you can pass None to
556 # setReference() which will return the AlignedVirtualTrajectory to
557 # the iterative method. Also note that the reference structure is
558 # copied into the AVT object as a deep copy (i.e. it does not share
559 # any atoms).
560 #
561 # See VirtualTrajectory for some basic examples in addition to
562 # below...
563 #
564 # Align using only C-alphas (the default)
565 # \code
566 # vtraj = loos.pyloos.AlignedVirtualTrajectory(traj1, traj2)
567 # \endcode
568 #
569 # Align using only backbone atoms
570 # \code
571 # vtraj = loos.pyloos.AlignedVirtualTrajectory(traj1, traj2, alignwith='name =~ "^(C|N|O|CA)$"')
572 # \endcode
573 #
574 # Add another trajectory
575 # \code
576 # vtraj.append(traj3)
577 # \endcode
578 #
579 # Align using only C-alphas and a reference structure
580 # \code
581 # refmodel = loos.createSystem('foo-ref.pdb')
582 # refsubset = loos.selectAtoms(refmodel, 'name == "CA"')
583 # vtraj = loos.pyloos.AlignedVirtualTrajectory(traj1, traj2, reference = refsubset)
584 # \endcode
585 
587  """
588  Combine loos.pyloos.Trajectory objects and align them
589 
590  >>> aligned = loos.pyloos.AlignedVirtualTrajectory(traj1, traj2, alignwith = 'backbone')
591  Supports the same keywords as VirtualTrajectory.
592  New keywords:
593  alignwith = Selection used for alignment (default is all C-alphas)
594  reference = AtomicGroup that all frames are aligned to (disables iterative alignment)
595 
596  See the Doxygen documentation for more details.
597  """
598 
599  def __init__(self, *trajs, **kwargs):
600  super(AlignedVirtualTrajectory, self).__init__(*trajs, **kwargs)
601  self._aligned = False
602  self._xformlist = []
603  self._rmsd = -1
604  self._iters = -1
605  if 'alignwith' in kwargs:
606  self._alignwith = kwargs['alignwith']
607  else:
608  self._alignwith = 'name == "CA"'
609 
610  if 'reference' in kwargs:
611  self._reference = copy.deepcopy(kwargs['reference'])
612  else:
613  self._reference = None
614 
615 
616  def append(self, *traj):
617  """
618  Add another trajectory at the end. Requires re-aligning
619  """
620  self._trajectories.extend(traj)
621  self._stale = True
622  self._aligned = False
623 
624  def alignWith(self, selection):
625  """
626  Change the selection used to align with. Requires re-aligning
627  """
628  self._alignwith = selection
629  self._aligned = False
630 
631 
632  def __iter__(self):
633  self._align()
634  self._index = 0
635  return(self)
636 
637 
638  def setReference(self, reference):
639  self._reference = copy.deepcopy(reference)
640  self._aligned = False
641 
642  def _align(self):
643  """
644  Align the frames (called implicitly on iterator or array access)
645  """
646  current_traj = None
647  current_subset = None
648  ensemble = []
649 
650  if self._stale:
651  self._initFrameList()
652 
653  if self._reference: # Align to a reference structure
654  self._xformlist = []
655  for i in range(len(self._framelist)):
656  t = self._trajectories[self._trajlist[i]]
657  if t != current_traj:
658  current_traj = t
659  current_subset = loos.selectAtoms(t.model(), self._alignwith)
660  t.readFrame(self._framelist[i])
661  m = current_subset.superposition(self._reference)
662  x = loos.XForm()
663  x.load(m)
664  self._xformlist.append(x)
665 
666  self._rmsd = 0.0
667  self._iters = 0
668 
669  else: # Iterative alignment
670 
671  ensemble = loos.DoubleVectorMatrix()
672 
673  for i in range(len(self._framelist)):
674  t = self._trajectories[self._trajlist[i]]
675  if t != current_traj:
676  current_traj = t
677  current_subset = loos.selectAtoms(t.model(), self._alignwith)
678  t.readFrame(self._framelist[i])
679  ensemble.push_back(current_subset.coordsAsVector())
680 
681  result = loos.iterativeAlignmentPy(ensemble)
682  (self._xformlist, self._rmsd, self._iters) = (loos.xformVectorToList(result.transforms), result.rmsd, result.iterations)
683 
684  self._aligned = True
685 
686 
687  def rmsd(self):
688  return(self._rmsd)
689 
690  def iters(self):
691  return(self._iters)
692 
693 
694  def _getSlice(self, s):
695  indices = list(range(*s.indices(self.__len__())))
696  ensemble = []
697  for i in indices:
698  frame = self._trajectories[self._trajlist[i]][self._framelist[i]].copy()
699  frame.applyTransform(self._xformlist[i])
700  ensemble.append(frame)
701  return(ensemble)
702 
703 
704  def __getitem__(self, i):
705  """
706  Returns the ith frame aligned. Supports Python slices. Negative indices are relative
707  to the end of the composite trajectory.
708  """
709  if not self._aligned:
710  self._align()
711 
712  if isinstance(i, slice):
713  return(self._getSlice(i))
714 
715  if (i < 0):
716  i += len(self._framelist)
717  if (i >= len(self._framelist)):
718  raise IndexError
719 
720  frame = self._trajectories[self._trajlist[i]][self._framelist[i]]
721  frame.applyTransform(self._xformlist[i])
722  return(frame)
Python-based wrapper for LOOS Trajectories This class turns a loos Trajectory into something more pyt...
Definition: trajectories.py:55
pTraj createTrajectory(const std::string &filename, const AtomicGroup &g)
Factory function for reading in a trajectory file.
Definition: sfactories.cpp:184
Matrix class for handling coordinate transforms...
Definition: XForm.hpp:91
Virtual trajectory composed of multiple Trajectory objects This class can combine multiple loos...
AtomicGroup selectAtoms(const AtomicGroup &source, const std::string selection)
Applies a string-based selection to an atomic group...
Definition: utils.cpp:195
A virtual trajectory that supports iterative alignment.
def frameLocation(self, i)
Returns information about the ith frame in the VirtualTrajectory The tuple returned has the following...
def __init__(self, fname, model, kwargs)
Instantiate a Trajectory object.
Definition: trajectories.py:71