Exemplo n.º 1
0
 def H(psit_xG):
     if self.keep_htpsit:
         result_xG = Htpsit_nG
     else:
         result_xG = reshape(self.operator.work1_xG, psit_xG.shape)
     wfs.apply_pseudo_hamiltonian(kpt, hamiltonian, psit_xG,
                                  result_xG)
     return result_xG
Exemplo n.º 2
0
 def H(psit_xG):
     if self.keep_htpsit:
         result_xG = Htpsit_nG
     else:
         result_xG = reshape(self.operator.work1_xG, psit_xG.shape)
     wfs.apply_pseudo_hamiltonian(kpt, hamiltonian, psit_xG,
                                  result_xG)
     hamiltonian.xc.apply_orbital_dependent_hamiltonian(
         kpt, psit_xG, result_xG, hamiltonian.dH_asp)
     return result_xG
Exemplo n.º 3
0
 def H(psit_xG):
     if self.keep_htpsit:
         result_xG = Htpsit_nG
     else:
         if use_mic:
             result_xG = self.operator.work1_xG_mic
         else:
             result_xG = reshape(self.operator.work1_xG, psit_xG.shape)
     if use_mic:
         psit_xG.update_device()
         wfs.apply_pseudo_hamiltonian(kpt, hamiltonian, psit_xG.array,
                                      result_xG.array)
         result_xG.update_device()
         stream.sync()
     else:
         wfs.apply_pseudo_hamiltonian(kpt, hamiltonian, psit_xG,
                                      result_xG)
     hamiltonian.xc.apply_orbital_dependent_hamiltonian(
         kpt, psit_xG, result_xG, hamiltonian.dH_asp)
     return result_xG
 def H(psit_xG):
     if self.keep_htpsit:
         result_xG = Htpsit_nG
     else:
         if use_mic:
             result_xG = self.operator.work1_xG_mic
         else:
             result_xG = reshape(self.operator.work1_xG, psit_xG.shape)
     if use_mic:
         psit_xG.update_device()
         wfs.apply_pseudo_hamiltonian(kpt, hamiltonian, psit_xG.array,
                                      result_xG.array)
         result_xG.update_device() 
         stream.sync()
     else:
         wfs.apply_pseudo_hamiltonian(kpt, hamiltonian, psit_xG,
                                      result_xG)
     hamiltonian.xc.apply_orbital_dependent_hamiltonian(
         kpt, psit_xG, result_xG, hamiltonian.dH_asp)
     return result_xG
Exemplo n.º 5
0
    def iterate_one_k_point(self, hamiltonian, wfs, kpt):
        """Do Davidson iterations for the kpoint"""
        niter = self.niter
        nbands = self.nbands

        gd = wfs.matrixoperator.gd

        psit_nG, Htpsit_nG = self.subspace_diagonalize(hamiltonian, wfs, kpt)
        # Note that psit_nG is now in self.operator.work1_nG and
        # Htpsit_nG is in kpt.psit_nG!

        H_2n2n = self.H_2n2n
        S_2n2n = self.S_2n2n
        eps_2n = self.eps_2n
        psit2_nG = reshape(self.Htpsit_nG, psit_nG.shape)

        self.timer.start("Davidson")
        R_nG = Htpsit_nG
        self.calculate_residuals(kpt, wfs, hamiltonian, psit_nG, kpt.P_ani, kpt.eps_n, R_nG)

        def integrate(a_G, b_G):
            return np.real(wfs.integrate(a_G, b_G, global_integral=False))

        for nit in range(niter):
            H_2n2n[:] = 0.0
            S_2n2n[:] = 0.0

            norm_n = np.zeros(nbands)
            error = 0.0
            for n in range(nbands):
                if kpt.f_n is None:
                    weight = kpt.weight
                else:
                    weight = kpt.f_n[n]
                if self.nbands_converge != "occupied":
                    if n < self.nbands_converge:
                        weight = kpt.weight
                    else:
                        weight = 0.0
                error += weight * integrate(R_nG[n], R_nG[n])

                ekin = self.preconditioner.calculate_kinetic_energy(psit_nG[n : n + 1], kpt)
                psit2_nG[n] = self.preconditioner(R_nG[n : n + 1], kpt, ekin)

                if self.normalize:
                    norm_n[n] = integrate(psit2_nG[n], psit2_nG[n])

                H_2n2n[n, n] = kpt.eps_n[n]
                S_2n2n[n, n] = 1.0

            if self.normalize:
                gd.comm.sum(norm_n)
                for norm, psit2_G in zip(norm_n, psit2_nG):
                    psit2_G *= norm ** -0.5

            # Calculate projections
            P2_ani = wfs.pt.dict(nbands)
            wfs.pt.integrate(psit2_nG, P2_ani, kpt.q)

            # Hamiltonian matrix
            # <psi2 | H | psi>
            wfs.apply_pseudo_hamiltonian(kpt, hamiltonian, psit2_nG, Htpsit_nG)
            gd.integrate(psit_nG, Htpsit_nG, global_integral=False, _transposed_result=self.H_nn)
            # gemm(1.0, psit_nG, Htpsit_nG, 0.0, self.H_nn, 'c')

            for a, P_ni in kpt.P_ani.items():
                P2_ni = P2_ani[a]
                dH_ii = unpack(hamiltonian.dH_asp[a][kpt.s])
                self.H_nn += np.dot(P2_ni, np.dot(dH_ii, P_ni.T.conj()))

            gd.comm.sum(self.H_nn, 0)
            H_2n2n[nbands:, :nbands] = self.H_nn

            # <psi2 | H | psi2>
            gd.integrate(psit2_nG, Htpsit_nG, global_integral=False, _transposed_result=self.H_nn)
            # r2k(0.5 * gd.dv, psit2_nG, Htpsit_nG, 0.0, self.H_nn)
            for a, P2_ni in P2_ani.items():
                dH_ii = unpack(hamiltonian.dH_asp[a][kpt.s])
                self.H_nn += np.dot(P2_ni, np.dot(dH_ii, P2_ni.T.conj()))

            gd.comm.sum(self.H_nn, 0)
            H_2n2n[nbands:, nbands:] = self.H_nn

            # Overlap matrix
            # <psi2 | S | psi>
            gd.integrate(psit_nG, psit2_nG, global_integral=False, _transposed_result=self.S_nn)
            # gemm(1.0, psit_nG, psit2_nG, 0.0, self.S_nn, 'c')

            for a, P_ni in kpt.P_ani.items():
                P2_ni = P2_ani[a]
                dO_ii = wfs.setups[a].dO_ii
                self.S_nn += np.dot(P2_ni, np.inner(dO_ii, P_ni.conj()))

            gd.comm.sum(self.S_nn, 0)
            S_2n2n[nbands:, :nbands] = self.S_nn

            # <psi2 | S | psi2>
            gd.integrate(psit2_nG, psit2_nG, global_integral=False, _transposed_result=self.S_nn)
            # rk(gd.dv, psit2_nG, 0.0, self.S_nn)
            for a, P2_ni in P2_ani.items():
                dO_ii = wfs.setups[a].dO_ii
                self.S_nn += np.dot(P2_ni, np.dot(dO_ii, P2_ni.T.conj()))

            gd.comm.sum(self.S_nn, 0)
            S_2n2n[nbands:, nbands:] = self.S_nn

            if gd.comm.rank == 0:
                m = 0
                if self.smin:
                    s_N, U_NN = np.linalg.eigh(S_2n2n)
                    m = int((s_N < self.smin).sum())

                if m == 0:
                    general_diagonalize(H_2n2n, eps_2n, S_2n2n)
                else:
                    T_Nn = np.dot(U_NN[:, m:], np.diag(s_N[m:] ** -0.5))
                    H_2n2n[:nbands, nbands:] = H_2n2n[nbands:, :nbands].conj().T
                    eps_2n[:-m], P_nn = np.linalg.eigh(np.dot(np.dot(T_Nn.T.conj(), H_2n2n), T_Nn))
                    H_2n2n[:-m] = np.dot(T_Nn, P_nn).T

            gd.comm.broadcast(H_2n2n, 0)
            gd.comm.broadcast(eps_2n, 0)

            kpt.eps_n[:] = eps_2n[:nbands]

            # Rotate psit_nG
            gd.gemm(1.0, psit_nG, H_2n2n[:nbands, :nbands], 0.0, Htpsit_nG)
            gd.gemm(1.0, psit2_nG, H_2n2n[:nbands, nbands:], 1.0, Htpsit_nG)
            psit_nG, Htpsit_nG = Htpsit_nG, psit_nG

            # Rotate P_uni:
            for a, P_ni in kpt.P_ani.items():
                P2_ni = P2_ani[a]
                gemm(1.0, P_ni.copy(), H_2n2n[:nbands, :nbands], 0.0, P_ni)
                gemm(1.0, P2_ni, H_2n2n[:nbands, nbands:], 1.0, P_ni)

            if nit < niter - 1:
                wfs.apply_pseudo_hamiltonian(kpt, hamiltonian, psit_nG, Htpsit_nG)
                R_nG = Htpsit_nG
                self.calculate_residuals(kpt, wfs, hamiltonian, psit_nG, kpt.P_ani, kpt.eps_n, R_nG)

        self.timer.stop("Davidson")
        error = gd.comm.sum(error)
        return error, psit_nG
Exemplo n.º 6
0
    def subspace_diagonalize(self, hamiltonian, wfs, kpt):
        """Diagonalize the Hamiltonian in the subspace of kpt.psit_nG

        *Htpsit_nG* is a work array of same size as psit_nG which contains
        the local part of the Hamiltonian times psit on exit

        First, the Hamiltonian (defined by *kin*, *vt_sG*, and
        *dH_asp*) is applied to the wave functions, then the *H_nn*
        matrix is calculated and diagonalized, and finally, the wave
        functions (and also Htpsit_nG are rotated.  Also the
        projections *P_ani* are rotated.

        It is assumed that the wave functions *psit_nG* are orthonormal
        and that the integrals of projector functions and wave functions
        *P_ani* are already calculated.

        Return ratated wave functions and H applied to the rotated
        wave functions if self.keep_htpsit is True.
        """

        if self.band_comm.size > 1 and wfs.bd.strided:
            raise NotImplementedError

        self.timer.start('Subspace diag')

        if use_mic:
            psit_nG = kpt.psit_nG_mic
            # psit_nG.update_device()
            # stream.sync()
        else:
            psit_nG = kpt.psit_nG
        P_ani = kpt.P_ani

        if self.keep_htpsit:
            if use_mic:
                Htpsit_nG = self.Htpsit_nG_mic
            else:
                Htpsit_nG = reshape(self.Htpsit_nG, psit_nG.shape)
        else:
            Htpsit_nG = None

        def H(psit_xG):
            if self.keep_htpsit:
                result_xG = Htpsit_nG
            else:
                if use_mic:
                    result_xG = self.operator.work1_xG_mic
                else:
                    result_xG = reshape(self.operator.work1_xG, psit_xG.shape)
            if use_mic:
                psit_xG.update_device()
                wfs.apply_pseudo_hamiltonian(kpt, hamiltonian, psit_xG.array,
                                             result_xG.array)
                result_xG.update_device()
                stream.sync()
            else:
                wfs.apply_pseudo_hamiltonian(kpt, hamiltonian, psit_xG,
                                             result_xG)
            hamiltonian.xc.apply_orbital_dependent_hamiltonian(
                kpt, psit_xG, result_xG, hamiltonian.dH_asp)
            return result_xG

        def dH(a, P_ni):
            return np.dot(P_ni, unpack(hamiltonian.dH_asp[a][kpt.s]))

        self.timer.start('calc_h_matrix')
        H_nn = self.operator.calculate_matrix_elements(psit_nG, P_ani, H, dH)
        hamiltonian.xc.correct_hamiltonian_matrix(kpt, H_nn)
        self.timer.stop('calc_h_matrix')

        diagonalization_string = repr(self.ksl)
        wfs.timer.start(diagonalization_string)
        self.ksl.diagonalize(H_nn, kpt.eps_n)
        # H_nn now contains the result of the diagonalization.
        wfs.timer.stop(diagonalization_string)

        self.timer.start('rotate_psi')
        psit_nG = self.operator.matrix_multiply(H_nn, psit_nG, P_ani)
        if self.keep_htpsit:
            if use_mic:
                Htpsit_nG = self.operator.matrix_multiply(
                    H_nn, Htpsit_nG, out_nG=kpt.psit_nG_mic)

            else:
                Htpsit_nG = self.operator.matrix_multiply(H_nn,
                                                          Htpsit_nG,
                                                          out_nG=kpt.psit_nG)

        # Rotate orbital dependent XC stuff:
        hamiltonian.xc.rotate(kpt, H_nn)

        self.timer.stop('rotate_psi')
        self.timer.stop('Subspace diag')

        if use_mic:
            psit_nG.update_host()
            stream.sync()
            if self.keep_htpsit:
                Htpsit_nG.update_host()
                stream.sync()
                return psit_nG.array, Htpsit_nG.array
            else:
                return psit_nG.array, Htpsit_nG
        else:
            return psit_nG, Htpsit_nG
    def subspace_diagonalize(self, hamiltonian, wfs, kpt):
        """Diagonalize the Hamiltonian in the subspace of kpt.psit_nG

        *Htpsit_nG* is a work array of same size as psit_nG which contains
        the local part of the Hamiltonian times psit on exit

        First, the Hamiltonian (defined by *kin*, *vt_sG*, and
        *dH_asp*) is applied to the wave functions, then the *H_nn*
        matrix is calculated and diagonalized, and finally, the wave
        functions (and also Htpsit_nG are rotated.  Also the
        projections *P_ani* are rotated.

        It is assumed that the wave functions *psit_nG* are orthonormal
        and that the integrals of projector functions and wave functions
        *P_ani* are already calculated.

        Return ratated wave functions and H applied to the rotated
        wave functions if self.keep_htpsit is True.
        """

        if self.band_comm.size > 1 and wfs.bd.strided:
            raise NotImplementedError

        self.timer.start('Subspace diag')

        if use_mic:
            psit_nG = kpt.psit_nG_mic
            # psit_nG.update_device()
            # stream.sync()
        else:
            psit_nG = kpt.psit_nG
        P_ani = kpt.P_ani

        if self.keep_htpsit:
            if use_mic:
                Htpsit_nG = self.Htpsit_nG_mic
            else:
                Htpsit_nG = reshape(self.Htpsit_nG, psit_nG.shape)
        else:
            Htpsit_nG = None

        def H(psit_xG):
            if self.keep_htpsit:
                result_xG = Htpsit_nG
            else:
                if use_mic:
                    result_xG = self.operator.work1_xG_mic
                else:
                    result_xG = reshape(self.operator.work1_xG, psit_xG.shape)
            if use_mic:
                psit_xG.update_device()
                wfs.apply_pseudo_hamiltonian(kpt, hamiltonian, psit_xG.array,
                                             result_xG.array)
                result_xG.update_device() 
                stream.sync()
            else:
                wfs.apply_pseudo_hamiltonian(kpt, hamiltonian, psit_xG,
                                             result_xG)
            hamiltonian.xc.apply_orbital_dependent_hamiltonian(
                kpt, psit_xG, result_xG, hamiltonian.dH_asp)
            return result_xG

        def dH(a, P_ni):
            return np.dot(P_ni, unpack(hamiltonian.dH_asp[a][kpt.s]))

        self.timer.start('calc_h_matrix')
        H_nn = self.operator.calculate_matrix_elements(psit_nG, P_ani,
                                                       H, dH)
        hamiltonian.xc.correct_hamiltonian_matrix(kpt, H_nn)
        self.timer.stop('calc_h_matrix')

        diagonalization_string = repr(self.ksl)
        wfs.timer.start(diagonalization_string)
        self.ksl.diagonalize(H_nn, kpt.eps_n)
        # H_nn now contains the result of the diagonalization.
        wfs.timer.stop(diagonalization_string)

        self.timer.start('rotate_psi')
        psit_nG = self.operator.matrix_multiply(H_nn, psit_nG, P_ani)
        if self.keep_htpsit:
            if use_mic:
                Htpsit_nG = self.operator.matrix_multiply(H_nn, Htpsit_nG,
                                                          out_nG=kpt.psit_nG_mic)
                 
            else:
                Htpsit_nG = self.operator.matrix_multiply(H_nn, Htpsit_nG,
                                                          out_nG=kpt.psit_nG)

        # Rotate orbital dependent XC stuff:
        hamiltonian.xc.rotate(kpt, H_nn)

        self.timer.stop('rotate_psi')
        self.timer.stop('Subspace diag')

        if use_mic:
            psit_nG.update_host()
            stream.sync()
            if self.keep_htpsit:
                Htpsit_nG.update_host()
                stream.sync()
                return psit_nG.array, Htpsit_nG.array
            else:
                return psit_nG.array, Htpsit_nG
        else:
            return psit_nG, Htpsit_nG
Exemplo n.º 8
0
    def iterate_one_k_point(self, hamiltonian, wfs, kpt):
        """Do conjugate gradient iterations for the k-point"""

        niter = self.niter

        phi_G = wfs.empty(q=kpt.q)
        phi_old_G = wfs.empty(q=kpt.q)

        comm = wfs.gd.comm

        psit_nG, Htpsit_nG = self.subspace_diagonalize(hamiltonian, wfs, kpt)
        # Note that psit_nG is now in self.operator.work1_nG and
        # Htpsit_nG is in kpt.psit_nG!

        R_nG = reshape(self.Htpsit_nG, psit_nG.shape)
        Htphi_G = R_nG[0]

        R_nG[:] = Htpsit_nG
        self.timer.start('Residuals')
        self.calculate_residuals(kpt, wfs, hamiltonian, psit_nG,
                                 kpt.P_ani, kpt.eps_n, R_nG)
        self.timer.stop('Residuals')

        self.timer.start('CG')

        total_error = 0.0
        for n in range(self.nbands):
            if extra_parameters.get('PK', False):
               N = n+1
            else:
               N = psit_nG.shape[0]+1
            R_G = R_nG[n]
            Htpsit_G = Htpsit_nG[n]
            gamma_old = 1.0
            phi_old_G[:] = 0.0
            error = np.real(wfs.integrate(R_G, R_G))
            for nit in range(niter):
                if (error * Hartree**2 < self.tolerance / self.nbands):
                    break

                ekin = self.preconditioner.calculate_kinetic_energy(
                    psit_nG[n:n + 1], kpt)

                pR_G = self.preconditioner(R_nG[n:n + 1], kpt, ekin)

                # New search direction
                gamma = comm.sum(np.vdot(pR_G, R_G).real)
                phi_G[:] = -pR_G - gamma / gamma_old * phi_old_G
                gamma_old = gamma
                phi_old_G[:] = phi_G[:]

                # Calculate projections
                P2_ai = wfs.pt.dict()
                wfs.pt.integrate(phi_G, P2_ai, kpt.q)

                # Orthonormalize phi_G to all bands
                self.timer.start('CG: orthonormalize')
                self.timer.start('CG: overlap')
                overlap_n = wfs.integrate(psit_nG[:N], phi_G,
                                        global_integral=False)
                self.timer.stop('CG: overlap')
                self.timer.start('CG: overlap2')
                for a, P2_i in P2_ai.items():
                    P_ni = kpt.P_ani[a]
                    dO_ii = wfs.setups[a].dO_ii
                    gemv(1.0, P_ni[:N].conjugate(), np.inner(dO_ii, P2_i), 
                         1.0, overlap_n)
                self.timer.stop('CG: overlap2')
                comm.sum(overlap_n)

                # phi_G -= overlap_n * kpt.psit_nG
                wfs.matrixoperator.gd.gemv(-1.0, psit_nG[:N], overlap_n, 
                                           1.0, phi_G, 'n')
                for a, P2_i in P2_ai.items():
                    P_ni = kpt.P_ani[a]
                    gemv(-1.0, P_ni[:N], overlap_n, 1.0, P2_i, 'n')

                norm = wfs.integrate(phi_G, phi_G, global_integral=False)
                for a, P2_i in P2_ai.items():
                    dO_ii = wfs.setups[a].dO_ii
                    norm += np.vdot(P2_i, np.inner(dO_ii, P2_i))
                norm = comm.sum(np.real(norm).item())
                phi_G /= sqrt(norm)
                for P2_i in P2_ai.values():
                    P2_i /= sqrt(norm)
                self.timer.stop('CG: orthonormalize')

                # find optimum linear combination of psit_G and phi_G
                an = kpt.eps_n[n]
                wfs.apply_pseudo_hamiltonian(kpt, hamiltonian,
                                             phi_G.reshape((1,) + phi_G.shape),
                                             Htphi_G.reshape((1,) +
                                                             Htphi_G.shape))
                b = wfs.integrate(phi_G, Htpsit_G, global_integral=False)
                c = wfs.integrate(phi_G, Htphi_G, global_integral=False)
                for a, P2_i in P2_ai.items():
                    P_i = kpt.P_ani[a][n]
                    dH_ii = unpack(hamiltonian.dH_asp[a][kpt.s])
                    b += dot(P2_i, dot(dH_ii, P_i.conj()))
                    c += dot(P2_i, dot(dH_ii, P2_i.conj()))
                b = comm.sum(np.real(b).item())
                c = comm.sum(np.real(c).item())

                theta = 0.5 * atan2(2 * b, an - c)
                enew = (an * cos(theta)**2 +
                        c * sin(theta)**2 +
                        b * sin(2.0 * theta))
                # theta can correspond either minimum or maximum
                if (enew - kpt.eps_n[n]) > 0.0:  # we were at maximum
                    theta += pi / 2.0
                    enew = (an * cos(theta)**2 +
                            c * sin(theta)**2 +
                            b * sin(2.0 * theta))

                kpt.eps_n[n] = enew
                psit_nG[n] *= cos(theta)
                # kpt.psit_nG[n] += sin(theta) * phi_G
                axpy(sin(theta), phi_G, psit_nG[n])
                for a, P2_i in P2_ai.items():
                    P_i = kpt.P_ani[a][n]
                    P_i *= cos(theta)
                    P_i += sin(theta) * P2_i

                if nit < niter - 1:
                    Htpsit_G *= cos(theta)
                    # Htpsit_G += sin(theta) * Htphi_G
                    axpy(sin(theta), Htphi_G, Htpsit_G)
                    #adjust residuals
                    R_G[:] = Htpsit_G - kpt.eps_n[n] * psit_nG[n]

                    coef_ai = wfs.pt.dict()
                    for a, coef_i in coef_ai.items():
                        P_i = kpt.P_ani[a][n]
                        dO_ii = wfs.setups[a].dO_ii
                        dH_ii = unpack(hamiltonian.dH_asp[a][kpt.s])
                        coef_i[:] = (dot(P_i, dH_ii) -
                                     dot(P_i * kpt.eps_n[n], dO_ii))
                    wfs.pt.add(R_G, coef_ai, kpt.q)
                    error_new = np.real(wfs.integrate(R_G, R_G))
                    if error_new / error < self.rtol:
                        # print >> self.f, "cg:iters", n, nit+1
                        break
                    if (self.nbands_converge == 'occupied' and
                        kpt.f_n is not None and kpt.f_n[n] == 0.0):
                        # print >> self.f, "cg:iters", n, nit+1
                        break
                    error = error_new

            if kpt.f_n is None:
                weight = 1.0
            else:
                weight = kpt.f_n[n]
            if self.nbands_converge != 'occupied':
                weight = kpt.weight * float(n < self.nbands_converge)
            total_error += weight * error
            # if nit == 3:
            #   print >> self.f, "cg:iters", n, nit+1

        self.timer.stop('CG')
        return total_error, psit_nG
Exemplo n.º 9
0
    def iterate_one_k_point(self, hamiltonian, wfs, kpt):
        """Do conjugate gradient iterations for the k-point"""

        niter = self.niter

        phi_G = wfs.empty(q=kpt.q)
        phi_old_G = wfs.empty(q=kpt.q)

        comm = wfs.gd.comm

        psit_nG, Htpsit_nG = self.subspace_diagonalize(hamiltonian, wfs, kpt)
        # Note that psit_nG is now in self.operator.work1_nG and
        # Htpsit_nG is in kpt.psit_nG!

        R_nG = reshape(self.Htpsit_nG, psit_nG.shape)
        Htphi_G = R_nG[0]

        R_nG[:] = Htpsit_nG
        self.timer.start('Residuals')
        self.calculate_residuals(kpt, wfs, hamiltonian, psit_nG, kpt.P_ani,
                                 kpt.eps_n, R_nG)
        self.timer.stop('Residuals')

        self.timer.start('CG')

        total_error = 0.0
        for n in range(self.nbands):
            if extra_parameters.get('PK', False):
                N = n + 1
            else:
                N = psit_nG.shape[0] + 1
            R_G = R_nG[n]
            Htpsit_G = Htpsit_nG[n]
            gamma_old = 1.0
            phi_old_G[:] = 0.0
            error = np.real(wfs.integrate(R_G, R_G))
            for nit in range(niter):
                if (error * Hartree**2 < self.tolerance / self.nbands):
                    break

                ekin = self.preconditioner.calculate_kinetic_energy(
                    psit_nG[n:n + 1], kpt)

                pR_G = self.preconditioner(R_nG[n:n + 1], kpt, ekin)

                # New search direction
                gamma = comm.sum(np.vdot(pR_G, R_G).real)
                phi_G[:] = -pR_G - gamma / gamma_old * phi_old_G
                gamma_old = gamma
                phi_old_G[:] = phi_G[:]

                # Calculate projections
                P2_ai = wfs.pt.dict()
                wfs.pt.integrate(phi_G, P2_ai, kpt.q)

                # Orthonormalize phi_G to all bands
                self.timer.start('CG: orthonormalize')
                self.timer.start('CG: overlap')
                overlap_n = wfs.integrate(psit_nG[:N],
                                          phi_G,
                                          global_integral=False)
                self.timer.stop('CG: overlap')
                self.timer.start('CG: overlap2')
                for a, P2_i in P2_ai.items():
                    P_ni = kpt.P_ani[a]
                    dO_ii = wfs.setups[a].dO_ii
                    gemv(1.0, P_ni[:N].conjugate(), np.inner(dO_ii, P2_i), 1.0,
                         overlap_n)
                self.timer.stop('CG: overlap2')
                comm.sum(overlap_n)

                # phi_G -= overlap_n * kpt.psit_nG
                wfs.matrixoperator.gd.gemv(-1.0, psit_nG[:N], overlap_n, 1.0,
                                           phi_G, 'n')
                for a, P2_i in P2_ai.items():
                    P_ni = kpt.P_ani[a]
                    gemv(-1.0, P_ni[:N], overlap_n, 1.0, P2_i, 'n')

                norm = wfs.integrate(phi_G, phi_G, global_integral=False)
                for a, P2_i in P2_ai.items():
                    dO_ii = wfs.setups[a].dO_ii
                    norm += np.vdot(P2_i, np.inner(dO_ii, P2_i))
                norm = comm.sum(np.real(norm).item())
                phi_G /= sqrt(norm)
                for P2_i in P2_ai.values():
                    P2_i /= sqrt(norm)
                self.timer.stop('CG: orthonormalize')

                # find optimum linear combination of psit_G and phi_G
                an = kpt.eps_n[n]
                wfs.apply_pseudo_hamiltonian(
                    kpt, hamiltonian, phi_G.reshape((1, ) + phi_G.shape),
                    Htphi_G.reshape((1, ) + Htphi_G.shape))
                b = wfs.integrate(phi_G, Htpsit_G, global_integral=False)
                c = wfs.integrate(phi_G, Htphi_G, global_integral=False)
                for a, P2_i in P2_ai.items():
                    P_i = kpt.P_ani[a][n]
                    dH_ii = unpack(hamiltonian.dH_asp[a][kpt.s])
                    b += dot(P2_i, dot(dH_ii, P_i.conj()))
                    c += dot(P2_i, dot(dH_ii, P2_i.conj()))
                b = comm.sum(np.real(b).item())
                c = comm.sum(np.real(c).item())

                theta = 0.5 * atan2(2 * b, an - c)
                enew = (an * cos(theta)**2 + c * sin(theta)**2 +
                        b * sin(2.0 * theta))
                # theta can correspond either minimum or maximum
                if (enew - kpt.eps_n[n]) > 0.0:  # we were at maximum
                    theta += pi / 2.0
                    enew = (an * cos(theta)**2 + c * sin(theta)**2 +
                            b * sin(2.0 * theta))

                kpt.eps_n[n] = enew
                psit_nG[n] *= cos(theta)
                # kpt.psit_nG[n] += sin(theta) * phi_G
                axpy(sin(theta), phi_G, psit_nG[n])
                for a, P2_i in P2_ai.items():
                    P_i = kpt.P_ani[a][n]
                    P_i *= cos(theta)
                    P_i += sin(theta) * P2_i

                if nit < niter - 1:
                    Htpsit_G *= cos(theta)
                    # Htpsit_G += sin(theta) * Htphi_G
                    axpy(sin(theta), Htphi_G, Htpsit_G)
                    #adjust residuals
                    R_G[:] = Htpsit_G - kpt.eps_n[n] * psit_nG[n]

                    coef_ai = wfs.pt.dict()
                    for a, coef_i in coef_ai.items():
                        P_i = kpt.P_ani[a][n]
                        dO_ii = wfs.setups[a].dO_ii
                        dH_ii = unpack(hamiltonian.dH_asp[a][kpt.s])
                        coef_i[:] = (dot(P_i, dH_ii) -
                                     dot(P_i * kpt.eps_n[n], dO_ii))
                    wfs.pt.add(R_G, coef_ai, kpt.q)
                    error_new = np.real(wfs.integrate(R_G, R_G))
                    if error_new / error < self.rtol:
                        # print >> self.f, "cg:iters", n, nit+1
                        break
                    if (self.nbands_converge == 'occupied'
                            and kpt.f_n is not None and kpt.f_n[n] == 0.0):
                        # print >> self.f, "cg:iters", n, nit+1
                        break
                    error = error_new

            if kpt.f_n is None:
                weight = 1.0
            else:
                weight = kpt.f_n[n]
            if self.nbands_converge != 'occupied':
                weight = kpt.weight * float(n < self.nbands_converge)
            total_error += weight * error
            # if nit == 3:
            #   print >> self.f, "cg:iters", n, nit+1

        self.timer.stop('CG')
        return total_error, psit_nG
Exemplo n.º 10
0
    def iterate_one_k_point(self, hamiltonian, wfs, kpt):
        """Do Davidson iterations for the kpoint"""
        niter = self.niter
        nbands = self.nbands
        mynbands = self.mynbands

        gd = wfs.matrixoperator.gd
        bd = self.operator.bd

        psit_nG, Htpsit_nG = self.subspace_diagonalize(hamiltonian, wfs, kpt)
        # Note that psit_nG is now in self.operator.work1_nG and
        # Htpsit_nG is in kpt.psit_nG!

        H_2n2n = self.H_2n2n
        S_2n2n = self.S_2n2n
        eps_2n = self.eps_2n

        self.timer.start('Davidson')

        if self.keep_htpsit:
            R_nG = Htpsit_nG
            psit2_nG = reshape(self.Htpsit_nG, psit_nG.shape)
        else:
            R_nG = wfs.empty(mynbands, q=kpt.q)
            psit2_nG = wfs.empty(mynbands, q=kpt.q)
            wfs.apply_pseudo_hamiltonian(kpt, hamiltonian, psit_nG, R_nG)
            wfs.pt.integrate(psit_nG, kpt.P_ani, kpt.q)
        
        self.calculate_residuals(kpt, wfs, hamiltonian, psit_nG,
                                 kpt.P_ani, kpt.eps_n, R_nG)

        def integrate(a_G, b_G):
            return np.real(wfs.integrate(a_G, b_G, global_integral=False))

        # Note on band parallelization
        # The "large" H_2n2n and S_2n2n matrices are at the moment
        # global and replicated over band communicator, and the
        # general diagonalization is performed in serial i.e. without
        # scalapack

        for nit in range(niter):
            H_2n2n[:] = 0.0
            S_2n2n[:] = 0.0

            norm_n = np.zeros(mynbands)
            error = 0.0
            for n in range(mynbands):
                if kpt.f_n is None:
                    weight = kpt.weight
                else:
                    weight = kpt.f_n[n]
                if self.nbands_converge != 'occupied':
                    if n < self.nbands_converge:
                        weight = kpt.weight
                    else:
                        weight = 0.0
                error += weight * integrate(R_nG[n], R_nG[n])

                ekin = self.preconditioner.calculate_kinetic_energy(
                    psit_nG[n:n + 1], kpt)
                psit2_nG[n] = self.preconditioner(R_nG[n:n + 1], kpt, ekin)

                if self.normalize:
                    norm_n[n] = integrate(psit2_nG[n], psit2_nG[n])

                N = bd.global_index(n)
                H_2n2n[N, N] = kpt.eps_n[n]
                S_2n2n[N, N] = 1.0

            bd.comm.sum(H_2n2n)
            bd.comm.sum(S_2n2n)

            if self.normalize:
                gd.comm.sum(norm_n)
                for norm, psit2_G in zip(norm_n, psit2_nG):
                    psit2_G *= norm**-0.5
        
            # Calculate projections
            P2_ani = wfs.pt.dict(mynbands)
            wfs.pt.integrate(psit2_nG, P2_ani, kpt.q)

            self.timer.start('calc. matrices')
            
            # Hamiltonian matrix
            # <psi2 | H | psi>

            def H(psit_xG):
                result_xG = R_nG
                wfs.apply_pseudo_hamiltonian(kpt, hamiltonian, psit_xG,
                                             result_xG)
                return result_xG

            def dH(a, P_ni):
                return np.dot(P_ni, unpack(hamiltonian.dH_asp[a][kpt.s]))

            H_nn = self.operator.calculate_matrix_elements(psit_nG, kpt.P_ani,
                                                           H, dH, psit2_nG,
                                                           P2_ani)

            H_2n2n[nbands:, :nbands] = H_nn

            # <psi2 | H | psi2>

            def H(psit_xG):
                # H | psi2 > already calculated in previous step
                result_xG = R_nG
                return result_xG

            def dH(a, P_ni):
                return np.dot(P_ni, unpack(hamiltonian.dH_asp[a][kpt.s]))

            H_nn = self.operator.calculate_matrix_elements(psit2_nG, P2_ani,
                                                           H, dH)

            H_2n2n[nbands:, nbands:] = H_nn

            # Overlap matrix
            # <psi2 | S | psi>

            def S(psit_G):
                return psit_G
            
            def dS(a, P_ni):
                return np.dot(P_ni, wfs.setups[a].dO_ii)

            S_nn = self.operator.calculate_matrix_elements(psit_nG, kpt.P_ani,
                                                           S, dS, psit2_nG,
                                                           P2_ani)

            S_2n2n[nbands:, :nbands] = S_nn

            # <psi2 | S | psi2>
            S_nn = self.operator.calculate_matrix_elements(psit2_nG, P2_ani,
                                                           S, dS)
            S_2n2n[nbands:, nbands:] = S_nn

            self.timer.stop('calc. matrices')

            self.timer.start('diagonalize')
            if gd.comm.rank == 0 and bd.comm.rank == 0:
                m = 0
                if self.smin:
                    s_N, U_NN = np.linalg.eigh(S_2n2n)
                    m = int((s_N < self.smin).sum())

                if m == 0:
                    general_diagonalize(H_2n2n, eps_2n, S_2n2n)
                else:
                    T_Nn = np.dot(U_NN[:, m:], np.diag(s_N[m:]**-0.5))
                    H_2n2n[:nbands, nbands:] = \
                        H_2n2n[nbands:, :nbands].conj().T
                    eps_2n[:-m], P_nn = np.linalg.eigh(
                        np.dot(np.dot(T_Nn.T.conj(), H_2n2n), T_Nn))
                    H_2n2n[:-m] = np.dot(T_Nn, P_nn).T

            gd.comm.broadcast(H_2n2n, 0)
            gd.comm.broadcast(eps_2n, 0)
            bd.comm.broadcast(H_2n2n, 0)
            bd.comm.broadcast(eps_2n, 0)

            self.operator.bd.distribute(eps_2n[:nbands], kpt.eps_n[:])

            self.timer.stop('diagonalize')

            self.timer.start('rotate_psi')
            # Rotate psit_nG

            # Memory references during rotate:
            # Case 1, no band parallelization:
            #   Before 1. matrix multiply: psit_nG -> operator.work1_xG
            #   After  1. matrix multiply: psit_nG -> R_nG
            #   After  2. matrix multiply: tmp_nG -> work1_xG
            #
            # Case 2, band parallelization
            # Work arrays used only in send/recv buffers,
            # psit_nG -> psit_nG
            # tmp_nG -> psit2_nG

            psit_nG = self.operator.matrix_multiply(H_2n2n[:nbands, :nbands],
                                                    psit_nG, kpt.P_ani,
                                                    out_nG=R_nG)

            tmp_nG = self.operator.matrix_multiply(H_2n2n[:nbands, nbands:],
                                                   psit2_nG, P2_ani)

            if bd.comm.size > 1:
                psit_nG += tmp_nG
            else:
                tmp_nG += psit_nG
                psit_nG, R_nG = tmp_nG, psit_nG
            
            for a, P_ni in kpt.P_ani.items():
                P2_ni = P2_ani[a]
                P_ni += P2_ni

            self.timer.stop('rotate_psi')

            if nit < niter - 1:
                wfs.apply_pseudo_hamiltonian(kpt, hamiltonian, psit_nG,
                                             R_nG)
                self.calculate_residuals(kpt, wfs, hamiltonian, psit_nG,
                                         kpt.P_ani, kpt.eps_n, R_nG)

        self.timer.stop('Davidson')
        error = gd.comm.sum(error)
        return error, psit_nG
Exemplo n.º 11
0
    def iterate_one_k_point(self, hamiltonian, wfs, kpt):
        """Do Davidson iterations for the kpoint"""
        niter = self.niter
        nbands = self.nbands

        gd = wfs.matrixoperator.gd

        psit_nG, Htpsit_nG = self.subspace_diagonalize(hamiltonian, wfs, kpt)
        # Note that psit_nG is now in self.operator.work1_nG and
        # Htpsit_nG is in kpt.psit_nG!

        H_2n2n = self.H_2n2n
        S_2n2n = self.S_2n2n
        eps_2n = self.eps_2n
        psit2_nG = reshape(self.Htpsit_nG, psit_nG.shape)

        self.timer.start('Davidson')
        R_nG = Htpsit_nG
        self.calculate_residuals(kpt, wfs, hamiltonian, psit_nG, kpt.P_ani,
                                 kpt.eps_n, R_nG)

        def integrate(a_G, b_G):
            return np.real(wfs.integrate(a_G, b_G, global_integral=False))

        for nit in range(niter):
            H_2n2n[:] = 0.0
            S_2n2n[:] = 0.0

            norm_n = np.zeros(nbands)
            error = 0.0
            for n in range(nbands):
                if kpt.f_n is None:
                    weight = kpt.weight
                else:
                    weight = kpt.f_n[n]
                if self.nbands_converge != 'occupied':
                    if n < self.nbands_converge:
                        weight = kpt.weight
                    else:
                        weight = 0.0
                error += weight * integrate(R_nG[n], R_nG[n])

                ekin = self.preconditioner.calculate_kinetic_energy(
                    psit_nG[n:n + 1], kpt)
                psit2_nG[n] = self.preconditioner(R_nG[n:n + 1], kpt, ekin)

                if self.normalize:
                    norm_n[n] = integrate(psit2_nG[n], psit2_nG[n])

                H_2n2n[n, n] = kpt.eps_n[n]
                S_2n2n[n, n] = 1.0

            if self.normalize:
                gd.comm.sum(norm_n)
                for norm, psit2_G in zip(norm_n, psit2_nG):
                    psit2_G *= norm**-0.5

            # Calculate projections
            P2_ani = wfs.pt.dict(nbands)
            wfs.pt.integrate(psit2_nG, P2_ani, kpt.q)

            # Hamiltonian matrix
            # <psi2 | H | psi>
            wfs.apply_pseudo_hamiltonian(kpt, hamiltonian, psit2_nG, Htpsit_nG)
            gd.integrate(psit_nG,
                         Htpsit_nG,
                         global_integral=False,
                         _transposed_result=self.H_nn)
            # gemm(1.0, psit_nG, Htpsit_nG, 0.0, self.H_nn, 'c')

            for a, P_ni in kpt.P_ani.items():
                P2_ni = P2_ani[a]
                dH_ii = unpack(hamiltonian.dH_asp[a][kpt.s])
                self.H_nn += np.dot(P2_ni, np.dot(dH_ii, P_ni.T.conj()))

            gd.comm.sum(self.H_nn, 0)
            H_2n2n[nbands:, :nbands] = self.H_nn

            # <psi2 | H | psi2>
            gd.integrate(psit2_nG,
                         Htpsit_nG,
                         global_integral=False,
                         _transposed_result=self.H_nn)
            # r2k(0.5 * gd.dv, psit2_nG, Htpsit_nG, 0.0, self.H_nn)
            for a, P2_ni in P2_ani.items():
                dH_ii = unpack(hamiltonian.dH_asp[a][kpt.s])
                self.H_nn += np.dot(P2_ni, np.dot(dH_ii, P2_ni.T.conj()))

            gd.comm.sum(self.H_nn, 0)
            H_2n2n[nbands:, nbands:] = self.H_nn

            # Overlap matrix
            # <psi2 | S | psi>
            gd.integrate(psit_nG,
                         psit2_nG,
                         global_integral=False,
                         _transposed_result=self.S_nn)
            # gemm(1.0, psit_nG, psit2_nG, 0.0, self.S_nn, 'c')

            for a, P_ni in kpt.P_ani.items():
                P2_ni = P2_ani[a]
                dO_ii = wfs.setups[a].dO_ii
                self.S_nn += np.dot(P2_ni, np.inner(dO_ii, P_ni.conj()))

            gd.comm.sum(self.S_nn, 0)
            S_2n2n[nbands:, :nbands] = self.S_nn

            # <psi2 | S | psi2>
            gd.integrate(psit2_nG,
                         psit2_nG,
                         global_integral=False,
                         _transposed_result=self.S_nn)
            # rk(gd.dv, psit2_nG, 0.0, self.S_nn)
            for a, P2_ni in P2_ani.items():
                dO_ii = wfs.setups[a].dO_ii
                self.S_nn += np.dot(P2_ni, np.dot(dO_ii, P2_ni.T.conj()))

            gd.comm.sum(self.S_nn, 0)
            S_2n2n[nbands:, nbands:] = self.S_nn

            if gd.comm.rank == 0:
                m = 0
                if self.smin:
                    s_N, U_NN = np.linalg.eigh(S_2n2n)
                    m = int((s_N < self.smin).sum())

                if m == 0:
                    general_diagonalize(H_2n2n, eps_2n, S_2n2n)
                else:
                    T_Nn = np.dot(U_NN[:, m:], np.diag(s_N[m:]**-0.5))
                    H_2n2n[:nbands, nbands:] = \
                        H_2n2n[nbands:, :nbands].conj().T
                    eps_2n[:-m], P_nn = np.linalg.eigh(
                        np.dot(np.dot(T_Nn.T.conj(), H_2n2n), T_Nn))
                    H_2n2n[:-m] = np.dot(T_Nn, P_nn).T

            gd.comm.broadcast(H_2n2n, 0)
            gd.comm.broadcast(eps_2n, 0)

            kpt.eps_n[:] = eps_2n[:nbands]

            # Rotate psit_nG
            gd.gemm(1.0, psit_nG, H_2n2n[:nbands, :nbands], 0.0, Htpsit_nG)
            gd.gemm(1.0, psit2_nG, H_2n2n[:nbands, nbands:], 1.0, Htpsit_nG)
            psit_nG, Htpsit_nG = Htpsit_nG, psit_nG

            # Rotate P_uni:
            for a, P_ni in kpt.P_ani.items():
                P2_ni = P2_ani[a]
                gemm(1.0, P_ni.copy(), H_2n2n[:nbands, :nbands], 0.0, P_ni)
                gemm(1.0, P2_ni, H_2n2n[:nbands, nbands:], 1.0, P_ni)

            if nit < niter - 1:
                wfs.apply_pseudo_hamiltonian(kpt, hamiltonian, psit_nG,
                                             Htpsit_nG)
                R_nG = Htpsit_nG
                self.calculate_residuals(kpt, wfs, hamiltonian, psit_nG,
                                         kpt.P_ani, kpt.eps_n, R_nG)

        self.timer.stop('Davidson')
        error = gd.comm.sum(error)
        return error, psit_nG