TRZReader check for n_atoms

package/CHANGELOG

@@ -16,8 +16,6 @@ The rules for this file:
 ??/??/?? richardjgowers, IAlibay, hmacdope, orbeckst, cbouy, lilyminium,
          daveminh, jbarnoud, yuxuanzhuang
 
-
-
   * 2.0.0
 
 Fixes
@@ -28,6 +26,9 @@ Fixes
   * Testsuite does not any more matplotlib.use('agg') (#2191)
   * In ChainReader, read_frame does not trigger change of iterating position. 
     (Issue #2723, PR #2815)
+  * TRZReader now checks `n_atoms` on reading. (Issue #2817, PR #2820)
+  * TRZWriter now writes `n_atoms` to the file. (Issue #2817, PR #2820)
+
 
 Enhancements
   * Added the RDKitParser which creates a `core.topology.Topology` object from

package/MDAnalysis/coordinates/TRZ.py

@@ -84,7 +84,6 @@
    :members:
 
 """
-import sys
 import warnings
 import numpy as np
 import os
@@ -141,7 +140,9 @@ class TRZReader(base.ReaderBase):
        Extra data (Temperature, Energies, Pressures, etc) now read
        into ts.data dictionary.
        Now passes a weakref of self to ts (ts._reader).
-
+    .. versionchanged:: 2.0.0
+       Now checks for the correct `n_atoms` on reading
+       and can raise :exc:`ValueError`.
     """
 
     format = "TRZ"
@@ -159,6 +160,12 @@ def __init__(self, trzfilename, n_atoms=None, **kwargs):
             number of atoms in trajectory, must be taken from topology file!
         convert_units : bool (optional)
             converts units to MDAnalysis defaults
+
+        Raises
+        ------
+        ValueError
+           If `n_atoms` or the number of atoms in the topology file do not
+           match the number of atoms in the trajectory.
         """
         super(TRZReader, self).__init__(trzfilename,  **kwargs)
 
@@ -244,6 +251,11 @@ def _read_next_timestep(self, ts=None):
 
         try:
             data = np.fromfile(self.trzfile, dtype=self._dtype, count=1)
+            if data['natoms'][0] != self.n_atoms:
+                raise ValueError("Supplied n_atoms {} is incompatible "
+                                 "with provided trajectory file. "
+                                 "Maybe `topology` is wrong?".format(
+                                                             self.n_atoms))
             ts.frame = data['nframe'][0] - 1  # 0 based for MDA
             ts._frame = data['ntrj'][0]
             ts.time = data['treal'][0]
@@ -588,6 +600,7 @@ def _write_next_frame(self, obj):
         out['p1a'], out['p1b'] = 20, 20
         out['nframe'] = ts.frame + 1  # TRZ wants 1 based
         out['ntrj'] = data['step']
+        out['natoms'] = self.n_atoms
         out['treal'] = data['time']
         out['p2a'], out['p2b'] = 72, 72
         out['box'] = self.convert_pos_to_native(unitcell, inplace=False)

testsuite/MDAnalysisTests/coordinates/test_trz.py

@@ -129,6 +129,10 @@ def test_get_writer_2(self, universe, tmpdir):
             assert_equal(isinstance(W, mda.coordinates.TRZ.TRZWriter), True)
             assert_equal(W.n_atoms, 100)
 
+    def test_get_wrong_n_atoms(self):
+        with pytest.raises(ValueError, match=r"Supplied n_atoms"):
+            mda.Universe(TRZ, n_atoms=8080)
+
 
 class TestTRZWriter(RefTRZ):
     prec = 3