Coverage for gpaw/lcao/eigensolver.py: 98%

1import numpy as np 

2 

3 

4class DirectLCAO: 

5 """Eigensolver for LCAO-basis calculation""" 

6 

7 def __init__(self, diagonalizer=None): 

8 self.diagonalizer = diagonalizer 

9 # ??? why should we be able to set 

10 # this diagonalizer in both constructor and initialize? 

11 self.has_initialized = False # XXX 

12 

13 def initialize(self, gd, dtype, nao, diagonalizer=None): 

14 self.gd = gd 

15 self.nao = nao 

16 if diagonalizer is not None: 

17 self.diagonalizer = diagonalizer 

18 assert self.diagonalizer is not None 

19 self.has_initialized = True # XXX 

20 

21 def reset(self): 

22 pass 

23 

24 @property 

25 def error(self): 

26 return 0.0 

27 

28 @error.setter 

29 def error(self, e): 

30 pass 

31 

32 def calculate_hamiltonian_matrix(self, hamiltonian, wfs, kpt, Vt_xMM=None, 

33 root=-1, add_kinetic=True): 

34 # XXX document parallel stuff, particularly root parameter 

35 assert self.has_initialized 

36 

37 bfs = wfs.basis_functions 

38 

39 # distributed_atomic_correction works with ScaLAPACK/BLACS in general. 

40 # If SL is not enabled, it will not work with band parallelization. 

41 # But no one would want that for a practical calculation anyway. 

42 # dH_asp = wfs.atomic_correction.redistribute(wfs, hamiltonian.dH_asp) 

43 # XXXXX fix atomic corrections 

44 dH_asp = hamiltonian.dH_asp 

45 

46 if Vt_xMM is None: 

47 wfs.timer.start('Potential matrix') 

48 vt_G = hamiltonian.vt_sG[kpt.s] 

49 Vt_xMM = bfs.calculate_potential_matrices(vt_G) 

50 wfs.timer.stop('Potential matrix') 

51 

52 if bfs.gamma and wfs.dtype == float: 

53 yy = 1.0 

54 H_MM = Vt_xMM[0] 

55 else: 

56 wfs.timer.start('Sum over cells') 

57 yy = 0.5 

58 k_c = wfs.kd.ibzk_qc[kpt.q] 

59 H_MM = (0.5 + 0.0j) * Vt_xMM[0] 

60 

61 H_MM += np.einsum('x,xMN->MN', 

62 np.exp(2j * np.pi * bfs.sdisp_xc[1:] @ k_c), 

63 Vt_xMM[1:], 

64 optimize=True) 

65 wfs.timer.stop('Sum over cells') 

66 

67 # Add atomic contribution 

68 # 

69 # -- a a a* 

70 # H += > P dH P 

71 # mu nu -- mu i ij nu j 

72 # aij 

73 # 

74 name = wfs.atomic_correction.__class__.__name__ 

75 wfs.timer.start(name) 

76 wfs.atomic_correction.calculate_hamiltonian(kpt, dH_asp, H_MM, yy) 

77 wfs.timer.stop(name) 

78 

79 wfs.timer.start('Distribute overlap matrix') 

80 H_MM = wfs.ksl.distribute_overlap_matrix( 

81 H_MM, root, add_hermitian_conjugate=(yy == 0.5)) 

82 wfs.timer.stop('Distribute overlap matrix') 

83 

84 if add_kinetic: 

85 H_MM += wfs.T_qMM[kpt.q] 

86 

87 return H_MM 

88 

89 def iterate(self, hamiltonian, wfs, occ=None): 

90 wfs.timer.start('LCAO eigensolver') 

91 

92 for s in set([kpt.s for kpt in wfs.kpt_u]): 

93 wfs.timer.start('Potential matrix') 

94 Vt_xMM = wfs.basis_functions.calculate_potential_matrices( 

95 hamiltonian.vt_sG[s]) 

96 wfs.timer.stop('Potential matrix') 

97 

98 for kpt in wfs.kpt_u: 

99 if kpt.s != s: 

100 continue 

101 self.iterate_one_k_point(hamiltonian, wfs, kpt, Vt_xMM) 

102 

103 wfs.set_orthonormalized(True) 

104 wfs.timer.stop('LCAO eigensolver') 

105 

106 def iterate_one_k_point(self, hamiltonian, wfs, kpt, Vt_xMM): 

107 if wfs.bd.comm.size > 1 and wfs.bd.strided: 

108 raise NotImplementedError 

109 

110 H_MM = self.calculate_hamiltonian_matrix(hamiltonian, wfs, kpt, Vt_xMM, 

111 root=0) 

112 

113 # Decomposition step of overlap matrix can be skipped if we have 

114 # cached the result and if the solver supports it (Elpa) 

115 may_decompose = self.diagonalizer.accepts_decomposed_overlap_matrix 

116 if may_decompose: 

117 S_MM = wfs.decomposed_S_qMM[kpt.q] 

118 is_already_decomposed = (S_MM is not None) 

119 if S_MM is None: 

120 # Contents of S_MM will be overwritten by eigensolver below. 

121 S_MM = wfs.decomposed_S_qMM[kpt.q] = wfs.S_qMM[kpt.q].copy() 

122 else: 

123 is_already_decomposed = False 

124 S_MM = wfs.S_qMM[kpt.q] 

125 

126 if kpt.eps_n is None: 

127 kpt.eps_n = np.empty(wfs.bd.mynbands) 

128 

129 diagonalization_string = repr(self.diagonalizer) 

130 wfs.timer.start(diagonalization_string) 

131 # May overwrite S_MM (so the results will be stored as decomposed) 

132 if kpt.C_nM is None: 

133 kpt.C_nM = wfs.bd.empty(wfs.setups.nao, dtype=wfs.dtype) 

134 

135 self.diagonalizer.diagonalize(H_MM, kpt.C_nM, kpt.eps_n, S_MM, 

136 is_already_decomposed) 

137 wfs.timer.stop(diagonalization_string) 

138 

139 wfs.timer.start('Calculate projections') 

140 # P_ani are not strictly necessary as required quantities can be 

141 # evaluated directly using P_aMi/Paaqim. We should perhaps get rid 

142 # of the places in the LCAO code using P_ani directly 

143 wfs.atomic_correction.calculate_projections(wfs, kpt) 

144 wfs.timer.stop('Calculate projections') 

145 

146 def __repr__(self): 

147 # The "diagonalizer" must be equal to the Kohn-Sham layout 

148 # object presently. That information will be printed in the 

149 # text output anyway so we do not need it here. 

150 # 

151 # Although maybe it may be better to print it anyway... 

152 return 'LCAO using direct dense diagonalizer' 

153 

154 def estimate_memory(self, mem, dtype): 

155 pass 

156 # self.diagonalizer.estimate_memory(mem, dtype) #XXX enable this