Coverage for /builds/alexhroom/ase/ase/md/switch_langevin.py: 92.16%

1from typing import Any, List, Optional 

2 

3import numpy as np 

4try: 

5 from numpy import trapezoid # NumPy 2.0.0 

6except ImportError: 

7 from numpy import trapz as trapezoid 

8 

9from ase import Atoms 

10from ase.calculators.mixing import MixedCalculator 

11from ase.md.langevin import Langevin 

12 

13 

14class SwitchLangevin(Langevin): 

15 """ 

16 MD class for carrying out thermodynamic integration between two 

17 hamiltonians 

18 

19 Read more at https://en.wikipedia.org/wiki/Thermodynamic_integration 

20 

21 The integration path is lambda 0 ---> 1 , with the switching hamiltonian 

22 H(lambda) = (1-lambda) * H1 + lambda * H2 

23 

24 Attributes 

25 ---------- 

26 path_data : numpy array 

27 col 1 (md-step), col 2 (lambda), col 3 (energy H1), col 4 (energy H2) 

28 

29 Parameters 

30 ---------- 

31 atoms : ASE Atoms object 

32 Atoms object for which MD will be run 

33 calc1 : ASE calculator object 

34 Calculator corresponding to first Hamiltonian 

35 calc2 : ASE calculator object 

36 Calculator corresponding to second Hamiltonian 

37 dt : float 

38 Timestep for MD simulation 

39 T : float 

40 Temperature in eV (deprecated) 

41 friction : float 

42 Friction for langevin dynamics 

43 n_eq : int 

44 Number of equilibration steps 

45 n_switch : int 

46 Number of switching steps 

47 """ 

48 

49 def __init__( 

50 self, 

51 atoms: Atoms, 

52 calc1, 

53 calc2, 

54 dt: float, 

55 T: Optional[float] = None, 

56 friction: Optional[float] = None, 

57 n_eq: Optional[int] = None, 

58 n_switch: Optional[int] = None, 

59 temperature_K: Optional[float] = None, 

60 **langevin_kwargs, 

61 ): 

62 super().__init__(atoms, dt, temperature=T, temperature_K=temperature_K, 

63 friction=friction, **langevin_kwargs) 

64 if friction is None: 

65 raise TypeError("Missing 'friction' argument.") 

66 if n_eq is None: 

67 raise TypeError("Missing 'n_eq' argument.") 

68 if n_switch is None: 

69 raise TypeError("Missing 'n_switch' argument.") 

70 self.n_eq = n_eq 

71 self.n_switch = n_switch 

72 self.lam = 0.0 

73 calc = MixedCalculator(calc1, calc2, weight1=1.0, weight2=0.0) 

74 self.atoms.calc = calc 

75 

76 self.path_data: List[Any] = [] 

77 

78 def run(self): 

79 """ Run the MD switching simulation """ 

80 forces = self.atoms.get_forces(md=True) 

81 

82 # run equilibration with calc1 

83 for _ in range(self.n_eq): 

84 forces = self.step(forces) 

85 self.nsteps += 1 

86 self.call_observers() 

87 

88 # run switch from calc1 to calc2 

89 self.path_data.append( 

90 [0, self.lam, 

91 *self.atoms.calc.get_energy_contributions(self.atoms)]) 

92 for step in range(1, self.n_switch): 

93 # update calculator 

94 self.lam = get_lambda(step, self.n_switch) 

95 self.atoms.calc.set_weights(1 - self.lam, self.lam) 

96 

97 # carry out md step 

98 forces = self.step(forces) 

99 self.nsteps += 1 

100 

101 # collect data 

102 self.call_observers() 

103 self.path_data.append( 

104 [step, self.lam, 

105 *self.atoms.calc.get_energy_contributions(self.atoms)]) 

106 

107 self.path_data = np.array(self.path_data) 

108 

109 def get_free_energy_difference(self): 

110 """ Return the free energy difference between calc2 and calc1, by 

111 integrating dH/dlam along the switching path 

112 

113 Returns 

114 ------- 

115 float 

116 Free energy difference, F2 - F1 

117 """ 

118 if len(self.path_data) == 0: 

119 raise ValueError('No free energy data found.') 

120 

121 lambdas = self.path_data[:, 1] 

122 U1 = self.path_data[:, 2] 

123 U2 = self.path_data[:, 3] 

124 delta_F = trapezoid(U2 - U1, lambdas) 

125 return delta_F 

126 

127 

128def get_lambda(step, n_switch): 

129 """ Return lambda value along the switching path """ 

130 assert step >= 0 and step <= n_switch 

131 t = step / (n_switch - 1) 

132 return t**5 * (70 * t**4 - 315 * t**3 + 540 * t**2 - 420 * t + 126)