Add timeout to minimiser and fix constraint asymmetry bug

src/somd2/config/_config.py

@@ -127,6 +127,7 @@ def __init__(
         somd1_compatibility=False,
         pert_file=None,
         save_energy_components=False,
+        timeout="300s",
     ):
         """
         Constructor.
@@ -301,6 +302,9 @@ def __init__(
         save_energy_components: bool
             Whether to save the energy contribution for each force when checkpointing.
             This is useful when debugging crashes.
+
+        timeout: str
+            Timeout for the minimiser.
         """
 
         # Setup logger before doing anything else
@@ -350,6 +354,7 @@ def __init__(
         self.somd1_compatibility = somd1_compatibility
         self.pert_file = pert_file
         self.save_energy_components = save_energy_components
+        self.timeout = timeout
 
         self.write_config = write_config
 
@@ -1267,6 +1272,29 @@ def save_energy_components(self, save_energy_components):
             raise ValueError("'save_energy_components' must be of type 'bool'")
         self._save_energy_components = save_energy_components
 
+    @property
+    def timeout(self):
+        return self._timeout
+
+    @timeout.setter
+    def timeout(self, timeout):
+        if not isinstance(timeout, str):
+            raise TypeError("'timeout' must be of type 'str'")
+
+        from sire.units import second
+
+        try:
+            t = _sr.u(timeout)
+        except:
+            raise ValueError(
+                f"Unable to parse 'timeout' as a Sire GeneralUnit: {timeout}"
+            )
+
+        if t.value() != 0 and not t.has_same_units(second):
+            raise ValueError("'timeout' units are invalid.")
+
+        self._timeout = t
+
     @property
     def output_directory(self):
         return self._output_directory

src/somd2/runner/_dynamics.py

@@ -184,6 +184,7 @@ def _setup_dynamics(self, equilibration=False):
         else:
             pressure = None
 
+        # Create the dynamics object.
         self._dyn = self._system.dynamics(
             integrator=self._config.integrator,
             temperature=self._config.temperature,
@@ -222,6 +223,18 @@ def _setup_dynamics(self, equilibration=False):
             map=self._config._extra_args,
         )
 
+        # We now need to re-initialize the context so that the constraints
+        # are updated correctly.
+
+        # Get the current positions.
+        positions = self._dyn._d._omm_mols.getState(getPositions=True).getPositions()
+
+        # Reinitialize the context to update the constraints.
+        self._dyn._d._omm_mols.reinitialize()
+
+        # Set the positions.
+        self._dyn._d._omm_mols.setPositions(positions)
+
     def _minimisation(
         self, lambda_min=None, constraint="none", perturbable_constraint="none"
     ):
@@ -239,6 +252,7 @@ def _minimisation(
         if lambda_min is None:
             _logger.info(f"Minimising at {_lam_sym} = {self._lambda_val}")
             try:
+                # Create a minimisation object.
                 m = self._system.minimisation(
                     cutoff_type=self._config.cutoff_type,
                     cutoff=self._config.cutoff,
@@ -256,13 +270,27 @@ def _minimisation(
                     shift_delta=self._config.shift_delta,
                     map=self._config._extra_args,
                 )
-                m.run()
+
+                # We now need to re-initialize the context so that the constraints
+                # are updated correctly.
+
+                # Get the current positions.
+                positions = m._d._omm_mols.getState(getPositions=True).getPositions()
+
+                # Reinitialize the context to update the constraints.
+                m._d._omm_mols.reinitialize()
+
+                # Set the positions.
+                m._d._omm_mols.setPositions(positions)
+
+                m.run(timeout=self._config.timeout)
                 self._system = m.commit()
             except:
                 raise
         else:
             _logger.info(f"Minimising at {_lam_sym} = {lambda_min}")
             try:
+                # Create a minimisation object.
                 m = self._system.minimisation(
                     cutoff_type=self._config.cutoff_type,
                     cutoff=self._config.cutoff,
@@ -280,7 +308,21 @@ def _minimisation(
                     shift_delta=self._config.shift_delta,
                     map=self._config._extra_args,
                 )
-                m.run()
+
+                # We now need to re-initialize the context so that the constraints
+                # are updated correctly.
+
+                # Get the current positions.
+                positions = m._d._omm_mols.getState(getPositions=True).getPositions()
+
+                # Reinitialize the context to update the constraints.
+                m._d._omm_mols.reinitialize()
+
+                # Set the positions.
+                m._d._omm_mols.setPositions(positions)
+
+                # Minimise and commit the changes.
+                m.run(timeout=self._config.timeout)
                 self._system = m.commit()
             except:
                 raise