Pexsi stage4 pattern reuse opt

source/source_esolver/esolver_ks_lcao.cpp

@@ -403,7 +403,13 @@ void ESolver_KS_LCAO<TK, TR>::hamilt2rho_single(UnitCell& ucell, int istep, int
     // 3) run Hsolver
     if (!skip_solve)
     {
+#ifdef __PEXSI
+        hsolver::HSolverLCAO<TK> hsolver_lcao_obj(&(this->pv),
+                                                   PARAM.inp.ks_solver,
+                                                   &this->pexsi_reuse_contexts_);
+#else
         hsolver::HSolverLCAO<TK> hsolver_lcao_obj(&(this->pv), PARAM.inp.ks_solver);
+#endif
         hsolver_lcao_obj.solve(static_cast<hamilt::Hamilt<TK>*>(this->p_hamilt), this->psi[0], this->pelec, *this->dmat.dm, 
           this->chr, PARAM.inp.nspin, skip_charge);
     }

source/source_esolver/esolver_ks_lcao.h

@@ -13,7 +13,11 @@
 #include "source_lcao/setup_dm.h" // mohan add 2025-10-30
 
 #include <memory>
+#include <vector>
 
+#ifdef __PEXSI
+#include "source_hsolver/module_pexsi/simple_pexsi.h"
+#endif
 
 // for Linear Response
 namespace LR
@@ -82,6 +86,9 @@ class ESolver_KS_LCAO : public ESolver_KS
     //! Add density matrix class, mohan add 2025-10-30
     LCAO_domain::Setup_DM<TK> dmat;
 
+#ifdef __PEXSI
+    std::vector<pexsi::PexsiComplexReuseContext> pexsi_reuse_contexts_;
+#endif
 
     // For deepks method, mohan add 2025-10-08
     Setup_DeePKS<TK> deepks;

source/source_esolver/esolver_ks_lcao_tddft.cpp

@@ -325,7 +325,13 @@ void ESolver_KS_LCAO_TDDFT<TR, Device>::hamilt2rho_single(UnitCell& ucell,
         if (this->psi != nullptr)
         {
             bool skip_charge = PARAM.inp.calculation == "nscf" ? true : false;
+#ifdef __PEXSI
+            hsolver::HSolverLCAO<std::complex<double>> hsolver_lcao_obj(&this->pv,
+                                                                         PARAM.inp.ks_solver,
+                                                                         &this->pexsi_reuse_contexts_);
+#else
             hsolver::HSolverLCAO<std::complex<double>> hsolver_lcao_obj(&this->pv, PARAM.inp.ks_solver);
+#endif
             hsolver_lcao_obj.solve(static_cast<hamilt::Hamilt<std::complex<double>>*>(this->p_hamilt),
                                    this->psi[0],
                                    this->pelec,

source/source_estate/elecstate_lcao.cpp

@@ -43,7 +43,7 @@ void ElecStateLCAO<double>::dm2rho(std::vector<double*> pexsi_DM,
     }
 
 #ifdef __PEXSI
-    dm->pexsi_EDM = pexsi_EDM;
+    dm->set_pexsi_EDM_pointer(pexsi_EDM);
 #endif
 
     for (int is = 0; is < nspin; is++)
@@ -80,7 +80,56 @@ void ElecStateLCAO<std::complex<double>>::dm2rho(std::vector<std::complex<double
 		std::vector<std::complex<double>*> pexsi_EDM,
 		DensityMatrix<std::complex<double>, double>* dm)
 {
-    ModuleBase::WARNING_QUIT("ElecStateLCAO", "pexsi is not completed for multi-k case");
+    ModuleBase::timer::start("ElecStateLCAO", "dm2rho");
+
+    const int dmk_size = dm->get_DMK_size();
+    if (static_cast<int>(pexsi_DM.size()) < dmk_size)
+    {
+        ModuleBase::WARNING_QUIT("ElecStateLCAO", "PEXSI multi-k density matrix buffer is smaller than DMK storage");
+    }
+
+    const int local_matrix_size = dm->get_DMK_nrow() * dm->get_DMK_ncol();
+    const double pexsi_spin_weight = PARAM.inp.nspin == 1 ? 0.5 : 1.0;
+    for (int ik = 0; ik < dmk_size; ++ik)
+    {
+        const double k_weight = ((this->klist != nullptr && ik < static_cast<int>(this->klist->wk.size()))
+                                     ? this->klist->wk[ik]
+                                     : 1.0)
+                                * pexsi_spin_weight;
+        for (int i = 0; i < local_matrix_size; ++i)
+        {
+            pexsi_DM[ik][i] *= k_weight;
+            pexsi_EDM[ik][i] *= k_weight;
+        }
+        dm->set_DMK_pointer(ik, pexsi_DM[ik]);
+    }
+
+#ifdef __PEXSI
+    dm->set_pexsi_EDM_pointer(pexsi_EDM);
+#endif
+
+    dm->cal_DMR();
+
+    for (int is = 0; is < PARAM.inp.nspin; is++)
+    {
+        ModuleBase::GlobalFunc::ZEROS(this->charge->rho[is], this->charge->nrxx);
+    }
+
+    ModuleBase::GlobalFunc::NOTE("Calculate the charge on real space grid!");
+    ModuleGint::cal_gint_rho(dm->get_DMR_vector(), PARAM.inp.nspin, this->charge->rho);
+    if (XC_Functional::get_ked_flag())
+    {
+        for (int is = 0; is < PARAM.inp.nspin; is++)
+        {
+            ModuleBase::GlobalFunc::ZEROS(this->charge->kin_r[0], this->charge->nrxx);
+        }
+        ModuleGint::cal_gint_tau(dm->get_DMR_vector(), PARAM.inp.nspin, this->charge->kin_r);
+    }
+
+    this->charge->renormalize_rho();
+
+    ModuleBase::timer::end("ElecStateLCAO", "dm2rho");
+    return;
 }
 
 

source/source_estate/module_dm/density_matrix.h

@@ -1,7 +1,10 @@
 #ifndef DENSITY_MATRIX_H
 #define DENSITY_MATRIX_H
 
+#include <algorithm>
+#include <cstddef>
 #include <string>
+#include <vector>
 
 #include "source_cell/module_neighbor/sltk_grid_driver.h"
 #include "source_lcao/record_adj.h"
@@ -267,6 +270,24 @@ class DensityMatrix
     std::vector<ModuleBase::ComplexMatrix> EDMK; // for TD-DFT
 
 #ifdef __PEXSI
+    void set_pexsi_EDM_pointer(const std::vector<TK*>& pexsi_EDM_in)
+    {
+        const int local_size = this->_paraV->nrow * this->_paraV->ncol;
+        this->pexsi_EDM_storage_.resize(pexsi_EDM_in.size());
+        this->pexsi_EDM.resize(pexsi_EDM_in.size());
+        for (std::size_t ik = 0; ik < pexsi_EDM_in.size(); ++ik)
+        {
+            this->pexsi_EDM_storage_[ik].resize(local_size);
+            if (local_size > 0)
+            {
+                std::copy(pexsi_EDM_in[ik],
+                          pexsi_EDM_in[ik] + local_size,
+                          this->pexsi_EDM_storage_[ik].begin());
+            }
+            this->pexsi_EDM[ik] = this->pexsi_EDM_storage_[ik].data();
+        }
+    }
+
     /**
      * @brief EDM storage for PEXSI
      * used in MD calculation
@@ -298,6 +319,10 @@ class DensityMatrix
     // std::vector<ModuleBase::ComplexMatrix> _DMK;
     std::vector<std::vector<TK>> _DMK;
 
+#ifdef __PEXSI
+    std::vector<std::vector<TK>> pexsi_EDM_storage_;
+#endif
+
     /**
      * @brief K_Vectors object, which is used to get k-point information
      */

source/source_estate/module_dm/test/test_dm_constructor.cpp

@@ -237,6 +237,28 @@ TEST_F(DMTest, DMConstructor_nspin2)
     delete kv;
 }
 
+#ifdef __PEXSI
+TEST_F(DMTest, PexsiEDMIsOwnedCopy)
+{
+    elecstate::DensityMatrix<double, double> DM(paraV, 1);
+    const int local_size = paraV->nrow * paraV->ncol;
+    std::vector<double> external_edm(local_size);
+    for (int i = 0; i < local_size; ++i)
+    {
+        external_edm[i] = static_cast<double>(i + 1);
+    }
+
+    DM.set_pexsi_EDM_pointer({external_edm.data()});
+    std::fill(external_edm.begin(), external_edm.end(), -1.0);
+
+    ASSERT_EQ(DM.pexsi_EDM.size(), 1);
+    for (int i = 0; i < local_size; ++i)
+    {
+        EXPECT_DOUBLE_EQ(DM.pexsi_EDM[0][i], static_cast<double>(i + 1));
+    }
+}
+#endif
+
 int main(int argc, char** argv)
 {
 #ifdef __MPI