# Source code for MDAnalysis.coordinates.memory

    # -*- Mode: python; tab-width: 4; indent-tabs-mode:nil; coding:utf-8 -*-
    # vim: tabstop=4 expandtab shiftwidth=4 softtabstop=4
    #
    # MDAnalysis --- https://www.mdanalysis.org
    # Copyright (c) 2006-2017 The MDAnalysis Development Team and contributors
    # (see the file AUTHORS for the full list of names)
    #
    # Released under the Lesser GNU Public Licence, v2.1 or any higher version
    #
    # Please cite your use of MDAnalysis in published work:
    #
    # R. J. Gowers, M. Linke, J. Barnoud, T. J. E. Reddy, M. N. Melo, S. L. Seyler,
    # D. L. Dotson, J. Domanski, S. Buchoux, I. M. Kenney, and O. Beckstein.
    # MDAnalysis: A Python package for the rapid analysis of molecular dynamics
    # simulations. In S. Benthall and S. Rostrup editors, Proceedings of the 15th
    # Python in Science Conference, pages 102-109, Austin, TX, 2016. SciPy.
    # doi: 10.25080/majora-629e541a-00e
    #
    # N. Michaud-Agrawal, E. J. Denning, T. B. Woolf, and O. Beckstein.
    # MDAnalysis: A Toolkit for the Analysis of Molecular Dynamics Simulations.
    # J. Comput. Chem. 32 (2011), 2319--2327, doi:10.1002/jcc.21787
    #
    """\
    =========================================================================
    Reading trajectories from memory --- :mod:`MDAnalysis.coordinates.memory`
    =========================================================================

    :Author: Wouter Boomsma
    :Year: 2016
    :Copyright: Lesser GNU Public License v2.1+
    :Maintainer: Wouter Boomsma <wb@di.ku.dk>, wouterboomsma on github

    .. versionadded:: 0.16.0

    The module contains a trajectory reader that operates on an array in
    memory, rather than reading from file. This makes it possible to
    operate on raw coordinates using existing MDAnalysis tools. In
    addition, it allows the user to make changes to the coordinates in a
    trajectory (e.g. through
    :attr:`MDAnalysis.core.groups.AtomGroup.positions`) without having
    to write the entire state to file.

    How to use the :class:`MemoryReader`
    ====================================

    The :class:`MemoryReader` can be used to either directly generate a
    trajectory as a numpy array or by transferring an existing trajectory
    to memory.

    In-memory representation of arbitrary trajectories
    ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~

    If sufficient memory is available to hold a whole trajectory in memory
    then analysis can be sped up substantially by transferring the
    trajectory to memory.

    The most straightforward use of the :class:`MemoryReader` is to simply
    use the ``in_memory=True`` flag for the
    :class:`~MDAnalysis.core.universe.Universe` class, which
    automatically transfers a trajectory to memory::

     import MDAnalysis as mda
     from MDAnalysisTests.datafiles import TPR, XTC

     universe = mda.Universe(TPR, XTC, in_memory=True)

    Of course, sufficient memory has to be available to hold the whole
    trajectory.

    Switching a trajectory to an in-memory representation
    ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~

    The decision to transfer the trajectory to memory can be made at any
    time with the
    :meth:`~MDAnalysis.core.universe.Universe.transfer_to_memory` method
    of a :class:`~MDAnalysis.core.universe.Universe`::

        universe = mda.Universe(TPR, XTC)
        universe.transfer_to_memory()

    This operation may take a while (with `verbose=True` a progress bar is
    displayed) but then subsequent operations on the trajectory directly
    operate on the in-memory array and will be very fast.

    Constructing a Reader from a numpy array
    ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~

    The :class:`MemoryReader` provides great flexibility because it
    becomes possible to create a :class:`~MDAnalysis.core.universe.Universe` directly
    from a numpy array.

    A simple example consists of a new universe created from the array
    extracted from a DCD
    :meth:`~MDAnalysis.coordinates.DCD.DCDReader.timeseries`::

        import MDAnalysis as mda
        from MDAnalysisTests.datafiles import DCD, PSF
        from MDAnalysis.coordinates.memory import MemoryReader

        universe = mda.Universe(PSF, DCD)

        coordinates = universe.trajectory.timeseries(universe.atoms)
        universe2 = mda.Universe(PSF, coordinates, format=MemoryReader, order='afc')

    .. _create-in-memory-trajectory-with-AnalysisFromFunction:

    .. rubric:: Creating an in-memory trajectory with
                :func:`~MDAnalysis.analysis.base.AnalysisFromFunction`

    The :meth:`~MDAnalysis.coordinates.DCD.DCDReader.timeseries` is
    currently only implemented for the
    :class:`~MDAnalysis.coordinates.DCD.DCDReader`. However, the
    :func:`MDAnalysis.analysis.base.AnalysisFromFunction` can provide the
    same functionality for any supported trajectory format::

      import MDAnalysis as mda
      from MDAnalysis.tests.datafiles import PDB, XTC

      from MDAnalysis.coordinates.memory import MemoryReader
      from MDAnalysis.analysis.base import AnalysisFromFunction

      u = mda.Universe(PDB, XTC)

      coordinates = AnalysisFromFunction(lambda ag: ag.positions.copy(),
                                         u.atoms).run().results['timeseries']
      u2 = mda.Universe(PDB, coordinates, format=MemoryReader)

    .. _creating-in-memory-trajectory-label:

    Creating an in-memory trajectory of a sub-system
    ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~

    Creating a trajectory for just a selection of an existing trajectory
    requires the transfer of the appropriate coordinates as well as
    creation of a topology of the sub-system. For the latter one can use
    the :func:`~MDAnalysis.core.universe.Merge` function, for the former
    the :meth:`~MDAnalysis.core.universe.Universe.load_new` method of a
    :class:`~MDAnalysis.core.universe.Universe` together with the
    :class:`MemoryReader`. In the following, an in-memory trajectory of
    only the protein is created::

      import MDAnalysis as mda
      from MDAnalysis.tests.datafiles import PDB, XTC

      from MDAnalysis.coordinates.memory import MemoryReader
      from MDAnalysis.analysis.base import AnalysisFromFunction

      u = mda.Universe(PDB, XTC)
      protein = u.select_atoms("protein")

      coordinates = AnalysisFromFunction(lambda ag: ag.positions.copy(),
                                         protein).run().results['timeseries']
      u2 = mda.Merge(protein)            # create the protein-only Universe
      u2.load_new(coordinates, format=MemoryReader)

    The protein coordinates are extracted into ``coordinates`` and then
    the in-memory trajectory is loaded from these coordinates. In
    principle, this could have all be done in one line::

      u2 = mda.Merge(protein).load_new(
               AnalysisFromFunction(lambda ag: ag.positions.copy(),
                                    protein).run().results['timeseries'],
               format=MemoryReader)

    The new :class:`~MDAnalysis.core.universe.Universe` ``u2`` can be used
    to, for instance, write out a new trajectory or perform fast analysis
    on the sub-system.

    Classes
    =======

    .. autoclass:: MemoryReader
       :members:
       :inherited-members:

    """
    import logging
    import errno
    import numpy as np
    import warnings
    import copy

    from . import base
    from .timestep import Timestep

    # These methods all pass in an existing *view* onto a larger array
    def _replace_positions_array(ts, new):
        """Replace the array of positions

        Replaces the array of positions by another array.

        Note
        ----
        The behavior of :meth:`_replace_positions_array` is different from the
        behavior of the :attr:`position` property that replaces the **content**
        of the array. The :meth:`_replace_positions_array` method should only be
        used to set the positions to a different frame in
        :meth:`MemoryReader._read_next_timestep`; there, the memory reader sets
        the positions to a view of the correct frame.  Modifying the positions
        for a given frame should be done with the :attr:`positions` attribute
        that does not break the link between the array of positions in the time
        step and the :attr:`MemoryReader.coordinate_array`.

        .. versionadded:: 0.19.0
        .. versionchanged:: 2.0.0
           This function, and the _repalace helper functions for velocities,
           forces, and dimensions, have been moved out of the now removed
           custom timestep object for :class:`MemoryReader`.
        """
        ts.has_positions = True
        ts._pos = new

    def _replace_velocities_array(ts, new):
        ts.has_velocities = True
        ts._velocities = new

    def _replace_forces_array(ts, new):
        ts.has_forces = True
        ts._forces = new

    def _replace_dimensions(ts, new):
        ts._unitcell = new

    class MemoryReader(base.ProtoReader):
        """
        MemoryReader works with trajectories represented as numpy arrays.

        A trajectory reader interface to a numpy array of the coordinates.
        For compatibility with the timeseries interface, support is provided for
        specifying the order of columns through the `order` keyword.

        .. versionadded:: 0.16.0
        .. versionchanged:: 1.0.0
           Support for the deprecated `format` keyword for
           :meth:`MemoryReader.timeseries` has now been removed.
        """

        format = "MEMORY"

        def __init__(
            self,
            coordinate_array,
            order="fac",
            dimensions=None,
            dt=1,
            filename=None,
            velocities=None,
            forces=None,
            **kwargs,
        ):
            """
            Parameters
            ----------
            coordinate_array : numpy.ndarray
                The underlying array of coordinates. The MemoryReader now
                necessarily requires a np.ndarray
            order : {"afc", "acf", "caf", "fac", "fca", "cfa"} (optional)
                the order/shape of the return data array, corresponding
                to (a)tom, (f)rame, (c)oordinates all six combinations
                of 'a', 'f', 'c' are allowed ie "fac" - return array
                where the shape is (frame, number of atoms,
                coordinates).
            dimensions: [A, B, C, alpha, beta, gamma] (optional)
                unitcell dimensions (*A*, *B*, *C*, *alpha*, *beta*, *gamma*)
                lengths *A*, *B*, *C* are in the MDAnalysis length unit (Å), and
                angles are in degrees. An array of dimensions can be given,
                which must then be shape (nframes, 6)
            dt: float (optional)
                The time difference between frames (ps).  If :attr:`time`
                is set, then `dt` will be ignored.
            filename: string (optional)
                The name of the file from which this instance is created. Set to ``None``
                when created from an array
            velocities : numpy.ndarray (optional)
                Atom velocities.  Must match shape of coordinate_array.  Will share order
                with coordinates.
            forces : numpy.ndarray (optional)
                Atom forces.  Must match shape of coordinate_array  Will share order
                with coordinates

            Raises
            ------
            TypeError if the coordinate array passed is not a np.ndarray

            Note
            ----
            At the moment, only a fixed `dimension` is supported, i.e., the same
            unit cell for all frames in `coordinate_array`. See issue `#1041`_.

            .. _`#1041`: https://github.com/MDAnalysis/mdanalysis/issues/1041

            .. versionchanged:: 0.19.0
                The input to the MemoryReader now must be a np.ndarray
                Added optional velocities and forces
            .. versionchanged:: 2.2.0
                Input kwargs are now stored under the :attr:`_kwargs` attribute,
                and are passed on class creation in :meth:`copy`.
            """

            super(MemoryReader, self).__init__()
            self.filename = filename
            self.stored_order = order
            self._kwargs = kwargs

            # See Issue #1685. The block below checks if the coordinate array
            # passed is of shape (N, 3) and if it is, the coordiante array is
            # reshaped to (1, N, 3)
            try:
                if coordinate_array.ndim == 2 and coordinate_array.shape[1] == 3:
                    coordinate_array = coordinate_array[np.newaxis, :, :]
            except AttributeError:
                errmsg = (
                    "The input has to be a numpy.ndarray that corresponds "
                    "to the layout specified by the 'order' keyword."
                )
                raise TypeError(errmsg) from None

            self.set_array(coordinate_array, order)
            self.n_frames = self.coordinate_array.shape[
                self.stored_order.find("f")
            ]
            self.n_atoms = self.coordinate_array.shape[self.stored_order.find("a")]

            if velocities is not None:
                try:
                    velocities = np.asarray(velocities, dtype=np.float32)
                except ValueError:
                    errmsg = (
                        f"'velocities' must be array-like got "
                        f"{type(velocities)}"
                    )
                    raise TypeError(errmsg) from None
                # if single frame, make into array of 1 frame
                if velocities.ndim == 2:
                    velocities = velocities[np.newaxis, :, :]
                if not velocities.shape == self.coordinate_array.shape:
                    raise ValueError(
                        "Velocities has wrong shape {} "
                        "to match coordinates {}"
                        "".format(velocities.shape, self.coordinate_array.shape)
                    )
                self.velocity_array = velocities.astype(np.float32, copy=False)
            else:
                self.velocity_array = None

            if forces is not None:
                try:
                    forces = np.asarray(forces, dtype=np.float32)
                except ValueError:
                    errmsg = f"'forces' must be array like got {type(forces)}"
                    raise TypeError(errmsg) from None
                if forces.ndim == 2:
                    forces = forces[np.newaxis, :, :]
                if not forces.shape == self.coordinate_array.shape:
                    raise ValueError(
                        "Forces has wrong shape {} "
                        "to match coordinates {}"
                        "".format(forces.shape, self.coordinate_array.shape)
                    )
                self.force_array = forces.astype(np.float32, copy=False)
            else:
                self.force_array = None

            provided_n_atoms = kwargs.pop("n_atoms", None)
            if provided_n_atoms is not None and provided_n_atoms != self.n_atoms:
                raise ValueError(
                    "The provided value for n_atoms ({}) "
                    "does not match the shape of the coordinate "
                    "array ({})".format(provided_n_atoms, self.n_atoms)
                )

            self.ts = self._Timestep(self.n_atoms, **kwargs)
            self.ts.dt = dt

            if dimensions is None:
                self.dimensions_array = np.zeros(
                    (self.n_frames, 6), dtype=np.float32
                )
            else:
                try:
                    dimensions = np.asarray(dimensions, dtype=np.float32)
                except ValueError:
                    errmsg = (
                        f"'dimensions' must be array-like got "
                        f"{type(dimensions)}"
                    )
                    raise TypeError(errmsg) from None
                if dimensions.shape == (6,):
                    # single box, tile this to trajectory length
                    # allows modifying the box of some frames
                    dimensions = np.tile(dimensions, (self.n_frames, 1))
                elif dimensions.shape != (self.n_frames, 6):
                    raise ValueError(
                        "Provided dimensions array has shape {}. "
                        "This must be a array of shape (6,) or "
                        "(n_frames, 6)".format(dimensions.shape)
                    )
                self.dimensions_array = dimensions

            self.ts.frame = -1
            self.ts.time = -1
            self._read_next_timestep()

        @staticmethod
        def _format_hint(thing):
            """For internal use: Check if MemoryReader can operate on *thing*

            .. versionadded:: 1.0.0
            """
            return isinstance(thing, np.ndarray)

        @staticmethod
        def parse_n_atoms(filename, order="fac", **kwargs):
            """Deduce number of atoms in a given array of coordinates

            Parameters
            ----------
            filename : numpy.ndarray
              data which will be used later in MemoryReader
            order : {"afc", "acf", "caf", "fac", "fca", "cfa"} (optional)
                the order/shape of the return data array, corresponding
                to (a)tom, (f)rame, (c)oordinates all six combinations
                of 'a', 'f', 'c' are allowed ie "fac" - return array
                where the shape is (frame, number of atoms,
                coordinates).

            Returns
            -------
            n_atoms : int
              number of atoms in system
            """
            # assume filename is a numpy array
            return filename.shape[order.find("a")]

        def copy(self):
            """Return a copy of this Memory Reader"""
            vels = (
                self.velocity_array.copy()
                if self.velocity_array is not None
                else None
            )
            fors = (
                self.force_array.copy() if self.force_array is not None else None
            )
            dims = self.dimensions_array.copy()

            new = self.__class__(
                self.coordinate_array.copy(),
                order=self.stored_order,
                dimensions=dims,
                velocities=vels,
                forces=fors,
                dt=self.ts.dt,
                filename=self.filename,
                **self._kwargs,
            )
            new[self.ts.frame]

            for auxname, auxread in self._auxs.items():
                new.add_auxiliary(auxname, auxread.copy())
            # since transformations are already applied to the whole trajectory
            # simply copy the property
            new.transformations = self.transformations

            return new

        def set_array(self, coordinate_array, order="fac"):
            """
            Set underlying array in desired column order.

            Parameters
            ----------
            coordinate_array : :class:`~numpy.ndarray` object
                The underlying array of coordinates
            order : {"afc", "acf", "caf", "fac", "fca", "cfa"} (optional)
                the order/shape of the return data array, corresponding
                to (a)tom, (f)rame, (c)oordinates all six combinations
                of 'a', 'f', 'c' are allowed ie "fac" - return array
                where the shape is (frame, number of atoms,
                coordinates).
            """
            # Only make copy if not already in float32 format
            self.coordinate_array = coordinate_array.astype("float32", copy=False)
            self.stored_format = order

        def get_array(self):
            """
            Return underlying array.
            """
            return self.coordinate_array

        def _reopen(self):
            """Reset iteration to first frame"""
            self.ts.frame = -1
            self.ts.time = -1

        def timeseries(
            self, asel=None, atomgroup=None, start=0, stop=-1, step=1, order="afc"
        ):
            """Return a subset of coordinate data for an AtomGroup in desired
            column order. If no selection is given, it will return a view of the
            underlying array, while a copy is returned otherwise.

            Parameters
            ---------
            asel : AtomGroup (optional)
                Atom selection. Defaults to ``None``, in which case the full set of
                coordinate data is returned. Note that in this case, a view
                of the underlying numpy array is returned, while a copy of the
                data is returned whenever `asel` is different from ``None``.

                .. deprecated:: 2.7.0
                   asel argument will be renamed to atomgroup in 3.0.0

            atomgroup: AtomGroup (optional)
                Same as `asel`, will replace `asel` in 3.0.0
            start : int (optional)
                the start trajectory frame
            stop : int (optional)
                the end trajectory frame

                .. deprecated:: 2.4.0
                   Note that `stop` is currently *inclusive* but will be
                   changed in favour of being *exclusive* in version 3.0.

            step : int (optional)
                the number of trajectory frames to skip
            order : {"afc", "acf", "caf", "fac", "fca", "cfa"} (optional)
                the order/shape of the return data array, corresponding
                to (a)tom, (f)rame, (c)oordinates all six combinations
                of 'a', 'f', 'c' are allowed ie "fac" - return array
                where the shape is (frame, number of atoms,
                coordinates).

            .. versionchanged:: 1.0.0
               Deprecated `format` keyword has been removed. Use `order` instead.
            .. versionchanged:: 2.4.0
                ValueError now raised instead of NoDataError for empty input
                AtomGroup
            """
            if asel is not None:
                warnings.warn(
                    "asel argument to timeseries will be renamed to"
                    "'atomgroup' in 3.0, see #3911",
                    category=DeprecationWarning,
                )
                if atomgroup:
                    raise ValueError(
                        "Cannot provide both asel and atomgroup kwargs"
                    )
                atomgroup = asel

            if stop != -1:
                warnings.warn(
                    "MemoryReader.timeseries inclusive `stop` "
                    "indexing will be removed in 3.0 in favour of exclusive "
                    "indexing",
                    category=DeprecationWarning,
                )

            array = self.get_array()
            if order == self.stored_order:
                pass
            elif order[0] == self.stored_order[0]:
                array = np.swapaxes(array, 1, 2)
            elif order[1] == self.stored_order[1]:
                array = np.swapaxes(array, 0, 2)
            elif order[2] == self.stored_order[2]:
                array = np.swapaxes(array, 0, 1)
            elif order[0] == self.stored_order[1]:
                array = np.swapaxes(array, 1, 0)
                array = np.swapaxes(array, 1, 2)
            elif order[0] == self.stored_order[2]:
                array = np.swapaxes(array, 2, 0)
                array = np.swapaxes(array, 1, 2)

            a_index = order.find("a")
            f_index = order.find("f")
            stop_index = stop + 1
            if stop_index == 0:
                stop_index = None
            basic_slice = (
                [slice(None)] * f_index
                + [slice(start, stop_index, step)]
                + [slice(None)] * (2 - f_index)
            )

            # Return a view if either:
            #   1) atomgroup is None
            #   2) atomgroup corresponds to the selection of all atoms.
            array = array[tuple(basic_slice)]
            if atomgroup is None or atomgroup is atomgroup.universe.atoms:
                return array
            else:
                if len(atomgroup) == 0:
                    raise ValueError(
                        "Timeseries requires at least one atom " "to analyze"
                    )
                # If selection is specified, return a copy
                return array.take(asel.indices, a_index)

        def _read_next_timestep(self, ts=None):
            """copy next frame into timestep"""

            if self.ts.frame >= self.n_frames - 1:
                raise IOError(errno.EIO, "trying to go over trajectory limit")
            if ts is None:
                ts = self.ts
            ts.frame += 1
            f_index = self.stored_order.find("f")
            basic_slice = (
                [slice(None)] * (f_index)
                + [self.ts.frame]
                + [slice(None)] * (2 - f_index)
            )
            _replace_positions_array(ts, self.coordinate_array[tuple(basic_slice)])
            _replace_dimensions(ts, self.dimensions_array[self.ts.frame])
            if self.velocity_array is not None:
                _replace_velocities_array(
                    ts, self.velocity_array[tuple(basic_slice)]
                )
            if self.force_array is not None:
                _replace_forces_array(ts, self.force_array[tuple(basic_slice)])

            ts.time = self.ts.frame * self.dt
            return ts

        def _read_frame(self, i):
            """read frame i"""
            # Frame number is incremented to zero by _read_next_timestep()
            self.ts.frame = i - 1
            return self._read_next_timestep()

        def __repr__(self):
            """String representation"""
            return "<{cls} with {nframes} frames of {natoms} atoms>" "".format(
                cls=self.__class__.__name__,
                nframes=self.n_frames,
                natoms=self.n_atoms,
            )

        def add_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(MemoryReader, self).add_transformations(*transformations)
            for i, ts in enumerate(self):
                for transform in self.transformations:
                    ts = transform(ts)

        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 frame

            return ts