Coverage for gpaw/new/lcao/hamiltonian.py: 84%

1from __future__ import annotations 

2 

3import numpy as np 

4from gpaw.core.matrix import Matrix 

5from gpaw.external import ExternalPotential 

6from gpaw.lfc import BasisFunctions 

7from gpaw.new import zips 

8from gpaw.new.calculation import DFTState 

9from gpaw.new.fd.pot_calc import FDPotentialCalculator 

10from gpaw.new.hamiltonian import Hamiltonian 

11from gpaw.new.lcao.wave_functions import LCAOWaveFunctions 

12from gpaw.typing import Array2D, Array3D 

13 

14 

15class HamiltonianMatrixCalculator: 

16 def calculate_matrix(self, 

17 wfs: LCAOWaveFunctions) -> Matrix: 

18 raise NotImplementedError 

19 

20 

21class CollinearHamiltonianMatrixCalculator(HamiltonianMatrixCalculator): 

22 def __init__(self, 

23 V_sxMM: list[np.ndarray], 

24 dH_saii: list[dict[int, Array2D]], 

25 basis: BasisFunctions, 

26 include_kinetic: bool = True): 

27 self.V_sxMM = V_sxMM 

28 self.dH_saii = dH_saii 

29 self.basis = basis 

30 self.include_kinetic = include_kinetic 

31 

32 def calculate_matrix(self, 

33 wfs: LCAOWaveFunctions) -> Matrix: 

34 if self.include_kinetic: 

35 return self._calculate_matrix_with_kinetic(wfs) 

36 return self._calculate_matrix_without_kinetic(wfs) 

37 

38 def _calculate_potential_matrix(self, 

39 wfs: LCAOWaveFunctions, 

40 V_xMM: Array3D) -> Matrix: 

41 data = V_xMM[0] 

42 _, M = data.shape 

43 if wfs.dtype == complex: 

44 data = data.astype(complex) 

45 V_MM = Matrix(M, M, data=data, dist=(wfs.band_comm, -1, 1)) 

46 if wfs.dtype == complex: 

47 phase_x = np.exp(-2j * np.pi * 

48 self.basis.sdisp_xc[1:] @ wfs.kpt_c) 

49 V_MM.data += np.einsum('x, xMN -> MN', 

50 2 * phase_x, V_xMM[1:], 

51 optimize=True) 

52 return V_MM 

53 

54 def _calculate_matrix_without_kinetic(self, 

55 wfs: LCAOWaveFunctions, 

56 V_xMM: Array3D = None, 

57 dH_aii: dict[int, Array2D] = None 

58 ) -> Matrix: 

59 if V_xMM is None: 

60 V_xMM = self.V_sxMM[wfs.spin] 

61 if dH_aii is None: 

62 dH_aii = self.dH_saii[wfs.spin] 

63 

64 V_MM = self._calculate_potential_matrix(wfs, V_xMM) 

65 

66 M1, M2 = V_MM.dist.my_row_range() 

67 for a, dH_ii in dH_aii.items(): 

68 P_Mi = wfs.P_aMi[a] 

69 V_MM.data += (P_Mi[M1:M2] @ dH_ii).conj() @ P_Mi.T # XXX use gemm 

70 

71 return V_MM 

72 

73 def _calculate_matrix_with_kinetic(self, 

74 wfs: LCAOWaveFunctions) -> Matrix: 

75 H_MM = self._calculate_matrix_without_kinetic(wfs) 

76 H_MM.data += wfs.T_MM.data 

77 

78 wfs.domain_comm.sum(H_MM.data, 0) 

79 

80 if wfs.domain_comm.rank == 0: 

81 if wfs.dtype == complex: 

82 H_MM.add_hermitian_conjugate(scale=0.5) 

83 else: 

84 H_MM.tril2full() 

85 return H_MM 

86 

87 

88class NonCollinearHamiltonianMatrixCalculator(HamiltonianMatrixCalculator): 

89 def __init__(self, matcalc: CollinearHamiltonianMatrixCalculator): 

90 self.matcalc = matcalc 

91 

92 def calculate_matrix(self, 

93 wfs: LCAOWaveFunctions) -> Matrix: 

94 V_sMM = [ 

95 self.matcalc._calculate_matrix_without_kinetic(wfs, V_xMM, dH_aii) 

96 for V_xMM, dH_aii in zips(self.matcalc.V_sxMM, 

97 self.matcalc.dH_saii)] 

98 

99 V_sMM[0] += wfs.T_MM 

100 

101 assert wfs.domain_comm.size == 1 

102 

103 for V_MM in V_sMM: 

104 wfs.domain_comm.sum(V_MM.data, 0) 

105 

106 if wfs.domain_comm.rank == 0: 

107 if wfs.dtype == complex: 

108 V_MM.add_hermitian_conjugate(scale=0.5) 

109 else: 

110 V_MM.tril2full() 

111 

112 _, M = V_MM.shape 

113 v_MM, x_MM, y_MM, z_MM = (V_MM.data for V_MM in V_sMM) 

114 H_sMsM = Matrix(2 * M, 2 * M, dtype=complex, dist=(wfs.band_comm,)) 

115 H_sMsM.data[:M, :M] = v_MM + z_MM 

116 H_sMsM.data[:M, M:] = x_MM + 1j * y_MM 

117 H_sMsM.data[M:, :M] = x_MM - 1j * y_MM 

118 H_sMsM.data[M:, M:] = v_MM - z_MM 

119 return H_sMsM 

120 

121 

122class LCAOHamiltonian(Hamiltonian): 

123 def __init__(self, 

124 basis: BasisFunctions): 

125 self.basis = basis 

126 

127 def create_hamiltonian_matrix_calculator(self, 

128 potential, 

129 ) -> HamiltonianMatrixCalculator: 

130 V_sxMM = [self.basis.calculate_potential_matrices(vt_R.data) 

131 for vt_R in potential.vt_sR.to_xp(np)] 

132 

133 dH_saii = [{a: dH_sii[s] 

134 for a, dH_sii 

135 in potential.dH_asii.to_xp(np).items()} 

136 for s in range(len(V_sxMM))] 

137 

138 matcalc = CollinearHamiltonianMatrixCalculator(V_sxMM, dH_saii, 

139 self.basis, 

140 include_kinetic=True) 

141 if len(V_sxMM) < 4: 

142 return matcalc 

143 

144 return NonCollinearHamiltonianMatrixCalculator(matcalc) 

145 

146 def create_kick_matrix_calculator(self, 

147 state: DFTState, 

148 ext: ExternalPotential, 

149 pot_calc: FDPotentialCalculator 

150 ) -> HamiltonianMatrixCalculator: 

151 from gpaw.utilities import unpack_hermitian 

152 vext_r = pot_calc.vbar_r.new() 

153 finegd = vext_r.desc._gd 

154 

155 vext_r.data = ext.get_potential(finegd) 

156 vext_R = pot_calc.restrict(vext_r) 

157 

158 nspins = state.ibzwfs.nspins 

159 

160 V_MM = self.basis.calculate_potential_matrices(vext_R.data) 

161 V_sxMM = [V_MM for s in range(nspins)] 

162 

163 W_aL = pot_calc.ghat_aLr.integrate(vext_r) 

164 

165 assert state.ibzwfs.ibz.bz.gamma_only 

166 setups_a = state.ibzwfs.wfs_qs[0][0].setups 

167 

168 dH_saii = [{a: unpack_hermitian(setups_a[a].Delta_pL @ W_L) 

169 for (a, W_L) in W_aL.items()} 

170 for s in range(nspins)] 

171 

172 return CollinearHamiltonianMatrixCalculator(V_sxMM, dH_saii, 

173 self.basis, 

174 include_kinetic=False)