reducing data copies

Becksteinlab · Oct 28, 2024 · a4e161b · a4e161b
1 parent c3ac798
commit a4e161b
Showing 1 changed file with 5 additions and 35 deletions.
diff --git a/imdclient/IMDClient.py b/imdclient/IMDClient.py
@@ -448,13 +448,7 @@ def _parse_imdframe(self):
             self._expect_header(
                 IMDHeaderType.IMD_FCOORDS, expected_value=self._n_atoms
             )
-            self._read(self._positions)
-            np.copyto(
-                self._imdf.positions,
-                np.frombuffer(
-                    self._positions, dtype=f"{self._imdsinfo.endianness}f"
-                ).reshape((self._n_atoms, 3)),
-            )
+            self._read(self._imdf.positions.ravel())
         elif (
             header.type == IMDHeaderType.IMD_FCOORDS
             and header.length == self._n_atoms
@@ -465,13 +459,7 @@ def _parse_imdframe(self):
                 self._imdf.energies = self._prev_energies
             else:
                 self._imdf.energies = None
-            self._read(self._positions)
-            np.copyto(
-                self._imdf.positions,
-                np.frombuffer(
-                    self._positions, dtype=f"{self._imdsinfo.endianness}f"
-                ).reshape((self._n_atoms, 3)),
-            )
+            self._read(self._imdf.positions.ravel())
         else:
             raise RuntimeError("IMDProducer: Unexpected packet type or length")
 
@@ -600,35 +588,17 @@ def _parse_imdframe(self):
             self._expect_header(
                 IMDHeaderType.IMD_FCOORDS, expected_value=self._n_atoms
             )
-            self._read(self._positions)
-            np.copyto(
-                self._imdf.positions,
-                np.frombuffer(
-                    self._positions, dtype=f"{self._imdsinfo.endianness}f"
-                ).reshape((self._n_atoms, 3)),
-            )
+            self._read(self._imdf.positions.ravel())
         if self._imdsinfo.velocities:
             self._expect_header(
                 IMDHeaderType.IMD_VELOCITIES, expected_value=self._n_atoms
             )
-            self._read(self._velocities)
-            np.copyto(
-                self._imdf.velocities,
-                np.frombuffer(
-                    self._velocities, dtype=f"{self._imdsinfo.endianness}f"
-                ).reshape((self._n_atoms, 3)),
-            )
+            self._read(self._imdf.velocities.ravel())
         if self._imdsinfo.forces:
             self._expect_header(
                 IMDHeaderType.IMD_FORCES, expected_value=self._n_atoms
             )
-            self._read(self._forces)
-            np.copyto(
-                self._imdf.forces,
-                np.frombuffer(
-                    self._forces, dtype=f"{self._imdsinfo.endianness}f"
-                ).reshape((self._n_atoms, 3)),
-            )
+            self._read(self._imdf.forces.ravel())
 
     def __del__(self):
         logger.debug("IMDProducer: I am being deleted")