Coverage for /builds/kinetik161/ase/ase/md/switch_langevin.py: 91.49%

1from typing import Any, List, Optional 

2 

3import numpy as np 

4 

5from ase import Atoms 

6from ase.calculators.mixing import MixedCalculator 

7from ase.md.langevin import Langevin 

8 

9 

10class SwitchLangevin(Langevin): 

11 """ 

12 MD class for carrying out thermodynamic integration between two 

13 hamiltonians 

14 

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

16 

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

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

19 

20 Attributes 

21 ---------- 

22 path_data : numpy array 

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

24 

25 Parameters 

26 ---------- 

27 atoms : ASE Atoms object 

28 Atoms object for which MD will be run 

29 calc1 : ASE calculator object 

30 Calculator corresponding to first Hamiltonian 

31 calc2 : ASE calculator object 

32 Calculator corresponding to second Hamiltonian 

33 dt : float 

34 Timestep for MD simulation 

35 T : float 

36 Temperature in eV (deprecated) 

37 friction : float 

38 Friction for langevin dynamics 

39 n_eq : int 

40 Number of equilibration steps 

41 n_switch : int 

42 Number of switching steps 

43 """ 

44 

45 def __init__( 

46 self, 

47 atoms: Atoms, 

48 calc1, 

49 calc2, 

50 dt: float, 

51 T: Optional[float] = None, 

52 friction: Optional[float] = None, 

53 n_eq: Optional[int] = None, 

54 n_switch: Optional[int] = None, 

55 temperature_K: Optional[float] = None, 

56 **langevin_kwargs, 

57 ): 

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

59 friction=friction, **langevin_kwargs) 

60 if friction is None: 

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

62 if n_eq is None: 

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

64 if n_switch is None: 

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

66 self.n_eq = n_eq 

67 self.n_switch = n_switch 

68 self.lam = 0.0 

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

70 self.atoms.calc = calc 

71 

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

73 

74 def run(self): 

75 """ Run the MD switching simulation """ 

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

77 

78 # run equilibration with calc1 

79 for _ in range(self.n_eq): 

80 forces = self.step(forces) 

81 self.nsteps += 1 

82 self.call_observers() 

83 

84 # run switch from calc1 to calc2 

85 self.path_data.append( 

86 [0, self.lam, 

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

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

89 # update calculator 

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

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

92 

93 # carry out md step 

94 forces = self.step(forces) 

95 self.nsteps += 1 

96 

97 # collect data 

98 self.call_observers() 

99 self.path_data.append( 

100 [step, self.lam, 

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

102 

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

104 

105 def get_free_energy_difference(self): 

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

107 integrating dH/dlam along the switching path 

108 

109 Returns 

110 ------- 

111 float 

112 Free energy difference, F2 - F1 

113 """ 

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

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

116 

117 lambdas = self.path_data[:, 1] 

118 U1 = self.path_data[:, 2] 

119 U2 = self.path_data[:, 3] 

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

121 return delta_F 

122 

123 

124def get_lambda(step, n_switch): 

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

126 assert step >= 0 and step <= n_switch 

127 t = step / (n_switch - 1) 

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