From e762c2f95993d7188338771e5bf420e81022196e Mon Sep 17 00:00:00 2001 From: anyschaefer <45199319+anyschaefer@users.noreply.github.com> Date: Wed, 26 Jul 2023 13:35:07 +0200 Subject: [PATCH 1/4] Update integrators.py Adding the LangevinSplitting Girsanov integrator to determine the Girsanov reweighting factors g and M. --- openmmtools/integrators.py | 241 +++++++++++++++++++++++++++++++++++++ 1 file changed, 241 insertions(+) diff --git a/openmmtools/integrators.py b/openmmtools/integrators.py index 4c0bada9..cece71c1 100644 --- a/openmmtools/integrators.py +++ b/openmmtools/integrators.py @@ -1557,6 +1557,247 @@ def _add_metropolize_finish(self): self.addComputeGlobal("naccept", "ntrials - nreject") self.addComputeGlobal("shadow_work", "0") +class LangevinSplittingGirsanov(LangevinIntegrator): + r"""Creates a Langevin splitting integrator method that allows you to + output the reweighting factors on-the-fly in addition to standard + trajectory properties. + This custom integrator is based on the :class:`openmmtools.integrators.LangevinIntegrator`. + For information about the splitting rules see the documentation there. + However, not all functionalities of the :class:`LangevinIntegrator` can + be used, since the theoretical background for this method is so far only + available for selective (ABO) splitting schemes (mts integrators bias + overlaps, etc.): + + A B O -> "R V O" + A B O B A -> "R V O V R" + A O B O A -> "R O V O R" + B O A O B -> "V O R O V" + O B A B O -> "O V R V O" + + Parameters + ---------- + temperature : np.unit.Quantity compatible with kelvin, + default: 298.0*unit.kelvin Fictitious "bath" temperature + + collision_rate : np.unit.Quantity compatible with 1/picoseconds, + default: 91.0/unit.picoseconds Collision rate + + timestep : np.unit.Quantity compatible with femtoseconds, + default: 1.0*unit.femtoseconds Integration timestep + + constraint_tolerance : float, default: 1.0e-8 + Tolerance for constraint solver + + measure_shadow_work : boolean, default: False + + + References + ---------- + [Kieninger, Ghysbrecht and Keller, 2023] Girsanov reweighting for simulations + of underdamped Langevin dynamics. Theory + + Example + ------- + # Import integratiom class + >>> import LangevinSplittingIntegratorWithGirsanovPathReweighting as LSIWGPR + + # Create a splitting integrator. + >>> integrator = LSIWGPR(nstxout=1, + >>> temperature=298.0 * unit.kelvin, + >>> collision_rate=1.0 / unit.picoseconds, + >>> timestep=1.0 * unit.femtoseconds, + >>> splitting="R O V O R" , + >>> constraint_tolerance=1e-8 + >>> ) + + # Prepare openMM simulation object (not shown in detail) + ... + >>> simulation = simulation.Simulation(top.topology, system, integrator, platform) + + # Run openMM simulation + >>> simulation.step(nsteps) + """ + + def __init__(self, nstxout, *args, **kwargs): + ''' + Langevin integrator with Girsanov reweighting + based on openmmtools.integrators.LangevinIntegrator + ''' + # Check splitting input keyword + if np.array([kwargs['splitting'].split()[i] in ['A', 'B'] for i in range(len(kwargs['splitting'].split()))]).any(): + raise Exception("Please use different notation for the splitting algorithm.\n A = R\n B = V") + if kwargs['splitting'] not in self._delta_eta_table.keys(): + raise Exception("Unfortunately, the splitting algorithm is not implemented.\n Modify '_delta_eta_table' according to\n \"string of splitting elements\" : (number of random variables, expression for delta eta, update rule)\n to include.") + if self._delta_eta_table[kwargs['splitting']] == 'nan': + raise Exception("Unfortunately, the splitting algorithm is not suitable for Girsanov reweighting.") + + # Set global variables + self._nstxout = nstxout + self._splitting = kwargs['splitting'] + self._timestep = kwargs['timestep'] + + # Initialize a new CustomIntegrator + super(LangevinSplittingGirsanov, self).__init__(*args, **kwargs) + + # Extend inherited functions + def _add_global_variables(self): + """Add global variables needed for reweighting. + """ + # Add global variables defined in LangevinIntegrator + super(LangevinSplittingGirsanov, self)._add_global_variables() + # Add variables for reweighting + self._init_rwght() + + def _add_integrator_steps(self): + """Add integrator steps, with updated difference of random variables + (forces -> minus bias in force group 1). The exponent of the Euler's + number in M reweighting factors expression is printed. + """ + # Integrate + self.addUpdateContextState() + self.addComputeTemperatureDependentConstants({"sigma": "sqrt(kT/m)"}) + + # Add random number + self._get_eta() + + # Create variable for reweighting factor M + U_idx = 0 + O_idx = 0 + for i, step in enumerate(self._delta_eta_table[self._splitting][2].split()): + if step == 'U': + # Because biase needs to be updated depending on position index + U_idx +=1 + self._get_delta_eta(U_idx-1) + elif step == 'O': + O_idx +=1 + function, _ = self._step_dispatch_table[step] + # Because _add_O_step needs extra input of eta_idx + function(O_idx-1) + else: + # Like in setup function but without possibility of mts + function, _ = self._step_dispatch_table[step] + function() + + # Trick to enable sumation over the path + # for n=0 and after tau steps of a path delta gives 0 + # so the integrals for the new path are recalculate + self.addComputeGlobal("ndivtau", "n/tau") + self.addComputeGlobal("onedelta","1 - delta(ndivtau-floor(ndivtau))") + + # Random number based reweighting factor logM(\eta) + self._get_logM() + # Increase timestep n for the next integration step + self.addComputeGlobal("n", "n + 1") + + def _add_O_step(self, eta_idx): + """Add an O step (stochastic velocity update) for reweighting (use etai). + """ + # update velocities with stored eta + self.addComputePerDof("v", "(a * v) + (b * sigma * Eta{idx})".format(idx=eta_idx)) + self.addConstrainVelocities() + + # Additional functions needed to output reweighting factors + @property + def _delta_eta_table(self): + """The dictionary provides information to set up the reweighting factors. + (number of random variables, expression for delta eta, update rule). + The expression for delta eta are derived in [Kieninger and Keller, 2023]. + The timestep dependent factors a and b (d, f or d', f' in [Kieninger and Keller, 2023]) + are provided in LangevinIntegrator with the inizialisation of global variables. + To set correct units for b (dimensionless in LangevinIntegrator) one needs to multiply + by the term 'sigma*m' (sigma=sqrt(kT/m)). + """ + dispatch_table = { + "R V O" : (1, ['a/(b*sigma*m) * timestep * ff0'], "R U V O"), + "R V O V R" : (1, ['1/(b*sigma*m) * (1 + a) * timestep/2 * ff0'], "R U V O V R"), + "V R O R V" : ('nan'), + "V R O R" : ('nan'), + "R O V O R" : (2, ['a/(b*sigma*m) * timestep * ff0'], "R U O V O R"), + "V O R O V" : (2, ['a/(b*sigma*m) * timestep/2 * ff0', '1/(b*sigma*m) * timestep/2 * ff1'], "U V O R U O V"), + "O V R V O" : (2, ['1/(b*sigma*m) * timestep/2 * ff0', 'a/(b*sigma*m) * timestep/2 * ff1'], "U O V R U V O") + } + return dispatch_table + + def _init_eta(self): + """Initialize random numbers according to the splitting scheme. + """ + for i in range(self._delta_eta_table[self._splitting][0]): + self.addPerDofVariable("Eta{}".format(i),0) + + def _init_delta_eta(self): + """Initialize difference in random numbers between biased and target + simulation according to the splitting scheme. + """ + for i in range(len(self._delta_eta_table[self._splitting][1])): + self.addPerDofVariable("DeltaEta{}".format(i),0) + self.addPerDofVariable("ff{}".format(i),0) + + + def _get_eta(self): + """Add random numbers drawn from Gaussian distribution + according to the splitting scheme. + """ + for i in range(self._delta_eta_table[self._splitting][0]): + self.addComputePerDof("Eta{}".format(i),"gaussian") + + def _get_delta_eta(self, idx): + """Add the difference in random numbers between the target and target + simulation according to the splitting scheme. The forces associated + with the bias must be updated for the respective time step. An update + rule is given in _delta_eta_table()). The force of the bias must be in + force group 1. + """ + # Update forces + self.addComputePerDof("ff{}".format(idx),"f1") + # Set delta eta + self.addComputePerDof("DeltaEta{}".format(idx), self._delta_eta_table[self._splitting][1][idx]) + + def _get_logM(self): + """Add the reweighting factor M in terms of the difference in random numbers + between the target and target simulation. + """ + # AOBOA has a combined random number + # ToDo: check + if self._splitting=='R O V O R': + self.addComputePerDof("Eta0","a*Eta0+Eta1") + SOP="(Eta0 * DeltaEta0)/(a*a+1) + 0.5 * (DeltaEta0 * DeltaEta0)/(a*a+1)" + + else: + # Sum over the path (SOP) of length \tau + SOP = str() + for i in range(len(self._delta_eta_table[self._splitting][1])): + SOP+="Eta{idx} * DeltaEta{idx} + 0.5 * (DeltaEta{idx} * DeltaEta{idx})".format(idx=i) + if i+1 < len(self._delta_eta_table[self._splitting][1]): + SOP+=" + " + self.addComputeSum("SumOverPath", SOP) + self.addComputeGlobal('M', "M * onedelta + SumOverPath") + + def _init_rwght(self): + """ Initialize + """ + ## Add a variable for the timestep n and size + self.addGlobalVariable("n", 0) + self.addGlobalVariable("timestep", self._timestep) + + ## Add a variable for \tau the length of a path \omega; + ## here given by the write-out freuquency nstxout + self.addGlobalVariable("tau", self._nstxout) + + ## Add variables to enable sumation over the path + ## cf. J. Chem. Phys. 146, 244112 (2017) EQ:(29) + self.addGlobalVariable("ndivtau", 0) + self.addGlobalVariable("onedelta", 0) + + ## Abb variable give the sum over the path + ## cf. J. Chem. Phys. 146, 244112 (2017) EQ:(25) + self.addGlobalVariable("SumOverPath", 0) + self.addGlobalVariable("M", 0) + + ## Add variable for \eta and \Delta\eta needed to give reweighting factor M(\eta) + ## cf. J. Chem. Phys. 154, 094102 (2021) EQ:(10) + self._init_eta() + self._init_delta_eta() + class NonequilibriumLangevinIntegrator(LangevinIntegrator): """Nonequilibrium integrator mix-in. From 2b828dc903c9b9bf54eec7b767e4ba3577cc6b3a Mon Sep 17 00:00:00 2001 From: anyschaefer <45199319+anyschaefer@users.noreply.github.com> Date: Mon, 4 Mar 2024 16:04:25 +0100 Subject: [PATCH 2/4] Update integrators.py --- openmmtools/integrators.py | 23 ++++++++++++----------- 1 file changed, 12 insertions(+), 11 deletions(-) diff --git a/openmmtools/integrators.py b/openmmtools/integrators.py index cece71c1..5395b19a 100644 --- a/openmmtools/integrators.py +++ b/openmmtools/integrators.py @@ -1576,6 +1576,8 @@ class LangevinSplittingGirsanov(LangevinIntegrator): Parameters ---------- + nstxout : int, write out frequency of simulation data + temperature : np.unit.Quantity compatible with kelvin, default: 298.0*unit.kelvin Fictitious "bath" temperature @@ -1587,19 +1589,18 @@ class LangevinSplittingGirsanov(LangevinIntegrator): constraint_tolerance : float, default: 1.0e-8 Tolerance for constraint solver - - measure_shadow_work : boolean, default: False References ---------- + [Schaefer and Keller, 2023] Implementation of Girsanov reweighting [Kieninger, Ghysbrecht and Keller, 2023] Girsanov reweighting for simulations of underdamped Langevin dynamics. Theory Example ------- - # Import integratiom class - >>> import LangevinSplittingIntegratorWithGirsanovPathReweighting as LSIWGPR + # Import integration class + >>> import LangevinSplittingGirsanov as LSIWGPR # Create a splitting integrator. >>> integrator = LSIWGPR(nstxout=1, @@ -1649,9 +1650,9 @@ def _add_global_variables(self): self._init_rwght() def _add_integrator_steps(self): - """Add integrator steps, with updated difference of random variables - (forces -> minus bias in force group 1). The exponent of the Euler's - number in M reweighting factors expression is printed. + """Add integrator steps, with update for random variable differences + (and bias forces -> in force group 1). The exponent of reweighting + factors M is computed. """ # Integrate self.addUpdateContextState() @@ -1690,7 +1691,7 @@ def _add_integrator_steps(self): self.addComputeGlobal("n", "n + 1") def _add_O_step(self, eta_idx): - """Add an O step (stochastic velocity update) for reweighting (use etai). + """Add a O step (stochastic velocity update) for reweighting (use Eta{idx}). """ # update velocities with stored eta self.addComputePerDof("v", "(a * v) + (b * sigma * Eta{idx})".format(idx=eta_idx)) @@ -1741,7 +1742,7 @@ def _get_eta(self): self.addComputePerDof("Eta{}".format(i),"gaussian") def _get_delta_eta(self, idx): - """Add the difference in random numbers between the target and target + """Add the difference in random numbers between the target and biased simulation according to the splitting scheme. The forces associated with the bias must be updated for the respective time step. An update rule is given in _delta_eta_table()). The force of the bias must be in @@ -1753,8 +1754,8 @@ def _get_delta_eta(self, idx): self.addComputePerDof("DeltaEta{}".format(idx), self._delta_eta_table[self._splitting][1][idx]) def _get_logM(self): - """Add the reweighting factor M in terms of the difference in random numbers - between the target and target simulation. + """Add the pre-reweighting factor M in terms of the difference in random numbers + between the target and biased simulation. """ # AOBOA has a combined random number # ToDo: check From bfbe2421d44cc28b94f2f51ff68da0ee7b3235c4 Mon Sep 17 00:00:00 2001 From: anyschaefer <45199319+anyschaefer@users.noreply.github.com> Date: Mon, 4 Mar 2024 16:42:14 +0100 Subject: [PATCH 3/4] Update integrators.py --- openmmtools/integrators.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/openmmtools/integrators.py b/openmmtools/integrators.py index 5395b19a..60d06cc7 100644 --- a/openmmtools/integrators.py +++ b/openmmtools/integrators.py @@ -1593,7 +1593,7 @@ class LangevinSplittingGirsanov(LangevinIntegrator): References ---------- - [Schaefer and Keller, 2023] Implementation of Girsanov reweighting + [Schaefer and Keller, 2024] Implementation of Girsanov reweighting [Kieninger, Ghysbrecht and Keller, 2023] Girsanov reweighting for simulations of underdamped Langevin dynamics. Theory From 28b569e1fdc90ce78033c3bde71feaf3ce32aa71 Mon Sep 17 00:00:00 2001 From: anyschaefer <45199319+anyschaefer@users.noreply.github.com> Date: Fri, 26 Apr 2024 15:49:33 +0200 Subject: [PATCH 4/4] Update integrators.py add reference to publication on implementation --- openmmtools/integrators.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/openmmtools/integrators.py b/openmmtools/integrators.py index 60d06cc7..07aaa0c3 100644 --- a/openmmtools/integrators.py +++ b/openmmtools/integrators.py @@ -1593,7 +1593,7 @@ class LangevinSplittingGirsanov(LangevinIntegrator): References ---------- - [Schaefer and Keller, 2024] Implementation of Girsanov reweighting + [Schaefer and Keller, 2024] Implementation of Girsanov reweighting in OpenMM and Deeptime [Kieninger, Ghysbrecht and Keller, 2023] Girsanov reweighting for simulations of underdamped Langevin dynamics. Theory