Fix: wrong band output when kpar > 1 (#5847)

* fix wrong band results when kpar>1

* fix compile

* fix: bug for init_wfc file when nspin = 2

* add new integral test

* format read_wfc_to_rho_test

* [pre-commit.ci lite] apply automatic fixes

---------

Co-authored-by: pre-commit-ci-lite[bot] <117423508+pre-commit-ci-lite[bot]@users.noreply.github.com>

Author: Qianruipku

source/module_cell/klist.cpp

@@ -13,34 +13,33 @@
 #include "module_cell/module_paw/paw_cell.h"
 #endif
 
-K_Vectors::K_Vectors()
+void K_Vectors::cal_ik_global()
 {
-
-    nspin = 0; // default spin.
-    kc_done = false;
-    kd_done = false;
-    nkstot_full = 0;
-    nks = 0;
-    nkstot = 0;
-    k_nkstot = 0; // LiuXh add 20180619
-}
-
-K_Vectors::~K_Vectors()
-{
-}
-
-int K_Vectors::get_ik_global(const int& ik, const int& nkstot)
-{
-    int nkp = nkstot / PARAM.inp.kpar;
-    int rem = nkstot % PARAM.inp.kpar;
-    if (GlobalV::MY_POOL < rem)
+    const int my_pool = this->para_k.my_pool;
+    this->ik2iktot.resize(this->nks);
+#ifdef __MPI
+    if(this->nspin == 2)
     {
-        return GlobalV::MY_POOL * nkp + GlobalV::MY_POOL + ik;
+        for (int ik = 0; ik < this->nks / 2; ++ik)
+        {
+            this->ik2iktot[ik] = this->para_k.startk_pool[my_pool] + ik;
+            this->ik2iktot[ik + this->nks / 2] = this->nkstot / 2 + this->para_k.startk_pool[my_pool] + ik;
+        }
     }
     else
     {
-        return GlobalV::MY_POOL * nkp + rem + ik;
+        for (int ik = 0; ik < this->nks; ++ik)
+        {
+            this->ik2iktot[ik] = this->para_k.startk_pool[my_pool] + ik;
+        }
     }
+#else
+    for (int ik = 0; ik < this->nks; ++ik)
+    {
+        this->ik2iktot[ik] = ik;
+    }
+#endif
+
 }
 
 void K_Vectors::set(const UnitCell& ucell,
@@ -162,6 +161,9 @@ void K_Vectors::set(const UnitCell& ucell,
     // set the k vectors for the up and down spin
     this->set_kup_and_kdw();
 
+    // get ik2iktot
+    this->cal_ik_global();
+
     this->print_klists(ofs);
 
     // std::cout << " NUMBER OF K-POINTS   : " << nkstot << std::endl;

source/module_cell/klist.h

@@ -26,8 +26,8 @@ class K_Vectors
     /// dim: [iks_ibz][(isym, kvec_d)]
     std::vector<std::map<int, ModuleBase::Vector3<double>>> kstars;
 
-    K_Vectors();
-    ~K_Vectors();
+    K_Vectors(){};
+    ~K_Vectors(){};
     K_Vectors& operator=(const K_Vectors&) = default;
     K_Vectors& operator=(K_Vectors&& rhs) = default;
 
@@ -106,23 +106,6 @@ class K_Vectors
      */
     void set_after_vc(const int& nspin, const ModuleBase::Matrix3& reciprocal_vec, const ModuleBase::Matrix3& latvec);
 
-    /**
-     * @brief Gets the global index of a k-point.
-     *
-     * This function gets the global index of a k-point based on its local index and the process pool ID.
-     * The global index is used when the k-points are distributed among multiple process pools.
-     *
-     * @param nkstot The total number of k-points.
-     * @param ik The local index of the k-point.
-     *
-     * @return int Returns the global index of the k-point.
-     *
-     * @note The function calculates the global index by dividing the total number of k-points (nkstot) by the number of
-     * process pools (KPAR), and adding the remainder if the process pool ID (MY_POOL) is less than the remainder.
-     * @note The function is declared as inline for efficiency.
-     */
-    static int get_ik_global(const int& ik, const int& nkstot);
-
     int get_nks() const
     {
         return this->nks;
@@ -157,19 +140,20 @@ class K_Vectors
     {
         this->nkstot_full = value;
     }
+    std::vector<int> ik2iktot; ///<[nks] map ik to the global index of k points
 
-private:
-    int nks;         // number of symmetry-reduced k points in this pool(processor, up+dw)
-    int nkstot;      /// number of symmetry-reduced k points in full k mesh
-    int nkstot_full; /// number of k points before symmetry reduction in full k mesh
+  private:
+    int nks = 0;         ///< number of symmetry-reduced k points in this pool(processor, up+dw)
+    int nkstot = 0;      ///< number of symmetry-reduced k points in full k mesh
+    int nkstot_full = 0; ///< number of k points before symmetry reduction in full k mesh
 
-    int nspin;
-    bool kc_done;
-    bool kd_done;
-    double koffset[3]={0.0};   // used only in automatic k-points.
-    std::string k_kword; // LiuXh add 20180619
-    int k_nkstot;        // LiuXh add 20180619
-    bool is_mp = false;  // Monkhorst-Pack
+    int nspin = 0;
+    bool kc_done = false;
+    bool kd_done = false;
+    double koffset[3] = {0.0}; // used only in automatic k-points.
+    std::string k_kword;       // LiuXh add 20180619
+    int k_nkstot = 0;          // LiuXh add 20180619
+    bool is_mp = false;        // Monkhorst-Pack
 
     /**
      * @brief Resize the k-point related vectors according to the new k-point number.
@@ -288,8 +272,8 @@ class K_Vectors
      * be recalculated.
      */
     void update_use_ibz(const int& nkstot_ibz,
-        const std::vector<ModuleBase::Vector3<double>>& kvec_d_ibz,
-        const std::vector<double>& wk_ibz);
+                        const std::vector<ModuleBase::Vector3<double>>& kvec_d_ibz,
+                        const std::vector<double>& wk_ibz);
 
     /**
      * @brief Sets both the direct and Cartesian k-vectors.
@@ -394,5 +378,11 @@ class K_Vectors
      * @note The function uses the FmtCore::format function to format the output.
      */
     void print_klists(std::ofstream& fn);
+
+    /**
+     * @brief Gets the global index of a k-point.
+     * @return this->ik2iktot[ik]
+     */
+    void cal_ik_global();
 };
 #endif // KVECT_H

source/module_cell/parallel_kpoints.h

@@ -68,14 +68,13 @@ class Parallel_Kpoints
         return *std::max_element(nks_pool.begin(), nks_pool.end());
     }
 
-  private:
-
+  public:
     int kpar = 0;         // number of pools
     int my_pool = 0;      // the pool index of the present processor
     int rank_in_pool = 0; // the rank in the present pool
     int nproc = 1;        // number of processors
     int nspin = 1;        // number of spins
-
+  private:
     std::vector<int> startpro_pool; // the first processor in each pool
 #ifdef __MPI
     void get_nks_pool(const int& nkstot);

source/module_elecstate/elecstate_print.cpp

@@ -2,6 +2,7 @@
 #include "elecstate_getters.h"
 #include "module_base/formatter.h"
 #include "module_base/global_variable.h"
+#include "module_base/parallel_common.h"
 #include "module_elecstate/potentials/H_Hartree_pw.h"
 #include "module_elecstate/potentials/efield.h"
 #include "module_elecstate/potentials/gatefield.h"
@@ -152,7 +153,9 @@ void print_scf_iterinfo(const std::string& ks_solver,
 void ElecState::print_eigenvalue(std::ofstream& ofs)
 {
     bool wrong = false;
-    for (int ik = 0; ik < this->klist->get_nks(); ++ik)
+    const int nks = this->klist->get_nks();
+    const int nkstot = this->klist->get_nkstot();
+    for (int ik = 0; ik < nks; ++ik)
     {
         for (int ib = 0; ib < this->ekb.nc; ++ib)
         {
@@ -164,76 +167,87 @@ void ElecState::print_eigenvalue(std::ofstream& ofs)
             }
         }
     }
+#ifdef __MPI
+    MPI_Allreduce(MPI_IN_PLACE, &wrong, 1, MPI_C_BOOL, MPI_LOR, MPI_COMM_WORLD);
+#endif
     if (wrong)
     {
         ModuleBase::WARNING_QUIT("print_eigenvalue", "Eigenvalues are too large!");
     }
+    std::stringstream ss;
+    if(PARAM.inp.out_alllog)
+    {
+        ss << PARAM.globalv.global_out_dir << "running_" << PARAM.inp.calculation << "_" << GlobalV::MY_RANK + 1 << ".log";
+    }
+    else
+    {
+        ss << PARAM.globalv.global_out_dir << "running_" << PARAM.inp.calculation << ".log";
+    }
+    std::string filename = ss.str();
+    std::vector<int> ngk_tot = this->klist->ngk;
 
-    if (GlobalV::MY_RANK != 0)
+#ifdef __MPI
+    if(!PARAM.inp.out_alllog)
     {
-        return;
+        Parallel_Common::bcast_string(filename);
     }
+    MPI_Allreduce(MPI_IN_PLACE, ngk_tot.data(), nks, MPI_INT, MPI_SUM, POOL_WORLD);
+#endif
 
     ModuleBase::TITLE("ESolver_KS_PW", "print_eigenvalue");
 
     ofs << "\n STATE ENERGY(eV) AND OCCUPATIONS ";
-    for (int ik = 0; ik < this->klist->get_nks(); ik++)
+    const int nk_fac = PARAM.inp.nspin == 2 ? 2 : 1;
+    const int nks_np = nks / nk_fac;
+    const int nkstot_np = nkstot / nk_fac;
+    ofs << "   NSPIN == " << PARAM.inp.nspin << std::endl;
+    for (int is = 0; is < nk_fac; ++is)
     {
-        ofs << std::setprecision(5);
-        ofs << std::setiosflags(std::ios::showpoint);
-        if (ik == 0)
+        if (is == 0 && nk_fac == 2)
         {
-            ofs << "   NSPIN == " << PARAM.inp.nspin << std::endl;
-            if (PARAM.inp.nspin == 2)
-            {
-                ofs << "SPIN UP : " << std::endl;
-            }
+            ofs << "SPIN UP : " << std::endl;
         }
-        else if (ik == this->klist->get_nks() / 2)
+        else if (is == 1 && nk_fac == 2)
         {
-            if (PARAM.inp.nspin == 2)
-            {
-                ofs << "SPIN DOWN : " << std::endl;
-            }
+            ofs << "SPIN DOWN : " << std::endl;
         }
 
-        if (PARAM.inp.nspin == 2)
+        for (int ip = 0; ip < GlobalV::KPAR; ++ip)
         {
-            if (this->klist->isk[ik] == 0)
-            {
-                ofs << " " << ik + 1 << "/" << this->klist->get_nks() / 2
-                    << " kpoint (Cartesian) = " << this->klist->kvec_c[ik].x << " " << this->klist->kvec_c[ik].y << " "
-                    << this->klist->kvec_c[ik].z << " (" << this->klist->ngk[ik] << " pws)" << std::endl;
-
-                ofs << std::setprecision(6);
-            }
-            if (this->klist->isk[ik] == 1)
+#ifdef __MPI
+            MPI_Barrier(MPI_COMM_WORLD);
+#endif
+            bool ip_flag = PARAM.inp.out_alllog || (GlobalV::RANK_IN_POOL == 0 && GlobalV::MY_STOGROUP == 0);
+            if (GlobalV::MY_POOL == ip && ip_flag)
             {
-                ofs << " " << ik + 1 - this->klist->get_nks() / 2 << "/" << this->klist->get_nks() / 2
-                    << " kpoint (Cartesian) = " << this->klist->kvec_c[ik].x << " " << this->klist->kvec_c[ik].y << " "
-                    << this->klist->kvec_c[ik].z << " (" << this->klist->ngk[ik] << " pws)" << std::endl;
+                const int start_ik = nks_np * is;
+                const int end_ik = nks_np * (is + 1);
+                for (int ik = start_ik; ik < end_ik; ++ik)
+                {
+                    std::ofstream ofs_eig(filename.c_str(), std::ios::app);
+                    ofs_eig << std::setprecision(5);
+                    ofs_eig << std::setiosflags(std::ios::showpoint);
+                    ofs_eig << " " << this->klist->ik2iktot[ik] + 1 - is * nkstot_np << "/" << nkstot_np
+                            << " kpoint (Cartesian) = " << this->klist->kvec_c[ik].x << " " << this->klist->kvec_c[ik].y
+                            << " " << this->klist->kvec_c[ik].z << " (" << ngk_tot[ik] << " pws)" << std::endl;
 
-                ofs << std::setprecision(6);
+                    ofs_eig << std::setprecision(6);
+                    ofs_eig << std::setiosflags(std::ios::showpoint);
+                    for (int ib = 0; ib < this->ekb.nc; ib++)
+                    {
+                        ofs_eig << std::setw(8) << ib + 1 << std::setw(15) << this->ekb(ik, ib) * ModuleBase::Ry_to_eV
+                                << std::setw(15) << this->wg(ik, ib) << std::endl;
+                    }
+                    ofs_eig << std::endl;
+                    ofs_eig.close();
+                }
             }
-        } // Pengfei Li  added  14-9-9
-        else
-        {
-            ofs << " " << ik + 1 << "/" << this->klist->get_nks()
-                << " kpoint (Cartesian) = " << this->klist->kvec_c[ik].x << " " << this->klist->kvec_c[ik].y << " "
-                << this->klist->kvec_c[ik].z << " (" << this->klist->ngk[ik] << " pws)" << std::endl;
-
-            ofs << std::setprecision(6);
         }
-
-        ofs << std::setprecision(6);
-        ofs << std::setiosflags(std::ios::showpoint);
-        for (int ib = 0; ib < this->ekb.nc; ib++)
-        {
-            ofs << std::setw(8) << ib + 1 << std::setw(15) << this->ekb(ik, ib) * ModuleBase::Ry_to_eV << std::setw(15)
-                << this->wg(ik, ib) << std::endl;
-        }
-        ofs << std::endl;
-    } // end ik
+#ifdef __MPI
+            MPI_Barrier(MPI_COMM_WORLD);
+#endif
+        ofs.seekp(0, std::ios::end);
+    }
     return;
 }
 

source/module_elecstate/module_charge/charge_init.cpp

@@ -248,7 +248,7 @@ void Charge::init_rho(elecstate::efermi& eferm_iout,
         const K_Vectors* kv = reinterpret_cast<const K_Vectors*>(klist);
         const int nkstot = kv->get_nkstot();
         const std::vector<int>& isk = kv->isk;
-        ModuleIO::read_wfc_to_rho(pw_wfc, symm, nkstot, isk, *this);
+        ModuleIO::read_wfc_to_rho(pw_wfc, symm, kv->ik2iktot.data(), nkstot, isk, *this);
     }
 }
 

source/module_elecstate/module_dm/test/test_cal_dm_R.cpp

@@ -6,13 +6,6 @@
 #include "module_hamilt_lcao/module_hcontainer/hcontainer.h"
 #include "module_cell/klist.h"
 
-K_Vectors::K_Vectors()
-{
-}
-
-K_Vectors::~K_Vectors()
-{
-}
 /************************************************
  *  unit test of DensityMatrix constructor
  ***********************************************/

source/module_elecstate/module_dm/test/test_dm_R_init.cpp

@@ -7,13 +7,6 @@
 #include "module_hamilt_lcao/module_hcontainer/hcontainer.h"
 #include "module_cell/klist.h"
 #undef private
-K_Vectors::K_Vectors()
-{
-}
-
-K_Vectors::~K_Vectors()
-{
-}
 /************************************************
  *  unit test of DensityMatrix constructor
  ***********************************************/

source/module_elecstate/module_dm/test/test_dm_constructor.cpp

@@ -5,12 +5,6 @@
 #include "module_elecstate/module_dm/density_matrix.h"
 #include "module_hamilt_lcao/module_hcontainer/hcontainer.h"
 #include "module_cell/klist.h"
-K_Vectors::K_Vectors()
-{
-}
-K_Vectors::~K_Vectors()
-{
-}
 /************************************************
  *  unit test of DensityMatrix constructor
  ***********************************************/

source/module_elecstate/module_dm/test/test_dm_io.cpp

@@ -33,15 +33,6 @@ Magnetism::~Magnetism()
 }
 
 #include "module_cell/klist.h"
-
-K_Vectors::K_Vectors()
-{
-}
-
-K_Vectors::~K_Vectors()
-{
-}
-
 #include "module_cell/module_neighbor/sltk_grid_driver.h"
 // mock find_atom() function
 void Grid_Driver::Find_atom(const UnitCell& ucell,