1import numpy as np 

2from ase.md.langevin import Langevin 

3from ase.calculators.mixing import MixedCalculator 

4 

5 

6class SwitchLangevin(Langevin): 

7 """ 

8 MD class for carrying out thermodynamic integration between two 

9 hamiltonians 

10 

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

12 

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

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

15 

16 Attributes 

17 ---------- 

18 path_data : numpy array 

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

20 

21 Parameters 

22 ---------- 

23 atoms : ASE Atoms object 

24 Atoms object for which MD will be run 

25 calc1 : ASE calculator object 

26 Calculator corresponding to first Hamiltonian 

27 calc2 : ASE calculator object 

28 Calculator corresponding to second Hamiltonian 

29 dt : float 

30 Timestep for MD simulation 

31 T : float 

32 Temperature in eV (deprecated) 

33 friction : float 

34 Friction for langevin dynamics 

35 n_eq : int 

36 Number of equilibration steps 

37 n_switch : int 

38 Number of switching steps 

39 """ 

40 

41 def __init__(self, atoms, calc1, calc2, dt, T=None, friction=None, 

42 n_eq=None, n_switch=None, temperature_K=None, 

43 **langevin_kwargs): 

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

45 friction=friction, **langevin_kwargs) 

46 if friction is None: 

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

48 if n_eq is None: 

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

50 if n_switch is None: 

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

52 self.n_eq = n_eq 

53 self.n_switch = n_switch 

54 self.lam = 0.0 

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

56 self.atoms.calc = calc 

57 

58 self.path_data = [] 

59 

60 def run(self): 

61 """ Run the MD switching simulation """ 

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

63 

64 # run equilibration with calc1 

65 for _ in range(self.n_eq): 

66 forces = self.step(forces) 

67 self.nsteps += 1 

68 self.call_observers() 

69 

70 # run switch from calc1 to calc2 

71 self.path_data.append( 

72 [0, self.lam, *self.atoms.calc.get_energy_contributions(self.atoms)]) 

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

74 # update calculator 

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

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

77 

78 # carry out md step 

79 forces = self.step(forces) 

80 self.nsteps += 1 

81 

82 # collect data 

83 self.call_observers() 

84 self.path_data.append( 

85 [step, self.lam, *self.atoms.calc.get_energy_contributions(self.atoms)]) 

86 

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

88 

89 def get_free_energy_difference(self): 

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

91 integrating dH/dlam along the switching path 

92 

93 Returns 

94 ------- 

95 float 

96 Free energy difference, F2 - F1 

97 """ 

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

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

100 

101 lambdas = self.path_data[:, 1] 

102 U1 = self.path_data[:, 2] 

103 U2 = self.path_data[:, 3] 

104 delta_F = np.trapz(U2 - U1, lambdas) 

105 return delta_F 

106 

107 

108def get_lambda(step, n_switch): 

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

110 assert step >= 0 and step <= n_switch 

111 t = step / (n_switch - 1) 

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