예제 #1
0
    def calculate(self, seperate_spin=None):
        """Calculate the non-interacting density response function. """
        calc = self.calc
        kd = self.kd
        gd = self.gd
        sdisp_cd = gd.sdisp_cd
        ibzk_kc = kd.ibzk_kc
        bzk_kc = kd.bzk_kc
        kq_k = self.kq_k
        f_skn = self.f_skn
        e_skn = self.e_skn

        # Matrix init
        chi0_wGG = np.zeros((self.Nw_local, self.npw, self.npw), dtype=complex)
        if self.hilbert_trans:
            specfunc_wGG = np.zeros((self.NwS_local, self.npw, self.npw), dtype=complex)

        # Prepare for the derivative of pseudo-wavefunction
        if self.optical_limit:
            d_c = [Gradient(gd, i, n=4, dtype=complex).apply for i in range(3)]
            dpsit_g = gd.empty(dtype=complex)
            tmp = np.zeros((3), dtype=complex)
            rhoG0_v = np.zeros(3, dtype=complex)

            self.chi0G0_wGv = np.zeros((self.Nw_local, self.npw, 3), dtype=complex)
            self.chi00G_wGv = np.zeros((self.Nw_local, self.npw, 3), dtype=complex)

            specfuncG0_wGv = np.zeros((self.NwS_local, self.npw, 3), dtype=complex)
            specfunc0G_wGv = np.zeros((self.NwS_local, self.npw, 3), dtype=complex)

        use_zher = False
        if self.eta < 1e-5:
            use_zher = True

        rho_G = np.zeros(self.npw, dtype=complex)
        t0 = time()

        if seperate_spin is None:
            spinlist = np.arange(self.nspins)
        else:
            spinlist = [seperate_spin]

        for spin in spinlist:
            if not (f_skn[spin] > self.ftol).any():
                self.chi0_wGG = chi0_wGG
                continue

            for k in range(self.kstart, self.kend):
                k_pad = False
                if k >= self.kd.nbzkpts:
                    k = 0
                    k_pad = True

                # Find corresponding kpoint in IBZ
                ibzkpt1 = kd.bz2ibz_k[k]
                if self.optical_limit:
                    ibzkpt2 = ibzkpt1
                else:
                    ibzkpt2 = kd.bz2ibz_k[kq_k[k]]

                if self.pwmode:
                    N_c = self.gd.N_c
                    k_c = self.kd.ibzk_kc[ibzkpt1]
                    eikr1_R = np.exp(2j * pi * np.dot(np.indices(N_c).T, k_c / N_c).T)
                    k_c = self.kd.ibzk_kc[ibzkpt2]
                    eikr2_R = np.exp(2j * pi * np.dot(np.indices(N_c).T, k_c / N_c).T)

                index1_g, phase1_g = kd.get_transform_wavefunction_index(self.gd.N_c - (self.pbc == False), k)
                index2_g, phase2_g = kd.get_transform_wavefunction_index(self.gd.N_c - (self.pbc == False), kq_k[k])

                for n in range(self.nvalbands):
                    if self.calc.wfs.world.size == 1:
                        if self.f_skn[spin][ibzkpt1, n] - self.ftol < 0:
                            continue

                    t1 = time()
                    if self.pwmode:
                        u = self.kd.get_rank_and_index(spin, ibzkpt1)[1]
                        psitold_g = calc.wfs._get_wave_function_array(u, n, realspace=True, phase=eikr1_R)
                    else:
                        u = None
                        psitold_g = self.get_wavefunction(ibzkpt1, n, True, spin=spin)

                    psit1new_g = kd.transform_wave_function(psitold_g, k, index1_g, phase1_g)

                    P1_ai = self.pawstuff(psit1new_g, k, n, spin, u, ibzkpt1)

                    psit1_g = psit1new_g.conj() * self.expqr_g

                    for m in self.mlist:
                        if self.nbands > 1000 and m % 200 == 0:
                            print("    ", k, n, m, time() - t0, file=self.txt)

                        check_focc = (f_skn[spin][ibzkpt1, n] - f_skn[spin][ibzkpt2, m]) > self.ftol

                        if not self.pwmode:
                            psitold_g = self.get_wavefunction(ibzkpt2, m, check_focc, spin=spin)

                        if check_focc:
                            if self.pwmode:
                                u = self.kd.get_rank_and_index(spin, ibzkpt2)[1]
                                psitold_g = calc.wfs._get_wave_function_array(u, m, realspace=True, phase=eikr2_R)

                            psit2_g = kd.transform_wave_function(psitold_g, kq_k[k], index2_g, phase2_g)

                            # zero padding is included through the FFT
                            rho_g = np.fft.fftn(psit2_g * psit1_g, s=self.nGrpad) * self.vol / self.nG0rpad
                            # Here, planewave cutoff is applied
                            rho_G = rho_g.ravel()[self.Gindex_G]

                            if self.optical_limit:
                                phase_cd = np.exp(2j * pi * sdisp_cd * kd.bzk_kc[kq_k[k], :, np.newaxis])
                                for ix in range(3):
                                    d_c[ix](psit2_g, dpsit_g, phase_cd)
                                    tmp[ix] = gd.integrate(psit1_g * dpsit_g)
                                rho_G[0] = -1j * np.dot(self.qq_v, tmp)

                                for ix in range(3):
                                    q2_c = np.diag((1, 1, 1))[ix] * self.qopt
                                    qq2_v = np.dot(q2_c, self.bcell_cv)  # summation over c
                                    rhoG0_v[ix] = -1j * np.dot(qq2_v, tmp)

                            P2_ai = self.pawstuff(psit2_g, kq_k[k], m, spin, u, ibzkpt2)

                            for a, id in enumerate(calc.wfs.setups.id_a):
                                P_p = np.outer(P1_ai[a].conj(), P2_ai[a]).ravel()
                                gemv(1.0, self.phi_aGp[a], P_p, 1.0, rho_G)

                                if self.optical_limit:
                                    gemv(1.0, self.phiG0_avp[a], P_p, 1.0, rhoG0_v)

                            if self.optical_limit:
                                if (
                                    np.abs(self.enoshift_skn[spin][ibzkpt2, m] - self.enoshift_skn[spin][ibzkpt1, n])
                                    > 0.1 / Hartree
                                ):
                                    rho_G[0] /= (
                                        self.enoshift_skn[spin][ibzkpt2, m] - self.enoshift_skn[spin][ibzkpt1, n]
                                    )
                                    rhoG0_v /= self.enoshift_skn[spin][ibzkpt2, m] - self.enoshift_skn[spin][ibzkpt1, n]
                                else:
                                    rho_G[0] = 0.0
                                    rhoG0_v[:] = 0.0

                            if k_pad:
                                rho_G[:] = 0.0

                            if self.optical_limit:
                                rho0G_Gv = np.outer(rho_G.conj(), rhoG0_v)
                                rhoG0_Gv = np.outer(rho_G, rhoG0_v.conj())
                                rho0G_Gv[0, :] = rhoG0_v * rhoG0_v.conj()
                                rhoG0_Gv[0, :] = rhoG0_v * rhoG0_v.conj()

                            if not self.hilbert_trans:
                                if not use_zher:
                                    rho_GG = np.outer(rho_G, rho_G.conj())

                                for iw in range(self.Nw_local):
                                    w = self.w_w[iw + self.wstart] / Hartree
                                    coef = 1.0 / (
                                        w + e_skn[spin][ibzkpt1, n] - e_skn[spin][ibzkpt2, m] + 1j * self.eta
                                    ) - 1.0 / (w - e_skn[spin][ibzkpt1, n] + e_skn[spin][ibzkpt2, m] + 1j * self.eta)
                                    C = (f_skn[spin][ibzkpt1, n] - f_skn[spin][ibzkpt2, m]) * coef

                                    if use_zher:
                                        czher(C.real, rho_G.conj(), chi0_wGG[iw])
                                    else:
                                        axpy(C, rho_GG, chi0_wGG[iw])

                                        if self.optical_limit:
                                            axpy(C, rho0G_Gv, self.chi00G_wGv[iw])
                                            axpy(C, rhoG0_Gv, self.chi0G0_wGv[iw])

                            else:
                                rho_GG = np.outer(rho_G, rho_G.conj())
                                focc = f_skn[spin][ibzkpt1, n] - f_skn[spin][ibzkpt2, m]
                                w0 = e_skn[spin][ibzkpt2, m] - e_skn[spin][ibzkpt1, n]
                                scal(focc, rho_GG)
                                if self.optical_limit:
                                    scal(focc, rhoG0_Gv)
                                    scal(focc, rho0G_Gv)

                                # calculate delta function
                                w0_id = int(w0 / self.dw)
                                if w0_id + 1 < self.NwS:
                                    # rely on the self.NwS_local is equal in each node!
                                    if self.wScomm.rank == w0_id // self.NwS_local:
                                        alpha = (w0_id + 1 - w0 / self.dw) / self.dw
                                        axpy(alpha, rho_GG, specfunc_wGG[w0_id % self.NwS_local])

                                        if self.optical_limit:
                                            axpy(alpha, rho0G_Gv, specfunc0G_wGv[w0_id % self.NwS_local])
                                            axpy(alpha, rhoG0_Gv, specfuncG0_wGv[w0_id % self.NwS_local])

                                    if self.wScomm.rank == (w0_id + 1) // self.NwS_local:
                                        alpha = (w0 / self.dw - w0_id) / self.dw
                                        axpy(alpha, rho_GG, specfunc_wGG[(w0_id + 1) % self.NwS_local])

                                        if self.optical_limit:
                                            axpy(alpha, rho0G_Gv, specfunc0G_wGv[(w0_id + 1) % self.NwS_local])
                                            axpy(alpha, rhoG0_Gv, specfuncG0_wGv[(w0_id + 1) % self.NwS_local])

                    #                            deltaw = delta_function(w0, self.dw, self.NwS, self.sigma)
                    #                            for wi in range(self.NwS_local):
                    #                                if deltaw[wi + self.wS1] > 1e-8:
                    #                                    specfunc_wGG[wi] += tmp_GG * deltaw[wi + self.wS1]
                    if self.kd.nbzkpts == 1:
                        if n == 0:
                            dt = time() - t0
                            totaltime = dt * self.nvalbands * self.nspins
                            self.printtxt("Finished n 0 in %d seconds, estimate %d seconds left." % (dt, totaltime))
                        if rank == 0 and self.nvalbands // 5 > 0:
                            if n > 0 and n % (self.nvalbands // 5) == 0:
                                dt = time() - t0
                                self.printtxt(
                                    "Finished n %d in %d seconds, estimate %d seconds left." % (n, dt, totaltime - dt)
                                )
                if calc.wfs.world.size != 1:
                    self.kcomm.barrier()
                if k == 0:
                    dt = time() - t0
                    totaltime = dt * self.nkpt_local * self.nspins
                    self.printtxt("Finished k 0 in %d seconds, estimate %d seconds left." % (dt, totaltime))

                if rank == 0 and self.nkpt_local // 5 > 0:
                    if k > 0 and k % (self.nkpt_local // 5) == 0:
                        dt = time() - t0
                        self.printtxt(
                            "Finished k %d in %d seconds, estimate %d seconds left.  " % (k, dt, totaltime - dt)
                        )
        self.printtxt("Finished summation over k")

        self.kcomm.barrier()

        # Hilbert Transform
        if not self.hilbert_trans:
            for iw in range(self.Nw_local):
                self.kcomm.sum(chi0_wGG[iw])
                if self.optical_limit:
                    self.kcomm.sum(self.chi0G0_wGv[iw])
                    self.kcomm.sum(self.chi00G_wGv[iw])

            if use_zher:
                assert (np.abs(chi0_wGG[0, 1:, 0]) < 1e-10).all()
                for iw in range(self.Nw_local):
                    chi0_wGG[iw] += chi0_wGG[iw].conj().T
                    for iG in range(self.npw):
                        chi0_wGG[iw, iG, iG] /= 2.0
                        assert np.abs(np.imag(chi0_wGG[iw, iG, iG])) < 1e-10
        else:
            for iw in range(self.NwS_local):
                self.kcomm.sum(specfunc_wGG[iw])
                if self.optical_limit:
                    self.kcomm.sum(specfuncG0_wGv[iw])
                    self.kcomm.sum(specfunc0G_wGv[iw])

            if self.wScomm.size == 1:
                chi0_wGG = hilbert_transform(
                    specfunc_wGG, self.w_w, self.Nw, self.dw, self.eta, self.full_hilbert_trans
                )[self.wstart : self.wend]
                self.printtxt("Finished hilbert transform !")
                del specfunc_wGG
            else:
                # redistribute specfunc_wGG to all nodes
                size = self.comm.size
                assert self.NwS % size == 0
                NwStmp1 = (rank % self.kcomm.size) * self.NwS // size
                NwStmp2 = (rank % self.kcomm.size + 1) * self.NwS // size
                specfuncnew_wGG = specfunc_wGG[NwStmp1:NwStmp2]
                del specfunc_wGG

                coords = np.zeros(self.wcomm.size, dtype=int)
                nG_local = self.npw ** 2 // self.wcomm.size
                if self.wcomm.rank == self.wcomm.size - 1:
                    nG_local = self.npw ** 2 - (self.wcomm.size - 1) * nG_local
                self.wcomm.all_gather(np.array([nG_local]), coords)

                specfunc_Wg = SliceAlongFrequency(specfuncnew_wGG, coords, self.wcomm)
                self.printtxt("Finished Slice Along Frequency !")
                chi0_Wg = hilbert_transform(specfunc_Wg, self.w_w, self.Nw, self.dw, self.eta, self.full_hilbert_trans)[
                    : self.Nw
                ]
                self.printtxt("Finished hilbert transform !")
                self.comm.barrier()
                del specfunc_Wg

                chi0_wGG = SliceAlongOrbitals(chi0_Wg, coords, self.wcomm)
                self.printtxt("Finished Slice along orbitals !")
                self.comm.barrier()
                del chi0_Wg

                if self.optical_limit:
                    specfuncG0_WGv = np.zeros((self.NwS, self.npw, 3), dtype=complex)
                    specfunc0G_WGv = np.zeros((self.NwS, self.npw, 3), dtype=complex)
                    self.wScomm.all_gather(specfunc0G_wGv, specfunc0G_WGv)
                    self.wScomm.all_gather(specfuncG0_wGv, specfuncG0_WGv)
                    specfunc0G_wGv = specfunc0G_WGv
                    specfuncG0_wGv = specfuncG0_WGv

            if self.optical_limit:
                self.chi00G_wGv = hilbert_transform(
                    specfunc0G_wGv, self.w_w, self.Nw, self.dw, self.eta, self.full_hilbert_trans
                )[self.wstart : self.wend]

                self.chi0G0_wGv = hilbert_transform(
                    specfuncG0_wGv, self.w_w, self.Nw, self.dw, self.eta, self.full_hilbert_trans
                )[self.wstart : self.wend]

        if self.optical_limit:
            self.chi00G_wGv /= self.vol
            self.chi0G0_wGv /= self.vol

        self.chi0_wGG = chi0_wGG
        self.chi0_wGG /= self.vol

        self.printtxt("")
        self.printtxt("Finished chi0 !")
예제 #2
0
    def calculate(self, seperate_spin=None):
        """Calculate the non-interacting density response function. """
        calc = self.calc
        kd = self.kd
        gd = self.gd
        sdisp_cd = gd.sdisp_cd
        ibzk_kc = kd.ibzk_kc
        bzk_kc = kd.bzk_kc
        kq_k = self.kq_k
        f_skn = self.f_skn
        e_skn = self.e_skn

        # Matrix init
        chi0_wGG = np.zeros((self.Nw_local, self.npw, self.npw), dtype=complex)
        if self.hilbert_trans:
            specfunc_wGG = np.zeros((self.NwS_local, self.npw, self.npw), dtype = complex)

        # Prepare for the derivative of pseudo-wavefunction
        if self.optical_limit:
            d_c = [Gradient(gd, i, n=4, dtype=complex).apply for i in range(3)]
            dpsit_g = gd.empty(dtype=complex)
            tmp = np.zeros((3), dtype=complex)
            rhoG0_v = np.zeros(3, dtype=complex)

            self.chi0G0_wGv = np.zeros((self.Nw_local, self.npw, 3), dtype=complex)
            self.chi00G_wGv = np.zeros((self.Nw_local, self.npw, 3), dtype=complex)

            specfuncG0_wGv = np.zeros((self.NwS_local, self.npw, 3), dtype=complex)
            specfunc0G_wGv = np.zeros((self.NwS_local, self.npw, 3), dtype=complex)
            
        use_zher = False
        if self.eta < 1e-5:
            use_zher = True

        rho_G = np.zeros(self.npw, dtype=complex)
        t0 = time()

        if seperate_spin is None:
            spinlist = np.arange(self.nspins)
        else:
            spinlist = [seperate_spin]
        
        for spin in spinlist:
            if not (f_skn[spin] > self.ftol).any():
                self.chi0_wGG = chi0_wGG
                continue

            for k in range(self.kstart, self.kend):
                k_pad = False
                if k >= self.kd.nbzkpts:
                    k = 0
                    k_pad = True
    
                # Find corresponding kpoint in IBZ
                ibzkpt1 = kd.bz2ibz_k[k]
                if self.optical_limit:
                    ibzkpt2 = ibzkpt1
                else:
                    ibzkpt2 = kd.bz2ibz_k[kq_k[k]]
    
                if self.pwmode:
                    N_c = self.gd.N_c
                    k_c = self.kd.ibzk_kc[ibzkpt1]
                    eikr1_R = np.exp(2j * pi * np.dot(np.indices(N_c).T, k_c / N_c).T)
                    k_c = self.kd.ibzk_kc[ibzkpt2]
                    eikr2_R = np.exp(2j * pi * np.dot(np.indices(N_c).T, k_c / N_c).T)
                    
                index1_g, phase1_g = kd.get_transform_wavefunction_index(self.gd.N_c - (self.pbc == False), k)
                index2_g, phase2_g = kd.get_transform_wavefunction_index(self.gd.N_c - (self.pbc == False), kq_k[k])
                
                for n in range(self.nvalbands):
                    if self.calc.wfs.world.size == 1:
                        if (self.f_skn[spin][ibzkpt1, n] - self.ftol < 0):
                            continue

                    t1 = time()
                    if self.pwmode:
                        u = self.kd.get_rank_and_index(spin, ibzkpt1)[1]
                        psitold_g = calc.wfs._get_wave_function_array(u, n, realspace=True, phase=eikr1_R)
                    else:
                        u = None
                        psitold_g = self.get_wavefunction(ibzkpt1, n, True, spin=spin)
    
                    psit1new_g = kd.transform_wave_function(psitold_g,k,index1_g,phase1_g)

                    P1_ai = self.pawstuff(psit1new_g, k, n, spin, u, ibzkpt1)

                    psit1_g = psit1new_g.conj() * self.expqr_g
    
                    for m in self.mlist:
                        if self.nbands > 1000 and m % 200 == 0:
                            print('    ', k, n, m, time() - t0, file=self.txt)
                    
                        check_focc = (f_skn[spin][ibzkpt1, n] - f_skn[spin][ibzkpt2, m]) > self.ftol
    
                        if not self.pwmode:
                            psitold_g = self.get_wavefunction(ibzkpt2, m, check_focc, spin=spin)
    
                        if check_focc:                            
                            if self.pwmode:
                                u = self.kd.get_rank_and_index(spin, ibzkpt2)[1]
                                psitold_g = calc.wfs._get_wave_function_array(u, m, realspace=True, phase=eikr2_R)
    
                            psit2_g = kd.transform_wave_function(psitold_g, kq_k[k], index2_g, phase2_g)
    
                            # zero padding is included through the FFT
                            rho_g = np.fft.fftn(psit2_g * psit1_g, s=self.nGrpad) * self.vol / self.nG0rpad
                            # Here, planewave cutoff is applied
                            rho_G = rho_g.ravel()[self.Gindex_G]
    
                            if self.optical_limit:
                                phase_cd = np.exp(2j * pi * sdisp_cd * kd.bzk_kc[kq_k[k], :, np.newaxis])
                                for ix in range(3):
                                    d_c[ix](psit2_g, dpsit_g, phase_cd)
                                    tmp[ix] = gd.integrate(psit1_g * dpsit_g)
                                rho_G[0] = -1j * np.dot(self.qq_v, tmp)

                                for ix in range(3):
                                    q2_c = np.diag((1,1,1))[ix] * self.qopt
                                    qq2_v = np.dot(q2_c, self.bcell_cv) # summation over c
                                    rhoG0_v[ix] = -1j * np.dot(qq2_v, tmp)

                            P2_ai = self.pawstuff(psit2_g, kq_k[k], m, spin, u, ibzkpt2)

                            for a, id in enumerate(calc.wfs.setups.id_a):
                                P_p = np.outer(P1_ai[a].conj(), P2_ai[a]).ravel()
                                gemv(1.0, self.phi_aGp[a], P_p, 1.0, rho_G)

                                if self.optical_limit:
                                    gemv(1.0, self.phiG0_avp[a], P_p, 1.0, rhoG0_v)

                            if self.optical_limit:
                                if np.abs(self.enoshift_skn[spin][ibzkpt2, m] -
                                          self.enoshift_skn[spin][ibzkpt1, n]) > 0.1/Hartree:
                                    rho_G[0] /= self.enoshift_skn[spin][ibzkpt2, m] \
                                                - self.enoshift_skn[spin][ibzkpt1, n]
                                    rhoG0_v /= self.enoshift_skn[spin][ibzkpt2, m] \
                                                - self.enoshift_skn[spin][ibzkpt1, n]
                                else:
                                    rho_G[0] = 0.
                                    rhoG0_v[:] = 0.
    
                            if k_pad:
                                rho_G[:] = 0.

                            if self.optical_limit:
                                rho0G_Gv = np.outer(rho_G.conj(), rhoG0_v)
                                rhoG0_Gv = np.outer(rho_G, rhoG0_v.conj())
                                rho0G_Gv[0,:] = rhoG0_v * rhoG0_v.conj()
                                rhoG0_Gv[0,:] = rhoG0_v * rhoG0_v.conj()

                            if not self.hilbert_trans:
                                if not use_zher:
                                    rho_GG = np.outer(rho_G, rho_G.conj())

                                for iw in range(self.Nw_local):
                                    w = self.w_w[iw + self.wstart] / Hartree
                                    coef = ( 1. / (w + e_skn[spin][ibzkpt1, n] - e_skn[spin][ibzkpt2, m]
                                                   + 1j * self.eta) 
                                           - 1. / (w - e_skn[spin][ibzkpt1, n] + e_skn[spin][ibzkpt2, m]
                                                   + 1j * self.eta) )
                                    C =  (f_skn[spin][ibzkpt1, n] - f_skn[spin][ibzkpt2, m]) * coef

                                    if use_zher:
                                        czher(C.real, rho_G.conj(), chi0_wGG[iw])
                                    else:
                                        axpy(C, rho_GG, chi0_wGG[iw])
                                        
                                        if self.optical_limit:
                                            axpy(C, rho0G_Gv, self.chi00G_wGv[iw])
                                            axpy(C, rhoG0_Gv, self.chi0G0_wGv[iw])
    
                            else:
                                rho_GG = np.outer(rho_G, rho_G.conj())
                                focc = f_skn[spin][ibzkpt1,n] - f_skn[spin][ibzkpt2,m]
                                w0 = e_skn[spin][ibzkpt2,m] - e_skn[spin][ibzkpt1,n]
                                scal(focc, rho_GG)
                                if self.optical_limit:
                                    scal(focc, rhoG0_Gv)
                                    scal(focc, rho0G_Gv)
    
                                # calculate delta function
                                w0_id = int(w0 / self.dw)
                                if w0_id + 1 < self.NwS:
                                    # rely on the self.NwS_local is equal in each node!
                                    if self.wScomm.rank == w0_id // self.NwS_local:
                                        alpha = (w0_id + 1 - w0/self.dw) / self.dw
                                        axpy(alpha, rho_GG, specfunc_wGG[w0_id % self.NwS_local] )

                                        if self.optical_limit:
                                            axpy(alpha, rho0G_Gv, specfunc0G_wGv[w0_id % self.NwS_local] )
                                            axpy(alpha, rhoG0_Gv, specfuncG0_wGv[w0_id % self.NwS_local] )
    
                                    if self.wScomm.rank == (w0_id+1) // self.NwS_local:
                                        alpha =  (w0 / self.dw - w0_id) / self.dw
                                        axpy(alpha, rho_GG, specfunc_wGG[(w0_id+1) % self.NwS_local] )

                                        if self.optical_limit:
                                            axpy(alpha, rho0G_Gv, specfunc0G_wGv[(w0_id+1) % self.NwS_local] )
                                            axpy(alpha, rhoG0_Gv, specfuncG0_wGv[(w0_id+1) % self.NwS_local] )

    #                            deltaw = delta_function(w0, self.dw, self.NwS, self.sigma)
    #                            for wi in range(self.NwS_local):
    #                                if deltaw[wi + self.wS1] > 1e-8:
    #                                    specfunc_wGG[wi] += tmp_GG * deltaw[wi + self.wS1]
                    if self.kd.nbzkpts == 1:
                        if n == 0:
                            dt = time() - t0
                            totaltime = dt * self.nvalbands * self.nspins
                            self.printtxt('Finished n 0 in %d seconds, estimate %d seconds left.' %(dt, totaltime) )
                        if rank == 0 and self.nvalbands // 5 > 0:
                            if n > 0 and n % (self.nvalbands // 5) == 0:
                                dt = time() - t0
                                self.printtxt('Finished n %d in %d seconds, estimate %d seconds left.'%(n, dt, totaltime-dt))
                if calc.wfs.world.size != 1:
                    self.kcomm.barrier()            
                if k == 0:
                    dt = time() - t0
                    totaltime = dt * self.nkpt_local * self.nspins
                    self.printtxt('Finished k 0 in %d seconds, estimate %d seconds left.' %(dt, totaltime))
                    
                if rank == 0 and self.nkpt_local // 5 > 0:            
                    if k > 0 and k % (self.nkpt_local // 5) == 0:
                        dt =  time() - t0
                        self.printtxt('Finished k %d in %d seconds, estimate %d seconds left.  '%(k, dt, totaltime - dt) )
        self.printtxt('Finished summation over k')

        self.kcomm.barrier()
        
        # Hilbert Transform
        if not self.hilbert_trans:
            for iw in range(self.Nw_local):
                self.kcomm.sum(chi0_wGG[iw])
                if self.optical_limit:
                    self.kcomm.sum(self.chi0G0_wGv[iw])
                    self.kcomm.sum(self.chi00G_wGv[iw])

            if use_zher:
                assert (np.abs(chi0_wGG[0,1:,0]) < 1e-10).all()
                for iw in range(self.Nw_local):
                    chi0_wGG[iw] += chi0_wGG[iw].conj().T
                    for iG in range(self.npw):
                        chi0_wGG[iw, iG, iG] /= 2.
                        assert np.abs(np.imag(chi0_wGG[iw, iG, iG])) < 1e-10 
        else:
            for iw in range(self.NwS_local):
                self.kcomm.sum(specfunc_wGG[iw])
                if self.optical_limit:
                    self.kcomm.sum(specfuncG0_wGv[iw])
                    self.kcomm.sum(specfunc0G_wGv[iw])

            if self.wScomm.size == 1:
                chi0_wGG = hilbert_transform(specfunc_wGG, self.w_w, self.Nw, self.dw, self.eta,
                                             self.full_hilbert_trans)[self.wstart:self.wend]
                self.printtxt('Finished hilbert transform !')
                del specfunc_wGG
            else:
                # redistribute specfunc_wGG to all nodes
                size = self.comm.size
                assert self.NwS % size == 0
                NwStmp1 = (rank % self.kcomm.size) * self.NwS // size
                NwStmp2 = (rank % self.kcomm.size + 1) * self.NwS // size 
                specfuncnew_wGG = specfunc_wGG[NwStmp1:NwStmp2]
                del specfunc_wGG
                
                coords = np.zeros(self.wcomm.size, dtype=int)
                nG_local = self.npw**2 // self.wcomm.size
                if self.wcomm.rank == self.wcomm.size - 1:
                    nG_local = self.npw**2 - (self.wcomm.size - 1) * nG_local
                self.wcomm.all_gather(np.array([nG_local]), coords)
        
                specfunc_Wg = SliceAlongFrequency(specfuncnew_wGG, coords, self.wcomm)
                self.printtxt('Finished Slice Along Frequency !')
                chi0_Wg = hilbert_transform(specfunc_Wg, self.w_w, self.Nw, self.dw, self.eta,
                                            self.full_hilbert_trans)[:self.Nw]
                self.printtxt('Finished hilbert transform !')
                self.comm.barrier()
                del specfunc_Wg
        
                chi0_wGG = SliceAlongOrbitals(chi0_Wg, coords, self.wcomm)
                self.printtxt('Finished Slice along orbitals !')
                self.comm.barrier()
                del chi0_Wg

                if self.optical_limit:
                    specfuncG0_WGv = np.zeros((self.NwS, self.npw, 3), dtype=complex)
                    specfunc0G_WGv = np.zeros((self.NwS, self.npw, 3), dtype=complex)
                    self.wScomm.all_gather(specfunc0G_wGv, specfunc0G_WGv)
                    self.wScomm.all_gather(specfuncG0_wGv, specfuncG0_WGv)
                    specfunc0G_wGv = specfunc0G_WGv
                    specfuncG0_wGv = specfuncG0_WGv

            if self.optical_limit:
                self.chi00G_wGv = hilbert_transform(specfunc0G_wGv, self.w_w, self.Nw, self.dw, self.eta,
                                             self.full_hilbert_trans)[self.wstart:self.wend]
                
                self.chi0G0_wGv = hilbert_transform(specfuncG0_wGv, self.w_w, self.Nw, self.dw, self.eta,
                                             self.full_hilbert_trans)[self.wstart:self.wend]

        if self.optical_limit:
            self.chi00G_wGv /= self.vol
            self.chi0G0_wGv /= self.vol

        
        self.chi0_wGG = chi0_wGG
        self.chi0_wGG /= self.vol

        self.printtxt('')
        self.printtxt('Finished chi0 !')
예제 #3
0
파일: chi.py 프로젝트: qsnake/gpaw
    def calculate(self, spin=0):
        """Calculate the non-interacting density response function. """

        calc = self.calc
        kd = self.kd
        gd = self.gd
        sdisp_cd = gd.sdisp_cd
        ibzk_kc = self.ibzk_kc
        bzk_kc = self.bzk_kc
        kq_k = self.kq_k
        pt = self.pt
        f_kn = self.f_kn
        e_kn = self.e_kn

        # Matrix init
        chi0_wGG = np.zeros((self.Nw_local, self.npw, self.npw), dtype=complex)
        if not (f_kn > self.ftol).any():
            self.chi0_wGG = chi0_wGG
            return

        if self.hilbert_trans:
            specfunc_wGG = np.zeros((self.NwS_local, self.npw, self.npw), dtype = complex)

        # Prepare for the derivative of pseudo-wavefunction
        if self.optical_limit:
            d_c = [Gradient(gd, i, n=4, dtype=complex).apply for i in range(3)]
            dpsit_g = gd.empty(dtype=complex)
            tmp = np.zeros((3), dtype=complex)

        rho_G = np.zeros(self.npw, dtype=complex)
        t0 = time()
        t_get_wfs = 0
        for k in range(self.kstart, self.kend):

            # Find corresponding kpoint in IBZ
            ibzkpt1 = kd.kibz_k[k]
            if self.optical_limit:
                ibzkpt2 = ibzkpt1
            else:
                ibzkpt2 = kd.kibz_k[kq_k[k]]
            
            for n in range(self.nstart, self.nend):
#                print >> self.txt, k, n, t_get_wfs, time() - t0
                t1 = time()
                psitold_g = self.get_wavefunction(ibzkpt1, n, True, spin=spin)
                t_get_wfs += time() - t1
                psit1new_g = kd.transform_wave_function(psitold_g, k)

                P1_ai = pt.dict()
                pt.integrate(psit1new_g, P1_ai, k)

                psit1_g = psit1new_g.conj() * self.expqr_g

                for m in range(self.nbands):

		    if self.hilbert_trans:
			check_focc = (f_kn[ibzkpt1, n] - f_kn[ibzkpt2, m]) > self.ftol
                    else:
                        check_focc = np.abs(f_kn[ibzkpt1, n] - f_kn[ibzkpt2, m]) > self.ftol

                    t1 = time()
                    psitold_g = self.get_wavefunction(ibzkpt2, m, check_focc, spin=spin)
                    t_get_wfs += time() - t1

                    if check_focc:
                        psit2_g = kd.transform_wave_function(psitold_g, kq_k[k])
                        P2_ai = pt.dict()
                        pt.integrate(psit2_g, P2_ai, kq_k[k])

                        # fft
                        tmp_g = np.fft.fftn(psit2_g*psit1_g) * self.vol / self.nG0

                        for iG in range(self.npw):
                            index = self.Gindex_G[iG]
                            rho_G[iG] = tmp_g[index[0], index[1], index[2]]

                        if self.optical_limit:
                            phase_cd = np.exp(2j * pi * sdisp_cd * bzk_kc[kq_k[k], :, np.newaxis])
                            for ix in range(3):
                                d_c[ix](psit2_g, dpsit_g, phase_cd)
                                tmp[ix] = gd.integrate(psit1_g * dpsit_g)
                            rho_G[0] = -1j * np.dot(self.qq_v, tmp)

                        # PAW correction
                        for a, id in enumerate(calc.wfs.setups.id_a):
                            P_p = np.outer(P1_ai[a].conj(), P2_ai[a]).ravel()
                            gemv(1.0, self.phi_aGp[a], P_p, 1.0, rho_G)

                        if self.optical_limit:
                            rho_G[0] /= e_kn[ibzkpt2, m] - e_kn[ibzkpt1, n]

                        rho_GG = np.outer(rho_G, rho_G.conj())
                        
                        if not self.hilbert_trans:
                            for iw in range(self.Nw_local):
                                w = self.w_w[iw + self.wstart] / Hartree
                                C =  (f_kn[ibzkpt1, n] - f_kn[ibzkpt2, m]) / (
                                     w + e_kn[ibzkpt1, n] - e_kn[ibzkpt2, m] + 1j * self.eta)
                                axpy(C, rho_GG, chi0_wGG[iw])
                        else:
                            focc = f_kn[ibzkpt1,n] - f_kn[ibzkpt2,m]
                            w0 = e_kn[ibzkpt2,m] - e_kn[ibzkpt1,n]
                            scal(focc, rho_GG)

                            # calculate delta function
                            w0_id = int(w0 / self.dw)
                            if w0_id + 1 < self.NwS:
                                # rely on the self.NwS_local is equal in each node!
                                if self.wScomm.rank == w0_id // self.NwS_local:
                                    alpha = (w0_id + 1 - w0/self.dw) / self.dw
                                    axpy(alpha, rho_GG, specfunc_wGG[w0_id % self.NwS_local] )

                                if self.wScomm.rank == (w0_id+1) // self.NwS_local:
                                    alpha =  (w0 / self.dw - w0_id) / self.dw
                                    axpy(alpha, rho_GG, specfunc_wGG[(w0_id+1) % self.NwS_local] )

#                            deltaw = delta_function(w0, self.dw, self.NwS, self.sigma)
#                            for wi in range(self.NwS_local):
#                                if deltaw[wi + self.wS1] > 1e-8:
#                                    specfunc_wGG[wi] += tmp_GG * deltaw[wi + self.wS1]
                if self.nkpt == 1:
                    if n == 0:
                        dt = time() - t0
                        totaltime = dt * self.nband_local
                        self.printtxt('Finished n 0 in %f seconds, estimated %f seconds left.' %(dt, totaltime) )
                    if rank == 0 and self.nband_local // 5 > 0:
                        if n > 0 and n % (self.nband_local // 5) == 0:
                            dt = time() - t0
                            self.printtxt('Finished n %d in %f seconds, estimated %f seconds left.'%(n, dt, totaltime-dt))
            if calc.wfs.world.size != 1:
                self.kcomm.barrier()            
            if k == 0:
                dt = time() - t0
                totaltime = dt * self.nkpt_local
                self.printtxt('Finished k 0 in %f seconds, estimated %f seconds left.' %(dt, totaltime))
                
            if rank == 0 and self.nkpt_local // 5 > 0:            
                if k > 0 and k % (self.nkpt_local // 5) == 0:
                    dt =  time() - t0
                    self.printtxt('Finished k %d in %f seconds, estimated %f seconds left.  '%(k, dt, totaltime - dt) )
        self.printtxt('Finished summation over k')

        self.kcomm.barrier()
        del rho_GG, rho_G
        # Hilbert Transform
        if not self.hilbert_trans:
            self.kcomm.sum(chi0_wGG)
        else:
            self.kcomm.sum(specfunc_wGG)
            if self.wScomm.size == 1:
                if not self.full_hilbert_trans:
                    chi0_wGG = hilbert_transform(specfunc_wGG, self.Nw, self.dw, self.eta)[self.wstart:self.wend]
                else:
                    chi0_wGG = full_hilbert_transform(specfunc_wGG, self.Nw, self.dw, self.eta)[self.wstart:self.wend]                
                self.printtxt('Finished hilbert transform !')
                del specfunc_wGG
            else:
                # redistribute specfunc_wGG to all nodes
                assert self.NwS % size == 0
                NwStmp1 = (rank % self.kcomm.size) * self.NwS // size
                NwStmp2 = (rank % self.kcomm.size + 1) * self.NwS // size 
                specfuncnew_wGG = specfunc_wGG[NwStmp1:NwStmp2]
                del specfunc_wGG
                
                coords = np.zeros(self.wcomm.size, dtype=int)
                nG_local = self.npw**2 // self.wcomm.size
                if self.wcomm.rank == self.wcomm.size - 1:
                    nG_local = self.npw**2 - (self.wcomm.size - 1) * nG_local
                self.wcomm.all_gather(np.array([nG_local]), coords)
        
                specfunc_Wg = SliceAlongFrequency(specfuncnew_wGG, coords, self.wcomm)
                self.printtxt('Finished Slice Along Frequency !')
                if not self.full_hilbert_trans:
                    chi0_Wg = hilbert_transform(specfunc_Wg, self.Nw, self.dw, self.eta)[:self.Nw]
                else:
                    chi0_Wg = full_hilbert_transform(specfunc_Wg, self.Nw, self.dw, self.eta)[:self.Nw]
                self.printtxt('Finished hilbert transform !')
                self.comm.barrier()
                del specfunc_Wg
        
                chi0_wGG = SliceAlongOrbitals(chi0_Wg, coords, self.wcomm)
                self.printtxt('Finished Slice along orbitals !')
                self.comm.barrier()
                del chi0_Wg
        
        self.chi0_wGG = chi0_wGG / self.vol

        self.printtxt('')
        self.printtxt('Finished chi0 !')

        return
예제 #4
0
    def calculate(self, spin=0):
        """Calculate the non-interacting density response function. """

        calc = self.calc
        kd = self.kd
        gd = self.gd
        sdisp_cd = gd.sdisp_cd
        ibzk_kc = self.ibzk_kc
        bzk_kc = self.bzk_kc
        kq_k = self.kq_k
        pt = self.pt
        f_kn = self.f_kn
        e_kn = self.e_kn

        # Matrix init
        chi0_wGG = np.zeros((self.Nw_local, self.npw, self.npw), dtype=complex)
        if not (f_kn > self.ftol).any():
            self.chi0_wGG = chi0_wGG
            return

        if self.hilbert_trans:
            specfunc_wGG = np.zeros((self.NwS_local, self.npw, self.npw),
                                    dtype=complex)

        # Prepare for the derivative of pseudo-wavefunction
        if self.optical_limit:
            d_c = [Gradient(gd, i, n=4, dtype=complex).apply for i in range(3)]
            dpsit_g = gd.empty(dtype=complex)
            tmp = np.zeros((3), dtype=complex)

        rho_G = np.zeros(self.npw, dtype=complex)
        t0 = time()
        t_get_wfs = 0
        for k in range(self.kstart, self.kend):

            # Find corresponding kpoint in IBZ
            ibzkpt1 = kd.kibz_k[k]
            if self.optical_limit:
                ibzkpt2 = ibzkpt1
            else:
                ibzkpt2 = kd.kibz_k[kq_k[k]]

            for n in range(self.nstart, self.nend):
                #                print >> self.txt, k, n, t_get_wfs, time() - t0
                t1 = time()
                psitold_g = self.get_wavefunction(ibzkpt1, n, True, spin=spin)
                t_get_wfs += time() - t1
                psit1new_g = kd.transform_wave_function(psitold_g, k)

                P1_ai = pt.dict()
                pt.integrate(psit1new_g, P1_ai, k)

                psit1_g = psit1new_g.conj() * self.expqr_g

                for m in range(self.nbands):

                    if self.hilbert_trans:
                        check_focc = (f_kn[ibzkpt1, n] -
                                      f_kn[ibzkpt2, m]) > self.ftol
                    else:
                        check_focc = np.abs(f_kn[ibzkpt1, n] -
                                            f_kn[ibzkpt2, m]) > self.ftol

                    t1 = time()
                    psitold_g = self.get_wavefunction(ibzkpt2,
                                                      m,
                                                      check_focc,
                                                      spin=spin)
                    t_get_wfs += time() - t1

                    if check_focc:
                        psit2_g = kd.transform_wave_function(
                            psitold_g, kq_k[k])
                        P2_ai = pt.dict()
                        pt.integrate(psit2_g, P2_ai, kq_k[k])

                        # fft
                        tmp_g = np.fft.fftn(
                            psit2_g * psit1_g) * self.vol / self.nG0

                        for iG in range(self.npw):
                            index = self.Gindex_G[iG]
                            rho_G[iG] = tmp_g[index[0], index[1], index[2]]

                        if self.optical_limit:
                            phase_cd = np.exp(2j * pi * sdisp_cd *
                                              bzk_kc[kq_k[k], :, np.newaxis])
                            for ix in range(3):
                                d_c[ix](psit2_g, dpsit_g, phase_cd)
                                tmp[ix] = gd.integrate(psit1_g * dpsit_g)
                            rho_G[0] = -1j * np.dot(self.qq_v, tmp)

                        # PAW correction
                        for a, id in enumerate(calc.wfs.setups.id_a):
                            P_p = np.outer(P1_ai[a].conj(), P2_ai[a]).ravel()
                            gemv(1.0, self.phi_aGp[a], P_p, 1.0, rho_G)

                        if self.optical_limit:
                            rho_G[0] /= e_kn[ibzkpt2, m] - e_kn[ibzkpt1, n]

                        rho_GG = np.outer(rho_G, rho_G.conj())

                        if not self.hilbert_trans:
                            for iw in range(self.Nw_local):
                                w = self.w_w[iw + self.wstart] / Hartree
                                C = (f_kn[ibzkpt1, n] - f_kn[ibzkpt2, m]) / (
                                    w + e_kn[ibzkpt1, n] - e_kn[ibzkpt2, m] +
                                    1j * self.eta)
                                axpy(C, rho_GG, chi0_wGG[iw])
                        else:
                            focc = f_kn[ibzkpt1, n] - f_kn[ibzkpt2, m]
                            w0 = e_kn[ibzkpt2, m] - e_kn[ibzkpt1, n]
                            scal(focc, rho_GG)

                            # calculate delta function
                            w0_id = int(w0 / self.dw)
                            if w0_id + 1 < self.NwS:
                                # rely on the self.NwS_local is equal in each node!
                                if self.wScomm.rank == w0_id // self.NwS_local:
                                    alpha = (w0_id + 1 -
                                             w0 / self.dw) / self.dw
                                    axpy(alpha, rho_GG,
                                         specfunc_wGG[w0_id % self.NwS_local])

                                if self.wScomm.rank == (w0_id +
                                                        1) // self.NwS_local:
                                    alpha = (w0 / self.dw - w0_id) / self.dw
                                    axpy(
                                        alpha, rho_GG,
                                        specfunc_wGG[(w0_id + 1) %
                                                     self.NwS_local])

#                            deltaw = delta_function(w0, self.dw, self.NwS, self.sigma)
#                            for wi in range(self.NwS_local):
#                                if deltaw[wi + self.wS1] > 1e-8:
#                                    specfunc_wGG[wi] += tmp_GG * deltaw[wi + self.wS1]
                if self.nkpt == 1:
                    if n == 0:
                        dt = time() - t0
                        totaltime = dt * self.nband_local
                        self.printtxt(
                            'Finished n 0 in %f seconds, estimated %f seconds left.'
                            % (dt, totaltime))
                    if rank == 0 and self.nband_local // 5 > 0:
                        if n > 0 and n % (self.nband_local // 5) == 0:
                            dt = time() - t0
                            self.printtxt(
                                'Finished n %d in %f seconds, estimated %f seconds left.'
                                % (n, dt, totaltime - dt))
            if calc.wfs.world.size != 1:
                self.kcomm.barrier()
            if k == 0:
                dt = time() - t0
                totaltime = dt * self.nkpt_local
                self.printtxt(
                    'Finished k 0 in %f seconds, estimated %f seconds left.' %
                    (dt, totaltime))

            if rank == 0 and self.nkpt_local // 5 > 0:
                if k > 0 and k % (self.nkpt_local // 5) == 0:
                    dt = time() - t0
                    self.printtxt(
                        'Finished k %d in %f seconds, estimated %f seconds left.  '
                        % (k, dt, totaltime - dt))
        self.printtxt('Finished summation over k')

        self.kcomm.barrier()
        del rho_GG, rho_G
        # Hilbert Transform
        if not self.hilbert_trans:
            self.kcomm.sum(chi0_wGG)
        else:
            self.kcomm.sum(specfunc_wGG)
            if self.wScomm.size == 1:
                if not self.full_hilbert_trans:
                    chi0_wGG = hilbert_transform(
                        specfunc_wGG, self.Nw, self.dw,
                        self.eta)[self.wstart:self.wend]
                else:
                    chi0_wGG = full_hilbert_transform(
                        specfunc_wGG, self.Nw, self.dw,
                        self.eta)[self.wstart:self.wend]
                self.printtxt('Finished hilbert transform !')
                del specfunc_wGG
            else:
                # redistribute specfunc_wGG to all nodes
                assert self.NwS % size == 0
                NwStmp1 = (rank % self.kcomm.size) * self.NwS // size
                NwStmp2 = (rank % self.kcomm.size + 1) * self.NwS // size
                specfuncnew_wGG = specfunc_wGG[NwStmp1:NwStmp2]
                del specfunc_wGG

                coords = np.zeros(self.wcomm.size, dtype=int)
                nG_local = self.npw**2 // self.wcomm.size
                if self.wcomm.rank == self.wcomm.size - 1:
                    nG_local = self.npw**2 - (self.wcomm.size - 1) * nG_local
                self.wcomm.all_gather(np.array([nG_local]), coords)

                specfunc_Wg = SliceAlongFrequency(specfuncnew_wGG, coords,
                                                  self.wcomm)
                self.printtxt('Finished Slice Along Frequency !')
                if not self.full_hilbert_trans:
                    chi0_Wg = hilbert_transform(specfunc_Wg, self.Nw, self.dw,
                                                self.eta)[:self.Nw]
                else:
                    chi0_Wg = full_hilbert_transform(specfunc_Wg, self.Nw,
                                                     self.dw,
                                                     self.eta)[:self.Nw]
                self.printtxt('Finished hilbert transform !')
                self.comm.barrier()
                del specfunc_Wg

                chi0_wGG = SliceAlongOrbitals(chi0_Wg, coords, self.wcomm)
                self.printtxt('Finished Slice along orbitals !')
                self.comm.barrier()
                del chi0_Wg

        self.chi0_wGG = chi0_wGG / self.vol

        self.printtxt('')
        self.printtxt('Finished chi0 !')

        return