Coverage for gpaw/directmin/lcao/etdm_helper_lcao.py: 94%

1""" 

2Helper class for LCAOETDM. 

3 

4Handles orbital initialization, setting reference orbitals, applying the 

5unitary transformation, calculating the gradient, and getting canonical 

6representation. 

7""" 

8 

9import numpy as np 

10from gpaw.lcao.eigensolver import DirectLCAO 

11from gpaw.directmin.functional.lcao import get_functional 

12from gpaw.utilities.tools import tri2full 

13 

14 

15class ETDMHelperLCAO(DirectLCAO): 

16 

17 def __init__(self, wfs, dens, ham, nkpts, functional, diagonalizer=None, 

18 orthonormalization='gramschmidt', 

19 need_init_orbs=True): 

20 

21 super(ETDMHelperLCAO, self).__init__(diagonalizer) 

22 super(ETDMHelperLCAO, self).initialize(wfs.gd, wfs.dtype, 

23 wfs.setups.nao, wfs.ksl) 

24 self.orthonormalization = orthonormalization 

25 self.need_init_orbs = need_init_orbs 

26 self.nkpts = nkpts 

27 self.func = get_functional(functional, wfs, dens, ham) 

28 self.reference_orbitals = {} 

29 self.initialize_orbitals(wfs, ham) 

30 

31 def __repr__(self): 

32 pass 

33 

34 def set_reference_orbitals(self, wfs, n_dim): 

35 for kpt in wfs.kpt_u: 

36 u = self.kpointval(kpt) 

37 self.reference_orbitals[u] = np.copy(kpt.C_nM[:n_dim[u]]) 

38 

39 def appy_transformation_kpt(self, wfs, u_mat, kpt, c_ref=None, 

40 broadcast=True, 

41 update_proj=True): 

42 """ 

43 If c_ref are not provided then 

44 kpt.C_nM <- u_mat kpt.C_nM 

45 otherwise kpt.C_nM <- u_mat c_ref 

46 """ 

47 

48 dimens1 = u_mat.shape[1] 

49 dimens2 = u_mat.shape[0] 

50 

51 if c_ref is None: 

52 kpt.C_nM[:dimens2] = u_mat @ kpt.C_nM[:dimens1] 

53 else: 

54 kpt.C_nM[:dimens2] = u_mat @ c_ref[:dimens1] 

55 

56 if broadcast: 

57 with wfs.timer('Broadcast coefficients'): 

58 wfs.gd.comm.broadcast(kpt.C_nM, 0) 

59 if update_proj: 

60 with wfs.timer('Calculate projections'): 

61 wfs.atomic_correction.calculate_projections(wfs, kpt) 

62 

63 def initialize_orbitals(self, wfs, ham): 

64 """ 

65 If it is the first use of the scf then initialize 

66 coefficient matrix using eigensolver 

67 """ 

68 

69 orthname = self.orthonormalization 

70 need_canon_coef = \ 

71 (not wfs.coefficients_read_from_file and self.need_init_orbs) 

72 if need_canon_coef or orthname == 'diag': 

73 super(ETDMHelperLCAO, self).iterate(ham, wfs) 

74 else: 

75 wfs.orthonormalize(type=orthname) 

76 wfs.coefficients_read_from_file = False 

77 self.need_init_orbs = False 

78 

79 def calc_grad(self, wfs, ham, kpt, evecs, evals, matrix_exp, 

80 representation, ind_up, constraints): 

81 """ 

82 Gradient w.r.t. skew-Hermitian matrices 

83 """ 

84 

85 h_mm = self.calculate_hamiltonian_matrix(ham, wfs, kpt) 

86 # make matrix hermitian 

87 tri2full(h_mm) 

88 # calc gradient and eigenstate error 

89 g_mat, error = self.func.get_gradients( 

90 h_mm, kpt.C_nM, kpt.f_n, evecs, evals, 

91 kpt, wfs, wfs.timer, matrix_exp, 

92 representation, ind_up, constraints) 

93 

94 return g_mat, error 

95 

96 def update_to_canonical_orbitals(self, wfs, ham, kpt, 

97 update_ref_orbs_canonical, restart): 

98 """ 

99 Choose canonical orbitals 

100 """ 

101 

102 h_mm = self.calculate_hamiltonian_matrix(ham, wfs, kpt) 

103 tri2full(h_mm) 

104 

105 if self.func.name == 'ks': 

106 if update_ref_orbs_canonical or restart: 

107 # Diagonalize entire Hamiltonian matrix 

108 with wfs.timer('Diagonalize and rotate'): 

109 kpt.C_nM, kpt.eps_n = rotate_subspace(h_mm, kpt.C_nM) 

110 else: 

111 # Diagonalize equally occupied subspaces separately 

112 f_unique = np.unique(kpt.f_n) 

113 for f in f_unique: 

114 with wfs.timer('Diagonalize and rotate'): 

115 kpt.C_nM[kpt.f_n == f, :], kpt.eps_n[kpt.f_n == f] = \ 

116 rotate_subspace(h_mm, kpt.C_nM[kpt.f_n == f, :]) 

117 elif self.func.name == 'PZ-SIC': 

118 self.func.get_lagrange_matrices( 

119 h_mm, kpt.C_nM, kpt.f_n, kpt, wfs, 

120 update_eigenvalues=True) 

121 

122 with wfs.timer('Calculate projections'): 

123 self.update_projections(wfs, kpt) 

124 

125 def sort_orbitals(self, wfs, kpt, ind): 

126 """ 

127 sort orbitals according to indices stored in ind 

128 """ 

129 kpt.C_nM[np.arange(len(ind)), :] = kpt.C_nM[ind, :] 

130 self.update_projections(wfs, kpt) 

131 

132 def update_projections(self, wfs, kpt): 

133 """ 

134 calculate projections kpt.P_ani 

135 """ 

136 

137 wfs.atomic_correction.calculate_projections(wfs, kpt) 

138 

139 def orbital_energies(self, wfs, ham, kpt): 

140 """ 

141 diagonal elements of hamiltonian matrix in orbital representation 

142 """ 

143 

144 if self.func.name == 'ks': 

145 h_mm = self.calculate_hamiltonian_matrix(ham, wfs, kpt) 

146 tri2full(h_mm) 

147 # you probably need only diagonal terms? 

148 # wouldn't "for" be faster? 

149 h_mm = kpt.C_nM.conj() @ h_mm.conj() @ kpt.C_nM.T 

150 energies = h_mm.diagonal().real.copy() 

151 elif self.func.name == 'PZ-SIC': 

152 energies = self.func.lagr_diag_s[wfs.eigensolver.kpointval(kpt)] 

153 

154 return energies 

155 

156 def kpointval(self, kpt): 

157 return self.nkpts * kpt.s + kpt.q 

158 

159 

160def rotate_subspace(h_mm, c_nm): 

161 """ 

162 choose canonical orbitals 

163 """ 

164 l_nn = c_nm @ h_mm @ c_nm.conj().T 

165 # check if diagonal then don't rotate? it could save a bit of time 

166 eps, w = np.linalg.eigh(l_nn) 

167 return w.T.conj() @ c_nm, eps