Coverage for /builds/ase/ase/ase/md/switch_langevin.py: 91.67%

1# fmt: off 

2 

3from typing import Any, List, Optional 

4 

5import numpy as np 

6from scipy.integrate import trapezoid 

7 

8from ase import Atoms 

9from ase.calculators.mixing import MixedCalculator 

10from ase.md.langevin import Langevin 

11 

12 

13class SwitchLangevin(Langevin): 

14 """ 

15 MD class for carrying out thermodynamic integration between two 

16 hamiltonians 

17 

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

19 

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

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

22 

23 Attributes 

24 ---------- 

25 path_data : numpy array 

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

27 

28 Parameters 

29 ---------- 

30 atoms : ASE Atoms object 

31 Atoms object for which MD will be run 

32 calc1 : ASE calculator object 

33 Calculator corresponding to first Hamiltonian 

34 calc2 : ASE calculator object 

35 Calculator corresponding to second Hamiltonian 

36 dt : float 

37 Timestep for MD simulation 

38 T : float 

39 Temperature in eV (deprecated) 

40 friction : float 

41 Friction for langevin dynamics 

42 n_eq : int 

43 Number of equilibration steps 

44 n_switch : int 

45 Number of switching steps 

46 """ 

47 

48 def __init__( 

49 self, 

50 atoms: Atoms, 

51 calc1, 

52 calc2, 

53 dt: float, 

54 T: Optional[float] = None, 

55 friction: Optional[float] = None, 

56 n_eq: Optional[int] = None, 

57 n_switch: Optional[int] = None, 

58 temperature_K: Optional[float] = None, 

59 **langevin_kwargs, 

60 ): 

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

62 friction=friction, **langevin_kwargs) 

63 if friction is None: 

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

65 if n_eq is None: 

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

67 if n_switch is None: 

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

69 self.n_eq = n_eq 

70 self.n_switch = n_switch 

71 self.lam = 0.0 

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

73 self.atoms.calc = calc 

74 

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

76 

77 def run(self): 

78 """ Run the MD switching simulation """ 

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

80 

81 # run equilibration with calc1 

82 for _ in range(self.n_eq): 

83 forces = self.step(forces) 

84 self.nsteps += 1 

85 self.call_observers() 

86 

87 # run switch from calc1 to calc2 

88 self.path_data.append( 

89 [0, self.lam, 

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

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

92 # update calculator 

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

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

95 

96 # carry out md step 

97 forces = self.step(forces) 

98 self.nsteps += 1 

99 

100 # collect data 

101 self.call_observers() 

102 self.path_data.append( 

103 [step, self.lam, 

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

105 

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

107 

108 def get_free_energy_difference(self): 

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

110 integrating dH/dlam along the switching path 

111 

112 Returns 

113 ------- 

114 float 

115 Free energy difference, F2 - F1 

116 """ 

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

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

119 

120 lambdas = self.path_data[:, 1] 

121 U1 = self.path_data[:, 2] 

122 U2 = self.path_data[:, 3] 

123 delta_F = trapezoid(U2 - U1, lambdas) 

124 return delta_F 

125 

126 

127def get_lambda(step, n_switch): 

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

129 assert step >= 0 and step <= n_switch 

130 t = step / (n_switch - 1) 

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