Coverage for gpaw/new/lcao/forces.py: 100%

1from __future__ import annotations 

2 

3from typing import Dict, Tuple 

4 

5import numpy as np 

6from gpaw.core.uniform_grid import UGArray 

7from gpaw.new import zips 

8from gpaw.new.lcao.wave_functions import LCAOWaveFunctions 

9from gpaw.new.potential import Potential 

10from gpaw.typing import Array2D, Array3D 

11from gpaw.utilities.blas import mmm 

12 

13Derivatives = Tuple[Array3D, 

14 Array3D, 

15 Dict[int, Array3D]] 

16 

17 

18class TCIDerivatives: 

19 def __init__(self, manytci, atomdist, nao: int): 

20 self.manytci = manytci 

21 self.atomdist = atomdist 

22 self.nao = nao 

23 

24 self._derivatives_q: dict[int, Derivatives] = {} 

25 

26 def calculate_derivatives(self, 

27 q: int) -> Derivatives: 

28 if not self._derivatives_q: 

29 dThetadR_qvMM, dTdR_qvMM = self.manytci.O_qMM_T_qMM( 

30 self.atomdist.comm, 

31 0, self.nao, 

32 False, derivative=True) 

33 

34 self.atomdist.comm.sum(dThetadR_qvMM) 

35 self.atomdist.comm.sum(dTdR_qvMM) 

36 

37 dPdR_aqvMi = self.manytci.P_aqMi( 

38 self.atomdist.indices, 

39 derivative=True) 

40 

41 dPdR_qavMi = [{a: dPdR_qvMi[q] 

42 for a, dPdR_qvMi in dPdR_aqvMi.items()} 

43 for q in range(len(dThetadR_qvMM))] 

44 

45 self._derivatives_q = { 

46 q: (dThetadR_vMM, dTdR_vMM, dPdR_avMi) 

47 for q, (dThetadR_vMM, dTdR_vMM, dPdR_avMi) 

48 in enumerate(zips(dThetadR_qvMM, dTdR_qvMM, dPdR_qavMi))} 

49 

50 return self._derivatives_q[q] 

51 

52 

53def add_force_contributions(wfs: LCAOWaveFunctions, 

54 potential: Potential, 

55 F_av: Array2D) -> None: 

56 (dThetadR_vMM, 

57 dTdR_vMM, 

58 dPdR_avMi) = wfs.tci_derivatives.calculate_derivatives(wfs.q) 

59 

60 indices = [] 

61 M1 = 0 

62 for a, sphere in enumerate(wfs.basis.sphere_a): 

63 M2 = M1 + sphere.Mmax 

64 indices.append((a, M1, M2)) 

65 M1 = M2 

66 

67 setups = wfs.setups 

68 

69 my_row_range = wfs.T_MM.dist.my_row_range() 

70 m1, m2 = my_row_range 

71 rhoT_MM = wfs.calculate_density_matrix(transposed=True) 

72 erhoT_MM = wfs.calculate_density_matrix(transposed=True, eigs=True) 

73 add_kinetic_term(rhoT_MM, dTdR_vMM[:, m1:m2], F_av, 

74 indices, wfs.atomdist.indices, my_row_range) 

75 add_pot_term(potential.vt_sR[wfs.spin], wfs.basis, wfs.q, rhoT_MM, F_av) 

76 add_den_mat_term(erhoT_MM, dThetadR_vMM[:, m1:m2], F_av, indices, 

77 wfs.atomdist.indices, my_row_range) 

78 

79 for b in wfs.atomdist.indices: 

80 add_den_mat_paw_term(b, 

81 setups[b].dO_ii, 

82 wfs.P_aMi[b], 

83 dPdR_avMi[b], 

84 erhoT_MM, 

85 indices, 

86 F_av, 

87 my_row_range) 

88 

89 add_atomic_density_term(b, 

90 potential.dH_asii[b][wfs.spin], 

91 wfs.P_aMi[b], 

92 dPdR_avMi[b], 

93 rhoT_MM, 

94 indices, 

95 F_av, 

96 my_row_range) 

97 

98 

99def add_kinetic_term(rhoT_MM, dTdR_vMM, F_av, indices, mya, my_row_range): 

100 """Calculate Kinetic energy term in LCAO 

101 

102 ::: 

103 

104 dT 

105 _a --- -- μν 

106 F += 2 Re > > ---- ρ 

107 --- -- _ νμ 

108 μ=a ν dR 

109 μν 

110 """ 

111 

112 for a, M1, M2 in indices: 

113 if a in mya and M2 >= my_row_range[0] and M1 <= my_row_range[1]: 

114 m1 = max(M1, my_row_range[0]) - my_row_range[0] 

115 m2 = min(M2, my_row_range[1]) - my_row_range[0] 

116 F_av[a, :] += 2 * np.einsum('vmM, mM -> v', 

117 dTdR_vMM[:, m1:m2], 

118 rhoT_MM[m1:m2]).real 

119 

120 

121def add_pot_term(vt_R: UGArray, 

122 basis, 

123 q: int, 

124 rhoT_MM, 

125 F_av) -> None: 

126 """Calculate potential term""" 

127 # Potential contribution 

128 # 

129 # ----- / d Phi (r) 

130 # a \ | mu ~ 

131 # F += -2 Re ) | ---------- v (r) Phi (r) dr rho 

132 # / | d R nu nu mu 

133 # ----- / a 

134 # mu in a; nu 

135 # 

136 F_av += basis.calculate_force_contribution(vt_R.data, 

137 rhoT_MM, 

138 q) 

139 

140 

141def add_den_mat_term(erhoT_MM, dThetadR_vMM, F_av, indices, mya, my_row_range): 

142 """Calculate density matrix term in LCAO""" 

143 # Density matrix contribution due to basis overlap 

144 # 

145 # ----- d Theta 

146 # a \ mu nu 

147 # F += -2 Re ) ------------ E 

148 # / d R nu mu 

149 # ----- mu nu 

150 # mu in a; nu 

151 # 

152 for a, M1, M2 in indices: 

153 if a in mya and M2 >= my_row_range[0] and M1 <= my_row_range[1]: 

154 m1 = max(M1, my_row_range[0]) - my_row_range[0] 

155 m2 = min(M2, my_row_range[1]) - my_row_range[0] 

156 F_av[a, :] -= 2 * np.einsum('vmM, mM -> v', 

157 dThetadR_vMM[:, m1:m2], 

158 erhoT_MM[m1:m2]).real 

159 

160 

161def add_den_mat_paw_term(b, dO_ii, P_Mi, dPdR_vMi, erhoT_MM, 

162 indices, F_av, my_row_range): 

163 """Calcualte PAW correction""" 

164 # TO DO: split this function into 

165 # _get_den_mat_paw_term (which calculate Frho_av) and 

166 # get_paw_correction (which calculate ZE_MM) 

167 # Density matrix contribution from PAW correction 

168 # 

169 # ----- ----- 

170 # a \ a \ b 

171 # F += 2 Re ) Z E - 2 Re ) Z E 

172 # / mu nu nu mu / mu nu nu mu 

173 # ----- ----- 

174 # mu nu b; mu in a; nu 

175 # 

176 # with 

177 # b* 

178 # ----- dP 

179 # b \ i mu b b 

180 # Z = ) -------- dS P 

181 # mu nu / dR ij j nu 

182 # ----- b mu 

183 # ij 

184 # 

185 m1, m2 = my_row_range 

186 dtype = P_Mi.dtype 

187 Z_MM = np.zeros((m2 - m1, len(P_Mi)), dtype) 

188 dOP_iM = np.zeros((len(dO_ii), len(P_Mi)), dtype) 

189 mmm(1.0, dO_ii.astype(dtype), 'N', P_Mi, 'C', 0.0, dOP_iM) 

190 for v in range(3): 

191 mmm(1.0, 

192 dPdR_vMi[v, m1:m2], 'N', 

193 dOP_iM, 'N', 

194 0.0, Z_MM) 

195 ZE_M = np.einsum('MN, MN -> M', Z_MM, erhoT_MM).real 

196 for a, M1, M2 in indices: 

197 if M2 >= my_row_range[0] and M1 <= my_row_range[1]: 

198 am1 = max(M1, my_row_range[0]) - my_row_range[0] 

199 am2 = min(M2, my_row_range[1]) - my_row_range[0] 

200 dE = 2 * ZE_M[am1:am2].sum() 

201 F_av[a, v] -= dE # the "b; mu in a; nu" term 

202 F_av[b, v] += dE # the "mu nu" term 

203 

204 

205def add_atomic_density_term(b, dH_ii, P_Mi, dPdR_vMi, rhoT_MM, 

206 indices, F_av, my_row_range): 

207 # Atomic density contribution 

208 # ----- ----- 

209 # a \ a \ b 

210 # F += -2 Re ) A rho + 2 Re ) A rho 

211 # / mu nu nu mu / mu nu nu mu 

212 # ----- ----- 

213 # mu nu b; mu in a; nu 

214 # 

215 # b* 

216 # ----- d P 

217 # b \ i mu b b 

218 # A = ) ------- dH P 

219 # mu nu / d R ij j nu 

220 # ----- b mu 

221 # ij 

222 # 

223 m1, m2 = my_row_range 

224 dtype = P_Mi.dtype 

225 A_MM = np.zeros((m2 - m1, len(P_Mi)), dtype) 

226 dHP_iM = np.zeros((len(dH_ii), len(P_Mi)), dtype) 

227 mmm(1.0, dH_ii.astype(dtype), 'N', P_Mi, 'C', 0.0, dHP_iM) 

228 for v in range(3): 

229 mmm(1.0, 

230 dPdR_vMi[v, m1:m2], 'N', 

231 dHP_iM, 'N', 

232 0.0, A_MM) 

233 AR_M = np.einsum('MN, MN -> M', A_MM, rhoT_MM).real 

234 for a, M1, M2 in indices: 

235 if M2 >= my_row_range[0] and M1 <= my_row_range[1]: 

236 am1 = max(M1, my_row_range[0]) - my_row_range[0] 

237 am2 = min(M2, my_row_range[1]) - my_row_range[0] 

238 dE = 2 * AR_M[am1:am2].sum() 

239 F_av[a, v] += dE # the "b; mu in a; nu" term 

240 F_av[b, v] -= dE # the "mu nu" term