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