Added on_step_finalize in MD simulation hooks (#390)

* added `on_step_finalize` hook for logging etc.

* switched to new hook for logging

src/schnetpack/md/simulation_hooks/basic_hooks.py

@@ -25,6 +25,9 @@ def on_step_middle(self, simulator: Simulator):
     def on_step_end(self, simulator: Simulator):
         pass
 
+    def on_step_finalize(self, simulator: Simulator):
+        pass
+
     def on_step_failed(self, simulator: Simulator):
         pass
 
@@ -51,7 +54,7 @@ def __init__(self, every_n_steps: int, remove_rotation: bool, wrap_positions: bo
         self.remove_rotation = remove_rotation
         self.wrap_positions = wrap_positions
 
-    def on_step_end(self, simulator: Simulator):
+    def on_step_finalize(self, simulator: Simulator):
         if simulator.step % self.every_n_steps == 0:
             simulator.system.remove_center_of_mass()
             simulator.system.remove_translation()

src/schnetpack/md/simulation_hooks/callback_hooks.py

@@ -35,7 +35,7 @@ def __init__(self, checkpoint_file: str, every_n_steps: int):
         self.every_n_steps = every_n_steps
         self.checkpoint_file = checkpoint_file
 
-    def on_step_end(self, simulator: Simulator):
+    def on_step_finalize(self, simulator: Simulator):
         """
         Store state_dict at specified intervals.
 
@@ -526,7 +526,7 @@ def on_simulation_start(self, simulator: Simulator):
         # Enable single writer, multiple reader flag
         self.file.swmr_mode = True
 
-    def on_step_end(self, simulator: Simulator):
+    def on_step_finalize(self, simulator: Simulator):
         """
         Update the buffer of each stream after each specified interval and flush the buffer to the main file if full.
 
@@ -604,7 +604,7 @@ def on_simulation_start(self, simulator):
         self.n_replicas = simulator.system.n_replicas
         self.n_molecules = simulator.system.n_molecules
 
-    def on_step_end(self, simulator):
+    def on_step_finalize(self, simulator: Simulator):
         """
         Routine for collecting and storing scalar properties of replicas and molecules during the simulation. Needs to
         be adapted based on the properties.
@@ -678,7 +678,7 @@ def __init__(self, log_file: str, properties: List, every_n_steps: int = 100):
 
         self.properties = properties
 
-    def on_step_end(self, simulator):
+    def on_step_finalize(self, simulator: Simulator):
         """
         Log the systems properties the given intervals.
 

src/schnetpack/md/simulator.py

@@ -148,6 +148,10 @@ def simulate(self, n_steps: int):
                 for hook in self.simulator_hooks[::-1]:
                     hook.on_step_end(self)
 
+                # Logging hooks etc
+                for hook in self.simulator_hooks:
+                    hook.on_step_finalize(self)
+
                 self.step += 1
                 self.effective_steps += 1