Merge pull request #366 from OpenBioSim/fix_amber_fep_mappings

Fix AMBER FEP mappings

Author: lohedges

python/BioSimSpace/Align/_align.py

@@ -163,7 +163,7 @@ def __init__(
 
         engine : str
             The molecular dynamics engine used to run the simulation. Available
-            options are "GROMACS", or "SOMD". If this argument is omitted then
+            options are "AMBER", "GROMACS", or "SOMD". If this argument is omitted then
             BioSimSpace will choose an appropriate engine for you.
 
         setup_only : bool
@@ -1158,8 +1158,8 @@ def _preprocess_data(data, estimator, **kwargs):
         # Assign defaults in case not passed via kwargs.
         auto_eq = False
         stat_ineff = False
-        truncate = False
-        truncate_keep = "start"
+        truncate_upper = 100
+        truncate_lower = 0
 
         # Parse kwargs.
         for key, value in kwargs.items():
@@ -1168,35 +1168,31 @@ def _preprocess_data(data, estimator, **kwargs):
                 auto_eq = value
             if key == "STATISTICALINEFFICIENCY":
                 stat_ineff = value
-            if key == "TRUNCATEPERCENTAGE":
-                truncate = value
-            if key == "TRUNCATEKEEP":
-                truncate_keep = value
+            if key == "TRUNCATEUPPER":
+                truncate_upper = value
+            if key == "TRUNCATELOWER":
+                truncate_lower = value
 
-        # First truncate data.
+        # Copy the data.
         raw_data = data
-        if truncate:
-            # Get the upper and lower bounds for truncate.
-            data_len = len(data[0])
-            data_step = round((data[0].index[-1][0] - data[0].index[-2][0]), 1)
-            data_kept = data_len * (truncate / 100)
-            data_time = data_kept * data_step
-            if truncate_keep == "start":
-                truncate_lower = 0
-                truncate_upper = data_time - data_step
-            if truncate_keep == "end":
-                truncate_lower = (data_len * data_step) - data_time
-                truncate_upper = (data_len * data_step) - data_step
 
-            try:
-                data = [
-                    _slicing(i, lower=truncate_lower, upper=truncate_upper)
-                    for i in raw_data
-                ]
-            except:
-                _warnings.warn("Could not truncate data.")
-                data = raw_data
-        else:
+        # Data length.
+        data_len = len(data[0])
+
+        # Step size.
+        data_step = round((data[0].index[-1][0] - data[0].index[-2][0]), 1)
+
+        # Get the upper and lower bounds for truncate.
+        truncate_lower = (data_len * (truncate_lower / 100)) * data_step
+        truncate_upper = (data_len * (truncate_upper / 100)) * data_step
+
+        try:
+            data = [
+                _slicing(i, lower=truncate_lower, upper=truncate_upper)
+                for i in raw_data
+            ]
+        except:
+            _warnings.warn("Could not truncate data.")
             data = raw_data
 
         # The decorrelate function calls either autoequilibration or statistical_inefficiency
@@ -1230,10 +1226,7 @@ def _preprocess_data(data, estimator, **kwargs):
                 )
             sampled_data = data
 
-        # Concatanate in alchemlyb format.
-        processed_data = _alchemlyb.concat(sampled_data)
-
-        return processed_data
+        return sampled_data
 
     @staticmethod
     def _analyse_internal(files, temperatures, lambdas, engine, estimator, **kwargs):
@@ -1319,6 +1312,7 @@ def _analyse_internal(files, temperatures, lambdas, engine, estimator, **kwargs)
         # Preprocess the data.
         try:
             processed_data = Relative._preprocess_data(data, estimator, **kwargs)
+            processed_data = _alchemlyb.concat(processed_data)
         except:
             _warnings.warn("Could not preprocess the data!")
             processed_data = _alchemlyb.concat(data)

python/BioSimSpace/FreeEnergy/_relative.py

@@ -277,6 +277,14 @@ def _setup(self, **kwargs):
                     "perturbable molecule!"
                 )
 
+            # Make sure the protocol is valid.
+            if self._protocol.getPerturbationType() != "full":
+                raise NotImplementedError(
+                    "AMBER currently only supports the 'full' perturbation "
+                    "type. Please use engine='SOMD' when running multistep "
+                    "perturbation types."
+                )
+
             # If this is vacuum simulation with pmemd.cuda then
             # we need to add a simulation box.
             if self._is_vacuum and self._is_pmemd_cuda:

python/BioSimSpace/Process/_amber.py

@@ -1488,13 +1488,12 @@ def getFrame(self, index):
 
         if not type(index) is int:
             raise TypeError("'index' must be of type 'int'")
-        max_index = (
-            int(
-                (self._protocol.getRunTime() / self._protocol.getTimeStep())
-                / self._protocol.getRestartInterval()
-            )
-            - 1
+
+        max_index = int(
+            (self._protocol.getRunTime() / self._protocol.getTimeStep())
+            / self._protocol.getRestartInterval()
         )
+
         if index < 0 or index > max_index:
             raise ValueError(f"'index' must be in range [0, {max_index}].")
 

python/BioSimSpace/Process/_openmm.py

@@ -42,14 +42,15 @@ class Equilibration(_Protocol, _PositionRestraintMixin):
     def __init__(
         self,
         timestep=_Types.Time(2, "femtosecond"),
-        runtime=_Types.Time(0.2, "nanoseconds"),
+        runtime=_Types.Time(1, "nanoseconds"),
         temperature_start=_Types.Temperature(300, "kelvin"),
         temperature_end=_Types.Temperature(300, "kelvin"),
         temperature=None,
         pressure=None,
         thermostat_time_constant=_Types.Time(1, "picosecond"),
-        report_interval=100,
-        restart_interval=500,
+        report_interval=200,
+        restart_interval=1000,
+        restart=False,
         restraint=None,
         force_constant=10 * _Units.Energy.kcal_per_mol / _Units.Area.angstrom2,
     ):
@@ -89,6 +90,9 @@ def __init__(
             The frequency at which restart configurations and trajectory
             frames are saved. (In integration steps.)
 
+        restart : bool
+            Whether this is a continuation of a previous simulation.
+
         restraint : str, [int]
             The type of restraint to perform. This should be one of the
             following options:
@@ -155,6 +159,9 @@ def __init__(
         # Set the restart interval.
         self.setRestartInterval(restart_interval)
 
+        # Set the restart flag.
+        self.setRestart(restart)
+
         # Set the posistion restraint.
         _PositionRestraintMixin.__init__(self, restraint, force_constant)
 
@@ -169,7 +176,7 @@ def _get_parm(self):
             f"thermostat_time_constant={self._thermostat_time_constant}, "
             f"report_interval={self._report_interval}, "
             f"restart_interval={self._restart_interval}, "
-            + _PositionRestraintMixin._get_parm(self)
+            f"restart={self._restart}, " + _PositionRestraintMixin._get_parm(self)
         )
 
     def __str__(self):
@@ -204,6 +211,7 @@ def __eq__(self, other):
             and self._thermostat_time_constant == other._thermostat_time_constant
             and self._report_interval == other._report_interval
             and self._restart_interval == other._restart_interval
+            and self._restart == other._restart
             and _PositionRestraintMixin.__eq__(self, other)
         )
 
@@ -484,6 +492,34 @@ def setRestartInterval(self, restart_interval):
 
         self._restart_interval = restart_interval
 
+    def isRestart(self):
+        """
+        Return whether this restart simulation.
+
+        Returns
+        -------
+
+        is_restart : bool
+            Whether this is a restart simulation.
+        """
+        return self._restart
+
+    def setRestart(self, restart):
+        """
+        Set the restart flag.
+
+        Parameters
+        ----------
+
+        restart : bool
+            Whether this is a restart simulation.
+        """
+        if isinstance(restart, bool):
+            self._restart = restart
+        else:
+            _warnings.warn("Non-boolean restart flag. Defaulting to False!")
+            self._restart = False
+
     def isConstantTemp(self):
         """
         Return whether the protocol has a constant temperature.

python/BioSimSpace/Protocol/_equilibration.py

@@ -44,14 +44,15 @@ def __init__(
         max_lam=1.0,
         num_lam=11,
         timestep=_Types.Time(2, "femtosecond"),
-        runtime=_Types.Time(0.2, "nanoseconds"),
+        runtime=_Types.Time(1, "nanoseconds"),
         temperature_start=_Types.Temperature(300, "kelvin"),
         temperature_end=_Types.Temperature(300, "kelvin"),
         temperature=None,
         pressure=None,
         thermostat_time_constant=_Types.Time(1, "picosecond"),
         report_interval=200,
         restart_interval=1000,
+        restart=False,
         restraint=None,
         force_constant=10 * _Units.Energy.kcal_per_mol / _Units.Area.angstrom2,
         perturbation_type="full",
@@ -106,6 +107,9 @@ def __init__(
             The frequency at which restart configurations and trajectory
             frames are saved. (In integration steps.)
 
+        restart : bool
+            Whether this is a continuation of a previous simulation.
+
         restraint : str, [int]
             The type of restraint to perform. This should be one of the
             following options:
@@ -217,4 +221,5 @@ def _to_regular_protocol(self):
             restart_interval=self.getRestartInterval(),
             restraint=self.getRestraint(),
             force_constant=self.getForceConstant(),
+            restart=self.isRestart(),
         )

python/BioSimSpace/Protocol/_free_energy_equilibration.py

@@ -46,8 +46,8 @@ def __init__(
         temperature=_Types.Temperature(300, "kelvin"),
         pressure=_Types.Pressure(1, "atmosphere"),
         thermostat_time_constant=_Types.Time(1, "picosecond"),
-        report_interval=100,
-        restart_interval=100,
+        report_interval=200,
+        restart_interval=1000,
         restart=False,
         restraint=None,
         force_constant=10 * _Units.Energy.kcal_per_mol / _Units.Area.angstrom2,

python/BioSimSpace/Protocol/_production.py

@@ -213,6 +213,6 @@ def steps(self):
             steps = self._protocol.getSteps()
         else:
             steps = _math.ceil(
-                self._protocol.getRunTime() / self._protocol.getTimeStep()
+                round(self._protocol.getRunTime() / self._protocol.getTimeStep(), 6)
             )
         return steps

python/BioSimSpace/_Config/_config.py

@@ -100,30 +100,21 @@ def createConfig(self, version=None, extra_options={}, extra_lines=[]):
             if not all(isinstance(line, str) for line in extra_lines):
                 raise TypeError("Lines in 'extra_lines' must be of type 'str'.")
 
-        # Make sure the report interval is a multiple of nstcalcenergy.
-        if isinstance(self._protocol, _FreeEnergyMixin):
-            nstcalcenergy = 250
-        else:
-            nstcalcenergy = 100
-        report_interval = self.reportInterval()
-        if report_interval % nstcalcenergy != 0:
-            report_interval = nstcalcenergy * _math.ceil(
-                report_interval / nstcalcenergy
-            )
-
         # Define some miscellaneous defaults.
         protocol_dict = {
             # Interval between writing to the log file.
-            "nstlog": report_interval,
+            "nstlog": self.reportInterval(),
             # Interval between writing to the energy file.
-            "nstenergy": report_interval,
+            "nstenergy": self.reportInterval(),
             # Interval between writing to the trajectory file.
             "nstxout-compressed": self.restartInterval(),
         }
 
         # Minimisation.
         if isinstance(self._protocol, _Protocol.Minimisation):
             protocol_dict["integrator"] = "steep"
+            # Maximum step size in nanometers.
+            protocol_dict["emstep"] = "0.001"
         else:
             # Timestep in picoseconds
             timestep = self._protocol.getTimeStep().picoseconds().value()
@@ -212,7 +203,7 @@ def createConfig(self, version=None, extra_options={}, extra_lines=[]):
         # Temperature control.
         if not isinstance(self._protocol, _Protocol.Minimisation):
             if isinstance(self._protocol, _FreeEnergyMixin):
-                # Langevin dynamics.
+                # Leap-frog stochastic dynamics integrator.
                 protocol_dict["integrator"] = "sd"
             else:
                 # Leap-frog molecular dynamics.
@@ -260,6 +251,24 @@ def createConfig(self, version=None, extra_options={}, extra_lines=[]):
                     "%.2f" % self._protocol.getTemperature().kelvin().value()
                 )
 
+            # Set as a continuation run.
+            if self.isRestart():
+                protocol_dict["continuation"] = "yes"
+                protocol_dict["gen-vel"] = "no"
+            # Generate velocities.
+            else:
+                protocol_dict["continuation"] = "no"
+                protocol_dict["gen-vel"] = "yes"
+                protocol_dict["gen-seed"] = "-1"
+                if isinstance(self._protocol, _Protocol.Equilibration):
+                    protocol_dict["gen-temp"] = (
+                        "%.2f" % self._protocol.getStartTemperature().kelvin().value()
+                    )
+                else:
+                    protocol_dict["gen-temp"] = (
+                        "%.2f" % self._protocol.getTemperature().kelvin().value()
+                    )
+
         # Free energies.
         if isinstance(self._protocol, _FreeEnergyMixin):
             # Extract the lambda array.
@@ -280,10 +289,14 @@ def createConfig(self, version=None, extra_options={}, extra_lines=[]):
             protocol_dict["couple-lambda1"] = "vdw-q"
             # Write all lambda values.
             protocol_dict["calc-lambda-neighbors"] = -1
-            # Calculate energies every 250 steps.
-            protocol_dict["nstcalcenergy"] = 250
-            # Write gradients every 250 steps.
-            protocol_dict["nstdhdl"] = 250
+            # Calculate energies at the report interval.
+            protocol_dict["nstcalcenergy"] = self.reportInterval()
+            # Write gradients at the report interval.
+            protocol_dict["nstdhdl"] = self.reportInterval()
+            # Soft-core parameters.
+            protocol_dict["sc-alpha"] = "0.30"
+            protocol_dict["sc-sigma"] = "0.25"
+            protocol_dict["sc-coul"] = "yes"
 
         # Put everything together in a line-by-line format.
         total_dict = {**protocol_dict, **extra_options}