[docs]classChainReader(base.ProtoReader):"""Reader that concatenates multiple trajectories on the fly. The :class:`ChainReader` is used by MDAnalysis internally to represent multiple trajectories as one virtual trajectory. Users typically do not need to use the :class:`ChainReader` explicitly. Chainreader can also handle a continuous trajectory split over several files. To use this pass the ``continuous == True`` keyword argument. Setting ``continuous=True`` will make the reader choose frames from the set of trajectories in such a way that the trajectory appears to be as continuous in time as possible, i.e. that time is strictly monotonically increasing. This means that there will be no duplicate time frames and no jumps backwards in time. However, there can be gaps in time (e.g., multiple time steps can appear to be missing). Ultimately, it is the user's responsibility to ensure that the input trajectories can be virtually stitched together in a meaningful manner. As an example take the following trajectory that is split into three parts. The column represents the time and the trajectory segments overlap. With the continuous chainreader only the frames marked with a + will be read. :: part01: ++++-- part02: ++++++- part03: ++++++++ .. warning:: The order in which trajectories are given to the chainreader can change what frames are used with the continuous option. The default chainreader will read all frames. The continuous option is currently only supported for XTC and TRR files. Notes ----- The trajectory API attributes exist but most of them only reflect the first trajectory in the list; :attr:`ChainReader.n_frames`, :attr:`ChainReader.n_atoms`, and :attr:`ChainReader.fixed` are properly set, though .. versionchanged:: 0.11.0 Frames now 0-based instead of 1-based .. versionchanged:: 0.13.0 :attr:`time` now reports the time summed over each trajectory's frames and individual :attr:`dt`. .. versionchanged:: 0.19.0 added ``continuous`` trajectory option .. versionchanged:: 0.19.0 limit output of __repr__ """format='CHAIN'def__init__(self,filenames,skip=1,dt=None,continuous=False,**kwargs):"""Set up the chain reader. Parameters ---------- filenames : str or list or sequence file name or list of file names; the reader will open all file names and provide frames in the order of trajectories from the list. Each trajectory must contain the same number of atoms in the same order (i.e. they all must belong to the same topology). The trajectory format is deduced from the extension of each file name. Extension: `filenames` are either a single file name or list of file names in either plain file names format or ``(filename, format)`` tuple combination. This allows explicit setting of the format for each individual trajectory file. skip : int (optional) skip step (also passed on to the individual trajectory readers); must be same for all trajectories dt : float (optional) Passed to individual trajectory readers to enforce a common time difference between frames, in MDAnalysis time units. If not set, each reader's `dt` will be used (either inferred from the trajectory files, or set to the reader's default) when reporting frame times; note that this might lead an inconsistent time difference between frames. continuous : bool (optional) treat all trajectories as one single long trajectory. Adds several checks; all trajectories have the same dt, they contain at least 2 frames, and they are all of the same file-type. Not implemented for all trajectory formats! This can be used to analyze GROMACS simulations without concatenating them prior to analysis. **kwargs : dict (optional) all other keyword arguments are passed on to each trajectory reader unchanged """super(ChainReader,self).__init__()filenames=asiterable(filenames)# Override here because single frame readers handle this argument as a# kwarg to a timestep which behaves differently if dt is present or not.ifdtisnotNone:kwargs['dt']=dtself.readers=[core.reader(filename,**kwargs)forfilenameinfilenames]self.filenames=np.array([fn[0]ifisinstance(fn,tuple)elsefnforfninfilenames])# pointer to "active" trajectory index into self.readersself.__active_reader_index=0self.skip=skipself.n_atoms=self._get_same('n_atoms')# Translation between virtual frames and frames in individual# trajectories. Assumes that individual trajectories i contain frames# that can be addressed with an index 0 <= f < n_frames[i]# Build a map of frames: ordered list of starting virtual frames; the# index i into this list corresponds to the index into self.readers## For virtual frame 0 <= k < sum(n_frames) find corresponding# trajectory i and local frame f (i.e. readers[i][f] will correspond to# ChainReader[k]).# build map 'start_frames', which is used by _get_local_frame()n_frames=self._get('n_frames')# [0]: frames are 0-indexed internally# (see Timestep.check_slice_indices())self._start_frames=np.cumsum([0]+n_frames)self.n_frames=np.sum(n_frames)self.dts=np.array(self._get('dt'))self.total_times=self.dts*n_frames#: source for trajectories frame (fakes trajectory)self.__chained_trajectories_iter=None# calculate new start_frames to have a time continuous trajectory.ifcontinuous:filetypes=np.unique([r.formatforrinself.readers])ifnotlen(filetypes)==1:raiseValueError("ChainReader: continuous=true only supported"" when all files are using the same format. ""found {}".format(filetypes))ifnp.any(np.array(n_frames)==1):raiseRuntimeError("ChainReader: Need at least two frames in ""every trajectory with continuous=True")iffiletypes[0]notin['XTC','TRR']:raiseNotImplementedError("ChainReader: continuous=True only ""supported for xtc and trr format")# TODO: allow floating point precision in dt checkdt=self._get_same('dt')n_frames=np.asarray(self._get('n_frames'))self.dts=np.ones(self.dts.shape)*dt# the sorting needs to happen on two levels. The first major level# is by start times and the second is by end times.# The second level of sorting is needed for cases like:# [0 1 2 3 4 5 6 7 8 9] [0 1 2 4]# to# [0 1 2 4] [0 1 2 3 4 5 6 7 8 9]# after that sort the chain reader will worktimes=[]forrinself.readers:r[0]start=r.ts.timer[-1]end=r.ts.timetimes.append((start,end))# sort stepsort_idx=multi_level_argsort(times)self.readers=[self.readers[i]foriinsort_idx]self.filenames=self.filenames[sort_idx]self.total_times=self.dts*n_frames[sort_idx]# filter step: remove indices if we have complete overlapiflen(self.readers)>1:used_idx=filter_times(np.array(times)[sort_idx],dt)self.readers=[self.readers[i]foriinused_idx]self.filenames=self.filenames[used_idx]self.total_times=self.dts[used_idx]*n_frames[used_idx]# rebuild lookup tablesf=[0,]n_frames=0forr1,r2inzip(self.readers[:-1],self.readers[1:]):r2[0],r1[0]r1_start_time=r1.timestart_time=r2.timer1[-1]ifr1.time<start_time:warnings.warn("Missing frame in continuous chain",UserWarning)# check for interleavingr1[1]ifr1_start_time<start_time<r1.time:raiseRuntimeError("ChainReader: Interleaving not supported with continuous=True.")# find end where trajectory was restarted fromfortsinr1[::-1]:ifts.time<start_time:breaksf.append(sf[-1]+ts.frame+1)n_frames+=ts.frame+1n_frames+=self.readers[-1].n_framesself._start_frames=sfself.n_frames=n_framesself._sf=sf# make sure that iteration always yields frame 0# rewind() also sets self.tsself.ts=Noneself.rewind()

[docs]def_get_local_frame(self,k):"""Find trajectory index and trajectory frame for chained frame `k`. Parameters ---------- k : int Frame `k` in the chained trajectory can be found in the trajectory at index *i* and frame index *f*. Frames are internally treated as 0-based indices into the trajectory. Returns ------- i : int trajectory f : int frame in trajectory i Raises ------ IndexError for `k<0` or `i<0`. Note ---- Does not check if `k` is larger than the maximum number of frames in the chained trajectory. """ifk<0:raiseIndexError("Virtual (chained) frames must be >= 0")# trajectory index ii=bisect.bisect_right(self._start_frames,k)-1ifi<0:raiseIndexError("Cannot find trajectory for virtual frame {0:d}".format(k))# local frame index f in trajectory i (frame indices are 0-based)f=k-self._start_frames[i]returni,f

defcopy(self):new=self.__class__(copy.copy(self.filenames))# seek the new reader to the same frame we started withnew[self.ts.frame]# then copy over the current Timestep in case it has# been modified since initial loadnew.ts=self.ts.copy()returnnew# attributes that can change with the current reader@propertydeffilename(self):"""Filename of the currently read trajectory"""returnself.active_reader.filename# TODO: check that skip_timestep is still supported in all readers# or should this be removed?@propertydefskip_timestep(self):returnself.active_reader.skip_timestep@propertydefdelta(self):returnself.active_reader.delta@propertydefperiodic(self):""":attr:`periodic` attribute of the currently read trajectory"""returnself.active_reader.periodic@propertydefunits(self):""":attr:`units` attribute of the currently read trajectory"""returnself.active_reader.units@propertydefcompressed(self):""":attr:`compressed` attribute of the currently read trajectory"""try:returnself.active_reader.compressedexceptAttributeError:returnNone@propertydefframe(self):"""Cumulative frame number of the current time step."""returnself.ts.frame@propertydeftime(self):"""Cumulative time of the current frame in MDAnalysis time units (typically ps)."""# Before 0.13 we had to distinguish between enforcing a common dt or# summing over each reader's times.# Now each reader is either already instantiated with a common dt, or# left at its default dt. In any case, we sum over individual times.trajindex,subframe=self._get_local_frame(self.frame)returnself.total_times[:trajindex].sum()+subframe*self.dts[trajindex]

[docs]def_apply(self,method,**kwargs):"""Execute `method` with `kwargs` for all readers."""return[reader.__getattribute__(method)(**kwargs)forreaderinself.readers]

[docs]def_get(self,attr):"""Get value of `attr` for all readers."""return[reader.__getattribute__(attr)forreaderinself.readers]

[docs]def_get_same(self,attr):"""Verify that `attr` has the same value for all readers and return value. Parameters ---------- attr : str attribute name Returns ------- value : int or float or str or object common value of the attribute Raises ------ ValueError if not all readers have the same value """values=np.array(self._get(attr))value=values[0]ifnotnp.allclose(values,value):bad_traj=np.array(self.filenames)[values!=value]raiseValueError("The following trajectories do not have the correct {0} "" ({1}):\n{2}".format(attr,value,bad_traj))returnvalue

def__activate_reader(self,i):"""Make reader `i` the active reader."""# private method, not to be used by user to avoid a total messifnot(0<=i<len(self.readers)):raiseIndexError("Reader index must be 0 <= i < {0:d}".format(len(self.readers)))self.__active_reader_index=i@propertydefactive_reader(self):"""Reader instance from which frames are currently being read."""returnself.readers[self.__active_reader_index]

[docs]def_read_frame(self,frame):"""Position trajectory at frame index `frame` and return :class:`~MDAnalysis.coordinates.base.Timestep`. The frame is translated to the corresponding reader and local frame index and the :class:`Timestep` instance in :attr:`ChainReader.ts` is updated. Notes ----- `frame` is 0-based, i.e. the first frame in the trajectory is accessed with ``frame = 0``. See Also -------- :meth:`~ChainReader._get_local_frame` """i,f=self._get_local_frame(frame)# seek to (1) reader i and (2) frame f in trajectory iself.__activate_reader(i)self.active_reader[f]# rely on reader to implement __getitem__()# update Timestepself.ts=self.active_reader.tsself.ts.frame=frame# continuous frames, 0-basedreturnself.ts

[docs]def_chained_iterator(self):"""Iterator that presents itself as a chained trajectory."""self._rewind()# must rewind all readersforiinrange(self.n_frames):j,f=self._get_local_frame(i)self.__activate_reader(j)self.ts=self.active_reader[f]self.ts.frame=iyieldself.ts

[docs]defrewind(self):"""Set current frame to the beginning."""self._rewind()self.__chained_trajectories_iter=self._chained_iterator()# set time step for frame 1self.ts=next(self.__chained_trajectories_iter)

def__iter__(self):"""Generator for all frames, starting at frame 1."""self._rewind()# start from first frameself.__chained_trajectories_iter=self._chained_iterator()fortsinself.__chained_trajectories_iter:yieldtsdef__repr__(self):iflen(self.filenames)>3:fnames="{fname} and {nfanmes} more".format(fname=os.path.basename(self.filenames[0]),nfanmes=len(self.filenames)-1)else:fnames=", ".join([os.path.basename(fn)forfninself.filenames])return("<{clsname} containing {fname} with {nframes} frames of {natoms} atoms>""".format(clsname=self.__class__.__name__,fname=fnames,nframes=self.n_frames,natoms=self.n_atoms))

[docs]defadd_transformations(self,*transformations):""" Add all transformations to be applied to the trajectory. This function take as list of transformations as an argument. These transformations are functions that will be called by the Reader and given a :class:`Timestep` object as argument, which will be transformed and returned to the Reader. The transformations can be part of the :mod:`~MDAnalysis.transformations` module, or created by the user, and are stored as a list `transformations`. This list can only be modified once, and further calls of this function will raise an exception. .. code-block:: python u = MDAnalysis.Universe(topology, coordinates) workflow = [some_transform, another_transform, this_transform] u.trajectory.add_transformations(*workflow) Parameters ---------- transform_list : list list of all the transformations that will be applied to the coordinates See Also -------- :mod:`MDAnalysis.transformations` """#Overrides :meth:`~MDAnalysis.coordinates.base.ProtoReader.add_transformations`#to avoid unintended behaviour where the coordinates of each frame are transformed#multiple times when iterating over the trajectory.#In this method, the trajectory is modified all at once and once only.super(ChainReader,self).add_transformations(*transformations)forrinself.readers:r.add_transformations(*transformations)

def_apply_transformations(self,ts):""" Applies the transformations to the timestep."""# Overrides :meth:`~MDAnalysis.coordinates.base.ProtoReader.add_transformations`# to avoid applying the same transformations multiple times on each framereturnts