FAHCustomNEQIntegrator

openmmtools/integrators.py

@@ -2007,6 +2007,9 @@ def _add_integrator_steps(self):
         self._add_alchemical_reset_step() # set alchemical parameters too
         self.endBlock()
 
+        #intermediate
+        self._add_antecedent()
+
         # Main body
         self.beginIfBlock("step >= 0")
         NonequilibriumLangevinIntegrator._add_integrator_steps(self) # includes H step which updates alchemical state and accumulates work
@@ -2025,6 +2028,12 @@ def _add_integrator_steps(self):
         self._add_alchemical_reset_step() # sets step to 0
         self.endBlock()
 
+    def _add_antecedent(self):
+        """
+        functionality to add stuff here
+        """
+        pass
+
     def _add_alchemical_perturbation_step(self):
         """
         Add alchemical perturbation step, accumulating protocol work.
@@ -2056,6 +2065,57 @@ def _add_alchemical_perturbation_step(self):
         self.addComputeGlobal("Enew", "energy")
         self.addComputeGlobal("protocol_work", "protocol_work + (Enew-Eold)")
 
+class FAHCustomNEQIntegrator(PeriodicNonequilibriumIntegrator):
+    """
+    Copy of the PeriodicNonequilibriumIntegrator that retains equilibrium MD at either endstate by caching equilibrium snapshots between annealing protocols
+    """
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs) #call the PeriodicNonequilibriumIntegrator init
+        self.addPerDofVariable("x0_eq_cache", 0) #initialize cache variables
+        self.addPerDofVariable("x1_eq_cache", 0)
+
+    def equip_equilibrium_cache(self, endstate, positions):
+        """
+        in order to run independent MD at the endstates, we need to equip each endstate cache with an i.i.d. snapshot
+
+        Parameters:
+        -----------
+
+        endstate : int
+            the lambda achemical endstate to which the positions will equip
+        positions : np.ndarray(N,3) (in nm implicitly)
+            array of the positions to be equipped
+        """
+        assert endstate in [0,1]
+        cache = 'x0_eq_cache' if endstate==0 else 'x1_eq_cache'
+        self.setPerDofVariableByName(cache, positions.tolist())
+
+    def _add_antecedent(self):
+        """
+        cache equilibrium snapshots and recover caches
+        """
+        #cache an equilibrium snapshot from lambda=0 (after eq)
+        self.beginIfBlock("step = n_steps_eq")
+        self.addComputePerDof("x0_eq_cache", "x")
+        self.endBlock()
+
+        #cache an equilibrium snapshot at lambda=1 (after eq)
+        self.beginIfBlock("step = n_steps_per_cycle")
+        self.addComputePerDof("x1_eq_cache", "x")
+        self.endBlock()
+
+        #recover the cache at lambda=0 (before eq)
+        self.beginIfBlock("step = 0")
+        self.addComputePerDof("x", "x0_eq_cache")
+        self.endBlock()
+
+        #recover the cache at lambda=1 (before eq)
+        self.beginIfBlock("step = n_steps_eq + n_steps_neq")
+        self.addComputePerDof("x", "x1_eq_cache")
+        self.endBlock()
+
+
+
 class ExternalPerturbationLangevinIntegrator(NonequilibriumLangevinIntegrator):
     """
     Create a LangevinSplittingIntegrator that accounts for external perturbations and tracks protocol work.

openmmtools/tests/test_integrators.py

@@ -795,7 +795,7 @@ def run_alchemical_langevin_integrator(nsteps=0, splitting="O { V R H R V } O"):
     if nsigma > NSIGMA_MAX:
         raise Exception("The free energy difference for the nonequilibrium switching for splitting '%s' and %d steps is not zero within statistical error." % (splitting, nsteps))
 
-def test_periodic_langevin_integrator(splitting="H V R O R V H", ncycles=40, nsteps_neq=1000, nsteps_eq=1000, write_trajectory=False):
+def test_periodic_langevin_integrator(splitting="H V R O R V H", ncycles=40, nsteps_neq=1000, nsteps_eq=1000, write_trajectory=False, test_FAH=False):
     """
     Test PeriodicNonequilibriumIntegrator
 
@@ -811,6 +811,8 @@ def test_periodic_langevin_integrator(splitting="H V R O R V H", ncycles=40, nst
         number of equilibration steps to run at endstates before annealing
     write_trajectory : bool, optional, default=True
         If True, will generate a PDB file that contains the harmonic oscillator trajectory
+    test_FAH : boolean, default False
+        whether to test the FAHCustomNEQIntegrator
     """
     #max deviation from the calculated free energy
     NSIGMA_MAX = 6
@@ -843,16 +845,34 @@ def test_periodic_langevin_integrator(splitting="H V R O R V H", ncycles=40, nst
     topology = testsystem.topology
 
     # Create integrator
-    from openmmtools.integrators import PeriodicNonequilibriumIntegrator
-    integrator = PeriodicNonequilibriumIntegrator(alchemical_functions=alchemical_functions,
-                                                  splitting=splitting,
-                                                  nsteps_eq=nsteps_eq,
-                                                  nsteps_neq=nsteps_neq,
-                                                  **integrator_kwargs)
+    from openmmtools.integrators import PeriodicNonequilibriumIntegrator, FAHCustomNEQIntegrator
+    if test_FAH:
+        from openmmtools.integrators import FAHCustomNEQIntegrator
+        integrator = FAHCustomNEQIntegrator(alchemical_functions=alchemical_functions,
+                                                      splitting=splitting,
+                                                      nsteps_eq=nsteps_eq,
+                                                      nsteps_neq=nsteps_neq,
+                                                      **integrator_kwargs)
+    else:
+        from openmmtools.integrators import PeriodicNonequilibriumIntegrator
+        integrator = PeriodicNonequilibriumIntegrator(alchemical_functions=alchemical_functions,
+                                                      splitting=splitting,
+                                                      nsteps_eq=nsteps_eq,
+                                                      nsteps_neq=nsteps_neq,
+                                                      **integrator_kwargs)
+
     platform = openmm.Platform.getPlatformByName("Reference")
     context = openmm.Context(system, integrator, platform)
     context.setPositions(positions)
 
+    if test_FAH: # we needto equip endstates
+        template = np.array([1., 1., 1.])
+        x0, x1 = (template * 0.*unit.angstroms).value_in_unit_system(unit.md_unit_system), (template * displacement).value_in_unit_system(unit.md_unit_system)
+        integrator.equip_equilibrium_cache(0, np.array([x0.tolist()]))
+        integrator.equip_equilibrium_cache(1, np.array([x1.tolist()]))
+
+
+
     nsteps_per_cycle = nsteps_eq + nsteps_neq + nsteps_eq + nsteps_neq
     assert integrator.getGlobalVariableByName("n_steps_per_cycle") == nsteps_per_cycle
 
@@ -887,9 +907,15 @@ def test_periodic_langevin_integrator(splitting="H V R O R V H", ncycles=40, nst
 
     step = 0
     for cycle in range(2):
+
         # eq (0)
         for i in range(nsteps_eq):
             integrator.step(1)
+            if test_FAH and i==0: #we have to assert that the equilibrium positions are appropriately set
+                state = context.getState(getPositions=True)
+                positions = state.getPositions()/unit.nanometers
+                assert np.allclose(integrator.getPerDofVariableByName('x0_eq_cache'), positions, atol=1e-1)
+
             step += 1
             assert integrator.getGlobalVariableByName("step") == (step % nsteps_per_cycle)
             assert np.isclose(integrator.getGlobalVariableByName("lambda"), 0.0)
@@ -902,6 +928,10 @@ def test_periodic_langevin_integrator(splitting="H V R O R V H", ncycles=40, nst
         # eq (1)
         for i in range(nsteps_eq):
             integrator.step(1)
+            if test_FAH and i==0: #we have to assert that the equilibrium positions are appropriately set for the
+                state = context.getState(getPositions=True)
+                positions = state.getPositions()/unit.nanometers
+                assert np.allclose(integrator.getPerDofVariableByName('x1_eq_cache'), positions, atol=1e-1)
             step += 1
             assert integrator.getGlobalVariableByName("step") == (step % nsteps_per_cycle)
             assert np.isclose(integrator.getGlobalVariableByName("lambda"), 1.0)
@@ -950,6 +980,9 @@ def test_periodic_langevin_integrator(splitting="H V R O R V H", ncycles=40, nst
     del context
     del integrator
 
+def test_FAHCustomNEQIntegrator():
+    test_periodic_langevin_integrator(splitting="H V R O R V H", ncycles=40, nsteps_neq=1000, nsteps_eq=1000, write_trajectory=False, test_FAH=True)
+
 def test_alchemical_langevin_integrator():
     for splitting in ["O V R H R V O", "H R V O V R H", "O { V R H R V } O"]:
         for nsteps in [0, 1, 10]: