From 7583a2efa1e45c95aee1c5205f503fa2f31ffa8c Mon Sep 17 00:00:00 2001
From: Tyler Reddy <tyler.je.reddy@gmail.com>
Date: Mon, 30 Dec 2024 16:35:57 -0700
Subject: [PATCH] ENH: PR 4873 revisions

* Expand `TPRReader()` support to include velocity handling,
and add tests/functionality for an additional tpx version (`133`).

[ci skip] [skip azp]
---
 package/MDAnalysis/coordinates/TPR.py         |  6 ++-
 package/MDAnalysis/topology/tpr/utils.py      |  2 +-
 .../MDAnalysisTests/coordinates/test_tpr.py   | 38 ++++++++++++++++---
 3 files changed, 38 insertions(+), 8 deletions(-)

diff --git a/package/MDAnalysis/coordinates/TPR.py b/package/MDAnalysis/coordinates/TPR.py
index 56fcbfd79b..ba9db02299 100644
--- a/package/MDAnalysis/coordinates/TPR.py
+++ b/package/MDAnalysis/coordinates/TPR.py
@@ -18,7 +18,7 @@ class TPRReader(base.SingleFrameReaderBase):
     # or perhaps combine the topology and coordinate reading
     # with some inheritance shenanigans?
     format = "TPR"
-    units = {"length": "nm"}
+    units = {"length": "nm", "velocity": "nm/ps"}
     _Timestep = Timestep
 
     def _read_first_frame(self):
@@ -71,3 +71,7 @@ def _read_first_frame(self):
         if th.bX:
             self.ts._pos = np.asarray(tpr_utils.ndo_rvec(data, th.natoms),
                                       dtype=np.float32)
+        if th.bV:
+            self.ts._velocities = np.asarray(tpr_utils.ndo_rvec(data, th.natoms),
+                                             dtype=np.float32)
+            self.ts.has_velocities = True
diff --git a/package/MDAnalysis/topology/tpr/utils.py b/package/MDAnalysis/topology/tpr/utils.py
index 4a4d03bc24..fe4c6bd5b4 100644
--- a/package/MDAnalysis/topology/tpr/utils.py
+++ b/package/MDAnalysis/topology/tpr/utils.py
@@ -433,7 +433,7 @@ def do_mtop(data, fver, tpr_resid_from_one=False):
     # src/gromacs/fileio/tpxio.cpp
     # TODO: expand tpx version support for striding to
     # the coordinates
-    if fver == 134:
+    if fver >= 133:
         # TODO: the following value is important, and not sure
         # how to access programmatically yet...
         # from GMX source code:
diff --git a/testsuite/MDAnalysisTests/coordinates/test_tpr.py b/testsuite/MDAnalysisTests/coordinates/test_tpr.py
index e3d2d2c28b..633815be94 100644
--- a/testsuite/MDAnalysisTests/coordinates/test_tpr.py
+++ b/testsuite/MDAnalysisTests/coordinates/test_tpr.py
@@ -1,34 +1,60 @@
 from MDAnalysisTests.datafiles import (TPR2024_4_bonded,
                                        TPR_EXTRA_2024_4,
-                                       TPR2024_4)
+                                       TPR2024_4,
+                                       TPR2024)
 import MDAnalysis as mda
 
 
 import pytest
+import numpy as np
 from numpy.testing import assert_allclose, assert_equal
 
 
-@pytest.mark.parametrize("tpr_file, exp_first_atom, exp_last_atom, exp_shape", [
-    (TPR2024_4_bonded,
+@pytest.mark.parametrize("tpr_file, exp_first_atom, exp_last_atom, exp_shape, exp_vel_first_atom, exp_vel_last_atom", [
+    (TPR2024_4_bonded, # tpx 134
      [4.446, 4.659, 2.384],
      [4.446, 4.659, 2.384],
      (14, 3),
+     np.zeros(3),
+     np.zeros(3),
     ),
     # same coordinates, different shape
-    (TPR_EXTRA_2024_4,
+    (TPR_EXTRA_2024_4, # tpx 134
      [4.446, 4.659, 2.384],
      [4.446, 4.659, 2.384],
      (18, 3),
+     np.zeros(3),
+     np.zeros(3),
     ),
     # different coordinates and different shape
-    (TPR2024_4,
+    (TPR2024_4, # tpx 134
      [3.25000e-01,  1.00400e+00,  1.03800e+00],
      [-2.56000e-01,  1.37300e+00,  3.59800e+00],
      (2263, 3),
+     np.zeros(3),
+     np.zeros(3),
+    ),
+    (TPR2024, # tpx 133
+     [3.25000e-01,  1.00400e+00,  1.03800e+00],
+     [-2.56000e-01,  1.37300e+00,  3.59800e+00],
+     (2263, 3),
+     np.zeros(3),
+     np.zeros(3),
     ),
 ])
-def test_basic_read_tpr(tpr_file, exp_first_atom, exp_last_atom, exp_shape):
+def test_basic_read_tpr(tpr_file,
+                        exp_first_atom,
+                        exp_last_atom,
+                        exp_shape,
+                        exp_vel_first_atom,
+                        exp_vel_last_atom):
+    # verify basic ability to read positions and
+    # velocities from GMX .tpr files
+    # expected values are from gmx dump
     u = mda.Universe(tpr_file)
     assert_allclose(u.atoms.positions[0, ...], exp_first_atom)
     assert_allclose(u.atoms.positions[-1, ...], exp_last_atom)
     assert_equal(u.atoms.positions.shape, exp_shape)
+    assert_allclose(u.atoms.velocities[0, ...], exp_vel_first_atom)
+    assert_allclose(u.atoms.velocities[-1, ...], exp_vel_last_atom)
+    assert_equal(u.atoms.velocities.shape, exp_shape)