Coverage for openxps/integrators/massive_ggmt.py: 81%

1""" 

2.. module:: openxps.integrators.massive_ggmt 

3 :platform: Linux, MacOS, Windows 

4 :synopsis: Massive Generalized Gaussian Moment Thermostat (GGMT) integrator. 

5 

6.. classauthor:: Charlles Abreu <craabreu@gmail.com> 

7 

8""" 

9 

10import numpy as np 

11import openmm as mm 

12from openmm import unit as mmunit 

13 

14from openxps.integrators.utils import IntegratorMixin, add_property 

15from openxps.utils import preprocess_args 

16 

17 

18@add_property("temperature", mmunit.kelvin) 

19@add_property("time constant", 1 / mmunit.picosecond) 

20@add_property("friction coefficient", 1 / mmunit.picosecond) 

21class MassiveGGMTIntegrator(IntegratorMixin, mm.CustomIntegrator): 

22 r"""A massive, second- and fourth-moment GGMT integrator :cite:`Liu2000`. 

23 

24 The massive GGMT integrator is effective at maintaining temperature control under 

25 steady-state conditions, such as in d-AFED/TAMD :cite:`Maragliano2006,Abrams2008` 

26 or UFED :cite:`Chen2012` simulations. 

27 

28 This implementation optionally introduces stochasticity by applying Langevin-type 

29 noise to the variables controlling the second- and fourth-moment of the velocity 

30 distribution. 

31 

32 .. note:: 

33 If stochasticity is enabled, the friction coefficient is initially set to the 

34 inverse of the time constant, but this value can be customized via 

35 :meth:`setFrictionCoefficient`. 

36 

37 .. warning:: 

38 This integrator is incompatible with systems subject to constraints. 

39 

40 Parameters 

41 ---------- 

42 temperature 

43 The temperature of the heat bath. 

44 timeConstant 

45 The time constant (a.k.a. coupling/damping/relaxation time) of the thermostat. 

46 stepSize 

47 The integration step size. 

48 

49 Keyword Arguments 

50 ----------------- 

51 stochastic 

52 If True, the integrator will apply stochasticity to the variables controlling 

53 the second- and fourth-moment of the velocity distribution. 

54 forceFirst 

55 If True, the integrator will apply a force-first scheme. Otherwise, it will 

56 apply a symmetric operator splitting scheme. 

57 

58 Example 

59 ------- 

60 >>> import openxps as xps 

61 >>> from openmm import unit 

62 >>> from openmmtools import testsystems 

63 

64 Symmetric scheme (default) 

65 

66 >>> integrator = xps.integrators.MassiveGGMTIntegrator( 

67 ... temperature=300 * unit.kelvin, 

68 ... timeConstant=100 * unit.femtoseconds, 

69 ... stepSize=2 * unit.femtoseconds, 

70 ... ) 

71 >>> integrator 

72 Per-dof variables: 

73 v1, v2, v2lb 

74 Global variables: 

75 kT = 2.494338... 

76 invQ = 40.090... 

77 invQ2 = 2.416... 

78 Computation steps: 

79 0: allow forces to update the context state 

80 1: v <- v + 0.5*dt*f/m 

81 2: x <- x + 0.5*dt*v 

82 3: v1 <- v1 + 0.5*dt*(m*v^2 - kT)*invQ 

83 4: v2lb <- -sqrt(v2*v2 + (mvv + e*e/mvv - 2*e)*invQ2); mvv = m*v*v; e=sqrt(3)*kT 

84 5: v2 <- select(step(new_v2 - v2lb), new_v2, 2*v2lb - new_v2); new_v2 = ... 

85 6: v <- v/sqrt(1 + 0.33...*dt*m*v^2*v2) 

86 7: v2 <- select(step(new_v2 - v2lb), new_v2, 2*v2lb - new_v2); new_v2 = ... 

87 8: v <- v*exp(-1*dt*(v1 + kT*v2)) 

88 9: v2lb <- -sqrt(v2*v2 + (mvv + e*e/mvv - 2*e)*invQ2); mvv = m*v*v; e=sqrt(3)*kT 

89 10: v2 <- select(step(new_v2 - v2lb), new_v2, 2*v2lb - new_v2); new_v2 = ... 

90 11: v <- v/sqrt(1 + 0.33...*dt*m*v^2*v2) 

91 12: v2 <- select(step(new_v2 - v2lb), new_v2, 2*v2lb - new_v2); new_v2 = ... 

92 13: v1 <- v1 + 0.5*dt*(m*v^2 - kT)*invQ 

93 14: x <- x + 0.5*dt*v 

94 15: v <- v + 0.5*dt*f/m 

95 

96 Force-first scheme 

97 

98 >>> integrator_ff = xps.integrators.MassiveGGMTIntegrator( 

99 ... temperature=300 * unit.kelvin, 

100 ... timeConstant=40 * unit.femtoseconds, 

101 ... stepSize=2 * unit.femtoseconds, 

102 ... forceFirst=True 

103 ... ) 

104 >>> model = testsystems.AlanineDipeptideVacuum() 

105 >>> try: 

106 ... integrator_ff.registerWithSystem(model.system, False) 

107 ... except ValueError as e: 

108 ... print(e) 

109 Massive GGMT integrators do not support constraints. 

110 """ 

111 

112 @preprocess_args 

113 def __init__( 

114 self, 

115 temperature: mmunit.Quantity, 

116 timeConstant: mmunit.Quantity, 

117 stepSize: mmunit.Quantity, 

118 *, 

119 stochastic: bool = False, 

120 forceFirst: bool = False, 

121 ) -> None: 

122 super().__init__(stepSize) 

123 self._forceFirst = forceFirst 

124 self._init_temperature(temperature) 

125 self._init_time_constant(timeConstant) 

126 self._init_friction_coefficient(1 / timeConstant if stochastic else 0.0) 

127 self._stochastic = stochastic 

128 self._add_variables() 

129 self.addUpdateContextState() 

130 self._add_boost(1 if forceFirst else 0.5) 

131 self._add_translation(0.5) 

132 self._add_thermostat(1) 

133 self._add_translation(0.5) 

134 if not forceFirst: 

135 self._add_boost(0.5) 

136 

137 def _add_variables(self) -> None: 

138 self.addGlobalVariable("kT", 0) 

139 self.addGlobalVariable("invQ", 0) 

140 self.addGlobalVariable("invQ2", 0) 

141 self.addPerDofVariable("v1", 0) 

142 self.addPerDofVariable("v2", 0) 

143 self.addPerDofVariable("v2lb", 0) 

144 if self._stochastic: 

145 self.addGlobalVariable("a", 0) 

146 self.addGlobalVariable("b1", 0) 

147 self.addGlobalVariable("b2", 0) 

148 self._update_global_variables() 

149 

150 def _update_global_variables(self) -> None: 

151 tau = self.getTimeConstant() 

152 kt = mmunit.MOLAR_GAS_CONSTANT_R * self.getTemperature() 

153 self.setGlobalVariableByName("kT", kt) 

154 self.setGlobalVariableByName("invQ", 1 / (kt * tau**2)) 

155 self.setGlobalVariableByName("invQ2", 3 / (8 * kt**3 * tau**2)) 

156 if self._stochastic: 

157 friction = self.getFrictionCoefficient() 

158 dt = self.getStepSize() 

159 a = np.exp(-friction * dt) 

160 b1 = np.sqrt(1 - a**2) / tau 

161 b2 = b1 * np.sqrt(3 / 8) / kt 

162 self.setGlobalVariableByName("a", a) 

163 self.setGlobalVariableByName("b1", b1) 

164 self.setGlobalVariableByName("b2", b2) 

165 

166 def _add_translation(self, fraction: float) -> None: 

167 self.addComputePerDof("x", f"x + {fraction}*dt*v") 

168 

169 def _add_boost(self, fraction: float) -> None: 

170 self.addComputePerDof("v", f"v + {fraction}*dt*f/m") 

171 

172 def _add_v1_boost(self, fraction: float) -> None: 

173 self.addComputePerDof("v1", f"v1 + {fraction}*dt*(m*v^2 - kT)*invQ") 

174 

175 def _add_v2_boost(self, fraction: float) -> None: 

176 self.addComputePerDof( 

177 "v2lb", 

178 "-sqrt(v2*v2 + (mvv + e*e/mvv - 2*e)*invQ2); mvv = m*v*v; e=sqrt(3)*kT", 

179 ) 

180 expression = ( 

181 "select(step(new_v2 - v2lb), new_v2, 2*v2lb - new_v2)" 

182 f"; new_v2 = v2 + {fraction / 2}*dt*((m*v^2)^2/3 - kT^2)*invQ2" 

183 ) 

184 self.addComputePerDof("v2", expression) 

185 self.addComputePerDof("v", f"v/sqrt(1 + {(2 * fraction) / 3}*dt*m*v^2*v2)") 

186 self.addComputePerDof("v2", expression) 

187 

188 def _add_v_scaling(self, fraction: float) -> None: 

189 self.addComputePerDof("v", f"v*exp(-{fraction}*dt*(v1 + kT*v2))") 

190 

191 def _add_thermostat(self, fraction: float) -> None: 

192 self._add_v1_boost(0.5 * fraction) 

193 self._add_v2_boost(0.5 * fraction) 

194 if self._stochastic: 

195 self._add_v_scaling(0.5 * fraction) 

196 self.addComputePerDof("v1", "a*v1 + b1*gaussian") 

197 self.addComputePerDof("v2", "a*v2 + b2*gaussian") 

198 self._add_v_scaling(0.5 * fraction) 

199 else: 

200 self._add_v_scaling(fraction) 

201 self._add_v2_boost(0.5 * fraction) 

202 self._add_v1_boost(0.5 * fraction) 

203 

204 def registerWithSystem(self, system: mm.System, isExtension: bool) -> None: 

205 if system.getNumConstraints() > 0: 

206 raise ValueError("Massive GGMT integrators do not support constraints.")