From dd310a162bc38933f69f94bce4aded1637a46cff Mon Sep 17 00:00:00 2001 From: mgastegger Date: Thu, 10 Feb 2022 17:27:13 +0100 Subject: [PATCH] Fixed energy logging for multiple molecules --- src/schnetpack/md/data/hdf5_data.py | 2 +- src/schnetpack/md/simulation_hooks/callback_hooks.py | 6 ++++-- 2 files changed, 5 insertions(+), 3 deletions(-) diff --git a/src/schnetpack/md/data/hdf5_data.py b/src/schnetpack/md/data/hdf5_data.py index 0caced13a..2efd4d15c 100644 --- a/src/schnetpack/md/data/hdf5_data.py +++ b/src/schnetpack/md/data/hdf5_data.py @@ -113,7 +113,7 @@ def _load_molecule_data(self): # Extract energies entry_start = 0 - entry_stop = 1 + entry_stop = self.n_molecules self.properties[f"{properties.energy}_system"] = raw_positions[ :, :, entry_start:entry_stop ].reshape(self.entries, self.n_replicas, self.n_molecules) diff --git a/src/schnetpack/md/simulation_hooks/callback_hooks.py b/src/schnetpack/md/simulation_hooks/callback_hooks.py index 921277c66..cad32ee3d 100644 --- a/src/schnetpack/md/simulation_hooks/callback_hooks.py +++ b/src/schnetpack/md/simulation_hooks/callback_hooks.py @@ -218,7 +218,9 @@ def _init_data_stream(self, simulator: Simulator): simulator (schnetpack.md.Simulator): Simulator class used in the molecular dynamics simulation. """ # Account for potential energy and positions - data_dimension = 1 + simulator.system.total_n_atoms * 3 + data_dimension = ( + simulator.system.n_molecules + simulator.system.total_n_atoms * 3 + ) # If requested, also store velocities if self.store_velocities: @@ -267,7 +269,7 @@ def update_buffer(self, buffer_position: int, simulator: Simulator): # Store energies start = 0 - stop = 1 + stop = simulator.system.n_molecules self.buffer[ buffer_position : buffer_position + 1, :, start:stop ] = simulator.system.energy.view(simulator.system.n_replicas, -1).detach()