Coverage for /builds/debichem-team/python-ase/ase/md/switch_langevin.py: 91.67%

1from typing import Any, List, Optional 

2 

3import numpy as np 

4from scipy.integrate import trapezoid 

5 

6from ase import Atoms 

7from ase.calculators.mixing import MixedCalculator 

8from ase.md.langevin import Langevin 

9 

10 

11class SwitchLangevin(Langevin): 

12 """ 

13 MD class for carrying out thermodynamic integration between two 

14 hamiltonians 

15 

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

17 

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

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

20 

21 Attributes 

22 ---------- 

23 path_data : numpy array 

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

25 

26 Parameters 

27 ---------- 

28 atoms : ASE Atoms object 

29 Atoms object for which MD will be run 

30 calc1 : ASE calculator object 

31 Calculator corresponding to first Hamiltonian 

32 calc2 : ASE calculator object 

33 Calculator corresponding to second Hamiltonian 

34 dt : float 

35 Timestep for MD simulation 

36 T : float 

37 Temperature in eV (deprecated) 

38 friction : float 

39 Friction for langevin dynamics 

40 n_eq : int 

41 Number of equilibration steps 

42 n_switch : int 

43 Number of switching steps 

44 """ 

45 

46 def __init__( 

47 self, 

48 atoms: Atoms, 

49 calc1, 

50 calc2, 

51 dt: float, 

52 T: Optional[float] = None, 

53 friction: Optional[float] = None, 

54 n_eq: Optional[int] = None, 

55 n_switch: Optional[int] = None, 

56 temperature_K: Optional[float] = None, 

57 **langevin_kwargs, 

58 ): 

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

60 friction=friction, **langevin_kwargs) 

61 if friction is None: 

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

63 if n_eq is None: 

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

65 if n_switch is None: 

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

67 self.n_eq = n_eq 

68 self.n_switch = n_switch 

69 self.lam = 0.0 

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

71 self.atoms.calc = calc 

72 

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

74 

75 def run(self): 

76 """ Run the MD switching simulation """ 

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

78 

79 # run equilibration with calc1 

80 for _ in range(self.n_eq): 

81 forces = self.step(forces) 

82 self.nsteps += 1 

83 self.call_observers() 

84 

85 # run switch from calc1 to calc2 

86 self.path_data.append( 

87 [0, self.lam, 

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

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

90 # update calculator 

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

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

93 

94 # carry out md step 

95 forces = self.step(forces) 

96 self.nsteps += 1 

97 

98 # collect data 

99 self.call_observers() 

100 self.path_data.append( 

101 [step, self.lam, 

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

103 

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

105 

106 def get_free_energy_difference(self): 

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

108 integrating dH/dlam along the switching path 

109 

110 Returns 

111 ------- 

112 float 

113 Free energy difference, F2 - F1 

114 """ 

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

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

117 

118 lambdas = self.path_data[:, 1] 

119 U1 = self.path_data[:, 2] 

120 U2 = self.path_data[:, 3] 

121 delta_F = trapezoid(U2 - U1, lambdas) 

122 return delta_F 

123 

124 

125def get_lambda(step, n_switch): 

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

127 assert step >= 0 and step <= n_switch 

128 t = step / (n_switch - 1) 

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