SepTop: Addressing issues raised in review

docs/guide/protocols/septop.rst

@@ -34,7 +34,10 @@ Orientational restraints
 
 Orientational, or Boresch-style, restraints are automaticallly (unless manually specified) applied between three protein and three ligand atoms using one bond,
 two angle, and three dihedral restraints. Reference atoms are picked based on different criteria, such as the root mean squared
-fluctuation of the atoms in a short MD simulation, the secondary structure of the protein, and the distance between atoms, based on heuristics from Baumann et al. [2]_.
+fluctuation of the atoms in a short MD simulation, the secondary structure of the protein, and the distance between atoms,
+based on heuristics from Baumann et al. [2]_ and Alibay et al. [3]_.
+Two strategies for selecting protein atoms are available, either picking atoms that are bonded to each other or that can span multiple residues.
+This can be specified using the ``restraint_settings.anchor_finding_strategy`` settings.
 
 Partial annihilation scheme
 ~~~~~~~~~~~~~~~~~~~~~~~~~~~
@@ -48,24 +51,30 @@ The lambda schedule
 Molecular interactions are modified along an alchemical path using a discrete set of lambda windows.
 For the transformation of ligand A to ligand B in the binding site, the following steps are carried out, starting with ligand A being fully interacting in the binding site while ligand B is decoupled.
 
-1. Insert the non-interacting dummy ligand B into the binding site and restrain it using orientational restraints. The contribution of the restraints is calculated analytically.
+1. Restrain the non-interacting dummy ligand B in the binding site. The contribution of the restraints is calculated analytically.
 2. Turn on the van der Waals (vdW) interactions of ligand B while also turning on orientational restraints on ligand A.
 3. Turn on the electrostatic interactions of ligand B while at the same time turning off the electrostatics of ligand A.
 4. Turn off vdW interactions of ligand A while simultaneously releasing restraints on ligand B.
-5. Release the restraints of the now dummy ligand A analytically and transfer the ligand into the solvent.
+5. Release the restraints of the now dummy ligand A analytically.
 
 The lambda schedule in the solvent phase is similar to the one in the complex, except that a single harmonic distance restraint is
 applied between the respective central atom in the two ligands to keep the ligands apart while doing the alchemical transformation.
-A soft-core potential from Beutler et al. [3]_ is applied to the Lennard-Jones potential to avoid instablilites in intermediate lambda windows.
+A soft-core potential from Beutler et al. [4]_ is applied to the Lennard-Jones potential to avoid instablilites in intermediate lambda windows.
 The lambda schedule is defined in the ``lambda_settings`` objects ``lambda_elec_A``, ``lambda_elec_B``,  ``lambda_vdw_A``, ``lambda_vdw_B``,
 ``lambda_restraints_A``, and ``lambda_restraints_B``.
 
 Simulation overview
 ~~~~~~~~~~~~~~~~~~~
 
-The :class:`.ProtocolDAG` of the :class:`SepTopProtocol <.SepTopProtocol>` contains :class:`.ProtocolUnit`\ s from both the complex and solvent transformations.
-This means that both legs of the thermodynamic cycle are constructed and run sequentially in the same :class:`.ProtocolDAG`. This is different from the :class:`.RelativeHybridTopologyProtocol` where the :class:`.ProtocolDAG` only runs a single leg of a thermodynamic cycle.
-If multiple ``protocol_repeats`` are run (default: ``protocol_repeats=3``), the :class:`.ProtocolDAG` contains multiple :class:`.ProtocolUnit`\ s of both complex and solvent transformations.
+The :class:`.ProtocolDAG` of the :class:`SepTopProtocol <.SepTopProtocol>` contains :class:`.ProtocolUnit`\ s from both
+the complex and solvent transformations.
+This means that both legs of the thermodynamic cycle are constructed and run sequentially in the same
+:class:`.ProtocolDAG`. This is different from the :class:`.RelativeHybridTopologyProtocol` where the
+:class:`.ProtocolDAG` only runs a single leg of a thermodynamic cycle.
+If multiple ``protocol_repeats`` are run (default: ``protocol_repeats=3``), the :class:`.ProtocolDAG`
+contains multiple :class:`.ProtocolUnit`\ s of both complex and solvent transformations.
+In that case, every :class:`.ProtocolUnit` would be run N times, where N is the number of ``protocol_repeats``. This means that also the
+selection of the atoms for restraints would be performed multiple times.
 
 Simulation steps
 """"""""""""""""
@@ -90,7 +99,7 @@ Simulation details
 
 Here are some details of how the simulation is carried out which are not detailed in the :class:`SepTopProtocol <.SepTopProtocol>`:
 
-* The protocol applies a `LangevinMiddleIntegrator <https://openmmtools.readthedocs.io/en/latest/api/generated/openmmtools.mcmc.LangevinDynamicsMove.html>`_ which uses Langevin dynamics, with the LFMiddle discretization [4]_.
+* The protocol applies a `LangevinMiddleIntegrator <https://openmmtools.readthedocs.io/en/latest/api/generated/openmmtools.mcmc.LangevinDynamicsMove.html>`_ which uses Langevin dynamics, with the LFMiddle discretization [5]_.
 * A `Monte Carlo barostat <https://docs.openmm.org/latest/api-python/generated/openmm.openmm.MonteCarloBarostat.html>`_ is used in the NPT ensemble to maintain constant pressure.
 
 Getting the free energy estimate
@@ -127,5 +136,6 @@ References
 
 .. [1] Separated topologies--a method for relative binding free energy calculations using orientational restraints, G. Rocklin, D. Mobley, K. Dill;  Chem Phys, 2013; 138(8):085104. doi: 10.1063/1.4792251.
 .. [2] Broadening the Scope of Binding Free Energy Calculations Using a Separated Topologies Approach, H. Baumann, E. Dybeck, C. McClendon, F. Pickard IV, V. Gapsys, L. Pérez-Benito, D. Hahn, G. Tresadern, A. Mathiowetz, D. Mobley, J. Chem. Theory Comput., 2023, 19, 15, 5058–5076
-.. [3] Avoiding singularities and numerical instabilities in free energy calculations based on molecular simulations, T.C. Beutler, A.E. Mark, R.C. van Schaik, P.R. Greber, and W.F. van Gunsteren, Chem. Phys. Lett., 222 529–539 (1994)
-.. [4] Unified Efficient Thermostat Scheme for the Canonical Ensemble with Holonomic or Isokinetic Constraints via Molecular Dynamics, Zhijun Zhang, Xinzijian Liu, Kangyu Yan, Mark E. Tuckerman, and Jian Liu, J. Phys. Chem. A 2019, 123, 28, 6056-6079
+.. [3] Evaluating the use of absolute binding free energy in the fragment optimisation process, I. Alibay, A. Magarkar, D. Seeliger, P. Biggin, Commun Chem 5, 105 (2022)
+.. [4] Avoiding singularities and numerical instabilities in free energy calculations based on molecular simulations, T.C. Beutler, A.E. Mark, R.C. van Schaik, P.R. Greber, and W.F. van Gunsteren, Chem. Phys. Lett., 222 529–539 (1994)
+.. [5] Unified Efficient Thermostat Scheme for the Canonical Ensemble with Holonomic or Isokinetic Constraints via Molecular Dynamics, Zhijun Zhang, Xinzijian Liu, Kangyu Yan, Mark E. Tuckerman, and Jian Liu, J. Phys. Chem. A 2019, 123, 28, 6056-6079

openfe/protocols/openmm_septop/base.py

@@ -47,7 +47,6 @@
 from openff.toolkit.topology import Molecule as OFFMolecule
 from openff.units import unit
 from openff.units.openmm import ensure_quantity, from_openmm, to_openmm
-from openmm import app
 from openmm import unit as omm_unit
 from openmmforcefields.generators import SystemGenerator
 from openmmtools import multistate
@@ -445,6 +444,41 @@ def _get_system_generator(
             )
         return system_generator
 
+    @staticmethod
+    def check_assign_velocities_with_virtual_site(
+        system: openmm.System,
+        integrator_settings: IntegratorSettings,
+    ):
+        """
+        Virtual sites sanity check - ensure we restart velocities when
+        there are virtual sites in the system
+
+        Parameters
+        ----------
+        system: openmm.System
+            The OpenMM System being checked.
+        integrator_settings: IntegratorSettings
+           Multistate sampler integrator settings.
+
+        Raises
+        ------
+        ValueError
+        * If the system has a virtual site but the integrator settings specify
+          that velocities are not reassigned.
+        """
+        # Virtual sites sanity check - ensure we restart velocities when
+        # there are virtual sites in the system
+        if integrator_settings.reassign_velocities:
+            return
+
+        for ix in range(system.getNumParticles()):
+            if system.isVirtualSite(ix):
+                errmsg = (
+                    "Simulations with virtual sites without velocity "
+                    "reassignments are unstable in openmmtools"
+                )
+                raise ValueError(errmsg)
+
     @staticmethod
     def _assign_partial_charges(
         partial_charge_settings: OpenFFPartialChargeSettings,
@@ -478,7 +512,7 @@ def _get_modeller(
         smc_components: dict[SmallMoleculeComponent, OFFMolecule],
         system_generator: SystemGenerator,
         solvation_settings: BaseSolvationSettings,
-    ) -> tuple[app.Modeller, dict[Component, npt.NDArray]]:
+    ) -> tuple[openmm.app.Modeller, dict[Component, npt.NDArray]]:
         """
         Get an OpenMM Modeller object and a list of residue indices
         for each component in the system.
@@ -502,7 +536,7 @@ def _get_modeller(
 
         Returns
         -------
-        system_modeller : app.Modeller
+        system_modeller : openmm.app.Modeller
           OpenMM Modeller object generated from ProteinComponent and
           OpenFF Molecules.
         comp_resids : dict[Component, npt.NDArray]
@@ -614,7 +648,6 @@ def get_smc_comps(
         dict[SmallMoleculeComponent, OFFMolecule],
         dict[SmallMoleculeComponent, OFFMolecule],
         dict[SmallMoleculeComponent, OFFMolecule],
-        dict[SmallMoleculeComponent, OFFMolecule],
     ]:
         # Get smcs for the different states and the common smcs
         smc_off_A = {m: m.to_openff() for m in alchem_comps["stateA"]}
@@ -629,7 +662,7 @@ def get_smc_comps(
         smc_comps_B = smc_off_B | smc_off_both
         smc_comps_AB = smc_off_A | smc_off_B | smc_off_both
 
-        return smc_comps_A, smc_comps_B, smc_comps_AB, smc_off_B
+        return smc_comps_A, smc_comps_B, smc_comps_AB
 
     def get_system(
         self,
@@ -652,10 +685,10 @@ def get_system(
 
         Returns
         -------
-        omm_system: app.System
-        omm_topology: app.Topology
+        omm_system: openmm.app.System
+        omm_topology: openmm.app.Topology
         positions: openmm.unit.Quantity
-        system_modeller: app.Modeller
+        system_modeller: openmm.app.Modeller
         comp_resids: dict[Component, npt.NDArray]
           A dictionary of residues for each component in the System.
         """
@@ -843,7 +876,7 @@ def _get_states(
 
     def _get_reporter(
         self,
-        topology: app.Topology,
+        topology: openmm.app.Topology,
         positions: openmm.unit.Quantity,
         simulation_settings: MultiStateSimulationSettings,
         output_settings: MultiStateOutputSettings,
@@ -853,7 +886,7 @@ def _get_reporter(
 
         Parameters
         ----------
-        topology : app.Topology
+        topology : openmm.app.Topology
           A Topology of the system being created.
         positions : openmm.unit.Quantity
           Positions of the pre-alchemical simulation system.

openfe/protocols/openmm_septop/equil_septop_method.py

@@ -14,16 +14,17 @@
 Current limitations
 -------------------
 
+* Transformations that involve net charge changes are currently not supported.
+  The ligands must have the same net charge.
 * Only small molecules are allowed to act as alchemical molecules.
   Alchemically changing protein or solvent components would induce
   perturbations which are too large to be handled by this Protocol.
 
 
 Acknowledgements
 ----------------
-This Protocol is based on, and leverages components originating from
-the SepTop implementation from the Mobleylab
-(https://github.com/MobleyLab/SeparatedTopologies) as well as
+This Protocol is based on and inspired by the SepTop implementation from
+the Mobleylab (https://github.com/MobleyLab/SeparatedTopologies) as well as
 femto (https://github.com/Psivant/femto).
 
 """
@@ -989,7 +990,7 @@ def _default_settings(cls):
             complex_equil_output_settings=SepTopEquilOutputSettings(
                 equil_nvt_structure=None,
                 equil_npt_structure="equil_npt",
-                production_trajectory_filename="equil_npt",
+                production_trajectory_filename="equil_production",
                 log_output="equil_simulation",
             ),
             complex_simulation_settings=MultiStateSimulationSettings(
@@ -1220,14 +1221,6 @@ def _create(
             self.settings.solvent_solvation_settings
         )
 
-        # Make sure that we have the full system for restraint trajectory analysis
-        if self.settings.complex_equil_output_settings.output_indices != "all":
-            errmsg = (
-                "Complex simulations need to output the full system "
-                "during equilibration simulations."
-            )
-            raise ValueError(errmsg)
-
         # Validate protein component
         system_validation.validate_protein(stateA)
 
@@ -1773,7 +1766,7 @@ def run(
         # 1. Get components
         self.logger.info("Creating and setting up the OpenMM systems")
         alchem_comps, solv_comp, prot_comp, smc_comps = self._get_components()
-        smc_comps_A, smc_comps_B, smc_comps_AB, smc_off_B = self.get_smc_comps(
+        smc_comps_A, smc_comps_B, smc_comps_AB = self.get_smc_comps(
             alchem_comps, smc_comps
         )
 
@@ -1802,13 +1795,18 @@ def run(
             )
         )
 
+        smc_B_unique_keys = smc_comps_B.keys() - smc_comps_A.keys()
+        smc_comp_B_unique = {key: smc_comps_B[key] for key in smc_B_unique_keys}
         omm_system_AB, omm_topology_AB, positions_AB, modeller_AB = self.get_system_AB(
             solv_comp,
             modeller_A,
             smc_comps_AB,
-            smc_off_B,
+            smc_comp_B_unique,
             settings,
         )
+        # Virtual sites sanity check - ensure we restart velocities when
+        # there are virtual sites in the system
+        self.check_assign_velocities_with_virtual_site(omm_system_AB, settings["integrator_settings"])
 
         # Get the comp_resids of the AB system
         resids_A = list(itertools.chain(*comp_resids_A.values()))
@@ -2124,7 +2122,7 @@ def run(
         # 1. Get components
         self.logger.info("Creating and setting up the OpenMM systems")
         alchem_comps, solv_comp, prot_comp, smc_comps = self._get_components()
-        smc_comps_A, smc_comps_B, smc_comps_AB, smc_off_B = self.get_smc_comps(
+        smc_comps_A, smc_comps_B, smc_comps_AB = self.get_smc_comps(
             alchem_comps, smc_comps
         )
 
@@ -2151,6 +2149,9 @@ def run(
                 settings,
             )
         )
+        # Virtual sites sanity check - ensure we restart velocities when
+        # there are virtual sites in the system
+        self.check_assign_velocities_with_virtual_site(omm_system_AB, settings["integrator_settings"])
 
         # 6. Get atom indices for ligand A and ligand B and the solvent in the
         # system AB

openfe/protocols/openmm_septop/equil_septop_settings.py

@@ -258,6 +258,12 @@ def must_be_monotonic(cls, v):
 class SepTopEquilOutputSettings(MDOutputSettings):
     # reporter settings
     output_indices = "all"
+    """
+    Selection string for which part of the system to write coordinates for.
+    The SepTop protocol enforces "all" since the full system output is
+    required in the complex leg.
+    Default "all". 
+    """
     production_trajectory_filename: Optional[str] = "simulation"
     """
     Basename for the path to the storage file for analysis. The protocol will
@@ -300,14 +306,22 @@ class SepTopEquilOutputSettings(MDOutputSettings):
     files of the respective endstates. Default 'simulation'.
     """
 
+    @validator("output_indices")
+    def must_be_all(cls, v):
+        if v != "all":
+            errmsg = ("Equilibration simulations need to output the full "
+                      f"system, got {v}.")
+            raise ValueError(errmsg)
+        return v
+
 
 class SepTopSettings(SettingsBaseModel):
     """
-    Configuration object for ``AbsoluteSolvationProtocol``.
+    Configuration object for ``SepTopProtocol``.
 
     See Also
     --------
-    openfe.protocols.openmm_afe.AbsoluteSolvationProtocol
+    openfe.protocols.openmm_septop.SepTopProtocol
     """
 
     protocol_repeats: int
@@ -414,3 +428,4 @@ def must_be_positive(cls, v):
     """
     Settings for the Boresch restraints in the complex
     """
+