Draft: Enable complex multi-k PEXSI path for LCAO

source/source_estate/elecstate_lcao.cpp

@@ -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->pexsi_EDM = 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_hsolver/diago_pexsi.cpp

@@ -6,10 +6,19 @@
 #include "diago_pexsi.h"
 #include "source_base/tool_title.h"
 #include "source_base/global_variable.h"
+#include "source_base/parallel_global.h"
 #include "source_base/tool_quit.h"
 #include "source_basis/module_ao/parallel_orbitals.h"
+#include "module_pexsi/simple_pexsi.h"
 #include "module_pexsi/pexsi_solver.h"
 
+#include <algorithm>
+#include <cfloat>
+#include <cmath>
+#include <utility>
+
+extern MPI_Comm DIAG_WORLD;
+
 typedef hamilt::MatrixBlock<double> matd;
 typedef hamilt::MatrixBlock<std::complex<double>> matcd;
 
@@ -19,46 +28,140 @@ template <typename T>
 std::vector<double> DiagoPexsi<T>::mu_buffer;
 
 template <typename T>
-DiagoPexsi<T>::DiagoPexsi(const Parallel_Orbitals* ParaV_in)
+DiagoPexsi<T>::DiagoPexsi(const Parallel_Orbitals* ParaV_in, std::unique_ptr<pexsi::IPexsiSolver> solver_in)
 {
+    this->ParaV = ParaV_in;
+    this->ps = std::move(solver_in);
+    if (this->ps == nullptr)
+    {
+        this->ps = std::make_unique<pexsi::PEXSI_Solver>();
+    }
+
     int nspin = PARAM.inp.nspin;
     if (PARAM.inp.nspin == 4)
     {
         nspin = 1;
     }
-    mu_buffer.resize(nspin);
-    for (int i = 0; i < nspin; i++)
+    if (static_cast<int>(mu_buffer.size()) < nspin)
     {
-        mu_buffer[i] = this->ps->pexsi_mu;
+        mu_buffer.resize(nspin, pexsi::PEXSI_Solver::pexsi_mu);
     }
+    this->resize_density_buffers(nspin);
 
-    this->ParaV = ParaV_in;
-    this->ps = std::make_unique<pexsi::PEXSI_Solver>();
+}
 
-    this->DM.resize(nspin);
-    this->EDM.resize(nspin);
-    for (int i = 0; i < nspin; i++)
+template <typename T>
+DiagoPexsi<T>::~DiagoPexsi()
+{
+}
+
+template <typename T>
+void DiagoPexsi<T>::resize_density_buffers(const int count)
+{
+    const int local_size = this->ParaV->nrow * this->ParaV->ncol;
+    this->DM.resize(count);
+    this->EDM.resize(count);
+    this->dm_buffer_.resize(count);
+    this->edm_buffer_.resize(count);
+    for (int i = 0; i < count; ++i)
     {
-        this->DM[i] = new T[ParaV->nrow * ParaV->ncol];
-        this->EDM[i] = new T[ParaV->nrow * ParaV->ncol];
+        if (static_cast<int>(this->dm_buffer_[i].size()) != local_size)
+        {
+            this->dm_buffer_[i].assign(local_size, T{});
+        }
+        if (static_cast<int>(this->edm_buffer_[i].size()) != local_size)
+        {
+            this->edm_buffer_[i].assign(local_size, T{});
+        }
+        this->DM[i] = this->dm_buffer_[i].data();
+        this->EDM[i] = this->edm_buffer_[i].data();
     }
+    if (static_cast<int>(mu_buffer.size()) < count)
+    {
+        mu_buffer.resize(count, pexsi::PEXSI_Solver::pexsi_mu);
+    }
+}
 
+template <typename T>
+void DiagoPexsi<T>::begin_mu_search()
+{
+    this->has_mu_lower_ = false;
+    this->has_mu_upper_ = false;
+    this->mu_lower_ = 0.0;
+    this->mu_upper_ = 0.0;
 }
 
 template <typename T>
-DiagoPexsi<T>::~DiagoPexsi()
+void DiagoPexsi<T>::begin_k_loop()
 {
-    int nspin = PARAM.inp.nspin;
-    if (PARAM.inp.nspin == 4)
+    this->num_electron_sum_ = 0.0;
+    this->num_electron_derivative_sum_ = 0.0;
+    this->totalEnergyH = 0.0;
+    this->totalEnergyS = 0.0;
+    this->totalFreeEnergy = 0.0;
+}
+
+template <typename T>
+void DiagoPexsi<T>::set_k_weight(const int ik, const double weight)
+{
+    if (ik >= static_cast<int>(this->k_weights_.size()))
     {
-        nspin = 1;
+        this->k_weights_.resize(ik + 1, 1.0);
     }
-    for (int i = 0; i < nspin; i++)
+    this->k_weights_[ik] = weight;
+}
+
+template <typename T>
+bool DiagoPexsi<T>::finish_k_loop(const double target_nelec)
+{
+    return true;
+}
+
+template <>
+bool DiagoPexsi<std::complex<double>>::finish_k_loop(const double target_nelec)
+{
+    const double residual = target_nelec - this->num_electron_sum_;
+    if (std::abs(residual) <= pexsi::PEXSI_Solver::pexsi_elec_thr)
+    {
+        return true;
+    }
+    if (residual > 0.0)
     {
-        delete[] this->DM[i];
-        delete[] this->EDM[i];
+        this->has_mu_lower_ = true;
+        this->mu_lower_ = mu_buffer[0];
+    }
+    else
+    {
+        this->has_mu_upper_ = true;
+        this->mu_upper_ = mu_buffer[0];
     }
 
+    double delta_mu = 0.0;
+    if (std::abs(this->num_electron_derivative_sum_) <= DBL_MIN)
+    {
+        if (this->has_mu_lower_ && this->has_mu_upper_)
+        {
+            mu_buffer[0] = 0.5 * (this->mu_lower_ + this->mu_upper_);
+            return false;
+        }
+        const double fallback_step = std::max(pexsi::PEXSI_Solver::pexsi_mu_guard, 0.5);
+        delta_mu = residual > 0.0 ? fallback_step : -fallback_step;
+    }
+    else
+    {
+        delta_mu = residual / this->num_electron_derivative_sum_;
+    }
+    if (pexsi::PEXSI_Solver::pexsi_mu_guard > 0.0 && std::abs(this->num_electron_derivative_sum_) > DBL_MIN)
+    {
+        delta_mu = std::max(-pexsi::PEXSI_Solver::pexsi_mu_guard,
+                            std::min(pexsi::PEXSI_Solver::pexsi_mu_guard, delta_mu));
+    }
+    if (mu_buffer.empty())
+    {
+        mu_buffer.push_back(pexsi::PEXSI_Solver::pexsi_mu);
+    }
+    mu_buffer[0] += delta_mu;
+    return false;
 }
 
 template <>
@@ -69,6 +172,12 @@ void DiagoPexsi<double>::diag(hamilt::Hamilt<double>* phm_in, psi::Psi<double>&
     phm_in->matrix(h_mat, s_mat);
     std::vector<double> eigen(PARAM.globalv.nlocal, 0.0);
     int ik = psi.get_current_k();
+    if (ik < 0 || ik >= static_cast<int>(this->DM.size()))
+    {
+        ModuleBase::WARNING_QUIT("DiagoPEXSI",
+                                 "PEXSI real path only has density buffers for Gamma/spin channels; multi-k requires "
+                                 "the complex expert-routine path");
+    }
     this->ps->prepare(this->ParaV->blacs_ctxt,
                       this->ParaV->nb,
                       this->ParaV->nrow,
@@ -90,11 +199,71 @@ void DiagoPexsi<std::complex<double>>::diag(hamilt::Hamilt<std::complex<double>>
                                             double* eigenvalue_in)
 {
     ModuleBase::TITLE("DiagoPEXSI", "diag");
-    ModuleBase::WARNING_QUIT("DiagoPEXSI", "PEXSI is not completed for multi-k case");
+    matcd h_mat, s_mat;
+    phm_in->matrix(h_mat, s_mat);
+    const int ik = psi.get_current_k();
+    if (ik < 0)
+    {
+        ModuleBase::WARNING_QUIT("DiagoPEXSI", "invalid k-point index for complex PEXSI path");
+    }
+    if (ik >= static_cast<int>(this->DM.size()))
+    {
+        this->resize_density_buffers(ik + 1);
+    }
+    if (mu_buffer.empty())
+    {
+        mu_buffer.push_back(pexsi::PEXSI_Solver::pexsi_mu);
+    }
+
+    MPI_Group grid_group;
+    MPI_Group world_group;
+    int grid_np = 0;
+    MPI_Comm_size(DIAG_WORLD, &grid_np);
+    MPI_Comm_group(DIAG_WORLD, &world_group);
+    int grid_proc_range[3] = {0, (GlobalV::NPROC / grid_np) * grid_np - 1, GlobalV::NPROC / grid_np};
+    MPI_Group_range_incl(world_group, 1, &grid_proc_range, &grid_group);
+
+    double num_electron = 0.0;
+    double num_electron_derivative = 0.0;
+    double total_energy_h = 0.0;
+    double total_energy_s = 0.0;
+    double total_free_energy = 0.0;
+    pexsi::simplePEXSIComplex(DIAG_WORLD,
+                              DIAG_WORLD,
+                              grid_group,
+                              this->ParaV->blacs_ctxt,
+                              PARAM.globalv.nlocal,
+                              this->ParaV->nb,
+                              this->ParaV->nrow,
+                              this->ParaV->ncol,
+                              'c',
+                              h_mat.p,
+                              s_mat.p,
+                              PARAM.inp.nelec,
+                              "PEXSIOPTION",
+                              this->DM[ik],
+                              this->EDM[ik],
+                              total_energy_h,
+                              total_energy_s,
+                              total_free_energy,
+                              mu_buffer[0],
+                              mu_buffer[0],
+                              &num_electron,
+                              &num_electron_derivative);
+
+    const double k_weight = ik < static_cast<int>(this->k_weights_.size()) ? this->k_weights_[ik] : 1.0;
+    this->num_electron_sum_ += k_weight * num_electron;
+    this->num_electron_derivative_sum_ += k_weight * num_electron_derivative;
+    this->totalEnergyH += k_weight * total_energy_h;
+    this->totalEnergyS += k_weight * total_energy_s;
+    this->totalFreeEnergy += k_weight * total_free_energy;
+
+    MPI_Group_free(&grid_group);
+    MPI_Group_free(&world_group);
 }
 
 template class DiagoPexsi<double>;
 template class DiagoPexsi<std::complex<double> >;
 
 } // namespace hsolver
-#endif
+#endif

source/source_hsolver/diago_pexsi.h

@@ -6,7 +6,7 @@
 #include "source_base/macros.h"   // GetRealType
 #include "source_hamilt/hamilt.h"
 #include "source_basis/module_ao/parallel_orbitals.h"
-#include "module_pexsi/pexsi_solver.h"
+#include "module_pexsi/pexsi_solver_interface.h"
 
 namespace hsolver
 {
@@ -19,16 +19,33 @@ class DiagoPexsi
     static std::vector<double> mu_buffer;
 
   public:
-    DiagoPexsi(const Parallel_Orbitals* ParaV_in);
+    explicit DiagoPexsi(const Parallel_Orbitals* ParaV_in,
+                        std::unique_ptr<pexsi::IPexsiSolver> solver_in = nullptr);
     void diag(hamilt::Hamilt<T>* phm_in, psi::Psi<T>& psi, Real* eigenvalue_in);
     const Parallel_Orbitals* ParaV = nullptr;
     std::vector<T*> DM;
     std::vector<T*> EDM;
-    double totalEnergyH;
-    double totalEnergyS;
-    double totalFreeEnergy;
-    std::unique_ptr<pexsi::PEXSI_Solver> ps;
+    double totalEnergyH = 0.0;
+    double totalEnergyS = 0.0;
+    double totalFreeEnergy = 0.0;
+    std::unique_ptr<pexsi::IPexsiSolver> ps;
+    void begin_mu_search();
+    void begin_k_loop();
+    void set_k_weight(const int ik, const double weight);
+    bool finish_k_loop(const double target_nelec);
     ~DiagoPexsi();
+
+  private:
+    std::vector<std::vector<T>> dm_buffer_;
+    std::vector<std::vector<T>> edm_buffer_;
+    std::vector<double> k_weights_;
+    double num_electron_sum_ = 0.0;
+    double num_electron_derivative_sum_ = 0.0;
+    bool has_mu_lower_ = false;
+    bool has_mu_upper_ = false;
+    double mu_lower_ = 0.0;
+    double mu_upper_ = 0.0;
+    void resize_density_buffers(const int count);
 };
 } // namespace hsolver