biotite 0.41.1__cp310-cp310-macosx_10_16_arm64.whl

biotite/structure/io/trajfile.py

@@ -0,0 +1,710 @@
+# This source code is part of the Biotite package and is distributed
+# under the 3-Clause BSD License. Please see 'LICENSE.rst' for further
+# information.
+
+__name__ = "biotite.structure.io"
+__author__ = "Patrick Kunzmann"
+__all__ = ["TrajectoryFile"]
+
+import itertools
+import abc
+import numpy as np
+from ..atoms import AtomArray, AtomArrayStack, stack, from_template
+from ...file import File
+
+
+class TrajectoryFile(File, metaclass=abc.ABCMeta):
+    """
+    This file class represents a trajectory file interfacing a
+    trajectory file class from `MDtraj`.
+    
+    A trajectory file stores atom coordinates over multiple (time)
+    frames. The file formats are usually binary and involve sometimes
+    heavy compression, so that a large number of frames can be stored
+    in relatively small space.
+    Since all :class:`TrajectoryFile` subclasses interface *MDtraj*
+    trajectory file classes, `MDtraj` must be installed to use any of
+    them.
+
+    Notes
+    -----
+    When extracting data from the file, only a reference to the
+    data arrays stored in this file are created.
+    The same is true, when setting data in the file.
+    Therefore, it is strongly recommended to make a copy of the
+    respective array, if the array is modified.
+    """
+    
+    def __init__(self):
+        super().__init__()
+        self._coord = None
+        self._time = None
+        self._box = None
+        self._model_count = None
+    
+
+    @classmethod
+    def read(cls, file_name, start=None, stop=None, step=None,
+             atom_i=None, chunk_size=None):
+        """
+        Read a trajectory file.
+        
+        A trajectory file can be seen as a file representation of an
+        :class:`AtomArrayStack`.
+        Therefore, `start`, `stop` and `step` represent slice parameters
+        of the index of the first dimension and
+        `atom_i` represents an index array for the second dimension.
+        
+        Parameters
+        ----------
+        file_name : str
+            The path of the file to be read.
+            A file-like-object cannot be used.
+        start : int, optional
+            The frame index, where file parsing is started. If no value
+            is given, parsing starts at the first frame.
+            The index starts at 0.
+        stop : int, optional
+            The exclusive frame index, where file parsing ends.
+            If no value is given, parsing stops after the last frame.
+            The index starts at 0.
+        step : int, optional
+            If this value is set, the function reads only every n-th
+            frame from the file.
+        atom_i : ndarray, dtype=int, optional
+            If this parameter is set, only the atoms at the given
+            indices are read from each frame.
+        chunk_size : int, optional
+            If this parameter is set, the trajectory is read in chunks:
+            Only the number of frames specified by this parameter are
+            read at once.
+            The resulting chunks of frames are automatically
+            concatenated, after all chunks are collected.
+            Use this parameter, if a :class:`MemoryError` is raised
+            when a trajectory file is read.
+            Although lower values can decrease the memory consumption of
+            reading trajectories, they also increase the computation
+            time.
+        
+        Returns
+        -------
+        file_object : TrajectoryFile
+            The parsed trajectory file.
+        """
+        file = cls()
+
+        if chunk_size is not None:
+            if chunk_size < 1:
+                raise ValueError("Chunk size must be greater than 0")
+            # Chunk size must be a multiple of step size to ensure that
+            # the step distance at the chunk border is the same as
+            # within a chunk
+            # -> round chunk size up to a multiple of step size
+            if step is not None and chunk_size % step != 0:
+                chunk_size = ((chunk_size // step) + 1) * step
+
+        traj_type = cls.traj_type()
+        with traj_type(file_name, "r") as f:
+            
+            if start is None:
+                start = 0
+            # Discard atoms before start
+            if start != 0:
+                if chunk_size is None or chunk_size > start:
+                    f.read(n_frames=start, stride=None, atom_indices=atom_i)
+                else:
+                    TrajectoryFile._read_chunk_wise(
+                        f, start, None, atom_i, chunk_size, discard=True
+                    )
+            
+            # The upcoming frames are saved
+            # Calculate the amount of frames to be read
+            if stop is None:
+                n_frames = None
+            else:
+                n_frames = stop-start
+            if step is not None and n_frames is not None:
+                # Divide number of frames by 'step' in order to convert
+                # 'step' into 'stride'
+                # Since the 0th frame is always included,
+                # the number of frames is decremented before division
+                # and incremented afterwards again
+                n_frames = ((n_frames - 1) // step) + 1
+            
+            # Read frames
+            if chunk_size is None:
+                result = f.read(n_frames, stride=step, atom_indices=atom_i)
+            else:
+                result = TrajectoryFile._read_chunk_wise(
+                    f, n_frames, step, atom_i, chunk_size, discard=False
+                )
+        
+        # nm to Angstrom
+        coord, box, time = cls.process_read_values(result)
+        file.set_coord(coord)
+        file.set_box(box)
+        file.set_time(time)
+
+        return file
+    
+
+    @classmethod
+    def read_iter(cls, file_name, start=None, stop=None, step=None,
+                  atom_i=None, stack_size=None):
+        """
+        Create an iterator over each frame of the given trajectory file
+        in the selected range.
+        
+        Parameters
+        ----------
+        file_name : str
+            The path of the file to be read.
+            A file-like-object cannot be used.
+        start : int, optional
+            The frame index, where file parsing is started. If no value
+            is given, parsing starts at the first frame.
+            The index starts at 0.
+        stop : int, optional
+            The exclusive frame index, where file parsing ends.
+            If no value is given, parsing stops at the end of file.
+            The index starts at 0.
+        step : int, optional
+            If this value is set, the function reads only every n-th
+            frame from the file.
+        atom_i : ndarray, dtype=int, optional
+            If this parameter is set, only the atoms at the given
+            indices are read from each frame.
+        stack_size : int, optional
+            If this parameter is set, the given number of frames *m* are
+            read at once.
+            As result an additional dimension is added to the return
+            values.
+            If the number of frames is not a multiple of `stack_size`,
+            the final stack is smaller than `stack_size`.
+        
+        Yields
+        ------
+        coord : ndarray, dtype=float32, shape=(n,3) or shape=(m,n,3)
+            The atom coordinates in the current frame or stack.
+        box : ndarray, dtype=float32, shape=(3,3) or shape=(m,3,3) or None
+            The box vectors of the current frame or stack.
+        time : float or ndarray, dtype=float32, shape=(n,) or None
+            The simulation time of the current frame or stack in *ps*.
+        
+        See also
+        --------
+        read_iter_structure
+        
+        Notes
+        -----
+        The `step` parameter does currently not work for *DCD* files.
+        """
+        traj_type = cls.traj_type()
+        with traj_type(file_name, "r") as f:
+            
+            if start is None:
+                start = 0
+            # Discard atoms before start
+            if start != 0:
+                f.read(n_frames=start, stride=None, atom_indices=atom_i)
+            
+            # The upcoming frames are read
+            # Calculate the amount of frames to be read
+            if stop is None:
+                n_frames = None
+            else:
+                n_frames = stop-start
+            if step is not None and n_frames is not None:
+                # Divide number of frames by 'step' in order to convert
+                # 'step' into 'stride'
+                # Since the 0th frame is always included,
+                # the number of frames is decremented before division
+                # and incremented afterwards again
+                n_frames = ((n_frames - 1) // step) + 1
+            
+
+            # Read frames
+            if stack_size is None:
+                remaining_frames = n_frames
+                while remaining_frames is None or remaining_frames > 0:
+                    result = f.read(1, stride=step, atom_indices=atom_i)
+                    if len(result[0]) == 0:
+                        # Empty array was read
+                        # -> no frames left -> stop iteration
+                        break
+                    coord, box, time = cls.process_read_values(result)
+                    # Only one frame
+                    # -> only one element in first dimension
+                    # -> remove first dimension
+                    coord = coord[0]
+                    box = box[0] if box is not None else None
+                    time = float(time[0]) if time is not None else None
+                    yield coord, box, time
+                    if remaining_frames is not None:
+                        remaining_frames -= 1
+            
+            else:
+                remaining_frames = n_frames
+                while remaining_frames is None or remaining_frames > 0:
+                    n_frames = (
+                        min(remaining_frames, stack_size)
+                        if remaining_frames is not None
+                        else stack_size
+                    )
+                    result = f.read(n_frames, stride=step, atom_indices=atom_i)
+                    if len(result[0]) == 0:
+                        # Empty array was read
+                        # -> no frames left -> stop iteration
+                        break
+                    coord, box, time = cls.process_read_values(result)
+                    yield coord, box, time
+                    if remaining_frames is not None:
+                        remaining_frames -= stack_size
+    
+
+    @classmethod
+    def read_iter_structure(cls, file_name, template, start=None, stop=None,
+                            step=None, atom_i=None, stack_size=None):
+        """
+        Create an iterator over each frame of the given trajectory file
+        in the selected range.
+
+        In contrast to :func:`read_iter()`, this function creates an
+        iterator over the structure as :class:`AtomArray`.
+        Since trajectory files usually only contain atom coordinate
+        information and no topology information, this method requires
+        a template atom array or stack. This template can be acquired
+        for example from a PDB file, which is associated with the
+        trajectory file. 
+        
+        Parameters
+        ----------
+        file_name : str
+            The path of the file to be read.
+            A file-like-object cannot be used.
+        template : AtomArray or AtomArrayStack
+            The template array or stack, where the atom annotation data
+            is taken from.
+        start : int, optional
+            The frame index, where file parsing is started. If no value
+            is given, parsing starts at the first frame.
+            The index starts at 0.
+        stop : int, optional
+            The exclusive frame index, where file parsing ends.
+            If no value is given, parsing stops at the end of file.
+            The index starts at 0.
+        step : int, optional
+            If this value is set, the function reads only every n-th
+            frame from the file.
+        atom_i : ndarray, dtype=int, optional
+            If this parameter is set, only the atoms at the given
+            indices are read from each frame in the file.
+        stack_size : int, optional
+            If this parameter is set, multiple frames are combined into
+            an :class:`AtomArrayStack`.
+            The number of frames in the :class:`AtomArrayStack` is
+            determined by this parameter.
+            If the number of frames is not a multiple of `stack_size`,
+            the final stack is smaller than `stack_size`.
+        
+        Yields
+        ------
+        structure : AtomArray or AtomArrayStack
+            The structure of the current frame as :class:`AtomArray`.
+            If `stack_size` is set, multiple frames are returned as
+            :class:`AtomArrayStack`.
+        
+        See also
+        --------
+        read_iter
+        
+        Notes
+        -----
+        This iterator creates a new copy of the given template for every
+        frame
+        (or stack of frames, if `stack_size` is set).
+        If a higher efficiency is required, please use the
+        :func:`read_iter()` function.
+
+        The `step` parameter does currently not work for *DCD* files.
+        """
+        if isinstance(template, AtomArrayStack):
+            template = template[0]
+        elif not isinstance(template, AtomArray):
+            raise TypeError(
+                f"An 'AtomArray' or 'AtomArrayStack' is expected as template, "
+                f"not '{type(template).__name__}'"
+            )
+        
+        for coord, box, _ in cls.read_iter(
+            file_name, start, stop, step, atom_i, stack_size
+        ):
+            if stack_size is None:
+                frame = template.copy()
+                frame.coord = coord
+                frame.box = box
+                yield frame
+            else:
+                yield from_template(template, coord, box)
+
+    
+    def write(self, file_name):
+        """
+        Write the content into a trajectory file.
+
+        Parameters
+        ----------
+        file_name : str
+            The path of the file to be written to.
+            A file-like-object cannot be used.
+        """
+        traj_type = self.traj_type()
+        param = self.prepare_write_values(self._coord, self._box, self._time)
+        with traj_type(file_name, 'w') as f:
+            f.write(**param)
+    
+
+    @classmethod
+    def write_iter(cls, file_name, coord, box=None, time=None):
+        """
+        Iterate over the given `coord` and write each item into
+        the file specified by `file_name`.
+
+        In contrast to :meth:`write()`, the data is not stored in an
+        intermediate :class:`TrajectoryFile`, but is directly written
+        to the file.
+        Hence, this class method may save a large amount of memory if
+        a large file should be written, if `coord` are provided as
+        generator.
+        
+        Parameters
+        ----------
+        file_name : str
+            The path of the file to be written to.
+            A file-like-object cannot be used.
+        coord : generator or array-like of ndarray, shape=(n,3), dtype=float
+            The atom coordinates for each frame.
+        box : generator or array-like of ndarray, shape=(3,3), dtype=float, optional
+            The three box vectors for each frame.
+        time : generator or array-like of float, optional
+            The simulation time in *ps* for each frame.
+
+        Notes
+        -----
+        The `time` parameter has no effect for *TNG* and *DCD* files.
+        """
+        if box is None:
+            box = itertools.repeat(None)
+        if time is None:
+            time = itertools.repeat(None)
+
+        traj_type = cls.traj_type()
+        with traj_type(file_name, 'w') as f:
+            for c, b, t in zip(coord, box, time):
+                if c.ndim != 2:
+                    raise IndexError(
+                        f"Expected ndarray with 2 dimensions, got {c.ndim}"
+                    )
+                # Add new dimension of length one
+                # for compatibility with 'prepare_write_values()'
+                c = c[np.newaxis, :]
+                if b is not None:
+                    b = np.expand_dims(b, axis=0)
+                if t is not None:
+                    t = np.expand_dims(t, axis=0)
+                param = cls.prepare_write_values(c, b, t)
+                f.write(**param)
+    
+
+    def get_coord(self):
+        """
+        Extract only the atom coordinates from the trajectory file.
+        
+        Returns
+        -------
+        coord : ndarray, dtype=float, shape=(m,n,3)
+            The coordinates stored in the trajectory file.
+        """
+        return self._coord
+    
+
+    def get_time(self):
+        """
+        Get the simlation time in *ps* values for each frame.
+        
+        Returns
+        -------
+        time : ndarray, dtype=float, shape=(m,)
+            A one dimensional array containing the time values for the
+            frames, that were read from the file.
+        """
+        return self._time
+    
+
+    def get_box(self):
+        """
+        Get the box vectors for each frame.
+        
+        Returns
+        -------
+        box : ndarray, dtype=float, shape=(m,3,3)
+            An array containing the box dimensions for the
+            frames, that were read from the file.
+        """
+        return self._box
+    
+
+    def set_coord(self, coord):
+        """
+        Set the atom coordinates in the trajectory file.
+        
+        Parameters
+        ----------
+        coord : ndarray, dtype=float, shape=(m,n,3)
+            The coordinates to be set.
+        """
+        self._check_model_count(coord)
+        self._coord = coord
+    
+
+    def set_time(self, time):
+        """
+        Set the simulation time of each frame in the trajectory file.
+        
+        Parameters
+        ----------
+        time : ndarray, dtype=float, shape=(m,)
+            The simulation time in *ps* to be set.
+        """
+        self._check_model_count(time)
+        self._time = time
+    
+
+    def set_box(self, box):
+        """
+        Set the periodic box vectors of each frame in the trajectory
+        file.
+        
+        Parameters
+        ----------
+        time : ndarray, dtype=float, shape=(m,3,3)
+            The box vectors to be set.
+        """
+        self._check_model_count(box)
+        self._box = box
+    
+
+    def get_structure(self, template):
+        """
+        Convert the trajectory file content into an
+        :class:`AtomArrayStack`.
+        
+        Since trajectory files usually only contain atom coordinate
+        information and no topology information, this method requires
+        a template atom array or stack. This template can be acquired
+        for example from a PDB file, which is associated with the
+        trajectory file. 
+        
+        Parameters
+        ----------
+        template : AtomArray or AtomArrayStack
+            The template array or stack, where the atom annotation data
+            is taken from.
+        
+        Returns
+        -------
+        array_stack : AtomArrayStack
+            A stack containing the annontation arrays from `template`
+            but the coordinates and the simulation boxes from the
+            trajectory file.
+        """
+        return from_template(template, self.get_coord(), self.get_box())
+    
+
+    def set_structure(self, structure, time=None):
+        """
+        Write an atom array (stack) into the trajectory file object.
+        
+        The topology information (chain, residue, etc.) is not saved in
+        the file.
+        
+        Parameters
+        ----------
+        structure : AtomArray or AtomArrayStack
+            The structure to be put into the trajectory file.
+        time : ndarray, dtype=float, shape=(n,), optional
+            The simulation time for each frame in `structure`.
+        """
+        coord = structure.coord
+        box = structure.box
+        if coord.ndim == 2:
+            coord = coord[np.newaxis, :, :]
+        if box is not None and box.ndim == 2:
+            box = box[np.newaxis, :, :]
+        self.set_coord(coord)
+        if box is not None:
+            self.set_box(box)
+        if time is not None:
+            self.set_time(time)
+
+    
+    def copy(self):
+        """
+        This operation is not implemented for trajectory files.
+        
+        Raises
+        ------
+        NotImplementedError
+        """
+        raise NotImplementedError("Copying is not implemented "
+                                  "for trajectory files")
+    
+
+    @classmethod
+    @abc.abstractmethod
+    def traj_type(cls):
+        """
+        The `MDtraj` files class to be used.
+        
+        PROTECTED: Override when inheriting.
+        
+        Returns
+        -------
+        class
+            An `MDtraj` subclass of :class:`TrajectoryFile`.
+        """
+        pass
+    
+
+    @classmethod
+    @abc.abstractmethod
+    def process_read_values(cls, read_values):
+        """
+        Convert the return value of the `read()` method of the
+        respective :class:`mdtraj.TrajectoryFile` into coordinates,
+        simulation box and simulation time.
+        
+        PROTECTED: Override when inheriting.
+        
+        Parameters
+        ----------
+        read_values : tuple
+            The return value of the respective
+            :func:`mdtraj.TrajectoryFile.read()` method.
+        
+        Returns
+        -------
+        coord : ndarray, dtype=float, shape=(m,n,3)
+            The atom coordinates in Å for each frame.
+        box : ndarray, dtype=float, shape=(m,3,3) or None
+            The box vectors in Å for each frame.
+        time : ndarray, dtype=float, shape=(m,) or None
+            The simulation time in ps for each frame.
+        """
+        pass
+    
+
+    @classmethod
+    @abc.abstractmethod
+    def prepare_write_values(cls, coord, box, time):
+        """
+        Convert the `coord`, `box` and `time` attribute into a
+        dictionary that is given as *kwargs* to the respective
+        :func:`mdtraj.TrajectoryFile.write()` method.
+
+        PROTECTED: Override when inheriting.
+
+        Parameters
+        ----------
+        coord : ndarray, dtype=float, shape=(m,n,3)
+            The atom coordinates in Å for each frame.
+        box : ndarray, dtype=float, shape=(m,3,3)
+            The box vectors in Å for each frame.
+        time : ndarray, dtype=float, shape=(m,)
+            The simulation time in ps for each frame.
+        
+        Returns
+        -------
+        parameters : dict
+            This dictionary is given as *kwargs* parameter to the
+            respective :func:`mdtraj.TrajectoryFile.write()` method.
+        """
+        pass
+
+
+    def _check_model_count(self, array):
+        """
+        Check if the amount of models in the given array is equal to
+        the amount of models in the file.
+        If not, raise an exception.
+        If the amount of models in the file is not set yet, set it with
+        the amount of models in the array.
+        """
+        if array is None:
+            return
+        if self._model_count is None:
+            self._model_count = len(array)
+        else:
+            if self._model_count != len(array):
+                raise ValueError(
+                    f"{len(array)} models were given, "
+                    f"but the file contains {self._model_count} models"
+                )
+    
+
+    @staticmethod
+    def _read_chunk_wise(file, n_frames, step, atom_i, chunk_size,
+                         discard=False):
+        """
+        Similar to :func:`read()`, just for chunk-wise reading of the
+        trajectory.
+
+        `n_frames` is already the actual number of frames in the output
+        arrays, i.e. the original number was divided by `step`.
+        """
+        chunks = []
+        remaining_frames = n_frames
+        # If n_frames is None, this is condition is never False
+        # -> break out of loop when read chunk is empty (see below)
+        while remaining_frames != 0:
+            if remaining_frames is not None:
+                n = min(remaining_frames, chunk_size)
+            else:
+                n = chunk_size
+            try:
+                chunk = file.read(n_frames=n, stride=step, atom_indices=atom_i)
+            except ValueError as e:
+                # MDTraj raises exception because no coordinates can be
+                # concatenated
+                # -> all frames have been read
+                # -> stop reading chunks
+                if str(e) != "need at least one array to concatenate":
+                    raise
+                else:
+                    break
+            if len(chunk[0]) == 0:
+                # Coordinates have a length of 0
+                # -> all frames have been read
+                # -> stop reading chunks
+                break
+            if not discard:
+                chunks.append(chunk)
+            if remaining_frames is not None:
+                remaining_frames -= n
+        
+        if not discard:
+            # Assemble the chunks into contiguous arrays
+            # for each value (coord, box, time)
+            result = [None] * len(chunks[0])
+            # Iterate over all values in the result tuple
+            # and concatenate the corresponding value from each chunk,
+            # if the value is not None
+            # The amount of values is determined from the first chunk
+            for i in range(len(chunks[0])):
+                if chunks[0][i] is not None:
+                    result[i] = np.concatenate([chunk[i] for chunk in chunks])
+                else:
+                    result[i] = None
+            return tuple(result)
+        else:
+            return None