diff --git a/openmmtools/integrators.py b/openmmtools/integrators.py index 4c0bada9..81cdfe40 100644 --- a/openmmtools/integrators.py +++ b/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. diff --git a/openmmtools/tests/test_integrators.py b/openmmtools/tests/test_integrators.py index 797f9dd9..f1adfdd8 100755 --- a/openmmtools/tests/test_integrators.py +++ b/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]: