Update NCDF Reader to read .ncrst Restart Files

package/AUTHORS

@@ -256,6 +256,7 @@ Chronological list of authors
   - James Rowe
   - Debasish Mohanty  
   - Abdulrahman Elbanna
+  - Jeremy M.G. Leung
 
 
 External code

package/CHANGELOG

@@ -15,7 +15,7 @@ The rules for this file:
 
 -------------------------------------------------------------------------------
 ??/??/?? IAlibay, orbeckst, BHM-Bob, TRY-ER, Abdulrahman-PROG, pbuslaev,
-         yuxuanzhuang, yuyuan871111
+         yuxuanzhuang, yuyuan871111, jpkrowe, jeremyleung521
 
  * 2.10.0
 
@@ -31,6 +31,8 @@ Fixes
    (Issue #4948, PR #5001)
 
 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/TRJ.py

@@ -359,7 +359,7 @@ class NCDFReader(base.ReaderBase):
     """Reader for `AMBER NETCDF format`_ (version 1.0).
 
     AMBER binary trajectories are automatically recognised by the
-    file extension ".ncdf".
+    file extension ".ncdf", ".nc", and '.ncrst'.
 
     The number of atoms (`n_atoms`) does not have to be provided as it can
     be read from the trajectory. The trajectory reader can randomly access
@@ -422,10 +422,12 @@ class NCDFReader(base.ReaderBase):
        the first two frames of the trajectory.
        :meth:`Writer` now also sets `convert_units`, `velocities`, `forces` and
        `scale_factor` information for the :class:`NCDFWriter`.
+    .. versionchanged:: 2.10.0
+       Now reads `AMBERRESTART` convention `.ncrst` files.
 
     """
 
-    format = ['NCDF', 'NC']
+    format = ['NCDF', 'NC', 'NCRST']
     multiframe = True
     version = "1.0"
     units = {'time': 'ps',
@@ -450,8 +452,14 @@ def __init__(self, filename, n_atoms=None, mmap=None, **kwargs):
         # AMBER NetCDF files should always have a convention
         try:
             conventions = self.trjfile.Conventions
-            if not ('AMBER' in conventions.decode('utf-8').split(',') or
-                    'AMBER' in conventions.decode('utf-8').split()):
+            check = [conventions.decode('utf-8').split(','),
+                     conventions.decode('utf-8').split()]
+
+            if 'AMBER' in check[0] or 'AMBER' in check[1]:
+                self.is_restart = False
+            elif 'AMBERRESTART' in check[0] or 'AMBERRESTART' in check[1]:
+                self.is_restart = True
+            else:
                 errmsg = ("NCDF trajectory {0} does not conform to AMBER "
                           "specifications, "
                           "http://ambermd.org/netcdf/nctraj.xhtml "
@@ -530,9 +538,12 @@ def __init__(self, filename, n_atoms=None, mmap=None, **kwargs):
             if self.n_frames is None:
                 self.n_frames = self.trjfile.variables['coordinates'].shape[0]
         except KeyError:
-            errmsg = (f"NCDF trajectory {self.filename} does not contain "
-                      f"frame information")
-            raise ValueError(errmsg) from None
+            if self.is_restart:
+                 self.n_frames = 1
+            else:
+                errmsg = (f"NCDF trajectory {self.filename} does not contain "
+                          f"frame information")
+                raise ValueError(errmsg) from None
 
         try:
             self.remarks = self.trjfile.title
@@ -645,21 +656,27 @@ def _read_frame(self, frame):
         if np.dtype(type(frame)) != np.dtype(int):
             # convention... for netcdf could also be a slice
             raise TypeError("frame must be a positive integer or zero")
-        if frame >= self.n_frames or frame < 0:
+        if (frame >= self.n_frames or frame < 0):
             raise IndexError("frame index must be 0 <= frame < {0}".format(
                 self.n_frames))
         # note: self.trjfile.variables['coordinates'].shape == (frames, n_atoms, 3)
-        ts._pos[:] = self._get_var_and_scale('coordinates', frame)
+
+        if self.is_restart:
+            check_frame = ()  # AMBERRESTART convention files have dimensionless datasets
+        else:
+            check_frame = int(frame)
+
+        ts._pos[:] = self._get_var_and_scale('coordinates', check_frame)
         if self.has_time:
-            ts.time = self._get_var_and_scale('time', frame)
+            ts.time = self._get_var_and_scale('time', check_frame)
         if self.has_velocities:
-            ts._velocities[:] = self._get_var_and_scale('velocities', frame)
+            ts._velocities[:] = self._get_var_and_scale('velocities', check_frame)
         if self.has_forces:
-            ts._forces[:] = self._get_var_and_scale('forces', frame)
+            ts._forces[:] = self._get_var_and_scale('forces', check_frame)
         if self.periodic:
             unitcell = np.zeros(6)
-            unitcell[:3] = self._get_var_and_scale('cell_lengths', frame)
-            unitcell[3:] = self._get_var_and_scale('cell_angles', frame)
+            unitcell[:3] = self._get_var_and_scale('cell_lengths', check_frame)
+            unitcell[3:] = self._get_var_and_scale('cell_angles', check_frame)
             ts.dimensions = unitcell
         if self.convert_units:
             self.convert_pos_from_native(ts._pos)  # in-place !
@@ -673,6 +690,8 @@ def _read_frame(self, frame):
             if self.periodic:
                 # in-place ! (only lengths)
                 self.convert_pos_from_native(ts.dimensions[:3])
+
+        del check_frame
         ts.frame = frame  # frame labels are 0-based
         self._current_frame = frame
         return ts

testsuite/MDAnalysisTests/coordinates/test_netcdf.py

@@ -41,6 +41,8 @@
     DLP_CONFIG,
     CPPTRAJ_TRAJ_TOP,
     CPPTRAJ_TRAJ,
+    PRMNCRST,
+    TRJNCRST,
 )
 from MDAnalysisTests.coordinates.test_trj import _TRJReaderTest
 from MDAnalysisTests.coordinates.reference import RefVGV, RefTZ2
@@ -389,6 +391,72 @@ def test_warn_user_no_time_information(self, u):
             u2 = mda.Universe(CPPTRAJ_TRAJ_TOP, CPPTRAJ_TRAJ)
 
 
+class TestNCDFReader5(object):
+    """NCRST Restart File with positions and forces, exported by CPPTRAJ.
+
+    Contributed by Jeremy M. G. Leung
+    """
+
+    prec = 6
+
+    @pytest.fixture(scope="class")
+    def u(self):
+        return mda.Universe(PRMNCRST, TRJNCRST)
+
+    def test_positions(self, u):
+        """Check positions on first frame"""
+        u.trajectory[0]
+        ref_1 = np.array(
+            [
+                [-1.1455358, -2.0177484, -0.55771565],
+                [-0.19042611, -2.2511053, -1.0282656],
+                [ 0.53238064, -1.5778863, -0.56737846],
+            ],
+            dtype=np.float64,
+        )
+        assert_almost_equal(ref_1, u.atoms.positions[:3], self.prec)
+
+    def test_velocities(self, u):
+        """Check forces on first frame"""
+        u.trajectory[0]
+        ref_1 = np.array(
+            [
+                [11.86471367, 31.22108269, -4.03538418],
+                [7.36676359, -4.68035316, 1.78124952],
+                [12.86675262, 1.39324546, -14.97190762],
+            ],
+            dtype=np.float64,
+        )
+        assert_almost_equal(ref_1, u.atoms.velocities[:3], self.prec)
+
+    def test_forces(self, u):
+        """Check forces on first frame"""
+        u.trajectory[0]
+        ref_1 = np.array(
+            [
+                [-9.7262249, -0.37627634, -24.79876137],
+                [-32.43384552, 37.47783279, -12.45422745],
+                [24.10939026, -19.80618095, -17.09523582],
+            ],
+            dtype=np.float64,
+        )
+        assert_almost_equal(ref_1, u.atoms.forces[:3], self.prec)
+
+    def test_time(self, u):
+        """Check time on first frame"""
+        ref = 5.0
+        assert_almost_equal(ref, u.trajectory[0].time, self.prec)
+
+    def test_dt(self, u):
+        """Default 1.0 fs"""
+        ref = 1.0
+        assert_almost_equal(ref, u.trajectory.dt, self.prec)
+        assert_almost_equal(ref, u.trajectory.ts.dt, self.prec)
+
+    def test_box(self, u):
+        assert u.trajectory[0].dimensions is None
+
+
 class _NCDFGenerator(object):
     """A class for generating abitrary ncdf files and exhaustively test
     edge cases which might not be found in the wild"""

testsuite/MDAnalysisTests/datafiles.py

@@ -191,7 +191,8 @@
     "CPPTRAJ_TRAJ_TOP",
     "CPPTRAJ_TRAJ",  # Amber ncdf extracted from CPPTRAJ without time variable
     "PRMcs",  # Amber (format, Issue 1331)
-    "PRMNCRST",  # Amber ncrst with positions/forces/velocities
+    "PRMNCRST",  # Amber topology for ncrst with positions/forces/velocities
+    "TRJNCRST",  # Amber ncrst with positions/forces/velocties
     "PRM_NCBOX",
     "TRJ_NCBOX",  # Amber parm7 + nc w/ pos/forces/vels/box
     "PRMNEGATIVE",  # Amber negative ATOMIC_NUMBER (Issue 2306)
@@ -662,6 +663,7 @@
 PRMcs = (_data_ref / "Amber/chitosan.prmtop").as_posix()
 
 PRMNCRST = (_data_ref / "Amber/ace_mbondi3.parm7").as_posix()
+TRJNCRST = (_data_ref / "Amber/ace_mbondi3.ncrst").as_posix()
 
 PRM_NCBOX = (_data_ref / "Amber/ace_tip3p.parm7").as_posix()
 TRJ_NCBOX = (_data_ref / "Amber/ace_tip3p.nc").as_posix()