Check consistency between multiple dump calls in NPYWriter

Author: tovrstra

CHANGELOG.md

@@ -16,6 +16,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 - By default, run only 100 optimization steps in `build_random_cell`.
 - Wrap atoms back into the cell when writing PDB trajectory files, for nicer visual.
+- Stricter consistency checking between multiple `dump` calls in `NPYWriter`.
 
 
 ## [0.0.0] - 2024-10-06

src/tinyff/trajectory.py

@@ -153,15 +153,24 @@ def dump(self, **kwargs):
 
     def dump_each(self, **kwargs):
         """Write data to NPY files without considering `self.stride`."""
-        for key, value in kwargs.items():
-            # Check array properties
-            shape, dtype = self.fields.get(key, (None, None))
-            arvalue = np.asarray(value)
-            if shape is None:
-                shape = arvalue.shape
-                dtype = arvalue.dtype
-                self.fields[key] = (shape, dtype)
-            else:
+        converted = {}
+        if len(self.fields) == 0:
+            # No checking, just record the given shapes and types
+            for key, value in kwargs.items():
+                arvalue = np.asarray(value)
+                converted[key] = arvalue
+                self.fields[key] = (arvalue.shape, arvalue.dtype)
+        else:
+            # Check kwargs
+            if set(self.fields) != set(kwargs):
+                raise TypeError(
+                    f"Received keys: {list(kwargs.keys())}. "
+                    f"Expected: {list(self.fields.keys())}"
+                )
+            for key, value in kwargs.items():
+                arvalue = np.asarray(value)
+                converted[key] = arvalue
+                shape, dtype = self.fields[key]
                 if shape != arvalue.shape:
                     raise TypeError(
                         f"The shape of {key}, {arvalue.shape}, differs from the first one, {shape}"
@@ -170,7 +179,9 @@ def dump_each(self, **kwargs):
                     raise TypeError(
                         f"The dtype of {key}, {arvalue.dtype}, differs from the first one, {dtype}"
                     )
-            # Append to NPY file
+
+        # Write only once all checks have passed
+        for key, value in converted.items():
             path = os.path.join(self.dir_out, f"{key}.npy")
             with NpyAppendArray(path, delete_if_exists=False) as npaa:
-                npaa.append(arvalue.reshape(1, *arvalue.shape))
+                npaa.append(value.reshape(1, *value.shape))

tests/test_trajectory.py

@@ -120,3 +120,24 @@ def test_npy_traj_stride(tmpdir):
     for foo in traj_foo:
         npy_writer.dump(foo=foo)
     assert (np.load(os.path.join(tmpdir, "out/foo.npy")) == [1, 4, 7]).all()
+
+
+def test_npy_traj_consistent_names(tmpdir):
+    npy_writer = NPYWriter(os.path.join(tmpdir, "out"))
+    npy_writer.dump(a=1, b=2)
+    with pytest.raises(TypeError):
+        npy_writer.dump(b=3, c=4)
+
+
+def test_npy_traj_consistent_shapes(tmpdir):
+    npy_writer = NPYWriter(os.path.join(tmpdir, "out"))
+    npy_writer.dump(a=[1, 2], b=3)
+    with pytest.raises(TypeError):
+        npy_writer.dump(a=3, b=4)
+
+
+def test_npy_traj_consistent_dtypes(tmpdir):
+    npy_writer = NPYWriter(os.path.join(tmpdir, "out"))
+    npy_writer.dump(a=1.0, b=3)
+    with pytest.raises(TypeError):
+        npy_writer.dump(a=3, b=4)