Update NCDF Reader to read .ncrst Restart Files

package/AUTHORS

@@ -258,7 +258,7 @@ Chronological list of authors
   - Abdulrahman Elbanna
   - Tulga-Erdene Sodjargal
   - Gareth Elliott
-
+  - Jeremy M.G. Leung
 
 External code
 -------------

package/CHANGELOG

@@ -15,9 +15,8 @@ The rules for this file:
 
 -------------------------------------------------------------------------------
 ??/??/?? IAlibay, orbeckst, BHM-Bob, TRY-ER, Abdulrahman-PROG, pbuslaev,
-         yuxuanzhuang, yuyuan871111, tanishy7777, tulga-rdn, Gareth-elliott,
-         hmacdope
-
+         yuxuanzhuang, yuyuan871111, tanishy7777, jpkrowe, tulga-rdn,
+         Gareth-elliott, hmacdope, jeremyleung521
 
  * 2.10.0
 
@@ -38,6 +37,8 @@ Fixes
    directly passing them. (Issue #3520, PR #5006)
 
 Enhancements
+ * Allow NCDF reader to read Amber Restart Files with extension `.ncrst`
+   (Issue #5030, PR #5031)
  * Improve parsing of topology information from LAMMPS dump files to allow
    reading of mass, charge and element attributes. (Issue #3449, PR #4995)
  * Added timestep support for XDR writers and readers (Issue #4905, PR #4908)

package/MDAnalysis/coordinates/INPCRD.py

@@ -22,29 +22,92 @@
 #
 
 
-"""INPCRD structure files in MDAnalysis --- :mod:`MDAnalysis.coordinates.INPCRD`
+"""AMBER restart files in MDAnalysis --- :mod:`MDAnalysis.coordinates.INPCRD`
 ================================================================================
 
-Read coordinates in Amber_ coordinate/restart file (suffix "inpcrd").
+AMBER_ can write :ref:`ASCII restart<ascii-restart>` ("inpcrd") and
+:ref:`binary restart<netcdf-restart>` ("ncrst") coordinate files. MDAnalysis
+supports reading of both file formats.
 
-.. _Amber: https://ambermd.org/FileFormats.php
+.. rubric:: Units
 
+AMBER restart files are assumed to be in the following units:
 
-Classes
--------
+* length in Angstrom (Å)
+* time in ps
+* velocity (NCRST only) in Å / ps
+* force (NCRST only) in kcal / (mol * Å)
+
+
+.. _ascii-restart:
+
+ASCII INPCRD restart files
+--------------------------
+
+ASCII AMBER_ INPCRD coordinate files (as defined in `AMBER INPCRD FORMAT`_)
+are handled by the :class:`INPReader`.
+
+AMBER ASICC restart files are recognised by the suffix '.inpcrd', '.restrt', or
+'.rst7'
 
 .. autoclass:: INPReader
    :members:
 
+
+.. _netcdf-restart:
+
+Binary NetCDF restart files
+---------------------------
+
+The `AMBER netcdf`_ restart format makes use of NetCDF_ (Network Common Data
+Form) format. Such binary restart files are recognised in MDAnalysis by the
+suffix '.ncrst', '.ncrestrt' or '.ncrst7' and read by the :class:`NCRSTReader`.
+
+Binary restart files can also contain velocity and force information, and can
+record the simulation time step. Whilst the `AMBER netcdf`_ format details
+default unit values of ångström and picoseconds, these can in theory occupy
+any unit type. However, at the moment MDAnalysis only supports the default
+types and will raise a :exc:`NotImplementedError` if anything else is detected.
+
+.. autoclass:: NCRSTReader
+   :members:
+
+
+.. Links
+
+.. _AMBER: http://ambermd.org
+.. _AMBER INPCRD FORMAT: http://ambermd.org/formats.html#restart
+.. _AMBER netcdf: http://ambermd.org/netcdf/nctraj.xhtml
+.. _NetCDF: http://www.unidata.ucar.edu/software/netcdf
+
 """
 
 from . import base
+import warnings
+import logging
+from scipy.io import netcdf_file
+
+from .timestep import Timestep
+from ..lib.util import store_init_arguments
+from .TRJ import NCDFPicklable, NCDFMixin
+
+logger = logging.getLogger("MDAnalysis.coordinates.AMBER")
 
 
 class INPReader(base.SingleFrameReaderBase):
-    """Reader for Amber restart files."""
+    """Reader for Amber restart files.
+
+    .. rubric:: Limitations
+
+    * Box information is not read (or checked for).
+    * Velocities are currently *not supported*.
+
+    .. versionchanged: 0.20.0
+       Now automatically detects files with .rst7 extension.
+
+    """
 
-    format = ["INPCRD", "RESTRT"]
+    format = ["INPCRD", "RESTRT", "RST7"]
     units = {"length": "Angstrom"}
 
     def _read_first_frame(self):
@@ -89,3 +152,116 @@
             f.readline()
             n_atoms = int(f.readline().split()[0])
         return n_atoms
+
+
+class NCRSTReader(base.SingleFrameReaderBase, NCDFMixin):
+    """Reader for `AMBER NETCDF format`_ (version 1.0 rev C) restart files.
+
+    This reader is a :class:`SingleFrameReaderBase` adaptation of the
+    :class:`NCDFReader` AMBER NETCDF trajectory reader.
+
+    AMBER binary restart files are automatically recognised by the file
+    extensions ".ncrst", ".ncrestrt", and ".ncrst7".
+
+    The number of atoms (`n_atoms`) does not have to be provided as it can
+    be read from the input NETCDF file.
+
+    Current simulation time is autodetected and if available is read into the
+    :attr:`Timestep.time` attribute.
+
+    Velocities are autodetected and read into the :attr:`Timestep._velocities`
+    attribute.
+
+    Forces are autodetected and read into the :attr:`Timestep._forces`
+    attribute.
+
+    Periodic unit cell information is detected and used to populate the
+    :attr:`Timestep.dimensions` attribute. (If no unit cell is available in
+    the restart file, then :attr:`Timestep.dimensions` will return
+    ``[0,0,0,0,0,0]``).
+
+    Support for the *mmap* keyword is available as detailed
+    in :class:`NCDFReader` and :mod:`scipy.io.netcdf.netcdf_file`. The use of
+    ``mmap=True`` leads to around a 2x read speed improvement in a ~ 1 million
+    atom system (AMBER STMV benchmark). As per the :class:`NCDFReader`, the
+    default behaviour is ``mmap=None``, which means that the default behaviour
+    of :class:`scipy.io.netcdf.netcdf_file` prevails.
+
+    The NCRST reader also uses a custom Timestep object with C-style memory
+    mapping in order to match the NCDFReader.
+
+    .. rubric:: Limitations
+
+    * Only NCRST files with time in ps, lengths in Angstroem and angles in
+      degree are processed.
+    * Restart files without coordinate information are not supported.
+    * Replica exchange variables are not supported.
+
+    .. _AMBER NETCDF format: http://ambermd.org/netcdf/nctraj.xhtml
+
+    See Also
+    --------
+    :class:`NCDFReader`
+    :class:`NCDFWriter`
+
+
+    .. versionadded: 2.10.0
+    """
+
+    format = ["NCRST", "NCRESTRT", "NCRST7"]
+    version = "1.0"
+    units = {
+        "time": "ps",
+        "length": "Angstrom",
+        "velocity": "Angstrom/ps",
+        "force": "kcal/(mol*Angstrom)",
+    }
+
+    _Timestep = Timestep
+
+    @store_init_arguments
+    def __init__(
+        self, filename, n_atoms=None, convert_units=None, mmap=None, **kwargs
+    ):
+        # Assign input mmap value
+        self._mmap = mmap
+
+        super(NCRSTReader, self).__init__(
+            filename, convert_units, n_atoms, **kwargs
+        )
+
+    @staticmethod
+    def parse_n_atoms(filename, **kwargs):
+        with scipy.io.netcdf_file(filename, mmap=None) as f:
+            n_atoms = f.dimensions["atom"]
+        return n_atoms
+
+    @staticmethod
+    def _verify_units(eval_units, expected_units):
+        if eval_units.decode("utf-8") != expected_units:
+            errmsg = (
+                "NCRSTReader currently assumes that the trajectory "
+                "was written in units of {0} instead of {1}".format(
+                    eval_units.decode("utf-8"), expected_units
+                )
+            )
+            raise NotImplementedError(errmsg)
+
+    def _read_first_frame(self):
+        """Function to read NetCDF restart file and fill timestep"""
+        # Open netcdf file via context manager
+        # ensure maskandscale is off so we don't end up double scaling
+
+        with NCDFPicklable(
+            self.filename, mode="r", mmap=self._mmap, maskandscale=False
+        ) as self.trjfile:
+
+            self._check_conventions(n_atoms=self.n_atoms)
+
+            self.n_frames = 1
+
+            self._read_values(
+                frame=()
+            )  # AMBERRESTART convention files have dimensionless datasets
+
+            self.ts.frame = 0  # 0-indexed