diff --git a/python/pyabacus/src/hsolver/py_diago_dav_subspace.hpp b/python/pyabacus/src/hsolver/py_diago_dav_subspace.hpp index 316bd9be74..e2314b961c 100644 --- a/python/pyabacus/src/hsolver/py_diago_dav_subspace.hpp +++ b/python/pyabacus/src/hsolver/py_diago_dav_subspace.hpp @@ -132,6 +132,11 @@ class PyDiagoDavSubspace std::copy(hpsi_ptr, hpsi_ptr + nvec * ld_psi, hpsi_out); }; + auto spsi_func = [this](const std::complex* psi_in, + std::complex* spsi_out, + const int ld_psi, + const int nvec) { syncmem_op()(spsi_out, psi_in, static_cast(ld_psi * nvec)); }; + obj = std::make_unique, base_device::DEVICE_CPU>>( precond_vec, nband, @@ -145,7 +150,7 @@ class PyDiagoDavSubspace nb2d ); - return obj->diag(hpsi_func, psi, nbasis, eigenvalue, diag_ethr, scf_type); + return obj->diag(hpsi_func, spsi_func, psi, nbasis, eigenvalue, diag_ethr, scf_type); } private: @@ -156,6 +161,10 @@ class PyDiagoDavSubspace int nband; std::unique_ptr, base_device::DEVICE_CPU>> obj; + + base_device::DEVICE_CPU* ctx = {}; + using syncmem_op = base_device::memory:: + synchronize_memory_op, base_device::DEVICE_CPU, base_device::DEVICE_CPU>; }; } // namespace py_hsolver diff --git a/source/source_hsolver/diago_dav_subspace.cpp b/source/source_hsolver/diago_dav_subspace.cpp index b98f23e696..4402fad9f8 100644 --- a/source/source_hsolver/diago_dav_subspace.cpp +++ b/source/source_hsolver/diago_dav_subspace.cpp @@ -54,6 +54,10 @@ Diago_DavSubspace::Diago_DavSubspace(const std::vector& precond resmem_complex_op()(this->hphi, this->nbase_x * this->dim, "DAV::hphi"); setmem_complex_op()(this->hphi, 0, this->nbase_x * this->dim); + // the product of S and psi in the reduced psi set + resmem_complex_op()(this->sphi, this->nbase_x * this->dim, "DAV::sphi"); + setmem_complex_op()(this->sphi, 0, this->nbase_x * this->dim); + // Hamiltonian on the reduced psi set resmem_complex_op()(this->hcc, this->nbase_x * this->nbase_x, "DAV::hcc"); setmem_complex_op()(this->hcc, 0, this->nbase_x * this->nbase_x); @@ -96,6 +100,7 @@ Diago_DavSubspace::~Diago_DavSubspace() template int Diago_DavSubspace::diag_once(const HPsiFunc& hpsi_func, + const HPsiFunc& spsi_func, T* psi_in, const int psi_in_dmax, Real* eigenvalue_in_hsolver, @@ -134,7 +139,11 @@ int Diago_DavSubspace::diag_once(const HPsiFunc& hpsi_func, // hphi[:, 0:nbase_x] = H * psi_in_iter[:, 0:nbase_x] hpsi_func(this->psi_in_iter, this->hphi, this->dim, this->notconv); - this->cal_elem(this->dim, nbase, this->notconv, this->psi_in_iter, this->hphi, this->hcc, this->scc); + // compute s*psi_in_iter + // sphi[:, 0:nbase_x] = S * psi_in_iter[:, 0:nbase_x] + spsi_func(this->psi_in_iter, this->sphi, this->dim, this->notconv); + + this->cal_elem(this->dim, nbase, this->notconv, this->psi_in_iter, this->sphi, this->hphi, this->hcc, this->scc); this->diag_zhegvx(nbase, this->notconv, this->hcc, this->scc, this->nbase_x, &eigenvalue_iter, this->vcc); @@ -152,16 +161,25 @@ int Diago_DavSubspace::diag_once(const HPsiFunc& hpsi_func, dav_iter++; this->cal_grad(hpsi_func, + spsi_func, this->dim, nbase, this->notconv, this->psi_in_iter, this->hphi, + this->sphi, this->vcc, unconv.data(), &eigenvalue_iter); - this->cal_elem(this->dim, nbase, this->notconv, this->psi_in_iter, this->hphi, this->hcc, this->scc); + this->cal_elem(this->dim, + nbase, + this->notconv, + this->psi_in_iter, + this->sphi, + this->hphi, + this->hcc, + this->scc); this->diag_zhegvx(nbase, this->n_band, this->hcc, this->scc, this->nbase_x, &eigenvalue_iter, this->vcc); @@ -238,6 +256,7 @@ int Diago_DavSubspace::diag_once(const HPsiFunc& hpsi_func, eigenvalue_in_hsolver, this->psi_in_iter, this->hphi, + this->sphi, this->hcc, this->scc, this->vcc); @@ -255,11 +274,13 @@ int Diago_DavSubspace::diag_once(const HPsiFunc& hpsi_func, template void Diago_DavSubspace::cal_grad(const HPsiFunc& hpsi_func, + const HPsiFunc& spsi_func, const int& dim, const int& nbase, const int& notconv, T* psi_iter, T* hphi, + T* spsi, T* vcc, const int* unconv, std::vector* eigenvalue_iter) @@ -331,7 +352,7 @@ void Diago_DavSubspace::cal_grad(const HPsiFunc& hpsi_func, notconv, nbase, this->one, - psi_iter, + sphi, this->dim, vcc, this->nbase_x, @@ -396,6 +417,7 @@ void Diago_DavSubspace::cal_grad(const HPsiFunc& hpsi_func, // update hpsi[:, nbase:nbase+notconv] // hpsi[:, nbase:nbase+notconv] = H * psi_iter[:, nbase:nbase+notconv] hpsi_func(psi_iter + nbase * dim, hphi + nbase * this->dim, this->dim, notconv); + spsi_func(psi_iter + nbase * dim, sphi + nbase * this->dim, this->dim, notconv); ModuleBase::timer::tick("Diago_DavSubspace", "cal_grad"); return; @@ -406,6 +428,7 @@ void Diago_DavSubspace::cal_elem(const int& dim, int& nbase, const int& notconv, const T* psi_iter, + const T* spsi, const T* hphi, T* hcc, T* scc) @@ -416,39 +439,39 @@ void Diago_DavSubspace::cal_elem(const int& dim, ModuleBase::gemm_op_mt() #else ModuleBase::gemm_op() -#endif - ('C', - 'N', - nbase + notconv, - notconv, - this->dim, - this->one, - psi_iter, - this->dim, - &hphi[nbase * this->dim], - this->dim, - this->zero, - &hcc[nbase * this->nbase_x], - this->nbase_x); +#endif + ('C', + 'N', + nbase + notconv, + notconv, + this->dim, + this->one, + psi_iter, + this->dim, + &hphi[nbase * this->dim], + this->dim, + this->zero, + &hcc[nbase * this->nbase_x], + this->nbase_x); #ifdef __DSP ModuleBase::gemm_op_mt() #else ModuleBase::gemm_op() #endif - ('C', - 'N', - nbase + notconv, - notconv, - this->dim, - this->one, - psi_iter, - this->dim, - psi_iter + nbase * this->dim, - this->dim, - this->zero, - &scc[nbase * this->nbase_x], - this->nbase_x); + ('C', + 'N', + nbase + notconv, + notconv, + this->dim, + this->one, + psi_iter, + this->dim, + spsi + nbase * this->dim, + this->dim, + this->zero, + &scc[nbase * this->nbase_x], + this->nbase_x); #ifdef __MPI if (this->diag_comm.nproc > 1) @@ -685,10 +708,11 @@ void Diago_DavSubspace::refresh(const int& dim, const Real* eigenvalue_in_hsolver, // const psi::Psi& psi, T* psi_iter, - T* hp, - T* sp, - T* hc, - T* vc) + T* hphi, + T* sphi, + T* hcc, + T* scc, + T* vcc) { ModuleBase::timer::tick("Diago_DavSubspace", "refresh"); @@ -714,6 +738,28 @@ void Diago_DavSubspace::refresh(const int& dim, // update hphi syncmem_complex_op()(hphi, psi_iter + nband * this->dim, this->dim * nband); +#ifdef __DSP + ModuleBase::gemm_op_mt() +#else + ModuleBase::gemm_op() +#endif + ('N', + 'N', + this->dim, + nband, + nbase, + this->one, + this->sphi, + this->dim, + this->vcc, + this->nbase_x, + this->zero, + psi_iter + nband * this->dim, + this->dim); + + // update sphi + syncmem_complex_op()(sphi, psi_iter + nband * this->dim, this->dim * nband); + nbase = nband; // set hcc/scc/vcc to 0 @@ -776,6 +822,7 @@ void Diago_DavSubspace::refresh(const int& dim, template int Diago_DavSubspace::diag(const HPsiFunc& hpsi_func, + const HPsiFunc& spsi_func, T* psi_in, const int psi_in_dmax, Real* eigenvalue_in_hsolver, @@ -791,7 +838,7 @@ int Diago_DavSubspace::diag(const HPsiFunc& hpsi_func, do { - sum_iter += this->diag_once(hpsi_func, psi_in, psi_in_dmax, eigenvalue_in_hsolver, ethr_band); + sum_iter += this->diag_once(hpsi_func, spsi_func, psi_in, psi_in_dmax, eigenvalue_in_hsolver, ethr_band); ++ntry; diff --git a/source/source_hsolver/diago_dav_subspace.h b/source/source_hsolver/diago_dav_subspace.h index fa16881d69..3b4c224ec8 100644 --- a/source/source_hsolver/diago_dav_subspace.h +++ b/source/source_hsolver/diago_dav_subspace.h @@ -41,6 +41,7 @@ class Diago_DavSubspace using HPsiFunc = std::function; int diag(const HPsiFunc& hpsi_func, + const HPsiFunc& spsi_func, T* psi_in, const int psi_in_dmax, Real* eigenvalue_in, @@ -81,6 +82,9 @@ class Diago_DavSubspace /// the product of H and psi in the reduced basis set T* hphi = nullptr; + /// the product of S and psi in the reduced basis set + T* sphi = nullptr; + /// Hamiltonian on the reduced basis T* hcc = nullptr; @@ -96,16 +100,25 @@ class Diago_DavSubspace base_device::AbacusDevice_t device = {}; void cal_grad(const HPsiFunc& hpsi_func, + const HPsiFunc& spsi_func, const int& dim, const int& nbase, const int& notconv, T* psi_iter, T* hphi, + T* spsi, T* vcc, const int* unconv, std::vector* eigenvalue_iter); - void cal_elem(const int& dim, int& nbase, const int& notconv, const T* psi_iter, const T* hphi, T* hcc, T* scc); + void cal_elem(const int& dim, + int& nbase, + const int& notconv, + const T* psi_iter, + const T* sphi, + const T* hphi, + T* hcc, + T* scc); void refresh(const int& dim, const int& nband, @@ -113,6 +126,7 @@ class Diago_DavSubspace const Real* eigenvalue, T* psi_iter, T* hphi, + T* sphi, T* hcc, T* scc, T* vcc); @@ -134,6 +148,7 @@ class Diago_DavSubspace T* vcc); int diag_once(const HPsiFunc& hpsi_func, + const HPsiFunc& spsi_func, T* psi_in, const int psi_in_dmax, Real* eigenvalue_in, diff --git a/source/source_hsolver/hsolver_pw.cpp b/source/source_hsolver/hsolver_pw.cpp index e36737848e..c061023546 100644 --- a/source/source_hsolver/hsolver_pw.cpp +++ b/source/source_hsolver/hsolver_pw.cpp @@ -380,6 +380,10 @@ void HSolverPW::hamiltSolvePsiK(hamilt::Hamilt* hm, }; bool scf = this->calculation_type == "nscf" ? false : true; + auto spsi_func = [hm](T* psi_in, T* spsi_out, const int ld_psi, const int nvec) { + hm->sPsi(psi_in, spsi_out, ld_psi, ld_psi, nvec); + }; + Diago_DavSubspace dav_subspace(pre_condition, psi.get_nbands(), psi.get_k_first() ? psi.get_current_ngk() @@ -393,7 +397,8 @@ void HSolverPW::hamiltSolvePsiK(hamilt::Hamilt* hm, PARAM.inp.nb2d); DiagoIterAssist::avg_iter += static_cast( - dav_subspace.diag(hpsi_func, psi.get_pointer(), psi.get_nbasis(), eigenvalue, this->ethr_band, scf)); + dav_subspace + .diag(hpsi_func, spsi_func, psi.get_pointer(), psi.get_nbasis(), eigenvalue, this->ethr_band, scf)); } else if (this->method == "dav") { diff --git a/source/source_lcao/module_lr/hsolver_lrtd.hpp b/source/source_lcao/module_lr/hsolver_lrtd.hpp index df49a9daa0..bbac18160a 100644 --- a/source/source_lcao/module_lr/hsolver_lrtd.hpp +++ b/source/source_lcao/module_lr/hsolver_lrtd.hpp @@ -105,13 +105,8 @@ namespace LR PARAM.inp.diag_subspace, PARAM.inp.nb2d); std::vector ethr_band(nband, diag_ethr); - hsolver::DiagoIterAssist::avg_iter - += static_cast(dav_subspace.diag( - hpsi_func, psi, - dim, - eigenvalue.data(), - ethr_band, - false /*scf*/)); + hsolver::DiagoIterAssist::avg_iter += static_cast( + dav_subspace.diag(hpsi_func, spsi_func, psi, dim, eigenvalue.data(), ethr_band, false /*scf*/)); } else if (method == "cg") { diff --git a/tools/stm/stm.py b/tools/stm/stm.py index a3a493367f..8bc3df7137 100644 --- a/tools/stm/stm.py +++ b/tools/stm/stm.py @@ -123,7 +123,7 @@ def scan2(self, bias, z, repeat=(1, 1)): # Returing scan with axes in Angstrom. return x, y, current - + def linescan(self, bias, current, p1, p2, npoints=50, z0=None): """Constant current line scan. @@ -225,6 +225,30 @@ def line_sts(self, bias0, bias1, biasstep, p1, p2, npoints=50): return biases, np.linspace(0, s, npoints), current, dIdV + def line_sts_ave(self, bias0, bias1, biasstep, npoints=50): + + biases = np.arange(bias0, bias1 + biasstep, biasstep) + current = np.zeros((npoints, len(biases))) + + for b in np.arange(len(biases)): + print(b, biases[b]) + self.read_ldos(biases[b]) + nz = self.ldos.shape[2] + + for i in range(npoints): + z = i / (npoints - 1) * nz + dz = z - np.floor(z) + z = int(z) % nz + current[i, b] = ((1 - dz) * self.ldos[:, :, z].mean() + + dz * self.ldos[:, :, (z + 1) % nz].mean()) + + dIdV = np.zeros((npoints, len(biases))) + for i in range(npoints): + dIdV[i, :] = np.gradient(current[i, :], biasstep) + + return biases, np.linspace(0, 1, npoints), current, dIdV + + def find_current(self, ldos, z): """ Finds current for given LDOS at height z.""" nz = self.ldos.shape[2] @@ -279,4 +303,4 @@ def find_height(ldos, current, h, z0=None): def delta(biases, bias, width): """Return a delta-function centered at 'bias'""" x = -((biases - bias) / width)**2 - return np.exp(x) / (np.sqrt(np.pi) * width) \ No newline at end of file + return np.exp(x) / (np.sqrt(np.pi) * width)