# Source code for MDAnalysis.auxiliary.EDR

    # -*- 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
    #

    """
    EDR auxiliary reader --- :mod:`MDAnalysis.auxiliary.EDR`
    ========================================================
    .. versionadded:: 2.4.0

    Background
    ----------

    `EDR files`_ are binary files following the XDR_ protocol. They are written by
    GROMACS during simulations and contain time-series non-trajectory data on the
    system, such as energies, temperature, or pressure.

    pyedr_ is a Python package that reads EDR binary files and returns them
    human-readable form as a dictionary of NumPy arrays. It is used by the EDR
    auxiliary reader to parse EDR files. As such, a dictionary with string keys and
    numpy array values is loaded into the :class:`EDRReader`. It is basically a
    Python-based version of the C++ code in GROMACS_.

    The classes in this module are based on the `pyedr`_ package. Pyedr is an
    optional dependency and must be installed to use this Reader. Use of the reader
    without pyedr installed will raise an `ImportError`. The variable `HAS_PYEDR`
    indicates whether this module has pyedr availble.

    The EDR auxiliary reader takes the output from pyedr and makes this data
    available within MDAnalysis. The usual workflow starts with creating an
    EDRReader and passing it the file to read as such::

        import MDAnalysis as mda
        aux = mda.auxiliary.EDR.EDRReader(some_edr_file)

    The newly created `aux` object contains all the data found in the EDR file. It
    is stored in the :attr:`.data_dict` dictionary, which maps the names GROMACS
    gave each data entry to a NumPy array that holds the relevant data. These
    GROMACS-given names are stored in and available through the :attr:`.terms`
    attribute. In addition to the numeric data, the new `EDRReader` also stores the
    units of each entry in the :attr:`.data_dict` dictionary in its
    :attr:`.unit_dict` dictionary.

    .. warning::
        Units are converted to `MDAnalysis base units`_ automatically unless
        otherwise specified. However, not all unit types have a defined base unit
        in MDAnalysis. (cf. :data:`MDAnalysis.units.MDANALYSIS_BASE_UNITS`).
        Pressure units, for example, are not currently defined, and
        will not be converted. This might cause inconsistencies between units!
        Conversion can be switched off by passing `convert_units=False` when the
        EDRReader is created::

            aux = mda.auxiliary.EDR.EDRReader(some_edr_file, convert_units=False)

    Standalone Usage of the EDRReader
    ---------------------------------

    The :class:`.EDRReader` can be used to access the data stored in EDR files on
    its own, without association of data to trajectory frames.
    This is useful, for example, for plotting. The data for a single term, a list
    of terms, or for all terms can be returned in dictionary form. "Time" is always
    returned in this dictionary to make plotting easier::

        temp = aux.get_data("Temperature")
        plt.plot(temp["Time"], temp["Temperature"])

        some_terms = aux.get_data(["Potential", "Kinetic En.", "Box-X"])
        plt.plot(some_terms["Time"], some_terms["Potential"])

        all_terms = aux.get_data()
        plt.plot(all_terms["Time"], all_terms["Pressure"])

    Adding EDR data to trajectories
    -------------------------------

    Like other AuxReaders, the :class:`.EDRReader` can attach its data to a
    trajectory by associating it to the appropriate time steps.
    In general, to add EDR data to a trajectory, one needs to provide two
    arguments.

    .. note::
        The following will change with the release of MDAnalysis 3.0. From then on,
        the order of these two arguments will be reversed.

    The first argument is the `aux_spec` dictionary. With it, users specify which
    entries from the EDR file they want to add, and they give it a more convenient
    name to be used in MDAnalysis (because GROMACS creates names like
    "#Surf*SurfTen" or "'Constr. rmsd'" which may be inconvenient to use.)
    This dictionary might look like this::

        aux_spec = {"epot": "Potential",
                    "surf_tension": "#Surf*SurfTen"}

    When provided as shown below, this would direct the :class:`.EDRReader` to take
    the data it finds under the "Potential" key in its :attr:`.data_dict`
    dictionary and attach it to the trajectory time steps under
    `u.trajectory.ts.aux.epot` (and the same for the surface tension).

    The second argument needed is the source of the EDR data itself. Here, either
    the path to the EDR file or a previously created :class:`.EDRReader` object
    can be provided.

    Examples::

        import MDAnalysis as mda
        from MDAnalysisTests.datafiles import AUX_EDR, AUX_EDR_TPR, AUX_EDR_XTC
        import matplotlib.pyplot as plt

    A :class:`Universe` and an :class:`.EDRReader` object are created and the data
    available in the EDR file is printed::

        In [1]: u = mda.Universe(AUX_EDR_TPR, AUX_EDR_XTC)
        In [2]: aux = mda.auxiliary.EDR.EDRReader(AUX_EDR)
        In [3]: aux.terms
        Out[3]: ['Time', 'Bond', 'Angle', ...]

    Data is associated with the trajectory, using an `aux_spec` dictionary to
    specify which data to add under which name. Any number of terms can be added
    using this method. The data is then accessible in the `ts.aux` namespace via
    both attribute and dictionary syntax::

        In [4]: u.trajectory.add_auxiliary({"epot": "Potential",
                                            "angle": "Angle"}, aux)
        In [5]: u.trajectory.ts.aux.epot
        Out[5]: -525164.0625
        In [6]: u.trajectory.ts.aux.Angle
        Out[6]: 3764.52734375
        In [7]: u.trajectory.ts.aux["epot"]
        Out[7]: -525164.0625

    .. note::
        Some GROMACS-provided :attr:`terms` have spaces. Unless an attribute name
        without a space is provided, these terms will not be accessible via the
        attribute syntax. Only the dictionary syntax will work in that case.

    Further, it is possible to add all data from the EDR file to the trajectory. To
    do this, the `aux_spec` dictionary is omitted, and the data source (the second
    argument as explained above) is provided explicitly as `auxdata`. When adding
    data this way, the terms in :attr:`.terms` become the names used in `ts.aux`::

        In [7]: u.trajectory.add_auxiliary(auxdata=aux)
        In [8]: u.trajectory.ts.aux["#Surf*SurfTen"]
        Out[8]: -1857.519287109375

    .. _EDR files: https://manual.gromacs.org/current/reference-manual/file-formats.html#edr
    .. _XDR: https://datatracker.ietf.org/doc/html/rfc1014
    .. _pyedr: https://github.com/mdanalysis/panedr
    .. _GROMACS: https://github.com/gromacs/gromacs/blob/main/src/gromacs/fileio/enxio.cpp
    .. _MDAnalysis base units: https://docs.mdanalysis.org/2.3.0/documentation_pages/units.html

    Classes
    -------

    .. autoclass:: EDRReader
       :members:

    The actual data for each step is stored by instances of EDRStep.

    .. autoclass:: EDRStep
       :members:

    """
    from pathlib import Path
    import warnings
    from typing import Optional, Union, Dict, List

    import numpy as np

    from . import base
    from .. import units

    try:
        import pyedr
    except ImportError:
        # Indicates whether pyedr is found
        HAS_PYEDR = False
    else:
        # Indicates whether pyedr is found
        HAS_PYEDR = True

    class EDRStep(base.AuxStep):
        """:class:`AuxStep` class for the .edr file format.

        Extends the base AuxStep class to allow selection of time and
        data-of-interest fields (by dictionary key) from the full set of data read
        each step.

        Parameters
        ----------
        time_selector : str, optional
            Name of the dictionary key that links to the time values (assumed to
            be in ps). Default value is "Time"
        data_selector : str | list of str | None, optional
            List of dictionary keys linking to data of interest in the EDR file to
            be stored in ``data``. Default value is ``None``.
        **kwargs
            Other AuxStep options.

        See Also
        --------
        :class:`MDAnalysis.auxiliary.base.AuxStep`
        """

        def __init__(
            self,
            time_selector: str = "Time",
            data_selector: Optional[str] = None,
            **kwargs,
        ):
            super(EDRStep, self).__init__(
                time_selector=time_selector, data_selector=data_selector, **kwargs
            )

        def _select_time(self, key: str) -> np.float64:
            """'Time' is one of the entries in the dict returned by pyedr.
            The base AuxStep Class uses the time_selector 'Time' to return the
            time value of each step."""
            return self._select_data(key)

        def _select_data(self, key: Union[str, None]) -> np.float64:
            if key is None:
                return
            try:
                return self._data[key]
            except KeyError:
                raise KeyError(
                    f"'{key}' is not a key in the data_dict dictionary."
                    " Check the EDRReader.terms attribute"
                )

    class EDRReader(base.AuxReader):
        """Auxiliary reader to read data from an .edr file.

        `EDR files`_
        are created by GROMACS during a simulation. They are binary files which
        contain time-series energy data and other data related to the simulation.

        Default reader for .edr files. All data from the file will be read and
        stored on initialisation.

        Parameters
        ----------
        filename : str
            Location of the file containing the auxiliary data.
        convert_units : bool, optional
            If True (default), units from the EDR file are automatically converted
            to MDAnalysis base units. If False, units are taken from the file
            as-is. Where no base unit is defined in MDAnalysis, no conversion takes
            place. Unit types in :data:`MDAnalysis.units.MDANALYSIS_BASE_UNITS`
            will be converted automatically by default.
        **kwargs
           Other AuxReader options.

        Attributes
        ----------
        _auxdata : pathlib.PosixPath
            path at which the auxiliary data file is located
        data_dict : dict
            dictionary that contains the auxiliary data, mapping the names GROMACS
            gave the entries in the EDR file to a NumPy array containing this data
        unit_dict : dict
            dictionary that contains the units of the auxiliary data, mapping the
            :attr:`data_selector` of the Reader (i.e. the name of the dataset in
            the EDR file) to its unit.
        _n_steps : int
            Number of steps for which auxdata is available
        terms : list
            Names of the auxiliary data entries available in `data_dict`. These are
            the names GROMACS set in the EDR file.

        See Also
        --------
        :class:`MDAnalysis.auxiliary.base.AuxReader`
        :meth:`MDAnalysis.coordinates.base.ReaderBase.add_auxiliary`

        Note
        ----
        The file is assumed to be of a size such that reading and storing the full
        contents is practical. A warning will be issued when memory usage exceeds
        1 GB. This warning limit can be changed via the ``memory_limit`` kwarg.
        """

        format = "EDR"
        _Auxstep = EDRStep

        def __init__(self, filename: str, convert_units: bool = True, **kwargs):
            if not HAS_PYEDR:
                raise ImportError(
                    "EDRReader: To read EDR files please install " "pyedr."
                )
            self._auxdata = Path(filename).resolve()
            self.data_dict = pyedr.edr_to_dict(filename)
            self.unit_dict = pyedr.get_unit_dictionary(filename)
            self.convert_units = convert_units
            if self.convert_units:
                self._convert_units()
            self._n_steps = len(self.data_dict["Time"])
            # attribute to communicate found energy terms to user
            self.terms = list(self.data_dict.keys())
            super(EDRReader, self).__init__(**kwargs)

        def _convert_units(self):
            """Called during :func:`__init__` to convert the units found in the EDR
            file to MDAnalysis base units"""
            unknown_units = []
            for term, unit in self.unit_dict.items():
                try:
                    unit_type = units.unit_types[unit]
                except KeyError:
                    if unit not in unknown_units:
                        unknown_units.append(unit)
                    continue  # skip conversion if unit not defined yet

                target_unit = units.MDANALYSIS_BASE_UNITS[unit_type]
                data = self.data_dict[term]
                self.data_dict[term] = units.convert(data, unit, target_unit)
                self.unit_dict[term] = units.MDANALYSIS_BASE_UNITS[unit_type]
            if unknown_units:
                warnings.warn(
                    "Could not find unit type for the following "
                    f"units: {unknown_units}"
                )

        def _memory_usage(self):
            size = 0
            for array in self.data_dict.values():
                size += array.nbytes
            return size

        def _read_next_step(self) -> EDRStep:
            """Read next auxiliary step and update ``auxstep``.

            Returns
            -------
            AuxStep object
                Updated with the data for the new step.

            Raises
            ------
            StopIteration
                When end of auxiliary data set is reached.
            """
            auxstep = self.auxstep
            new_step = self.step + 1
            if new_step < self.n_steps:
                auxstep._data = {
                    term: self.data_dict[term][self.step + 1]
                    for term in self.terms
                }
                auxstep.step = new_step
                return auxstep
            else:
                self.rewind()
                if self.n_steps > 1:
                    raise StopIteration

        def _go_to_step(self, i: int) -> EDRStep:
            """Move to and read i-th auxiliary step.

            Parameters
            ----------
            i: int
                Step number (0-indexed) to move to

            Returns
            -------
            :class:`EDRStep`

            Raises
            ------
            ValueError
                If step index not in valid range.
            """
            if i >= self.n_steps or i < 0:
                raise ValueError(
                    "Step index {0} is not valid for auxiliary "
                    "(num. steps {1})".format(i, self.n_steps)
                )
            self.auxstep.step = i - 1
            self.next()
            return self.auxstep

        def read_all_times(self) -> np.ndarray:
            """Get list of time at each step.

            Returns
            -------
            NumPy array of float
                Time at each step.
            """
            return self.data_dict[self.time_selector]

        def get_data(
            self, data_selector: Union[str, List[str], None] = None
        ) -> Dict[str, np.ndarray]:
            """Returns the auxiliary data contained in the :class:`EDRReader`.
            Returns either all data or data specified as `data_selector` in form
            of a str or a list of any of :attr:`EDRReader.terms`. `Time` is
            always returned to allow easy plotting.

            Parameters
            ----------
            data_selector: str, List[str], None
                Keys to be extracted from the auxiliary reader's data dictionary.
                If ``None``, returns all data found in :attr:`.data_dict`.
            Returns
            -------
            data_dict : dict
                Dictionary mapping `data_selector` keys to NumPy arrays of the
                auxiliary data.

            Raises
            ------
            KeyError
                if an invalid data_selector key is passed.
            """
            if data_selector is None:
                return self.data_dict

            def _get_data_term(term, datadict):
                try:
                    return datadict[term]
                except KeyError:
                    raise KeyError(
                        f"data selector {term} is invalid. Check the "
                        "EDRReader's `terms` attribute."
                    )

            data_dict = {"Time": self.data_dict["Time"]}

            if isinstance(data_selector, list):
                for term in data_selector:
                    data_dict[term] = _get_data_term(term, self.data_dict)
            else:
                term = data_selector
                data_dict[term] = _get_data_term(term, self.data_dict)

            return data_dict