Coverage for ufedmm/integrators.py: 76%

1""" 

2.. module:: integrators 

3 :platform: Unix, Windows 

4 :synopsis: Unified Free Energy Dynamics Integrators 

5 

6.. moduleauthor:: Charlles Abreu <abreu@eq.ufrj.br> 

7""" 

8 

9import numpy as np 

10import openmm 

11from openmm import unit 

12 

13from ufedmm.ufedmm import _standardized 

14 

15 

16def add_inner_nonbonded_force(system, inner_switch, inner_cutoff, force_group_index): 

17 """ 

18 Adds a new force group to an :OpenMM:`System` containing a :OpenMM:`NonbondedForce` 

19 object, with the purpose of performing multiple time-scale integration via the 

20 RESPA2 splitting scheme of Morrone, Zhou, and Berne :cite:`Morrone_2010`. Also, 

21 assigns the provided `force_group_index` to this new group and `force_group_index+1` 

22 to the original :OpenMM:`NonbondedForce`. When used in an instance of any subclass 

23 of :class:`AbstractMiddleRespaIntegrator`, the new force group must be set as 

24 embodied by the :OpenMM:`NonbondedForce` as opposed to being complimentary to it. 

25 

26 .. warning: 

27 The new force group is not intended to contribute to the system energy. Its sole 

28 purpose is to provide a smooth, short-range force calculator for some 

29 intermediary time scale in a RESPA-type integration. 

30 

31 Parameters 

32 ---------- 

33 system : openmm.System 

34 The system the inner force will be added to, which must contain a 

35 :OpenMM:`NonbondedForce`. 

36 inner_switch : float or unit.Quantity 

37 The inner switching distance, where the interaction of an atom pair begins 

38 to switch off to zero. 

39 inner_cutoff : float or unit.Quantity 

40 The inner cutoff distance, where the interaction of an atom pairs completely 

41 switches off. 

42 force_group_index : int 

43 The force group the new interactions will belong to. The old 

44 :OpenMM:`NonbondedForce` will be automatically assigned to 

45 `force_group_index+1`. 

46 

47 Example 

48 ------- 

49 >>> import ufedmm 

50 >>> from openmm import unit 

51 >>> dt = 2*unit.femtoseconds 

52 >>> temp = 300*unit.kelvin 

53 >>> tau = 10*unit.femtoseconds 

54 >>> gamma = 10/unit.picoseconds 

55 >>> model = ufedmm.AlanineDipeptideModel() 

56 >>> ufedmm.add_inner_nonbonded_force( 

57 ... model.system, 5*unit.angstroms, 8*unit.angstroms, 1 

58 ... ) 

59 """ 

60 if openmm.__version__ < "7.5": 

61 raise Exception("add_inner_nonbonded_force requires OpenMM version >= 7.5") 

62 try: 

63 nonbonded_force = next( 

64 filter(lambda f: isinstance(f, openmm.NonbondedForce), system.getForces()) 

65 ) 

66 except StopIteration: 

67 raise Exception("add_inner_nonbonded_force requires system with NonbondedForce") 

68 if ( 

69 nonbonded_force.getNumParticleParameterOffsets() > 0 

70 or nonbonded_force.getNumExceptionParameterOffsets() > 0 

71 ): 

72 raise Exception("add_inner_nonbonded_force does not support parameter offsets") 

73 periodic = nonbonded_force.usesPeriodicBoundaryConditions() 

74 rs = _standardized(inner_switch) 

75 rc = _standardized(inner_cutoff) 

76 a = rc + rs 

77 b = rc * rs 

78 c = (30 / (rc - rs) ** 5) * np.array([b**2, -2 * a * b, a**2 + 2 * b, -2 * a, 1]) 

79 f0s = sum([c[n] * rs ** (n + 1) / (n + 1) for n in range(5)]) 

80 

81 def coeff(n, m): 

82 return c[m - 1] if m == n else c[m - 1] / (m - n) 

83 

84 def func(n, m): 

85 return "*log(r)" if m == n else (f"*r^{m-n}" if m > n else f"/r^{n-m}") 

86 

87 def val(n, m): 

88 return f0s if m == 0 else (coeff(n, m) - coeff(0, m) if n != m else coeff(n, m)) 

89 

90 def sgn(n, m): 

91 return "+" if m > 0 and val(n, m) >= 0 else "" 

92 

93 def S(n): 

94 return "".join(f"{sgn(n, m)}{val(n, m)}{func(n, m)}" for m in range(6)) 

95 

96 potential = "eps4*((sigma/r)^12-(sigma/r)^6)+Qprod/r" 

97 potential += ( 

98 f"+step(r-{rs})*(eps4*(sigma^12*({S(12)})-sigma^6*({S(6)}))+Qprod*({S(1)}))" 

99 ) 

100 mixing_rules = "; Qprod=Q1*Q2" 

101 mixing_rules += "; sigma=halfsig1+halfsig2" 

102 mixing_rules += "; eps4=sqrt4eps1*sqrt4eps2" 

103 force = openmm.CustomNonbondedForce(potential + mixing_rules) 

104 for parameter in ["Q", "halfsig", "sqrt4eps"]: 

105 force.addPerParticleParameter(parameter) 

106 force.setNonbondedMethod( 

107 force.CutoffPeriodic if periodic else force.CutoffNonPeriodic 

108 ) 

109 force.setCutoffDistance(inner_cutoff) 

110 force.setUseLongRangeCorrection(False) 

111 ONE_4PI_EPS0 = 138.93545764438198 

112 for index in range(nonbonded_force.getNumParticles()): 

113 charge, sigma, epsilon = map( 

114 _standardized, nonbonded_force.getParticleParameters(index) 

115 ) 

116 force.addParticle( 

117 [charge * np.sqrt(ONE_4PI_EPS0), sigma / 2, np.sqrt(4 * epsilon)] 

118 ) 

119 non_exclusion_exceptions = [] 

120 for index in range(nonbonded_force.getNumExceptions()): 

121 i, j, q1q2, sigma, epsilon = nonbonded_force.getExceptionParameters(index) 

122 q1q2, sigma, epsilon = map(_standardized, [q1q2, sigma, epsilon]) 

123 force.addExclusion(i, j) 

124 if q1q2 != 0.0 or epsilon != 0.0: 

125 non_exclusion_exceptions.append( 

126 (i, j, q1q2 * ONE_4PI_EPS0, sigma, 4 * epsilon) 

127 ) 

128 force.setForceGroup(force_group_index) 

129 system.addForce(force) 

130 if non_exclusion_exceptions: 

131 exceptions = openmm.CustomBondForce(f"step({rc}-r)*({potential})") 

132 for parameter in ["Qprod", "sigma", "eps4"]: 

133 exceptions.addPerBondParameter(parameter) 

134 for i, j, Qprod, sigma, eps4 in non_exclusion_exceptions: 

135 exceptions.addBond(i, j, [Qprod, sigma, eps4]) 

136 exceptions.setForceGroup(force_group_index) 

137 system.addForce(exceptions) 

138 nonbonded_force.setForceGroup(force_group_index + 1) 

139 

140 

141class CustomIntegrator(openmm.CustomIntegrator): 

142 """An extension of :OpenMM:`CustomIntegrator` class with an extra per-dof variable 

143 named `temperature`, whose content is the temperature of the heat bath associated to 

144 each degree of freedom. A per-dof temperature is necessary if the extended-space 

145 variables and the physical system are coupled adiabatically to thermostats at 

146 different temperatures. Otherwise, any other OpenMM integrator can be used. 

147 

148 Parameters 

149 ---------- 

150 temperature : float or unit.Quantity 

151 The temperature. 

152 step_size : float or unit.Quantity 

153 The step size with which to integrate the equations of motion. 

154 """ 

155 

156 def __init__(self, temperature, step_size): 

157 super().__init__(step_size) 

158 self.temperature = temperature 

159 self.addPerDofVariable("kT", unit.MOLAR_GAS_CONSTANT_R * temperature) 

160 self._up_to_date = False 

161 

162 def __repr__(self): 

163 """A human-readable version of each integrator step (adapted from openmmtools) 

164 

165 Returns 

166 ------- 

167 readable_lines : str 

168 A list of human-readable versions of each step of the integrator 

169 """ 

170 readable_lines = [] 

171 

172 self.getNumPerDofVariables() > 0 and readable_lines.append("Per-dof variables:") 

173 per_dof = [] 

174 for index in range(self.getNumPerDofVariables()): 

175 per_dof.append(self.getPerDofVariableName(index)) 

176 readable_lines.append(" " + ", ".join(per_dof)) 

177 

178 self.getNumGlobalVariables() > 0 and readable_lines.append("Global variables:") 

179 for index in range(self.getNumGlobalVariables()): 

180 name = self.getGlobalVariableName(index) 

181 value = self.getGlobalVariable(index) 

182 readable_lines.append(f" {name} = {value}") 

183 

184 readable_lines.append("Computation steps:") 

185 

186 step_type_str = [ 

187 "{target} <- {expr}", 

188 "{target} <- {expr}", 

189 "{target} <- sum({expr})", 

190 "constrain positions", 

191 "constrain velocities", 

192 "allow forces to update the context state", 

193 "if ({expr}):", 

194 "while ({expr}):", 

195 "end", 

196 ] 

197 indent_level = 0 

198 for step in range(self.getNumComputations()): 

199 line = "" 

200 step_type, target, expr = self.getComputationStep(step) 

201 if step_type == 8: 

202 indent_level -= 1 

203 command = step_type_str[step_type].format(target=target, expr=expr) 

204 line += "{:4d}: ".format(step) + " " * indent_level + command 

205 if step_type in [6, 7]: 

206 indent_level += 1 

207 readable_lines.append(line) 

208 return "\n".join(readable_lines) 

209 

210 def update_temperatures(self, system_temperature, extended_space_temperatures): 

211 nparticles = len(self.getPerDofVariableByName("kT")) - len( 

212 extended_space_temperatures 

213 ) 

214 temperatures = [system_temperature] * nparticles + extended_space_temperatures 

215 kT = [ 

216 unit.MOLAR_GAS_CONSTANT_R * T * openmm.Vec3(1, 1, 1) for T in temperatures 

217 ] 

218 self.setPerDofVariableByName("kT", kT) 

219 self._up_to_date = True 

220 

221 def step(self, steps): 

222 if not self._up_to_date: 

223 self.update_temperatures(self.temperature, []) 

224 super().step(steps) 

225 

226 

227class AbstractMiddleRespaIntegrator(CustomIntegrator): 

228 """An abstract class for middle-type, multiple time-scale integrators. 

229 

230 .. warning:: 

231 This class is meant for inheritance only and does not actually include 

232 thermostatting. Concrete subclasses are available, such as 

233 :class:`MiddleMassiveNHCIntegrator` and :class:`GeodesicLangevinIntegrator`, 

234 for instance. 

235 

236 Child classes will differ by the thermostat algorithm, which must be implemented 

237 by overriding the `_bath` method (see the example below). 

238 Temperature is treated as a per-dof parameter so as to allow adiabatic simulations. 

239 

240 The following :term:`ODE` system is solved for every degree of freedom in the 

241 system, with possibly :math:`n_c` holonomic constraints and with forces possibly 

242 split into :math:`m` parts according to their characteristic time scales: 

243 

244 .. math:: 

245 & \\dot{r}_i = v_i \\\\ 

246 & \\dot{v}_i = \\frac{\\sum_{k=1}^m F_i^{[k]}}{m_i} 

247 + \\sum_{k=1}^{n_c} \\lambda_k \\nabla_{r_i} \\sigma_k 

248 + \\mathrm{bath}(T_i, v_i) \\\\ 

249 & \\sigma_k(\\mathbf{r}) = 0 

250 

251 An approximate solution is obtained by applying the Trotter-Suzuki splitting 

252 formula. In the particular case of two time scales, the default splitting scheme 

253 goes as follows: 

254 

255 .. math:: 

256 e^{\\Delta t\\mathcal{L}} = 

257 e^{\\frac{\\Delta t}{2}\\mathcal{L}^{[1]}_v} 

258 \\left[ 

259 e^{\\frac{\\Delta t}{2 n_0}\\mathcal{L}^{[0]}_v} 

260 \\left( 

261 e^{\\frac{\\Delta t}{2 n_0 n_b}\\mathcal{L}_r} 

262 e^{\\frac{\\Delta t}{n_0 n_b}\\mathcal{L}_\\mathrm{bath}} 

263 e^{\\frac{\\Delta t}{2 n_0 n_b}\\mathcal{L}_r} 

264 \\right)^{n_b} 

265 e^{\\frac{\\Delta t}{2 n_0}\\mathcal{L}^{[0]}_v} 

266 \\right]^{n_0} 

267 e^{\\frac{\\Delta t}{2}\\mathcal{L}^{[1]}_v} 

268 

269 Each exponential operator is the solution of a particular subsystem of equations. 

270 If :math:`\\mathrm{bath}(T_i, v_i) = 0`, the scheme above is time-reversible, 

271 measure-preserving, and symplectic. It is referred to as the ``VV-Middle`` scheme 

272 :cite:`Zhang_2019`, where VV stands for Velocity Verlet. An alternative approach 

273 is also available, which is: 

274 

275 .. math:: 

276 e^{\\Delta t\\mathcal{L}} = 

277 \\left[ 

278 \\left( 

279 e^{\\frac{\\Delta t}{2 n_0 n_b}\\mathcal{L}_r} 

280 e^{\\frac{\\Delta t}{n_0 n_b}\\mathcal{L}_\\mathrm{bath}} 

281 e^{\\frac{\\Delta t}{2 n_0 n_b}\\mathcal{L}_r} 

282 \\right)^{n_b} 

283 e^{\\frac{\\Delta t}{n_0}\\mathcal{L}^{[0]}_v} 

284 \\right]^{n_0} 

285 e^{\\Delta t \\mathcal{L}^{[1]}_v} 

286 

287 This is referred to as the ``LF-Middle`` scheme :cite:`Zhang_2019`, where LF stands 

288 for Leap-Frog. In contrast to the previous scheme, it is not time-reversible. 

289 However, in single time-scale simulations, the two approaches result in equivalent 

290 coordinate trajectories, while the latter provides a velocity trajectory more 

291 consistent with the Maxwell-Boltzmann distribution at the specified temperature 

292 :cite:`Zhang_2019`. 

293 

294 Parameters 

295 ---------- 

296 temperature : unit.Quantity 

297 The temperature of the heat bath. 

298 step_size : float or unit.Quantity 

299 The outer step size with which to integrate the equations of motion. 

300 

301 Keyword Args 

302 ------------ 

303 num_rattles : int, default=0 

304 The number of RATTLE computations for geodesic integration 

305 :cite:`Leimkuhler_2016`. If ``num_rattles=0``, then no constraints are 

306 considered at all. 

307 scheme : str, default='VV-Middle' 

308 Which splitting scheme will be used. Valid options are 'VV-Middle' and 

309 'LF-Middle'. 

310 respa_loops : list(int), default=[1] 

311 A list of `m` integers, where `respa_loops[k]` determines how many substeps 

312 with force group `k` are internally executed for every step with force group 

313 `k+1`. 

314 bath_loops : int, default=1 

315 The number of iterations of the bath operator per each step at time scale 

316 `0`. This is useful when the bath operator is not exact, but derived from a 

317 splitting solution. 

318 embodied_force_groups : list(int), default=[] 

319 A list of indices of force groups. The presence of an index `i` is this list 

320 means that the contribution of force group `i` is embodied in force group 

321 `i+1`. Therefore, such contribution must be properly subtracted during the 

322 integration at time scale `i+1`. This feature requires OpenMM 7.5 or a newer 

323 version. 

324 unroll_loops : bool, default=True 

325 Whether the integrator loops should be unrolled for improving efficiency. 

326 Using ``unroll_loops=False`` can be useful for printing the integrator 

327 steps. 

328 

329 Example 

330 ------- 

331 >>> from ufedmm import integrators 

332 >>> from openmm import unit 

333 >>> class MiddleNoseHooverIntegrator(integrators.AbstractMiddleRespaIntegrator): 

334 ... def __init__(self, ndof, tau, temperature, step_size, num_rattles=1): 

335 ... super().__init__( 

336 ... temperature, step_size, num_rattles, 'VV-Middle', [1], 1 

337 ... ) 

338 ... kB = 8.3144626E-3*unit.kilojoules_per_mole/unit.kelvin 

339 ... gkT = ndof*unit.MOLAR_GAS_CONSTANT_R*temperature 

340 ... self.addGlobalVariable('gkT', gkT) 

341 ... self.addGlobalVariable('Q', gkT*tau**2) 

342 ... self.addGlobalVariable('v_eta', 0) 

343 ... self.addGlobalVariable('twoK', 0) 

344 ... self.addGlobalVariable('scaling', 1) 

345 ... def _bath(self, fraction): 

346 ... self.addComputeSum('twoK', 'm*v*v') 

347 ... self.addComputeGlobal( 

348 ... 'v_eta', f'v_eta + {0.5*fraction}*dt*(twoK - gkT)/Q' 

349 ... ) 

350 ... self.addComputeGlobal('scaling', f'exp(-{fraction}*dt*v_eta)') 

351 ... self.addComputePerDof('v', f'v*scaling') 

352 ... self.addComputeGlobal( 

353 ... 'v_eta', f'v_eta + {0.5*fraction}*dt*(scaling^2*twoK - gkT)/Q' 

354 ... ) 

355 >>> integrator = MiddleNoseHooverIntegrator( 

356 ... 500, 10*unit.femtoseconds, 300*unit.kelvin, 

357 ... 1*unit.femtoseconds, num_rattles=0 

358 ... ) 

359 >>> print(integrator) 

360 Per-dof variables: 

361 kT 

362 Global variables: 

363 gkT = 1247.169392722986 

364 Q = 0.1247169392722986 

365 v_eta = 0.0 

366 twoK = 0.0 

367 scaling = 1.0 

368 Computation steps: 

369 0: allow forces to update the context state 

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

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

372 3: twoK <- sum(m*v*v) 

373 4: v_eta <- v_eta + 0.5*dt*(twoK - gkT)/Q 

374 5: scaling <- exp(-1.0*dt*v_eta) 

375 6: v <- v*scaling 

376 7: v_eta <- v_eta + 0.5*dt*(scaling^2*twoK - gkT)/Q 

377 8: x <- x + 0.5*dt*v 

378 9: v <- v + 0.5*dt*f/m 

379 """ 

380 

381 def __init__( 

382 self, 

383 temperature, 

384 step_size, 

385 num_rattles=0, 

386 scheme="VV-Middle", 

387 respa_loops=[1], 

388 bath_loops=1, 

389 intertwine=True, 

390 embodied_force_groups=[], 

391 unroll_loops=True, 

392 ): 

393 if scheme not in ["LF-Middle", "VV-Middle"]: 

394 raise Exception(f"Invalid value {scheme} for keyword scheme") 

395 super().__init__(temperature, step_size) 

396 self._num_rattles = num_rattles 

397 self._scheme = scheme 

398 self._respa_loops = respa_loops 

399 self._bath_loops = bath_loops 

400 self._intertwine = intertwine 

401 self._subtractive_groups = embodied_force_groups 

402 num_rattles > 0 and self.addPerDofVariable("x0", 0) 

403 num_rattles > 1 and self.addGlobalVariable("irattle", 0) 

404 if not unroll_loops: 

405 for scale, n in enumerate(respa_loops): 

406 n > 1 and self.addGlobalVariable(f"irespa{scale}", 0) 

407 bath_loops > 1 and self.addGlobalVariable("ibath", 0) 

408 if embodied_force_groups: 

409 if openmm.__version__ < "7.5": 

410 raise Exception( 

411 "Use of `embodied_force_groups` option requires OpenMM >= 7.5" 

412 ) 

413 self.addPerDofVariable("f_emb", 0) 

414 integration_groups = set(range(len(respa_loops))) - set( 

415 embodied_force_groups 

416 ) 

417 self.setIntegrationForceGroups(integration_groups) 

418 

419 self.addUpdateContextState() 

420 self._step_initialization() 

421 if unroll_loops: 

422 self._integrate_respa_unrolled(1, len(respa_loops) - 1) 

423 else: 

424 self._integrate_respa(1, len(respa_loops) - 1) 

425 

426 def _step_initialization(self): 

427 pass 

428 

429 def _integrate_respa(self, fraction, scale): 

430 if scale >= 0: 

431 n = self._respa_loops[scale] 

432 if n > 1: 

433 self.addComputeGlobal(f"irespa{scale}", "0") 

434 self.beginWhileBlock(f"irespa{scale} < {n-1/2}") 

435 self._boost(fraction / (2 * n if self._scheme == "VV-Middle" else n), scale) 

436 self._integrate_respa(fraction / n, scale - 1) 

437 self._scheme == "VV-Middle" and self._boost(fraction / (2 * n), scale) 

438 if n > 1: 

439 self.addComputeGlobal(f"irespa{scale}", f"irespa{scale} + 1") 

440 self.endBlock() 

441 else: 

442 self._intertwine or self._translation(0.5 * fraction) 

443 n = self._bath_loops 

444 if n > 1: 

445 self.addComputeGlobal("ibath", "0") 

446 self.beginWhileBlock(f"ibath < {n-1/2}") 

447 self._intertwine and self._translation(0.5 * fraction / n) 

448 self._bath(fraction / n) 

449 self._num_rattles > 0 and self.addConstrainVelocities() 

450 self._intertwine and self._translation(0.5 * fraction / n) 

451 if n > 1: 

452 self.addComputeGlobal("ibath", "ibath + 1") 

453 self.endBlock() 

454 self._intertwine or self._translation(0.5 * fraction) 

455 

456 def _integrate_respa_unrolled(self, fraction, scale): 

457 if scale >= 0: 

458 n = self._respa_loops[scale] 

459 for i in range(n): 

460 self._boost( 

461 fraction / (2 * n if self._scheme == "VV-Middle" and i == 0 else n), 

462 scale, 

463 ) 

464 self._integrate_respa_unrolled(fraction / n, scale - 1) 

465 self._scheme == "VV-Middle" and i == n - 1 and self._boost( 

466 fraction / (2 * n), scale 

467 ) 

468 else: 

469 n = self._bath_loops 

470 self._intertwine or self._translation(0.5 * fraction) 

471 for i in range(n): 

472 self._intertwine and self._translation( 

473 fraction / (2 * n if i == 0 else n) 

474 ) 

475 self._bath(fraction / n) 

476 self._num_rattles > 0 and self.addConstrainVelocities() 

477 i == n - 1 and self._intertwine and self._translation( 

478 fraction / (2 * n) 

479 ) 

480 self._intertwine or self._translation(0.5 * fraction) 

481 

482 def _translation(self, fraction): 

483 if self._num_rattles > 1: 

484 self.addComputeGlobal("irattle", "0") 

485 self.beginWhileBlock(f"irattle < {self._num_rattles-1/2}") 

486 self.addComputePerDof("x", f"x + {fraction/max(1, self._num_rattles)}*dt*v") 

487 if self._num_rattles > 0: 

488 self.addComputePerDof("x0", "x") 

489 self.addConstrainPositions() 

490 self.addComputePerDof( 

491 "v", f"v + (x - x0)/({fraction/self._num_rattles}*dt)" 

492 ) 

493 self.addConstrainVelocities() 

494 if self._num_rattles > 1: 

495 self.addComputeGlobal("irattle", "irattle + 1") 

496 self.endBlock() 

497 

498 def _boost(self, fraction, scale): 

499 if len(self._respa_loops) > 1: 

500 if scale - 1 in self._subtractive_groups: 

501 self.addComputePerDof("f_emb", f"f{scale-1}") 

502 self.addComputePerDof("v", f"v + {fraction}*dt*(f{scale}-f_emb)/m") 

503 else: 

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

505 else: 

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

507 self._num_rattles > 0 and self.addConstrainVelocities() 

508 

509 def _bath(self, fraction): 

510 return 

511 

512 

513class GeodesicLangevinIntegrator(AbstractMiddleRespaIntegrator): 

514 """A geodesic Langevin integrator :cite:`Leimkuhler_2016`, which can be integrated 

515 by using either the LF-Middle or the VV-Middle scheme :cite:`Zhang_2019`. 

516 

517 .. note: 

518 The VV-Middle scheme is also known as the BAOAB :cite:`Leimkuhler_2016` method. 

519 

520 Parameters 

521 ---------- 

522 temperature : float or unit.Quantity 

523 The temperature. 

524 friction_coefficient : float or unit.Quantity 

525 The friction coefficient. 

526 step_size : float or unit.Quantity 

527 The time-step size. 

528 

529 Keyword Args 

530 ------------ 

531 num_rattles : int, default=1 

532 The number of RATTLE computations for geodesic integration 

533 :cite:`Leimkuhler_2016`. If ``num_rattles=0``, then no constraints are 

534 considered at all. 

535 scheme : str, default='LF-Middle' 

536 Which splitting scheme will be used. Valid options are 'VV-Middle' and 

537 'LF-Middle'. 

538 **kwargs 

539 All other keyword arguments in :class:`AbstractMiddleRespaIntegrator`. 

540 

541 Example 

542 ------- 

543 >>> import ufedmm 

544 >>> dt = 2*unit.femtoseconds 

545 >>> temp = 300*unit.kelvin 

546 >>> gamma = 10/unit.picoseconds 

547 >>> ufedmm.GeodesicLangevinIntegrator( 

548 ... temp, gamma, dt, num_rattles=1, scheme='VV-Middle' 

549 ... ) 

550 Per-dof variables: 

551 kT, x0 

552 Global variables: 

553 friction = 10.0 

554 Computation steps: 

555 0: allow forces to update the context state 

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

557 2: constrain velocities 

558 3: x <- x + 0.5*dt*v 

559 4: x0 <- x 

560 5: constrain positions 

561 6: v <- v + (x - x0)/(0.5*dt) 

562 7: constrain velocities 

563 8: v <- z*v + sqrt((1 - z*z)*kT/m)*gaussian; z = exp(-friction*1.0*dt) 

564 9: constrain velocities 

565 10: x <- x + 0.5*dt*v 

566 11: x0 <- x 

567 12: constrain positions 

568 13: v <- v + (x - x0)/(0.5*dt) 

569 14: constrain velocities 

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

571 16: constrain velocities 

572 """ 

573 

574 def __init__( 

575 self, 

576 temperature, 

577 friction_coefficient, 

578 step_size, 

579 num_rattles=1, 

580 scheme="LF-Middle", 

581 **kwargs, 

582 ): 

583 super().__init__( 

584 temperature, step_size, num_rattles=num_rattles, scheme=scheme, **kwargs 

585 ) 

586 self.addGlobalVariable("friction", friction_coefficient) 

587 

588 def _bath(self, fraction): 

589 expression = ( 

590 f"z*v + sqrt((1 - z*z)*kT/m)*gaussian; z = exp(-friction*{fraction}*dt)" 

591 ) 

592 self.addComputePerDof("v", expression) 

593 

594 

595class MiddleMassiveNHCIntegrator(AbstractMiddleRespaIntegrator): 

596 """A massive, middle-type Nose-Hoover Chain Thermostat solver :cite:`Martyna_1992` 

597 with optional multiple time-scale integration via RESPA. 

598 

599 To enable RESPA, the forces in OpenMM system must be split into distinct force 

600 groups and the keyword ``respa_loop`` (see below) must be a list with multiple 

601 entries. 

602 

603 Parameters 

604 ---------- 

605 temperature : float or unit.Quantity 

606 The temperature. 

607 time_constant : float or unit.Quantity 

608 The characteristic time constant. 

609 step_size : float or unit.Quantity 

610 The time-step size. 

611 

612 Keyword Args 

613 ------------ 

614 nchain : int, default=2 

615 The number of thermostats in each Nose-Hoover chain. 

616 track_energy : bool, default=False 

617 Whether to track the thermostat energy term. 

618 **kwargs 

619 All keyword arguments in :class:`AbstractMiddleRespaIntegrator`, except 

620 ``num_rattles``. 

621 

622 Example 

623 ------- 

624 >>> import ufedmm 

625 >>> temp, tau, dt = 300*unit.kelvin, 10*unit.femtoseconds, 2*unit.femtoseconds 

626 >>> integrator = ufedmm.MiddleMassiveNHCIntegrator( 

627 ... temp, tau, dt, respa_loops=[4, 1], unroll_loops=False 

628 ... ) 

629 >>> print(integrator) 

630 Per-dof variables: 

631 kT, Q, v1, v2 

632 Global variables: 

633 irespa0 = 0.0 

634 Computation steps: 

635 0: allow forces to update the context state 

636 1: v <- v + 0.5*dt*f1/m 

637 2: irespa0 <- 0 

638 3: while (irespa0 < 3.5): 

639 4: v <- v + 0.125*dt*f0/m 

640 5: x <- x + 0.125*dt*v 

641 6: v2 <- v2 + 0.125*dt*(Q*v1^2 - kT)/Q 

642 7: v1 <- (v1*z + 0.125*dt*(m*v^2 - kT)/Q)*z; z=exp(-0.0625*dt*v2) 

643 8: v <- v*exp(-0.25*dt*v1) 

644 9: v1 <- (v1*z + 0.125*dt*(m*v^2 - kT)/Q)*z; z=exp(-0.0625*dt*v2) 

645 10: v2 <- v2 + 0.125*dt*(Q*v1^2 - kT)/Q 

646 11: x <- x + 0.125*dt*v 

647 12: v <- v + 0.125*dt*f0/m 

648 13: irespa0 <- irespa0 + 1 

649 14: end 

650 15: v <- v + 0.5*dt*f1/m 

651 """ 

652 

653 def __init__( 

654 self, 

655 temperature, 

656 time_constant, 

657 step_size, 

658 nchain=2, 

659 track_energy=False, 

660 **kwargs, 

661 ): 

662 if "num_rattles" in kwargs.keys() and kwargs["num_rattles"] != 0: 

663 raise ValueError(f"{self.__class__.__name__} cannot handle constraints") 

664 self._tau = _standardized(time_constant) 

665 self._nchain = nchain 

666 self._track_energy = track_energy 

667 super().__init__(temperature, step_size, **kwargs) 

668 self.addPerDofVariable("Q", 0) 

669 for i in range(nchain): 

670 self.addPerDofVariable(f"v{i+1}", 0) 

671 if track_energy: 

672 self.addPerDofVariable(f"eta{i+1}", 0) 

673 

674 def update_temperatures(self, system_temperature, extended_space_temperatures): 

675 super().update_temperatures(system_temperature, extended_space_temperatures) 

676 Q = [self._tau**2 * kT for kT in self.getPerDofVariableByName("kT")] 

677 self.setPerDofVariableByName("Q", Q) 

678 

679 def _bath(self, fraction): 

680 n = self._nchain 

681 

682 def a(i): 

683 return f"(Q*v{i-1}^2 - kT)/Q" if i > 1 else "(m*v^2 - kT)/Q" 

684 

685 def z(i): 

686 return f"exp(-{fraction/4}*dt*v{i+1})" 

687 

688 self.addComputePerDof(f"v{n}", f"v{n} + {fraction/2}*dt*{a(n)}") 

689 for i in reversed(range(1, n)): 

690 self.addComputePerDof( 

691 f"v{i}", f"(v{i}*z + {fraction/2}*dt*{a(i)})*z; z={z(i)}" 

692 ) 

693 self.addComputePerDof("v", f"v*exp(-{fraction}*dt*v1)") 

694 for i in range(1, n): 

695 self.addComputePerDof( 

696 f"v{i}", f"(v{i}*z + {fraction/2}*dt*{a(i)})*z; z={z(i)}" 

697 ) 

698 self.addComputePerDof(f"v{n}", f"v{n} + {fraction/2}*dt*{a(n)}") 

699 

700 

701class MiddleMassiveGGMTIntegrator(AbstractMiddleRespaIntegrator): 

702 """A massive, middle-type Generalized Gaussian Moment Thermostat :cite:`Liu_2000` 

703 solver with optional multiple time-scale integration via RESPA. 

704 

705 To enable RESPA, the forces in OpenMM system must be split into distinct force 

706 groups and the keyword ``respa_loop`` (see below) must be a list with multiple 

707 entries. 

708 

709 Parameters 

710 ---------- 

711 temperature : float or unit.Quantity 

712 The temperature. 

713 time_constant : float or unit.Quantity 

714 The characteristic time constant. 

715 step_size : float or unit.Quantity 

716 The time-step size. 

717 

718 Keyword Args 

719 ------------ 

720 **kwargs 

721 All keyword arguments in :class:`AbstractMiddleRespaIntegrator`, except 

722 ``num_rattles``. 

723 

724 Example 

725 ------- 

726 >>> import ufedmm 

727 >>> temp, tau, dt = 300*unit.kelvin, 10*unit.femtoseconds, 2*unit.femtoseconds 

728 >>> integrator = ufedmm.MiddleMassiveGGMTIntegrator(temp, tau, dt) 

729 >>> print(integrator) 

730 Per-dof variables: 

731 kT, Q1, Q2, v1, v2 

732 Computation steps: 

733 0: allow forces to update the context state 

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

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

736 3: v1 <- v1 + 0.5*dt*(m*v^2 - kT)/Q1 

737 4: v2 <- v2 + 0.5*dt*((m*v^2)^2/3 - kT^2)/Q2 

738 5: v <- v*exp(-1.0*dt*(v1 + kT*v2))/sqrt(1 + 2.0*dt*m*v^2*v2/3) 

739 6: v1 <- v1 + 0.5*dt*(m*v^2 - kT)/Q1 

740 7: v2 <- v2 + 0.5*dt*((m*v^2)^2/3 - kT^2)/Q2 

741 8: x <- x + 0.5*dt*v 

742 9: v <- v + 0.5*dt*f/m 

743 """ 

744 

745 def __init__(self, temperature, time_constant, step_size, **kwargs): 

746 if "num_rattles" in kwargs.keys() and kwargs["num_rattles"] != 0: 

747 raise ValueError(f"{self.__class__.__name__} cannot handle constraints") 

748 self._tau = _standardized(time_constant) 

749 super().__init__(temperature, step_size, **kwargs) 

750 self.addPerDofVariable("Q1", 0) 

751 self.addPerDofVariable("Q2", 0) 

752 self.addPerDofVariable("v1", 0) 

753 self.addPerDofVariable("v2", 0) 

754 

755 def set_extended_space_time_constants(self, time_constants): 

756 self._xs_taus = [_standardized(tau) for tau in time_constants] 

757 

758 def update_temperatures(self, system_temperature, extended_space_temperatures): 

759 super().update_temperatures(system_temperature, extended_space_temperatures) 

760 kT_vectors = self.getPerDofVariableByName("kT") 

761 kT3_vectors = [openmm.Vec3(kT.x**3, kT.y**3, kT.z**3) for kT in kT_vectors] 

762 if hasattr(self, "_xs_taus"): 

763 num_particles = len(kT_vectors) - len(extended_space_temperatures) 

764 taus = [self._tau] * num_particles + self._xs_taus 

765 Q1 = [kT * tau**2 for kT, tau in zip(kT_vectors, taus)] 

766 Q2 = [8 / 3 * kT3 * tau**2 for kT3, tau in zip(kT3_vectors, taus)] 

767 else: 

768 Q1 = [kT * self._tau**2 for kT in kT_vectors] 

769 Q2 = [8 / 3 * kT3 * self._tau**2 for kT3 in kT3_vectors] 

770 self.setPerDofVariableByName("Q1", Q1) 

771 self.setPerDofVariableByName("Q2", Q2) 

772 

773 def _bath(self, fraction): 

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

775 self.addComputePerDof("v2", f"v2 + {fraction/2}*dt*((m*v^2)^2/3 - kT^2)/Q2") 

776 self.addComputePerDof( 

777 "v", 

778 f"v*exp(-{fraction}*dt*(v1 + kT*v2))/sqrt(1 + {2*fraction}*dt*m*v^2*v2/3)", 

779 ) 

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

781 self.addComputePerDof("v2", f"v2 + {fraction/2}*dt*((m*v^2)^2/3 - kT^2)/Q2") 

782 

783 

784class RegulatedNHLIntegrator(AbstractMiddleRespaIntegrator): 

785 """A regulated version of the massive Nose-Hoover-Langevin 

786 :cite:`Samoletov_2007,Leimkuhler_2009` method. Regulation means that the system 

787 Hamiltonian is modified so that velocities remain below a temperature-dependent 

788 limit. This is closely related to the SIN(R) method :cite:`Leimkuhler_2013` and 

789 allows multiple time-scale integration with very large outer time steps, without 

790 resonance. 

791 

792 .. info: 

793 If `regulation_parameter = 1` (default), this method is equivalent to SIN(R) 

794 with a single thermostat per degree of freedom (that is, `L=1`). 

795 

796 The following :term:`SDE` system is solved for every degree of freedom in the 

797 system: 

798 

799 .. math:: 

800 & dr_i = v_i dt \\\\ 

801 & dp_i = F_i dt - v_{\\eta_i} m_i v_i dt \\\\ 

802 & dv_{\\eta_i} = \\frac{1}{Q}\\left(\\frac{n+1}{n} m_i v_i^2 - k_B T\\right) dt 

803 - \\gamma v_{\\eta_i} dt + \\sqrt{\\frac{2\\gamma k_B T}{Q}} dW_i, 

804 

805 where: 

806 

807 .. math:: 

808 v_i = c_i \\tanh\\left(\\frac{p_i}{m_i c_i}\\right). 

809 

810 Here, :math:`n` is the regulation parameter and :math:`c_i` is the maximum speed for 

811 degree of freedom `i`, given by :math:`c_i = \\sqrt{\\frac{n k T}{m_i}}`. The 

812 inertial parameter :math:`Q` is defined as :math:`Q = n k_B T \\tau^2`, with 

813 :math:`\\tau` being a relaxation time :cite:`Tuckerman_1992`. An approximate 

814 solution is obtained by applying the Trotter-Suzuki splitting formula: 

815 

816 .. math:: 

817 e^{\\Delta t\\mathcal{L}} = 

818 e^{\\frac{\\Delta t}{2}\\mathcal{L}^1_p} 

819 \\left[e^{\\frac{\\delta t}{2}\\mathcal{L}^0_p} 

820 e^{\\frac{\\delta t}{2}\\mathcal{L}_r} 

821 e^{\\delta t \\mathcal{L}_\\mathrm{bath}} 

822 e^{\\frac{\\delta t}{2}\\mathcal{L}_r} 

823 e^{\\frac{\\delta t}{2}\\mathcal{L}^0_p}\\right]^m 

824 e^{\\frac{\\Delta t}{2}\\mathcal{L}^1_p} 

825 

826 where :math:`\\delta t = \\frac{\\Delta t}{m}`. Each exponential operator above is 

827 the solution of a differential equation. 

828 

829 The exact solution for the physical-system part is: 

830 

831 .. math:: 

832 r_i(t) = r_i^0 + c_i \\mathrm{tanh}\\left(\\frac{p_i}{m c_i}\\right) t 

833 

834 .. math:: 

835 p_i(t) = p_i^0 + F_i t 

836 

837 The bath propagator is further split as: 

838 

839 .. math:: 

840 e^{\\delta t \\mathcal{L}_\\mathrm{bath}} = 

841 e^{\\frac{\\delta t}{2m}\\mathcal{L}_B} 

842 e^{\\frac{\\delta t}{2m}\\mathcal{L}_S} 

843 e^{\\frac{\\delta t}{m}\\mathcal{L}_O} 

844 e^{\\frac{\\delta t}{2m}\\mathcal{L}_S} 

845 e^{\\frac{\\delta t}{2m}\\mathcal{L}_B} 

846 

847 Part 'B' is a boost, whose solution is: 

848 

849 .. math:: 

850 v_{\\eta_i}(t) = v_{\\eta_i}^0 + 

851 \\frac{1}{Q}\\left(\\frac{n+1}{n} m_i v_i^2 - k_B T\\right) t 

852 

853 Part 'S' is a scaling, whose solution is: 

854 

855 .. math:: 

856 p_i(t) = m_i c_i \\mathrm{arcsinh}\\left[ 

857 \\sinh\\left(\\frac{p_i^0}{m_i c_i}\\right) e^{- v_{\\eta_i} t} 

858 \\right] 

859 

860 Part 'O' is an Ornstein–Uhlenbeck process, whose solution is: 

861 

862 .. math:: 

863 v_{\\eta_i}(t) = v_{\\eta_i}^0 e^{-\\gamma t} 

864 + \\sqrt{\\frac{k_B T}{Q}(1-e^{-2\\gamma t})} R_N 

865 

866 where :math:`R_N` is a normally distributed random number. 

867 

868 Parameters 

869 ---------- 

870 step_size : float or unit.Quantity 

871 The outer step size with which to integrate the equations of motion. 

872 loops : int 

873 The number of internal substeps at each time step. 

874 temperature : unit.Quantity 

875 The temperature of the heat bath. 

876 time_scale : unit.Quantity (time) 

877 The relaxation time (:math:`\\tau`) of the Nose-Hoover thermostat. 

878 friction_coefficient : unit.Quantity (1/time) 

879 The friction coefficient (:math:`\\gamma`) of the Langevin thermostat. 

880 regulation_parameter : int or float 

881 The regulation parameter n. 

882 

883 Keyword Args 

884 ------------ 

885 semi_regulated : bool, default=True 

886 Whether to use the semi-regulated NHL version of the method instead of its 

887 fully-regulated version. 

888 split_ornstein_uhlenbeck : bool, default=True 

889 Whether to split the drifted Ornstein-Uhlenbeck operator. 

890 **kwargs 

891 All keyword arguments in :class:`AbstractMiddleRespaIntegrator`, except 

892 ``num_rattles``. 

893 """ 

894 

895 def __init__( 

896 self, 

897 temperature, 

898 time_constant, 

899 friction_coefficient, 

900 step_size, 

901 regulation_parameter, 

902 semi_regulated=True, 

903 split_ornstein_uhlenbeck=True, 

904 **kwargs, 

905 ): 

906 if "num_rattles" in kwargs.keys() and kwargs["num_rattles"] != 0: 

907 raise ValueError(f"{self.__class__.__name__} cannot handle constraints") 

908 self._tau = np.sqrt(regulation_parameter) * time_constant 

909 self._n = regulation_parameter 

910 self._split = split_ornstein_uhlenbeck 

911 self._semi_regulated = semi_regulated 

912 super().__init__(temperature, step_size, **kwargs) 

913 self.addPerDofVariable("invQ", 0) 

914 self.addPerDofVariable("v_eta", 0) 

915 self.addPerDofVariable("c", 0) 

916 self.addGlobalVariable("friction", friction_coefficient) 

917 self.addGlobalVariable("omega", 1.0 / self._tau) 

918 self.addGlobalVariable("aa", 0) 

919 self.addGlobalVariable("bb", 0) 

920 

921 def update_temperatures(self, system_temperature, extended_space_temperatures): 

922 super().update_temperatures(system_temperature, extended_space_temperatures) 

923 kT_vectors = self.getPerDofVariableByName("kT") 

924 tauSq = _standardized(self._tau) ** 2 

925 Q = [tauSq * kT for kT in kT_vectors] 

926 invQ = [openmm.Vec3(*map(lambda x: 1 / x if x > 0.0 else 0.0, q)) for q in Q] 

927 self.setPerDofVariableByName("invQ", invQ) 

928 

929 def _step_initialization(self): 

930 self.addComputePerDof("c", f"sqrt({self._n}*kT/m)") 

931 n = np.prod(self._respa_loops) * self._bath_loops 

932 self.addComputeGlobal("aa", f"exp(-friction*dt/{n})") 

933 self.addComputeGlobal("bb", "omega*sqrt(1-aa^2)") 

934 

935 def _translation(self, fraction):