Coverage for gpaw/tddft/ehrenfest.py: 93%

84 statements  

« prev     ^ index     » next       coverage.py v7.7.1, created at 2025-07-08 00:17 +0000

1from ase.units import Bohr, AUT, _me, _amu 

2from gpaw.tddft.units import attosec_to_autime 

3from gpaw.forces import calculate_forces 

4 

5############################################################################### 

6# EHRENFEST DYNAMICS WITHIN THE PAW METHOD 

7# WORKS WITH PAW AS LONG THERE ISN'T TOO 

8# TOO MUCH OVERLAP BETWEEN THE SPHERES 

9# SUPPORTS ALSO HGH PSEUDOPOTENTIALS 

10############################################################################### 

11 

12# m a(t+dt) = F[psi(t),x(t)] 

13# x(t+dt/2) = x(t) + v(t) dt/2 + .5 a(t) (dt/2)^2 

14# vh(t+dt/2) = v(t) + .5 a(t) dt/2 

15# m a(t+dt/2) = F[psi(t),x(t+dt/2)] 

16# v(t+dt/2) = vh(t+dt/2) + .5 a(t+dt/2) dt/2 

17# 

18# psi(t+dt) = U(t,t+dt) psi(t) 

19# 

20# m a(t+dt/2) = F[psi(t+dt),x(t+dt/2)] 

21# x(t+dt) = x(t+dt/2) + v(t+dt/2) dt/2 + .5 a(t+dt/2) (dt/2)^2 

22# vh(t+dt) = v(t+dt/2) + .5 a(t+dt/2) dt/2 

23# m a(t+dt) = F[psi(t+dt),x(t+dt)] 

24# v(t+dt) = vh(t+dt/2) + .5 a(t+dt/2) dt/2 

25 

26# TODO: move force corrections from forces.py to this module, as well as 

27# the cg method for calculating the inverse of S from overlap.py 

28 

29 

30class EhrenfestVelocityVerlet: 

31 

32 def __init__(self, calc, mass_scale=1.0, setups='paw'): 

33 """Initializes the Ehrenfest MD calculator. 

34 

35 Parameters 

36 ---------- 

37 

38 calc: TDDFT Object 

39 

40 mass_scale: 1.0 

41 Scaling coefficient for atomic masses 

42 

43 setups: {'paw', 'hgh'} 

44 Type of setups to use for the calculation 

45 

46 Note 

47 ------ 

48 

49 Use propagator = 'EFSICN' for when creating the TDDFT object from a 

50 PAW ground state calculator and propagator = 'EFSICN_HGH' for HGH 

51 pseudopotentials 

52 

53 """ 

54 self.calc = calc 

55 self.setups = setups 

56 self.x = self.calc.atoms.positions.copy() / Bohr 

57 self.xn = self.x.copy() 

58 self.v = self.x.copy() 

59 amu_to_aumass = _amu / _me 

60 if self.calc.atoms.get_velocities() is not None: 

61 self.v = self.calc.atoms.get_velocities() / (Bohr / AUT) 

62 else: 

63 self.v[:] = 0.0 

64 self.calc.atoms.set_velocities(self.v) 

65 

66 self.vt = self.v.copy() 

67 self.vh = self.v.copy() 

68 self.time = 0.0 

69 

70 self.M = calc.atoms.get_masses() * amu_to_aumass * mass_scale 

71 

72 self.a = self.v.copy() 

73 self.ah = self.a.copy() 

74 self.an = self.a.copy() 

75 self.F = self.a.copy() 

76 

77 self.calc.get_td_energy() 

78 self.F = self.get_forces() 

79 

80 for i in range(len(self.F)): 

81 self.a[i] = self.F[i] / self.M[i] 

82 

83 def get_forces(self): 

84 return calculate_forces(self.calc.wfs, 

85 self.calc.td_density.get_density(), 

86 self.calc.td_hamiltonian.hamiltonian, 

87 self.calc.log) 

88 

89 def propagate(self, dt): 

90 """Performs one Ehrenfest MD propagation step 

91 

92 Parameters 

93 ---------- 

94 

95 dt: scalar 

96 Time step (in attoseconds) used for the Ehrenfest MD step 

97 

98 """ 

99 self.x = self.calc.atoms.positions.copy() / Bohr 

100 self.v = self.calc.atoms.get_velocities() / (Bohr / AUT) 

101 

102 dt = dt * attosec_to_autime 

103 

104 # m a(t+dt) = F[psi(t),x(t)] 

105 self.calc.atoms.positions = self.x * Bohr 

106 self.calc.set_positions(self.calc.atoms) 

107 self.calc.get_td_energy() 

108 self.F = self.get_forces() 

109 

110 for i in range(len(self.F)): 

111 self.a[i] = self.F[i] / self.M[i] 

112 

113 # x(t+dt/2) = x(t) + v(t) dt/2 + .5 a(t) (dt/2)^2 

114 # vh(t+dt/2) = v(t) + .5 a(t) dt/2 

115 self.xh = self.x + self.v * dt / 2 + .5 * self.a * dt / 2 * dt / 2 

116 self.vhh = self.v + .5 * self.a * dt / 2 

117 

118 # m a(t+dt/2) = F[psi(t),x(t+dt/2)a] 

119 self.calc.atoms.positions = self.xh * Bohr 

120 self.calc.set_positions(self.calc.atoms) 

121 self.calc.get_td_energy() 

122 self.F = self.get_forces() 

123 

124 for i in range(len(self.F)): 

125 self.ah[i] = self.F[i] / self.M[i] 

126 

127 # v(t+dt/2) = vh(t+dt/2) + .5 a(t+dt/2) dt/2 

128 self.vh = self.vhh + 0.5 * self.ah * dt / 2 

129 

130 # Propagate wf 

131 # psi(t+dt) = U(t,t+dt) psi(t) 

132 self.propagate_single(dt) 

133 

134 # m a(t+dt/2) = F[psi(t+dt),x(t+dt/2)] 

135 self.calc.atoms.positions = self.xh * Bohr 

136 self.calc.set_positions(self.calc.atoms) 

137 self.calc.get_td_energy() 

138 self.F = self.get_forces() 

139 

140 for i in range(len(self.F)): 

141 self.ah[i] = self.F[i] / self.M[i] 

142 

143 # x(t+dt) = x(t+dt/2) + v(t+dt/2) dt/2 + .5 a(t+dt/2) (dt/2)^2 

144 # vh(t+dt) = v(t+dt/2) + .5 a(t+dt/2) dt/2 

145 self.xn = self.xh + self.vh * dt / 2 + 0.5 * self.ah * dt / 2 * dt / 2 

146 self.vhh = self.vh + .5 * self.ah * dt / 2 

147 

148 # m a(t+dt) = F[psi(t+dt),x(t+dt)] 

149 self.calc.atoms.positions = self.xn * Bohr 

150 self.calc.set_positions(self.calc.atoms) 

151 self.calc.get_td_energy() 

152 self.calc.update_eigenvalues() 

153 self.F = self.get_forces() 

154 

155 for i in range(len(self.F)): 

156 self.an[i] = self.F[i] / self.M[i] 

157 

158 # v(t+dt) = vh(t+dt/2) + .5 a(t+dt/2) dt/2 

159 self.vn = self.vhh + .5 * self.an * dt / 2 

160 

161 # update 

162 self.x[:] = self.xn 

163 self.v[:] = self.vn 

164 self.a[:] = self.an 

165 

166 # update atoms 

167 self.calc.atoms.set_positions(self.x * Bohr) 

168 self.calc.atoms.set_velocities(self.v * Bohr / AUT) 

169 

170 def propagate_single(self, dt): 

171 if self.setups == 'paw': 

172 self.calc.propagator.propagate(self.time, dt, self.vh) 

173 else: 

174 self.calc.propagator.propagate(self.time, dt) 

175 

176 def get_energy(self): 

177 """Updates kinetic, electronic and total energies""" 

178 

179 self.Ekin = 0.5 * (self.M * (self.v**2).sum(axis=1)).sum() 

180 self.e_coulomb = self.calc.get_td_energy() 

181 self.Etot = self.Ekin + self.e_coulomb 

182 return self.Etot 

183 

184 def get_velocities_in_au(self): 

185 return self.v 

186 

187 def set_velocities_in_au(self, v): 

188 self.v[:] = v 

189 self.calc.atoms.set_velocities(v * Bohr / AUT)