Fix: bug of out_band when kpar > 1 (#3963)

Qianruipku · WHUweiqingzhou · mohanchen · web-flow · commit ef53ec1f1810 · 2024-04-17T11:03:35.000+08:00
* Fix: bug of out_band when kpar &gt; 1

* fix bug in UT of Gatherkvec

* add ndim variables in parallel_kpoints

* initialize variables in parallel_kpoints

---------

Co-authored-by: wqzhou &lt;33364058+WHUweiqingzhou@users.noreply.github.com&gt;
Co-authored-by: Mohan Chen &lt;mohan.chen.chen.mohan@gmail.com&gt;
diff --git a/source/module_cell/parallel_kpoints.cpp b/source/module_cell/parallel_kpoints.cpp
@@ -25,6 +25,13 @@ void Parallel_Kpoints::kinfo(int &nkstot)
     this->get_nks_pool(nkstot);
     this->get_startk_pool(nkstot);
     this->get_whichpool(nkstot);
+    this->kpar = GlobalV::KPAR;
+    this->nkstot_np = nkstot;
+    this->nks_np = this->nks_pool[GlobalV::MY_POOL];
+#else
+    this->kpar = 1;
+    this->nkstot_np = nkstot;
+    this->nks_np = nkstot;
 #endif
     return;
 }
@@ -91,6 +98,19 @@ void Parallel_Kpoints::get_startk_pool(const int &nkstot)
     }
     return;
 }
+
+void Parallel_Kpoints::gatherkvec(const std::vector<ModuleBase::Vector3<double>>& vec_local,
+                                  std::vector<ModuleBase::Vector3<double>>& vec_global) const
+{
+    vec_global.resize(this->nkstot_np, ModuleBase::Vector3<double>(0.0, 0.0, 0.0));
+    for (int i = 0; i < this->nks_np; ++i)
+    {
+        vec_global[i + startk_pool[GlobalV::MY_POOL]] = vec_local[i] / double(GlobalV::NPROC_IN_POOL);
+    }
+
+    MPI_Allreduce(MPI_IN_PLACE, &vec_global[0], 3 * this->nkstot_np, MPI_DOUBLE, MPI_SUM, MPI_COMM_WORLD);
+    return;
+}
 #endif
 
 
diff --git a/source/module_cell/parallel_kpoints.h b/source/module_cell/parallel_kpoints.h
@@ -1,45 +1,61 @@
 #ifndef PARALLEL_KPOINTS_H
 #define PARALLEL_KPOINTS_H
 
+#include "module_base/complexarray.h"
 #include "module_base/global_function.h"
 #include "module_base/global_variable.h"
-#include "module_base/complexarray.h"
 #include "module_base/realarray.h"
+#include "module_base/vector3.h"
 
 class Parallel_Kpoints
 {
-	public:
-	
-	Parallel_Kpoints();
-	~Parallel_Kpoints();
-
-	void kinfo(int &nkstot);
-	
-	// collect value from each pool to wk.
-	void pool_collection(double &value, const double *wk, const int &ik);
-
-	// collect value from each pool to overlap.
-	void pool_collection(double *valuea, double *valueb, const ModuleBase::realArray &a, const ModuleBase::realArray &b, const int &ik);
-	void pool_collection(std::complex<double> *value, const ModuleBase::ComplexArray &w, const int &ik);
+  public:
+    Parallel_Kpoints();
+    ~Parallel_Kpoints();
+
+    void kinfo(int& nkstot);
+
+    // collect value from each pool to wk.
+    void pool_collection(double& value, const double* wk, const int& ik);
+
+    // collect value from each pool to overlap.
+    void pool_collection(double* valuea,
+                         double* valueb,
+                         const ModuleBase::realArray& a,
+                         const ModuleBase::realArray& b,
+                         const int& ik);
+    void pool_collection(std::complex<double>* value, const ModuleBase::ComplexArray& w, const int& ik);
+#ifdef __MPI
+    /**
+     * @brief gather kpoints from all processors
+     *
+     * @param vec_local kpoint vector in local processor
+     * @param vec_global kpoint vector in all processors
+     */
+    void gatherkvec(const std::vector<ModuleBase::Vector3<double>>& vec_local,
+                    std::vector<ModuleBase::Vector3<double>>& vec_global) const;
+#endif
 
-	// information about pool
-	int *nproc_pool;
-	int *startpro_pool;
+    // information about pool, dim: GlobalV::KPAR
+    int* nproc_pool = nullptr;
+    int* startpro_pool = nullptr;
 
-	// inforamation about kpoints //qianrui add comment
-	int* nks_pool; //number of k-points in each pool
-	int* startk_pool; //the first k-point in each pool
-	int* whichpool; //whichpool[k] : the pool which k belongs to
+    // inforamation about kpoints, dim: GlobalV::KPAR
+    int* nks_pool = nullptr;    // number of k-points in each pool
+    int* startk_pool = nullptr; // the first k-point in each pool
+    int kpar = 0;               // number of pools
 
-	private:
+    // information about which pool each k-point belongs to,
+    int* whichpool = nullptr; // whichpool[k] : the pool which k belongs to, dim: nkstot_np
+    int nkstot_np = 0;        // number of k-points without spin, kv.nkstot = nkstot_np * nspin(1 or 2)
+    int nks_np = 0;           // number of k-points without spin in the present pool
 
+  private:
 #ifdef __MPI
-	void get_nks_pool(const int &nkstot);
-	void get_startk_pool(const int &nkstot);
-	void get_whichpool(const int &nkstot);
+    void get_nks_pool(const int& nkstot);
+    void get_startk_pool(const int& nkstot);
+    void get_whichpool(const int& nkstot);
 #endif
-
-
 };
 
 #endif
diff --git a/source/module_cell/test/parallel_kpoints_test.cpp b/source/module_cell/test/parallel_kpoints_test.cpp
@@ -23,6 +23,8 @@
  *      to call another three functions: get_nks_pool(),
  *      get_startk_pool(), get_whichpool(), which divide all kpoints
  *      into KPAR groups.
+ *   iii.Parallel_Kpoints::gatherkvec() is an interface to gather kpoints
+ *      vectors from all processors.
  * The default number of processes is set to 4 in parallel_kpoints_test.sh.
  * One may modify it to do more tests, or adapt this unittest to local
  * environment.
@@ -120,6 +122,88 @@ class ParaKpoints : public ::testing::TestWithParam<ParaPrepare>
 {
 };
 
+TEST(Parallel_KpointsTest, GatherkvecTest) {
+    // Initialize Parallel_Kpoints object
+    Parallel_Kpoints parallel_kpoints;
+    
+    // Initialize local and global vectors
+    std::vector<ModuleBase::Vector3<double>> vec_local;
+    std::vector<ModuleBase::Vector3<double>> vec_global;
+
+    // Populate vec_local with some data
+	int npool = 1;
+	if(GlobalV::NPROC > 2)
+	{
+		npool = 3;
+	}
+	else if(GlobalV::NPROC == 2)
+	{
+		npool = 2;
+	}
+	GlobalV::KPAR = npool;
+
+	if(GlobalV::MY_RANK == 0)
+	{
+    	vec_local.push_back(ModuleBase::Vector3<double>(1.0, 1.0, 1.0));
+		GlobalV::NPROC_IN_POOL = 1;
+		GlobalV::MY_POOL = 0;
+		parallel_kpoints.nks_np = 1;
+	}
+	else if(GlobalV::MY_RANK == 1)
+	{
+		vec_local.push_back(ModuleBase::Vector3<double>(2.0, 2.0, 2.0));
+		vec_local.push_back(ModuleBase::Vector3<double>(3.0, 4.0, 5.0));
+		GlobalV::NPROC_IN_POOL = 1;
+		GlobalV::MY_POOL = 1;
+		parallel_kpoints.nks_np = 2;
+	}
+	else
+	{
+		vec_local.push_back(ModuleBase::Vector3<double>(3.0, 3.0, 3.0));
+		GlobalV::NPROC_IN_POOL = GlobalV::NPROC - 2;
+		GlobalV::MY_POOL = 2;
+		parallel_kpoints.nks_np = 1;
+	}
+
+	parallel_kpoints.startk_pool = new int[npool];
+	parallel_kpoints.nkstot_np = 1;
+	parallel_kpoints.startk_pool[0] = 0;
+	if (npool >= 2)
+	{
+		parallel_kpoints.nkstot_np += 2;
+		parallel_kpoints.startk_pool[1] = 1;
+	}
+	if(npool >= 3)
+	{
+		parallel_kpoints.nkstot_np += 1;
+		parallel_kpoints.startk_pool[2] = 3;
+	}
+	
+    // Call gatherkvec method
+    parallel_kpoints.gatherkvec(vec_local, vec_global);
+
+    // Check the values of vec_global
+    EXPECT_EQ(vec_global[0].x, 1.0);
+	EXPECT_EQ(vec_global[0].y, 1.0);
+	EXPECT_EQ(vec_global[0].z, 1.0);
+
+	if(npool >= 2)
+	{
+		EXPECT_EQ(vec_global[1].x, 2.0);
+		EXPECT_EQ(vec_global[1].y, 2.0);
+		EXPECT_EQ(vec_global[1].z, 2.0);
+		EXPECT_EQ(vec_global[2].x, 3.0);
+		EXPECT_EQ(vec_global[2].y, 4.0);
+		EXPECT_EQ(vec_global[2].z, 5.0);
+	}
+	if(npool >= 3)
+	{
+		EXPECT_EQ(vec_global[3].x, 3.0);
+		EXPECT_EQ(vec_global[3].y, 3.0);
+		EXPECT_EQ(vec_global[3].z, 3.0);
+	}
+}
+
 TEST_P(ParaKpoints,DividePools)
 {
 	ParaPrepare pp = GetParam();
diff --git a/source/module_esolver/esolver_ks_lcao.cpp b/source/module_esolver/esolver_ks_lcao.cpp
@@ -393,24 +393,13 @@ void ESolver_KS_LCAO<TK, TR>::post_process(void)
 
     if (INPUT.out_band[0]) // pengfei 2014-10-13
     {
-        int nks = 0;
-        if (nspin0 == 1)
-        {
-            nks = this->kv.nkstot;
-        }
-        else if (nspin0 == 2)
-        {
-            nks = this->kv.nkstot / 2;
-        }
-
         for (int is = 0; is < nspin0; is++)
         {
             std::stringstream ss2;
             ss2 << GlobalV::global_out_dir << "BANDS_" << is + 1 << ".dat";
             GlobalV::ofs_running << "\n Output bands in file: " << ss2.str() << std::endl;
             ModuleIO::nscf_band(is, 
                                 ss2.str(), 
-                                nks, 
                                 GlobalV::NBANDS, 
                                 0.0, 
                                 INPUT.out_band[1],
diff --git a/source/module_esolver/esolver_ks_pw.cpp b/source/module_esolver/esolver_ks_pw.cpp
@@ -1257,23 +1257,13 @@ void ESolver_KS_PW<T, Device>::post_process(void)
 
     if (INPUT.out_band[0]) // pengfei 2014-10-13
     {
-        int nks = 0;
-        if (nspin0 == 1)
-        {
-            nks = this->kv.nkstot;
-        }
-        else if (nspin0 == 2)
-        {
-            nks = this->kv.nkstot / 2;
-        }
         for (int is = 0; is < nspin0; is++)
         {
             std::stringstream ss2;
             ss2 << GlobalV::global_out_dir << "BANDS_" << is + 1 << ".dat";
             GlobalV::ofs_running << "\n Output bands in file: " << ss2.str() << std::endl;
             ModuleIO::nscf_band(is,
                                 ss2.str(),
-                                nks,
                                 GlobalV::NBANDS,
                                 0.0,
                                 INPUT.out_band[1],
diff --git a/source/module_io/nscf_band.cpp b/source/module_io/nscf_band.cpp
@@ -8,7 +8,6 @@
 void ModuleIO::nscf_band(
 	const int &is,
 	const std::string &out_band_dir, 
-	const int &nks, 
 	const int &nband,
 	const double &fermie,
 	const int &precision,
@@ -18,6 +17,10 @@ void ModuleIO::nscf_band(
 {
 	ModuleBase::TITLE("ModuleIO","nscf_band");
 	ModuleBase::timer::tick("ModuleIO", "nscf_band");
+	// number of k points without spin; nspin = 1,2, nkstot = nkstot_np * nspin; 
+	//                                  nspin = 4, nkstot = nkstot_np
+	const int nkstot_np = Pkpoints->nkstot_np;
+	const int nks_np = Pkpoints->nks_np;
 
 #ifdef __MPI
 	if(GlobalV::MY_RANK==0)
@@ -26,15 +29,16 @@ void ModuleIO::nscf_band(
 		ofs.close();
 	}
 	MPI_Barrier(MPI_COMM_WORLD);
-
 	std::vector<double> klength;
-	klength.resize(nks);
+	klength.resize(nkstot_np);
 	klength[0] = 0.0;
-	for(int ik=0; ik<nks; ik++)
+	std::vector<ModuleBase::Vector3<double>> kvec_c_global;
+	Pkpoints->gatherkvec(kv.kvec_c, kvec_c_global);
+	for(int ik=0; ik<nkstot_np; ik++)
 	{
 		if (ik>0)
 		{
-			auto delta=kv.kvec_c[ik]-kv.kvec_c[ik-1];
+			auto delta=kvec_c_global[ik]-kvec_c_global[ik-1];
 			klength[ik] = klength[ik-1];
 			klength[ik] += (kv.kl_segids[ik] == kv.kl_segids[ik-1]) ? delta.norm() : 0.0;
 		}
@@ -44,23 +48,21 @@ void ModuleIO::nscf_band(
 			/* then get the local kpoint index, which starts definitly from 0 */
 			const int ik_now = ik - Pkpoints->startk_pool[GlobalV::MY_POOL];
 			/* if present kpoint corresponds the spin of the present one */
-			if( kv.isk[ik_now+is*nks] == is )
-			{ 
-				if ( GlobalV::RANK_IN_POOL == 0)
+			assert( kv.isk[ik_now+is*nks_np] == is );
+			if ( GlobalV::RANK_IN_POOL == 0)
+			{
+				formatter::PhysicalFmt physfmt; // create a physical formatter temporarily
+				std::ofstream ofs(out_band_dir.c_str(), std::ios::app);
+				physfmt.adjust_formatter_flexible(4, 0, false); // for integer
+				ofs << physfmt.get_p_formatter()->format(ik+1);
+				physfmt.adjust_formatter_flexible(precision, 4.0/double(precision), false); // for decimal
+				ofs << physfmt.get_p_formatter()->format(klength[ik]);
+				for(int ib = 0; ib < nband; ib++)
 				{
-					formatter::PhysicalFmt physfmt; // create a physical formatter temporarily
-					std::ofstream ofs(out_band_dir.c_str(), std::ios::app);
-					physfmt.adjust_formatter_flexible(4, 0, false); // for integer
-					ofs << physfmt.get_p_formatter()->format(ik+1);
-					physfmt.adjust_formatter_flexible(precision, 4.0/double(precision), false); // for decimal
-					ofs << physfmt.get_p_formatter()->format(klength[ik]);
-					for(int ib = 0; ib < nband; ib++)
-					{
-						ofs << physfmt.get_p_formatter()->format((ekb(ik_now+is*nks, ib)-fermie) * ModuleBase::Ry_to_eV);
-					}
-					ofs << std::endl;
-					ofs.close();	
+					ofs << physfmt.get_p_formatter()->format((ekb(ik_now+is*nks_np, ib)-fermie) * ModuleBase::Ry_to_eV);
 				}
+				ofs << std::endl;
+				ofs.close();	
 			}
 		}
 		MPI_Barrier(MPI_COMM_WORLD);
@@ -73,7 +75,7 @@ void ModuleIO::nscf_band(
 		if(GlobalV::MY_POOL == ip && GlobalV::RANK_IN_POOL == 0)
 		{
 			std::ofstream ofs(out_band_dir.c_str(),ios::app);
-			for(int ik=0;ik<nks;ik++)
+			for(int ik=0;ik<nkstot_np;ik++)
 			{
 				ofs<<std::setw(12)<<ik;
 				for(int ib = 0; ib < nband; ib++)
@@ -92,10 +94,10 @@ void ModuleIO::nscf_band(
 //	std::cout<<out_band_dir<<std::endl;
 	formatter::PhysicalFmt physfmt; // create a physical formatter temporarily
 	std::vector<double> klength;
-	klength.resize(nks);
+	klength.resize(nkstot_np);
 	klength[0] = 0.0;
 	std::ofstream ofs(out_band_dir.c_str());
-	for(int ik=0;ik<nks;ik++)
+	for(int ik=0;ik<nkstot_np;ik++)
 	{
 		if (ik>0)
 		{
diff --git a/source/module_io/nscf_band.h b/source/module_io/nscf_band.h
diff --git a/source/module_io/test_serial/nscf_band_test.cpp b/source/module_io/test_serial/nscf_band_test.cpp
diff --git a/tests/integrate/107_PW_OB_outputbands/INPUT b/tests/integrate/107_PW_OB_outputbands/INPUT
