Slicing trajectories
MDAnalysis trajectories can be indexed to return a Timestep
, or sliced to give a FrameIterator
.
In [1]: import MDAnalysis as mda
In [2]: from MDAnalysis.tests.datafiles import PSF, DCD
In [3]: u = mda.Universe(PSF, DCD)
In [4]: u.trajectory[4]
Out[4]: < Timestep 4 >
Indexing a trajectory shifts the Universe
to point towards that particular frame, updating dynamic data such as Universe.atoms.positions
.
Note
The trajectory frame is not read from the MD data. It is the internal index assigned by MDAnalysis.
In [5]: u.trajectory.frame
Out[5]: 4
Creating a FrameIterator
by slicing a trajectory does not shift the Universe
to a new frame, but iterating over the sliced trajectory will rewind the trajectory back to the first frame.
In [6]: fiter = u.trajectory[10::10]
In [7]: frames = [ts.frame for ts in fiter]
In [8]: print(frames, u.trajectory.frame)
[10, 20, 30, 40, 50, 60, 70, 80, 90] 0
You can also create a sliced trajectory with boolean indexing and fancy indexing. Boolean indexing allows you to select only frames that meet a certain condition, by passing a ndarray
with the same length as the original trajectory. Only frames that have a boolean value of True
will be in the resulting FrameIterator
. For example, to select only the frames of the trajectory with an RMSD under 2 angstrom:
In [9]: from MDAnalysis.analysis import rms
In [10]: protein = u.select_atoms('protein')
In [11]: rmsd = rms.RMSD(protein, protein).run()
In [12]: bools = rmsd.results.rmsd.T[-1] < 2
In [13]: print(bools)
[ True True True True True True True True True True True True
True True False False False False False False False False False False
False False False False False False False False False False False False
False False False False False False False False False False False False
False False False False False False False False False False False False
False False False False False False False False False False False False
False False False False False False False False False False False False
False False False False False False False False False False False False
False False]
In [14]: fiter = u.trajectory[bools]
In [15]: print([ts.frame for ts in fiter])
[0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13]
You can also use fancy indexing to control the order of specific frames.
In [16]: indices = [10, 2, 3, 9, 4, 55, 2]
In [17]: print([ts.frame for ts in u.trajectory[indices]])
[10, 2, 3, 9, 4, 55, 2]
You can even slice a FrameIterator
to create a new FrameIterator
.
In [18]: print([ts.frame for ts in fiter[::3]])
[0, 3, 6, 9, 12]