1"""Module defining ``Eigensolver`` classes.""" 2 3from math import pi, sqrt, sin, cos, atan2 4 5import numpy as np 6from numpy import dot 7from ase.units import Hartree 8 9from gpaw.utilities.blas import axpy 10from gpaw.utilities import unpack 11from gpaw.eigensolvers.eigensolver import Eigensolver 12 13 14class CG(Eigensolver): 15 """Conjugate gardient eigensolver 16 17 It is expected that the trial wave functions are orthonormal 18 and the integrals of projector functions and wave functions 19 are already calculated. 20 21 Solution steps are: 22 23 * Subspace diagonalization 24 * Calculate all residuals 25 * Conjugate gradient steps 26 """ 27 28 def __init__(self, niter=4, rtol=0.30, tw_coeff=False): 29 """Construct conjugate gradient eigen solver. 30 31 parameters: 32 33 niter: int 34 Maximum number of conjugate gradient iterations per band 35 rtol: float 36 If change in residual is less than rtol, iteration for band is 37 not continued 38 39 """ 40 Eigensolver.__init__(self) 41 self.niter = niter 42 self.rtol = rtol 43 self.orthonormalization_required = False 44 self.tw_coeff = tw_coeff 45 46 self.tolerance = None 47 48 def __repr__(self): 49 return 'CG(niter=%d, rtol=%5.1e)' % (self.niter, self.rtol) 50 51 def todict(self): 52 return {'name': 'cg', 'niter': self.niter} 53 54 def initialize(self, wfs): 55 if wfs.bd.comm.size > 1: 56 raise ValueError('CG eigensolver does not support band ' 57 'parallelization. This calculation parallelizes ' 58 'over %d band groups.' % wfs.bd.comm.size) 59 Eigensolver.initialize(self, wfs) 60 61 def iterate_one_k_point(self, ham, wfs, kpt, weights): 62 """Do conjugate gradient iterations for the k-point""" 63 self.timer.start('CG') 64 65 niter = self.niter 66 67 phi_G, phi_old_G, Htphi_G = wfs.empty(3, q=kpt.q) 68 69 comm = wfs.gd.comm 70 if self.tw_coeff: 71 # Wait! What business does the eigensolver have changing 72 # the properties of the Hamiltonian? We are not updating 73 # the Hamiltonian here. Moreover, what is supposed to 74 # happen if this function is called multiple times per 75 # iteration? Then we keep dividing the potential by the 76 # same number. What on earth is the meaning of this? 77 # 78 # Also the parameter tw_coeff is undocumented. What is it? 79 ham.vt_sG /= self.tw_coeff 80 # Assuming the ordering in dH_asp and wfs is the same 81 for a in ham.dH_asp.keys(): 82 ham.dH_asp[a] /= self.tw_coeff 83 84 psit = kpt.psit 85 R = psit.new(buf=wfs.work_array) 86 P = kpt.projections 87 P2 = P.new() 88 89 self.subspace_diagonalize(ham, wfs, kpt) 90 91 Htpsit = psit.new(buf=self.Htpsit_nG) 92 93 R.array[:] = Htpsit.array 94 self.calculate_residuals(kpt, wfs, ham, psit, 95 P, kpt.eps_n, R, P2) 96 97 total_error = 0.0 98 for n in range(self.nbands): 99 N = self.nbands 100 R_G = R.array[n] 101 Htpsit_G = Htpsit.array[n] 102 psit_G = psit.array[n] 103 gamma_old = 1.0 104 phi_old_G[:] = 0.0 105 error = np.real(wfs.integrate(R_G, R_G)) 106 for nit in range(niter): 107 if (error * Hartree**2 < self.tolerance / self.nbands): 108 break 109 110 ekin = self.preconditioner.calculate_kinetic_energy(psit_G, 111 kpt) 112 113 pR_G = self.preconditioner(R_G, kpt, ekin) 114 115 # New search direction 116 gamma = comm.sum(np.vdot(pR_G, R_G).real) 117 phi_G[:] = -pR_G - gamma / gamma_old * phi_old_G 118 gamma_old = gamma 119 phi_old_G[:] = phi_G[:] 120 121 # Calculate projections 122 P2_ai = wfs.pt.dict() 123 wfs.pt.integrate(phi_G, P2_ai, kpt.q) 124 125 # Orthonormalize phi_G to all bands 126 self.timer.start('CG: orthonormalize') 127 self.timer.start('CG: overlap') 128 overlap_n = wfs.integrate(psit.array[:N], phi_G, 129 global_integral=False) 130 self.timer.stop('CG: overlap') 131 self.timer.start('CG: overlap2') 132 for a, P2_i in P2_ai.items(): 133 P_ni = kpt.P_ani[a] 134 dO_ii = wfs.setups[a].dO_ii 135 overlap_n += np.dot(P_ni[:N].conjugate(), 136 np.dot(dO_ii, P2_i)) 137 self.timer.stop('CG: overlap2') 138 comm.sum(overlap_n) 139 140 phi_G -= psit.array[:N].T.dot(overlap_n).T 141 142 for a, P2_i in P2_ai.items(): 143 P_ni = kpt.P_ani[a] 144 P2_i -= np.dot(overlap_n, P_ni[:N]) 145 146 norm = wfs.integrate(phi_G, phi_G, global_integral=False) 147 for a, P2_i in P2_ai.items(): 148 dO_ii = wfs.setups[a].dO_ii 149 norm += np.vdot(P2_i, np.dot(dO_ii, P2_i)) 150 norm = comm.sum(float(np.real(norm))) 151 phi_G /= sqrt(norm) 152 for P2_i in P2_ai.values(): 153 P2_i /= sqrt(norm) 154 self.timer.stop('CG: orthonormalize') 155 156 # find optimum linear combination of psit_G and phi_G 157 an = kpt.eps_n[n] 158 wfs.apply_pseudo_hamiltonian(kpt, ham, 159 phi_G.reshape((1,) + phi_G.shape), 160 Htphi_G.reshape((1,) + 161 Htphi_G.shape)) 162 b = wfs.integrate(phi_G, Htpsit_G, global_integral=False) 163 c = wfs.integrate(phi_G, Htphi_G, global_integral=False) 164 for a, P2_i in P2_ai.items(): 165 P_i = kpt.P_ani[a][n] 166 dH_ii = unpack(ham.dH_asp[a][kpt.s]) 167 b += dot(P2_i, dot(dH_ii, P_i.conj())) 168 c += dot(P2_i, dot(dH_ii, P2_i.conj())) 169 b = comm.sum(float(np.real(b))) 170 c = comm.sum(float(np.real(c))) 171 172 theta = 0.5 * atan2(2 * b, an - c) 173 enew = (an * cos(theta)**2 + 174 c * sin(theta)**2 + 175 b * sin(2.0 * theta)) 176 # theta can correspond either minimum or maximum 177 if (enew - kpt.eps_n[n]) > 0.0: # we were at maximum 178 theta += pi / 2.0 179 enew = (an * cos(theta)**2 + 180 c * sin(theta)**2 + 181 b * sin(2.0 * theta)) 182 183 kpt.eps_n[n] = enew 184 psit_G *= cos(theta) 185 # kpt.psit_nG[n] += sin(theta) * phi_G 186 axpy(sin(theta), phi_G, psit_G) 187 for a, P2_i in P2_ai.items(): 188 P_i = kpt.P_ani[a][n] 189 P_i *= cos(theta) 190 P_i += sin(theta) * P2_i 191 192 if nit < niter - 1: 193 Htpsit_G *= cos(theta) 194 # Htpsit_G += sin(theta) * Htphi_G 195 axpy(sin(theta), Htphi_G, Htpsit_G) 196 # adjust residuals 197 R_G[:] = Htpsit_G - kpt.eps_n[n] * psit_G 198 199 coef_ai = wfs.pt.dict() 200 for a, coef_i in coef_ai.items(): 201 P_i = kpt.P_ani[a][n] 202 dO_ii = wfs.setups[a].dO_ii 203 dH_ii = unpack(ham.dH_asp[a][kpt.s]) 204 coef_i[:] = (dot(P_i, dH_ii) - 205 dot(P_i * kpt.eps_n[n], dO_ii)) 206 wfs.pt.add(R_G, coef_ai, kpt.q) 207 error_new = np.real(wfs.integrate(R_G, R_G)) 208 if error_new / error < self.rtol: 209 # print >> self.f, "cg:iters", n, nit+1 210 break 211 if (self.nbands_converge == 'occupied' and 212 kpt.f_n is not None and kpt.f_n[n] == 0.0): 213 # print >> self.f, "cg:iters", n, nit+1 214 break 215 error = error_new 216 217 total_error += weights[n] * error 218 # if nit == 3: 219 # print >> self.f, "cg:iters", n, nit+1 220 if self.tw_coeff: # undo the scaling for calculating energies 221 for i in range(len(kpt.eps_n)): 222 kpt.eps_n[i] *= self.tw_coeff 223 ham.vt_sG *= self.tw_coeff 224 # Assuming the ordering in dH_asp and wfs is the same 225 for a in ham.dH_asp.keys(): 226 ham.dH_asp[a] *= self.tw_coeff 227 228 self.timer.stop('CG') 229 return total_error 230