fix the parallel copy of eigenvectors

Author: maki49

source/module_lr/hamilt_casida.cpp

@@ -8,7 +8,7 @@ namespace LR
         const int no = this->nocc[0];
         const int nv = this->nvirt[0];
         const auto& px = this->pX[0];
-        const int nloc_per_band = nk * px.get_local_size();
+        const int ldim = nk * px.get_local_size();
         int npairs = no * nv;
         std::vector<T> Amat_full(this->nk * npairs * this->nk * npairs, 0.0);
         for (int ik = 0;ik < this->nk;++ik)
@@ -30,7 +30,7 @@ namespace LR
                     hamilt::Operator<T>* node(this->ops);
                     while (node != nullptr)
                     {   // act() on and return the k1-first type of psi
-                        node->act(1, nloc_per_band, /*npol=*/1, X_bj.get_pointer(), A_aibj.get_pointer());
+                        node->act(1, ldim, /*npol=*/1, X_bj.get_pointer(), A_aibj.get_pointer());
                         node = (hamilt::Operator<T>*)(node->next_op);
                     }
                     // reduce ai for a fixed bj

source/module_lr/hamilt_casida.h

@@ -128,17 +128,11 @@ namespace LR
                     for (int ik = 0;ik < nk;++ik) { this->DM_trans->set_DMK_pointer(ik, dm_trans_2d[ik].data<T>()); }
                 };
         }
-        ~HamiltLR()
-        {
-            if (this->ops != nullptr)
-            {
-                delete this->ops;
-            }
-        };
+        ~HamiltLR() { delete this->ops; }
 
-        virtual std::vector<T> matrix()const;
+        std::vector<T> matrix()const;
 
-        virtual void hPsi(const T* psi_in, T* hpsi, const int ld_psi, const int& nband) const
+        void hPsi(const T* psi_in, T* hpsi, const int ld_psi, const int& nband) const
         {
             assert(ld_psi == nk * pX[0].get_local_size());
             for (int ib = 0;ib < nband;++ib)
@@ -154,6 +148,30 @@ namespace LR
             }
         }
 
+        void global2local(T* lvec, const T* gvec, const int& nband) const
+        {
+            const int npairs = nocc[0] * nvirt[0];
+            for (int ib = 0;ib < nband;++ib)
+            {
+                const int loffset_b = ib * nk * pX[0].get_local_size();
+                const int goffset_b = ib * nk * npairs;
+                for (int ik = 0;ik < nk;++ik)
+                {
+                    const int loffset = loffset_b + ik * pX[0].get_local_size();
+                    const int goffset = goffset_b + ik * npairs;
+                    for (int lo = 0;lo < pX[0].get_col_size();++lo)
+                    {
+                        const int go = pX[0].local2global_col(lo);
+                        for (int lv = 0;lv < pX[0].get_row_size();++lv)
+                        {
+                            const int gv = pX[0].local2global_row(lv);
+                            lvec[loffset + lo * pX[0].get_row_size() + lv] = gvec[goffset + go * nvirt[0] + gv];
+                        }
+                    }
+                }
+            }
+        }
+
     private:
         const std::vector<int>& nocc;
         const std::vector<int>& nvirt;

source/module_lr/hamilt_ulr.hpp

@@ -35,7 +35,8 @@ namespace LR
             const std::vector<Parallel_2D>& pX_in,   ///< {up, down}
             const Parallel_2D& pc_in,
             const Parallel_Orbitals& pmat_in) :nocc(nocc), nvirt(nvirt), pX(pX_in), nk(kv_in.get_nks() / nspin),
-            nloc_per_band(nk* pX[0].get_local_size() + nk * pX[1].get_local_size())
+            ldim(nk* pX[0].get_local_size() + nk * pX[1].get_local_size()),
+            gdim(nk* std::inner_product(nocc.begin(), nocc.end(), nvirt.begin(), 0))
         {
             ModuleBase::TITLE("HamiltULR", "HamiltULR");
             this->DM_trans = LR_Util::make_unique<elecstate::DensityMatrix<T, T>>(&pmat_in, 1, kv_in.kvec_d, nk);
@@ -92,7 +93,7 @@ namespace LR
         void hPsi(const T* psi_in, T* hpsi, const int ld_psi, const int& nband) const
         {
             ModuleBase::TITLE("HamiltULR", "hPsi");
-            assert(ld_psi == this->nloc_per_band);
+            assert(ld_psi == this->ldim);
             const std::vector<int64_t> xdim_is = { nk * pX[0].get_local_size(), nk * pX[1].get_local_size() };
             /// band-wise act (also works for close-shell, but not efficient)
             for (int ib = 0;ib < nband;++ib)
@@ -121,8 +122,7 @@ namespace LR
             const std::vector<int> npairs = { this->nocc[0] * this->nvirt[0], this->nocc[1] * this->nvirt[1] };
             const std::vector<int64_t> ldim_is = { nk * pX[0].get_local_size(), nk * pX[1].get_local_size() };
             const std::vector<int> gdim_is = { nk * npairs[0], nk * npairs[1] };
-            const int global_size = this->nk * (npairs[0] + npairs[1]);
-            std::vector<T> Amat_full(global_size * global_size);
+            std::vector<T> Amat_full(gdim * gdim);
             for (int is_bj : {0, 1})
             {
                 const int no = this->nocc[is_bj];
@@ -137,13 +137,13 @@ namespace LR
                         for (int b = 0;b < nv;++b)
                         {
                             const int gcol = goffset_bj + ik_bj * npairs[is_bj] + j * nv + b;//global
-                            std::vector<T> X_bj(this->nloc_per_band, T(0));
+                            std::vector<T> X_bj(this->ldim, T(0));
                             const int lj = px.global2local_col(j);
                             const int lb = px.global2local_row(b);
                             const int lcol = loffset_bj + ik_bj * px.get_local_size() + lj * px.get_row_size() + lb;//local
                             if (px.in_this_processor(b, j)) { X_bj[lcol] = T(1); }
                             this->cal_dm_trans(is_bj, X_bj.data() + loffset_bj);
-                            std::vector<T> Aloc_col(this->nloc_per_band, T(0)); // a col of A matrix (local)
+                            std::vector<T> Aloc_col(this->ldim, T(0)); // a col of A matrix (local)
                             for (int is_ai : {0, 1})
                             {
                                 const int goffset_ai = is_ai * gdim_is[0];
@@ -159,22 +159,54 @@ namespace LR
                                 for (int ik_ai = 0;ik_ai < this->nk;++ik_ai)
                                 {
                                     LR_Util::gather_2d_to_full(pax, Aloc_col.data() + loffset_ai + ik_ai * pax.get_local_size(),
-                                        Amat_full.data() + gcol * global_size /*col, bj*/ + goffset_ai + ik_ai * npairs[is_ai]/*row, ai*/,
+                                        Amat_full.data() + gcol * gdim /*col, bj*/ + goffset_ai + ik_ai * npairs[is_ai]/*row, ai*/,
                                         false, nv, no);
                                 }
 #else
-                                std::memcpy(Amat_full.data() + gcol * global_size + goffset_ai, Aloc_col.data() + goffset_ai, gdim_is[is_ai] * sizeof(T));
+                                std::memcpy(Amat_full.data() + gcol * gdim + goffset_ai, Aloc_col.data() + goffset_ai, gdim_is[is_ai] * sizeof(T));
 #endif
                             }
                         }
                     }
                 }
             }
             std::cout << "Full A matrix:" << std::endl;
-            LR_Util::print_value(Amat_full.data(), global_size, global_size);
+            LR_Util::print_value(Amat_full.data(), gdim, gdim);
             return Amat_full;
         }
 
+        /// copy global data (eigenvectors) to local memory
+        void global2local(T* lvec, const T* gvec, const int& nband) const
+        {
+            const std::vector<int> npairs = { this->nocc[0] * this->nvirt[0], this->nocc[1] * this->nvirt[1] };
+            const std::vector<int64_t> ldim_is = { nk * pX[0].get_local_size(), nk * pX[1].get_local_size() };
+            const std::vector<int> gdim_is = { nk * npairs[0], nk * npairs[1] };
+            for (int ib = 0;ib < nband;++ib)
+            {
+                const int loffset_b = ib * this->ldim;
+                const int goffset_b = ib * this->gdim;
+                for (int is : {0, 1})
+                {
+                    const int loffset_bs = loffset_b + is * ldim_is[0];
+                    const int goffset_bs = goffset_b + is * gdim_is[0];
+                    for (int ik = 0;ik < nk;++ik)
+                    {
+                        const int loffset = loffset_bs + ik * pX[is].get_local_size();
+                        const int goffset = goffset_bs + ik * npairs[is];
+                        for (int lo = 0;lo < pX[is].get_col_size();++lo)
+                        {
+                            const int go = pX[is].local2global_col(lo);
+                            for (int lv = 0;lv < pX[is].get_row_size();++lv)
+                            {
+                                const int gv = pX[is].local2global_row(lv);
+                                lvec[loffset + lo * pX[is].get_row_size() + lv] = gvec[goffset + go * nvirt[is] + gv];
+                            }
+                        }
+                    }
+                }
+            }
+        }
+
     private:
         const std::vector<int>& nocc;
         const std::vector<int>& nvirt;
@@ -183,7 +215,8 @@ namespace LR
 
 
         const int nk = 1;
-        const int nloc_per_band = 1;
+        const int ldim = 1;
+        const int gdim = 1;
 
         /// 4 operator lists: uu, ud, du, dd
         std::vector<hamilt::Operator<T>*> ops;

source/module_lr/hsolver_lrtd.hpp

@@ -50,17 +50,18 @@ namespace LR
             if (method == "lapack")
             {
                 std::vector<T> Amat_full = hm.matrix();
-                eigenvalue.resize(dim);
-                if (hermitian) { LR_Util::diag_lapack(dim, Amat_full.data(), eigenvalue.data()); }
+                const int gdim = std::sqrt(Amat_full.size());
+                eigenvalue.resize(gdim);
+                if (hermitian) { LR_Util::diag_lapack(gdim, Amat_full.data(), eigenvalue.data()); }
                 else
                 {
-                    std::vector<std::complex<double>> eig_complex(dim);
-                    LR_Util::diag_lapack_nh(dim, Amat_full.data(), eig_complex.data());
+                    std::vector<std::complex<double>> eig_complex(gdim);
+                    LR_Util::diag_lapack_nh(gdim, Amat_full.data(), eig_complex.data());
                     print_eigs(eig_complex, "Right eigenvalues: of the non-Hermitian matrix: (Ry)");
-                    for (int i = 0; i < dim; i++) { eigenvalue[i] = eig_complex[i].real(); }
+                    for (int i = 0; i < gdim; i++) { eigenvalue[i] = eig_complex[i].real(); }
                 }
                 // copy eigenvectors
-                std::memcpy(psi, Amat_full.data(), sizeof(T) * dim * nband);
+                hm.global2local(psi, Amat_full.data(), nband);
             }
             else
             {