Added tmpdir functionality for MDs, fixed calculator bug

src/schnetpack/md/calculators/base_calculator.py

@@ -2,6 +2,7 @@
 from typing import List, Union, Dict, Optional, Tuple
 
 from typing import TYPE_CHECKING
+from contextlib import nullcontext
 
 import numpy as np
 
@@ -43,6 +44,7 @@ class MDCalculator(nn.Module):
         stress_label (str, optional): Name of the property corresponding to the stress.
         property_conversion (dict(float, str)): Optional dictionary of conversion factors for other properties predicted
                              by the model. Only changes the units used for logging the various outputs.
+        gradients_rquired (bool): If set to true, enable accumulation of computational graph in calculator.
     """
 
     def __init__(
@@ -54,6 +56,7 @@ def __init__(
         energy_label: Optional[str] = None,
         stress_label: Optional[str] = None,
         property_conversion: Dict[str, Union[str, float]] = {},
+        gradients_required: bool = False,
     ):
         super(MDCalculator, self).__init__()
         # Get required properties and filter non-unique entries
@@ -93,6 +96,12 @@ def __init__(
         self.force_conversion = self.energy_conversion / self.position_conversion
         self.stress_conversion = self.energy_conversion / self.position_conversion**3
 
+        # set up gradient context for passing results
+        if gradients_required:
+            self.grad_context = torch.no_grad()
+        else:
+            self.grad_context = nullcontext()
+
     def calculate(self, system: System):
         """
         Main calculator routine, which needs to be implemented individually.
@@ -116,31 +125,31 @@ def _update_system(self, system: System):
         Args:
             system (schnetpack.md.System): System object containing current state of the simulation.
         """
-
-        # Collect all requested properties (including forces)
-        for p in self.required_properties:
-            if p not in self.results:
-                raise MDCalculatorError(
-                    "Requested property {:s} not in " "results".format(p)
-                )
-            else:
-                dim = self.results[p].shape
-                # Bring to general structure of MD code. Second dimension can be n_mol or n_mol x n_atoms.
-                system.properties[p] = (
-                    self.results[p].view(system.n_replicas, -1, *dim[1:])
-                    * self.property_conversion[p]
-                )
-
-        # Set the forces for the system (at this point, already detached)
-        self._set_system_forces(system)
-
-        # Store potential energy to system if requested:
-        if self.energy_label is not None:
-            self._set_system_energy(system)
-
-        # Set stress of the system if requested:
-        if self.stress_label is not None:
-            self._set_system_stress(system)
+        with self.grad_context:
+            # Collect all requested properties (including forces)
+            for p in self.required_properties:
+                if p not in self.results:
+                    raise MDCalculatorError(
+                        "Requested property {:s} not in " "results".format(p)
+                    )
+                else:
+                    dim = self.results[p].shape
+                    # Bring to general structure of MD code. Second dimension can be n_mol or n_mol x n_atoms.
+                    system.properties[p] = (
+                        self.results[p].view(system.n_replicas, -1, *dim[1:])
+                        * self.property_conversion[p]
+                    )
+
+            # Set the forces for the system (at this point, already detached)
+            self._set_system_forces(system)
+
+            # Store potential energy to system if requested:
+            if self.energy_label is not None:
+                self._set_system_energy(system)
+
+            # Set stress of the system if requested:
+            if self.stress_label is not None:
+                self._set_system_stress(system)
 
     def _get_system_molecules(self, system: System):
         """

src/schnetpack/md/cli.py

@@ -2,12 +2,15 @@
 import uuid
 import torch
 import os
+import shutil
 
 from datetime import datetime
 
 import hydra
 from omegaconf import DictConfig, OmegaConf
 
+import tempfile
+
 import schnetpack.md
 from schnetpack.utils import str2class, int2precision
 from schnetpack.utils.script import print_config
@@ -20,6 +23,7 @@
 log = logging.getLogger(__name__)
 
 OmegaConf.register_new_resolver("uuid", lambda x: str(uuid.uuid1()))
+OmegaConf.register_new_resolver("tmpdir", tempfile.mkdtemp, use_cache=True)
 
 
 class MDSetupError(Exception):
@@ -302,10 +306,24 @@ def simulate(config: DictConfig):
     #   Finally run simulation
     # ===========================================
 
-    log.info("Running simulation...")
+    log.info("Running simulation in {:s}...".format(hydra_wd))
 
     start = datetime.now()
     simulator.simulate(config.dynamics.n_steps)
     stop = datetime.now()
 
+    final_dir = hydra.utils.to_absolute_path(config.simulation_dir)
+    if final_dir != hydra_wd:
+        logging.info(
+            "Moving simulation output from {:s} to {:s}...".format(hydra_wd, final_dir)
+        )
+
+        if os.path.exists(final_dir):
+            logging.info(
+                "Destination {:s} already exists, moving data to subdirectory...".format(
+                    final_dir
+                )
+            )
+        shutil.move(hydra_wd, final_dir)
+
     log.info("Finished after: {:s}".format(str(stop - start)))