Merge pull request #390 from QCMM/gro_update

Update for load_one info in gro files

Author: tovrstra

iodata/formats/gromacs.py

@@ -90,6 +90,7 @@ def _helper_read_frame(lit: LineIterator) -> tuple:
     attypes = []
     pos = np.zeros((natoms, 3), np.float32)
     vel = np.zeros((natoms, 3), np.float32)
+    has_velocities = True
     for i in range(natoms):
         line = next(lit)
         resnums.append(int(line[:5]))
@@ -99,11 +100,18 @@ def _helper_read_frame(lit: LineIterator) -> tuple:
         pos[i, 0] = float(words[0])
         pos[i, 1] = float(words[1])
         pos[i, 2] = float(words[2])
-        vel[i, 0] = float(words[3])
-        vel[i, 1] = float(words[4])
-        vel[i, 2] = float(words[5])
+        # velocities are optional in gro files; only parse if present
+        if len(words) >= 6:
+            vel[i, 0] = float(words[3])
+            vel[i, 1] = float(words[4])
+            vel[i, 2] = float(words[5])
+        else:
+            has_velocities = False
     pos *= nanometer  # atom coordinates are in nanometers
-    vel *= nanometer / picosecond
+    if has_velocities:
+        vel *= nanometer / picosecond
+    else:
+        vel = None
     # Read the cell line
     cell = np.zeros((3, 3), np.float32)
     words = next(lit).split()

iodata/periodic.py

@@ -144,7 +144,6 @@
 
 sym2num: dict[str, int] = {value: key for key, value in num2sym.items()}
 
-
 # Labels used for bond types.
 
 num2bond = {

iodata/test/data/water_no_vel.gro

@@ -0,0 +1,9 @@
+MD of 2 waters, t= 0.0
+    6
+    1WATER  OW1    1   0.126   1.624   1.679
+    1WATER  HW2    2   0.190   1.661   1.747
+    1WATER  HW3    3   0.177   1.568   1.613
+    2WATER  OW1    4   1.275   0.053   0.622
+    2WATER  HW2    5   1.337   0.002   0.680
+    2WATER  HW3    6   1.326   0.120   0.568
+   1.82060   1.82060   1.82060

iodata/test/test_gromacs.py

@@ -27,7 +27,7 @@
 
 
 def test_load_water():
-    # test gro file of one water
+    """Test gro file of one water."""
     with as_file(files("iodata.test.data").joinpath("water.gro")) as fn_gro:
         mol = load_one(str(fn_gro))
     check_water(mol)
@@ -48,10 +48,28 @@ def check_water(mol):
 
 
 def test_load_many():
+    """Test water gro file with two frames."""
     with as_file(files("iodata.test.data").joinpath("water2.gro")) as fn_gro:
         mols = list(load_many(str(fn_gro)))
     assert len(mols) == 2
     assert mols[0].extra["time"] == 0.0 * picosecond
     assert mols[1].extra["time"] == 1.0 * picosecond
     for mol in mols:
         check_water(mol)
+
+
+def test_load_without_velocities():
+    """A GRO file without velocity columns should set velocities to None."""
+    with as_file(files("iodata.test.data").joinpath("water_no_vel.gro")) as fn_gro:
+        mol = load_one(str(fn_gro))
+
+    # basic checks
+    assert mol.title == "MD of 2 waters"
+    assert mol.atcoords.shape == (6, 3)
+
+    # coordinates should match the ones in the file (compare in nm)
+    assert_allclose(mol.atcoords[-1] / nanometer, [1.326, 0.120, 0.568])
+
+    # velocities should be present in extra but None
+    assert "velocities" in mol.extra
+    assert mol.extra["velocities"] is None