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