Merge pull request #26 from OpenBioSim/feature_somd1_compatibility2

Further SOMD1 compatibility updates

src/somd2/config/_config.py

@@ -74,6 +74,7 @@ def __init__(
         timestep="4fs",
         temperature="300K",
         pressure="1 atm",
+        barostat_frequency=25,
         integrator="langevin_middle",
         cutoff_type="pme",
         cutoff="7.5A",
@@ -125,6 +126,9 @@ def __init__(
         pressure: str
             Simulation pressure. (Simulations will run in the NVT ensemble unless a pressure is specified.)
 
+        barostat_frequency: int
+            The number of integration steps between barostat updates.
+
         integrator: str
             Integrator to use for simulation.
 
@@ -249,6 +253,7 @@ def __init__(
         self.runtime = runtime
         self.temperature = temperature
         self.pressure = pressure
+        self.barostat_frequency = barostat_frequency
         self.integrator = integrator
         self.cutoff_type = cutoff_type
         self.cutoff = cutoff
@@ -432,6 +437,10 @@ def pressure(self, pressure):
             try:
                 p = _sr.u(pressure)
             except:
+                # Handle special case of pressure = "none"
+                if pressure.lower().replace(" ", "") == "none":
+                    self._pressure = None
+                    return
                 raise ValueError(
                     f"Unable to parse 'pressure' as a Sire GeneralUnit: {pressure}"
                 )
@@ -443,6 +452,20 @@ def pressure(self, pressure):
         else:
             self._pressure = pressure
 
+    @property
+    def barostat_frequency(self):
+        return self._barostat_frequency
+
+    @barostat_frequency.setter
+    def barostat_frequency(self, barostat_frequency):
+        if not isinstance(barostat_frequency, int):
+            raise TypeError("'barostat_frequency' must be of type 'int'")
+
+        if barostat_frequency <= 0:
+            raise ValueError("'barostat_frequency' must be a positive integer")
+
+        self._barostat_frequency = barostat_frequency
+
     @property
     def integrator(self):
         return self._integrator

src/somd2/runner/_dynamics.py

@@ -158,14 +158,10 @@ def _setup_dynamics(self, equilibration=False):
         else:
             pressure = None
 
-        try:
-            map = self._config._extra_args
-        except:
-            map = None
-
         self._dyn = self._system.dynamics(
             temperature=self._config.temperature,
             pressure=pressure,
+            barostat_frequency=self._config.barostat_frequency,
             timestep=(
                 self._config.equilibration_timestep
                 if equilibration
@@ -186,7 +182,7 @@ def _setup_dynamics(self, equilibration=False):
             swap_end_states=self._config.swap_end_states,
             com_reset_frequency=self._config.com_reset_frequency,
             vacuum=not self._has_space,
-            map=map,
+            map=self._config._extra_args,
         )
 
     def _minimisation(self, lambda_min=None, perturbable_constraint="none"):
@@ -200,6 +196,7 @@ def _minimisation(self, lambda_min=None, perturbable_constraint="none"):
             Lambda value at which to run minimisation, if None run at pre-set
             lambda_val.
         """
+
         if lambda_min is None:
             _logger.info(f"Minimising at {_lam_sym} = {self._lambda_val}")
             try:
@@ -210,6 +207,7 @@ def _minimisation(self, lambda_min=None, perturbable_constraint="none"):
                     lambda_value=self._lambda_val,
                     platform=self._config.platform,
                     vacuum=not self._has_space,
+                    constraint=self._config.constraint,
                     perturbable_constraint=perturbable_constraint,
                     include_constrained_energies=self._config.include_constrained_energies,
                     dynamic_constraints=self._config.dynamic_constraints,
@@ -230,6 +228,7 @@ def _minimisation(self, lambda_min=None, perturbable_constraint="none"):
                     lambda_value=lambda_min,
                     platform=self._config.platform,
                     vacuum=not self._has_space,
+                    constraint=self._config.constraint,
                     perturbable_constraint=perturbable_constraint,
                     include_constrained_energies=self._config.include_constrained_energies,
                     dynamic_constraints=self._config.dynamic_constraints,

src/somd2/runner/_runner.py

@@ -94,6 +94,7 @@ def __init__(self, system, config):
         if not isinstance(config, _Config):
             raise TypeError("'config' must be of type 'somd2.config.Config'")
         self._config = config
+        self._config._extra_args = {}
 
         # We're running in SOMD1 compatibility mode.
         if self._config.somd1_compatibility:
@@ -138,9 +139,21 @@ def __init__(self, system, config):
                             "Unable to convert water topology to AMBER format for SOMD1 compatibility."
                         )
 
+            # Only check for light atoms by the maxium end state mass if running
+            # in SOMD1 compatibility mode.
+            if self._config.somd1_compatibility:
+                self._config._extra_args["check_for_h_by_max_mass"] = True
+                self._config._extra_args["check_for_h_by_mass"] = False
+                self._config._extra_args["check_for_h_by_mass"] = False
+                self._config._extra_args["check_for_h_by_element"] = False
+                self._config._extra_args["check_for_h_by_ambertype"] = False
+
         # Check for a periodic space.
         self._check_space()
 
+        # Check the end state contraints.
+        self._check_end_state_constraints()
+
         # Set the lambda values.
         self._lambda_values = [
             round(i / (self._config.num_lambda - 1), 5)
@@ -235,6 +248,41 @@ def _check_space(self):
                 )
                 self._config.cutoff_type = "rf"
 
+    def _check_end_state_constraints(self):
+        """
+        Internal function to check whether the constrants are the same at the two
+        end states.
+        """
+
+        # Find all perturbable molecules in the system..
+        pert_mols = self._system.molecules("property is_perturbable")
+
+        # Check constraints at lambda = 0 and lambda = 1 for each perturbable molecule.
+        for mol in pert_mols:
+            # Create a dynamics object.
+            d = mol.dynamics(
+                constraint=self._config.constraint,
+                perturbable_constraint=self._config.perturbable_constraint,
+                platform="cpu",
+                map=self._config._extra_args,
+            )
+
+            # Get the constraints at lambda = 0.
+            constraints0 = d.get_constraints()
+
+            # Update to lambda = 1.
+            d.set_lambda(1)
+
+            # Get the constraints at lambda = 1.
+            constraints1 = d.get_constraints()
+
+            # Check for equivalence.
+            for c0, c1 in zip(constraints0, constraints1):
+                if c0 != c1:
+                    _logger.info(
+                        f"Constraints are at not the same at {_lam_sym} = 0 and {_lam_sym} = 1."
+                    )
+
     def _check_directory(self):
         """
         Find the name of the file from which simulations will be started.
@@ -657,17 +705,15 @@ def run(self):
                     threads_per_worker = (
                         self._config.max_threads // self._config.num_lambda
                     )
-                self._config._extra_args = {"threads": threads_per_worker}
+                self._config._extra_args["threads"] = threads_per_worker
 
             # (Multi-)GPU platform.
             elif self._is_gpu:
                 self.max_workers = len(self._gpu_pool)
-                self._config._extra_args = {}
 
             # All other platforms.
             else:
                 self._max_workers = 1
-                self._config._extra_args = {}
 
             import concurrent.futures as _futures
 
@@ -696,7 +742,7 @@ def run(self):
         # Serial configuration.
         elif self._config.num_lambda is not None:
             if self._config.platform == "cpu":
-                self._config.extra_args = {"threads": self._config.max_threads}
+                self._config._extra_args = {"threads": self._config.max_threads}
             self._lambda_values = [
                 round(i / (self._config.num_lambda - 1), 5)
                 for i in range(0, self._config.num_lambda)