Coverage for openxps/integrators/symmetric_langevin.py: 95%

1""" 

2.. module:: openxps.integrators.symmetric_langevin 

3 :platform: Linux, MacOS, Windows 

4 :synopsis: BAOAB 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.utils import preprocess_args 

15 

16from .utils import IntegratorMixin, add_property 

17 

18 

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

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

21class SymmetricLangevinIntegrator(IntegratorMixin, mm.CustomIntegrator): 

22 """A symmetric Langevin integrator using the BAOAB algorithm :cite:`Leimkuhler2013`. 

23 

24 Parameters 

25 ---------- 

26 temperature 

27 The temperature of the heat bath. 

28 frictionCoeff 

29 The friction coefficient which couples the system to the heat bath. 

30 stepSize 

31 The step size with which to integrate the system. 

32 

33 Example 

34 ------- 

35 >>> import openxps as xps 

36 >>> from openmm import unit 

37 >>> integrator = xps.integrators.SymmetricLangevinIntegrator( 

38 ... 300 * unit.kelvin, 1 / unit.picosecond, 1 * unit.femtosecond 

39 ... ) 

40 >>> integrator 

41 Per-dof variables: 

42 x1 

43 Global variables: 

44 kT = 2.49433... 

45 a = 0.999000... 

46 b = 0.044699... 

47 Computation steps: 

48 0: allow forces to update the context state 

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

50 2: constrain velocities 

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

52 4: v <- a*v + b*sqrt(kT/m)*gaussian 

53 5: x <- x + 0.5*dt*v 

54 6: x1 <- x 

55 7: constrain positions 

56 8: v <- v + (x-x1)/dt + 0.5*dt*f/m 

57 9: constrain velocities 

58 """ 

59 

60 @preprocess_args 

61 def __init__( 

62 self, 

63 temperature: mmunit.Quantity, 

64 frictionCoeff: mmunit.Quantity, 

65 stepSize: mmunit.Quantity, 

66 ) -> None: 

67 super().__init__(stepSize) 

68 self._init_temperature(temperature) 

69 self._init_friction_coefficient(frictionCoeff) 

70 self._add_variables() 

71 self.addUpdateContextState() 

72 self.addComputePerDof("v", "v + 0.5*dt*f/m") 

73 self.addConstrainVelocities() 

74 self.addComputePerDof("x", "x + 0.5*dt*v") 

75 self.addComputePerDof("v", "a*v + b*sqrt(kT/m)*gaussian") 

76 self.addComputePerDof("x", "x + 0.5*dt*v") 

77 self.addComputePerDof("x1", "x") 

78 self.addConstrainPositions() 

79 self.addComputePerDof("v", "v + (x-x1)/dt + 0.5*dt*f/m") 

80 self.addConstrainVelocities() 

81 

82 def _add_variables(self) -> None: 

83 self.addGlobalVariable("kT", 0) 

84 self.addGlobalVariable("a", 0) 

85 self.addGlobalVariable("b", 0) 

86 self.addPerDofVariable("x1", 0) 

87 self._update_global_variables() 

88 

89 def _update_global_variables(self) -> None: 

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

91 friction = self.getFrictionCoefficient() 

92 dt = self.getStepSize() 

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

94 b = np.sqrt(1 - a**2) 

95 self.setGlobalVariableByName("kT", kt) 

96 self.setGlobalVariableByName("a", a) 

97 self.setGlobalVariableByName("b", b) 

98 

99 def setStepSize(self, size: mmunit.Quantity) -> None: 

100 super().setStepSize(size) 

101 self._update_global_variables()