fix write_istate_info test

mohanchen · mohanchen · commit e78b052298d2 · 2025-06-02T13:01:21.000+08:00
diff --git a/source/module_cell/klist.cpp b/source/module_cell/klist.cpp
@@ -178,12 +178,6 @@ void K_Vectors::renew(const int& kpoint_number)
     isk.resize(kpoint_number);
     ngk.resize(kpoint_number);
 
-    /*ModuleBase::Memory::record("KV::kvec_c",sizeof(double) * kpoint_number*3);
-    ModuleBase::Memory::record("KV::kvec_d",sizeof(double) * kpoint_number*3);
-    ModuleBase::Memory::record("KV::wk",sizeof(double) * kpoint_number*3);
-    ModuleBase::Memory::record("KV::isk",sizeof(int) * kpoint_number*3);
-    ModuleBase::Memory::record("KV::ngk",sizeof(int) * kpoint_number*3);*/
-
     return;
 }
 
diff --git a/source/module_cell/parallel_kpoints.cpp b/source/module_cell/parallel_kpoints.cpp
@@ -48,17 +48,12 @@ void Parallel_Kpoints::get_whichpool(const int& nkstot)
 {
     this->whichpool.resize(nkstot, 0);
 
-    // std::cout << " calculate : whichpool" << std::endl;
-    // std::cout << " nkstot is " << nkstot << std::endl;
-
-
     for (int i = 0; i < this->kpar; i++)
     {
         for (int ik = 0; ik < this->nks_pool[i]; ik++)
         {
             const int k_now = ik + startk_pool[i];
             this->whichpool[k_now] = i;
-            // ofs_running << "\n whichpool[" << k_now <<"] = " << whichpool[k_now];
         }
     }
 
@@ -72,19 +67,13 @@ void Parallel_Kpoints::get_nks_pool(const int& nkstot)
     const int nks_ave = nkstot / this->kpar;
     const int remain = nkstot % this->kpar;
 
-    // ofs_running << "\n nkstot = " << nkstot;
-    // ofs_running << "\n this->kpar = " << this->kpar;
-    // ofs_running << "\n nks_ave = " << nks_ave;
-
     for (int i = 0; i < this->kpar; i++)
-
     {
         this->nks_pool[i] = nks_ave;
         if (i < remain)
         {
             nks_pool[i]++;
         }
-        // ofs_running << "\n nks_pool[i] = " << nks_pool[i];
     }
     return;
 }
@@ -93,14 +82,10 @@ void Parallel_Kpoints::get_startk_pool(const int& nkstot)
 {
     startk_pool.resize(this->kpar, 0);
 
-    // const int remain = nkstot%this->kpar;
-
     startk_pool[0] = 0;
     for (int i = 1; i < this->kpar; i++)
-
     {
         startk_pool[i] = startk_pool[i - 1] + nks_pool[i - 1];
-        // ofs_running << "\n startk_pool[i] = " << startk_pool[i];
     }
     return;
 }
@@ -120,7 +105,6 @@ void Parallel_Kpoints::set_startpro_pool()
         {
             startpro_pool[i]++;
         }
-        // ofs_running << "\n startpro_pool[i] = " << startpro_pool[i];
     }
     return;
 }
@@ -138,8 +122,6 @@ void Parallel_Kpoints::gatherkvec(const std::vector<ModuleBase::Vector3<double>>
         {
             vec_global[i + startk_pool[this->my_pool]] = vec_local[i];
         }
-        // vec_global[i + startk_pool[MY_POOL]] = vec_local[i] / double(NPROC_IN_POOL);
-
     }
 
     MPI_Allreduce(MPI_IN_PLACE, &vec_global[0], 3 * this->nkstot_np, MPI_DOUBLE, MPI_SUM, MPI_COMM_WORLD);
@@ -152,7 +134,6 @@ void Parallel_Kpoints::pool_collection(double& value, const double* wk, const in
 #ifdef __MPI
 
     const int ik_now = ik - this->startk_pool[this->my_pool];
-    // ofs_running << "\n\n ik=" << ik << " ik_now=" << ik_now;
 
     const int pool = this->whichpool[ik];
 
@@ -167,8 +148,6 @@ void Parallel_Kpoints::pool_collection(double& value, const double* wk, const in
             }
             else
             {
-
-                // ofs_running << " receive data.";
                 MPI_Status ierror;
                 MPI_Recv(&value, 1, MPI_DOUBLE, this->startpro_pool[pool], ik, MPI_COMM_WORLD, &ierror);
 
@@ -178,18 +157,12 @@ void Parallel_Kpoints::pool_collection(double& value, const double* wk, const in
         {
             if (this->my_pool == pool)
             {
-
-                // ofs_running << " send data.";
-
                 MPI_Send(&wk[ik_now], 1, MPI_DOUBLE, 0, ik, MPI_COMM_WORLD);
             }
         }
     }
     else
     {
-
-        // ofs_running << "\n do nothing.";
-
     }
 
     MPI_Barrier(MPI_COMM_WORLD);
@@ -238,13 +211,10 @@ void Parallel_Kpoints::pool_collection_aux(T* value, const V& w, const int& dim,
     T* p = &w.ptr[begin];
     // temprary restrict kpar=1 for NSPIN=2 case for generating_orbitals
     int pool = 0;
-    if (this->nspin != 2) {
-        pool = this->whichpool[ik];
-}
-
-
-    // ofs_running << "\n ik=" << ik;
-
+	if (this->nspin != 2) 
+	{
+		pool = this->whichpool[ik];
+	}
 
     if (this->rank_in_pool == 0)
     {
@@ -261,7 +231,6 @@ void Parallel_Kpoints::pool_collection_aux(T* value, const V& w, const int& dim,
             }
             else
             {
-                // ofs_running << " receive data.";
                 MPI_Status ierror;
                 MPI_Recv(value, dim, MPI_DOUBLE, this->startpro_pool[pool], ik * 2 + 0, MPI_COMM_WORLD, &ierror);
             }
@@ -270,14 +239,12 @@ void Parallel_Kpoints::pool_collection_aux(T* value, const V& w, const int& dim,
         {
             if (this->my_pool == pool)
             {
-                // ofs_running << " send data.";
                 MPI_Send(p, dim, MPI_DOUBLE, 0, ik * 2 + 0, MPI_COMM_WORLD);
             }
         }
     }
     else
     {
-        // ofs_running << "\n do nothing.";
     }
     MPI_Barrier(MPI_COMM_WORLD);
 
@@ -292,4 +259,4 @@ void Parallel_Kpoints::pool_collection_aux(T* value, const V& w, const int& dim,
     }
     // data transfer ends.
 #endif
-}
+}
diff --git a/source/module_io/test/write_istate_info_test.cpp b/source/module_io/test/write_istate_info_test.cpp
@@ -43,13 +43,14 @@ class IstateInfoTest : public ::testing::Test
 
 TEST_F(IstateInfoTest, OutIstateInfoS1)
 {
-    // preconditions
+    // Global variables 
     GlobalV::KPAR = 1;
     PARAM.input.nbands = 4;
     PARAM.sys.nbands_l = 4;
     PARAM.input.nspin = 1;
     PARAM.sys.global_out_dir = "./";
-    // mpi setting
+
+    // MPI setting
     Parallel_Global::init_pools(GlobalV::NPROC,
                                 GlobalV::MY_RANK,
                                 PARAM.input.bndpar,
@@ -60,90 +61,54 @@ TEST_F(IstateInfoTest, OutIstateInfoS1)
                                 GlobalV::NPROC_IN_POOL,
                                 GlobalV::RANK_IN_POOL,
                                 GlobalV::MY_POOL);
-    kv->set_nkstot(100);
+
+    const int nkstot_init = 10;
+    kv->set_nkstot(nkstot_init);
     int nkstot = kv->get_nkstot();
-    kv->para_k.kinfo(nkstot, GlobalV::KPAR, GlobalV::MY_POOL, GlobalV::RANK_IN_POOL, GlobalV::NPROC_IN_POOL, PARAM.input.nspin);
-    // std::cout<<"my_rank "<<GlobalV::MY_RANK<<" pool rank/size: "
-    //	<<GlobalV::RANK_IN_POOL<<"/"<<GlobalV::NPROC_IN_POOL<<std::endl;
-    // std::cout<<"MY_POOL "<<GlobalV::MY_POOL<<std::endl;
+    kv->para_k.kinfo(nkstot, GlobalV::KPAR, GlobalV::MY_POOL, GlobalV::RANK_IN_POOL, 
+    GlobalV::NPROC_IN_POOL, PARAM.input.nspin);
     kv->set_nks(kv->para_k.nks_pool[GlobalV::MY_POOL]);
-    // std::cout<<"nks "<<kv->get_nks()<<std::endl;
-    ekb.create(kv->get_nks(), PARAM.input.nbands);
-    wg.create(kv->get_nks(), PARAM.input.nbands);
-    ekb.fill_out(0.15);
-    wg.fill_out(0.0);
-    kv->kvec_d.resize(kv->get_nkstot());
-    int i = 0;
-    for (auto& kd: kv->kvec_d)
+
+    // The number of plane waves for each k point
+    kv->ngk.resize(nkstot);
+    kv->ik2iktot.resize(nkstot);
+    for(int i=0; i<nkstot; ++i)
     {
-        kd.set(0.01 * i, 0.01 * i, 0.01 * i);
-        ++i;
+        kv->ngk[i]=299;
+        kv->ik2iktot[i]=i;
     }
-    ModuleIO::write_istate_info(ekb, wg, *kv);
-    std::ifstream ifs;
-    ifs.open("istate.info");
-    std::string str((std::istreambuf_iterator<char>(ifs)), std::istreambuf_iterator<char>());
-    EXPECT_THAT(
-        str,
-        testing::HasSubstr("BAND               Energy(ev)               Occupation                Kpoint = 100"));
-    EXPECT_THAT(str, testing::HasSubstr("(0.99 0.99 0.99)"));
-    EXPECT_THAT(str, testing::HasSubstr("4                2.0408547                        0"));
-    ifs.close();
-    remove("istate.info");
-}
 
-TEST_F(IstateInfoTest, OutIstateInfoS2)
-{
-    // preconditions
-    GlobalV::KPAR = 1;
-    PARAM.input.nbands = 4;
-    PARAM.sys.nbands_l = 4;
-    PARAM.input.nspin = 2;
-    PARAM.sys.global_out_dir = "./";
-    // mpi setting
-    Parallel_Global::init_pools(GlobalV::NPROC,
-                                GlobalV::MY_RANK,
-                                PARAM.input.bndpar,
-                                GlobalV::KPAR,
-                                GlobalV::NPROC_IN_BNDGROUP,
-                                GlobalV::RANK_IN_BPGROUP,
-                                GlobalV::MY_BNDGROUP,
-                                GlobalV::NPROC_IN_POOL,
-                                GlobalV::RANK_IN_POOL,
-                                GlobalV::MY_POOL);
-    kv->set_nkstot(100);
-    int nkstot = kv->get_nkstot();
-    kv->para_k.kinfo(nkstot, GlobalV::KPAR, GlobalV::MY_POOL, GlobalV::RANK_IN_POOL, GlobalV::NPROC_IN_POOL, PARAM.input.nspin);
-    // std::cout<<"my_rank "<<GlobalV::MY_RANK<<" pool rank/size: "
-    //	<<GlobalV::RANK_IN_POOL<<"/"<<GlobalV::NPROC_IN_POOL<<std::endl;
-    // std::cout<<"MY_POOL "<<GlobalV::MY_POOL<<std::endl;
-    kv->set_nks(kv->para_k.nks_pool[GlobalV::MY_POOL]);
-    // std::cout<<"nks "<<kv->get_nks()<<std::endl;
+    // Initialize the number of bands
     ekb.create(kv->get_nks(), PARAM.input.nbands);
     wg.create(kv->get_nks(), PARAM.input.nbands);
+
+    // fill the eigenvalues
     ekb.fill_out(0.15);
+    
+    // fill the weights
     wg.fill_out(0.0);
-    kv->kvec_d.resize(kv->get_nkstot());
+
+    // setup coordinates of k-points
+    kv->kvec_c.resize(kv->get_nkstot());
     int i = 0;
-    for (auto& kd: kv->kvec_d)
+    for (auto& kd: kv->kvec_c)
     {
         kd.set(0.01 * i, 0.01 * i, 0.01 * i);
         ++i;
     }
+   
+    // write eigenvalues and occupations
     ModuleIO::write_istate_info(ekb, wg, *kv);
+
+    // check the output files
     std::ifstream ifs;
-    ifs.open("istate.info");
+    ifs.open("eig.txt");
     std::string str((std::istreambuf_iterator<char>(ifs)), std::istreambuf_iterator<char>());
-    EXPECT_THAT(str,
-                testing::HasSubstr("BAND       Spin up Energy(ev)               Occupation     Spin down Energy(ev)    "
-                                   "           Occupation                Kpoint = 50"));
-    EXPECT_THAT(
-        str,
-        testing::HasSubstr(
-            "4                  2.04085                        0                  2.04085                        0"));
-    EXPECT_THAT(str, testing::HasSubstr("(0.49 0.49 0.49)"));
+    EXPECT_THAT(str, testing::HasSubstr("Electronic state energy (eV) and occupations"));
+    EXPECT_THAT(str, testing::HasSubstr("spin=1 k-point=1/10 Cartesian=0.0000000 0.0000000 0.0000000 (299 plane wave)"));
+    EXPECT_THAT(str, testing::HasSubstr("1 2.040854700000000 0.000000000000000"));
     ifs.close();
-    remove("istate.info");
+    remove("eig.txt");
 }
 
 #ifdef __MPI
diff --git a/source/module_io/write_istate_info.cpp b/source/module_io/write_istate_info.cpp
@@ -42,7 +42,6 @@ void ModuleIO::write_istate_info(const ModuleBase::matrix &ekb,const ModuleBase:
         ModuleBase::WARNING_QUIT("ModuleIO::write_istate_info", "Eigenvalues are too large!");
     }
 
-
     std::vector<int> ngk_tot = kv.ngk;
 
 #ifdef __MPI
@@ -73,7 +72,9 @@ void ModuleIO::write_istate_info(const ModuleBase::matrix &ekb,const ModuleBase:
 #ifdef __MPI
             MPI_Barrier(MPI_COMM_WORLD);
 #endif
+
             bool ip_flag = PARAM.inp.out_alllog || (GlobalV::RANK_IN_POOL == 0 && GlobalV::MY_BNDGROUP == 0);
+
             if (GlobalV::MY_POOL == ip && ip_flag)
             {
                 std::ofstream ofs_eig(filename.c_str(), std::ios::app);

Original file line number	Diff line number	Diff line change
`@@ -48,17 +48,12 @@ void Parallel_Kpoints::get_whichpool(const int& nkstot)`
`48`	`48`	`{`
`49`	`49`	`this->whichpool.resize(nkstot, 0);`
`50`	`50`
`51`		`- // std::cout << " calculate : whichpool" << std::endl;`
`52`		`- // std::cout << " nkstot is " << nkstot << std::endl;`
`53`		`-`
`54`		`-`
`55`	`51`	`for (int i = 0; i < this->kpar; i++)`
`56`	`52`	`{`
`57`	`53`	`for (int ik = 0; ik < this->nks_pool[i]; ik++)`
`58`	`54`	`{`
`59`	`55`	`const int k_now = ik + startk_pool[i];`
`60`	`56`	`this->whichpool[k_now] = i;`
`61`		`- // ofs_running << "\n whichpool[" << k_now <<"] = " << whichpool[k_now];`
`62`	`57`	`}`
`63`	`58`	`}`
`64`	`59`
`@@ -72,19 +67,13 @@ void Parallel_Kpoints::get_nks_pool(const int& nkstot)`
`72`	`67`	`const int nks_ave = nkstot / this->kpar;`
`73`	`68`	`const int remain = nkstot % this->kpar;`
`74`	`69`
`75`		`- // ofs_running << "\n nkstot = " << nkstot;`
`76`		`- // ofs_running << "\n this->kpar = " << this->kpar;`
`77`		`- // ofs_running << "\n nks_ave = " << nks_ave;`
`78`		`-`
`79`	`70`	`for (int i = 0; i < this->kpar; i++)`
`80`		`-`
`81`	`71`	`{`
`82`	`72`	`this->nks_pool[i] = nks_ave;`
`83`	`73`	`if (i < remain)`
`84`	`74`	`{`
`85`	`75`	`nks_pool[i]++;`
`86`	`76`	`}`
`87`		`- // ofs_running << "\n nks_pool[i] = " << nks_pool[i];`
`88`	`77`	`}`
`89`	`78`	`return;`
`90`	`79`	`}`
`@@ -93,14 +82,10 @@ void Parallel_Kpoints::get_startk_pool(const int& nkstot)`
`93`	`82`	`{`
`94`	`83`	`startk_pool.resize(this->kpar, 0);`
`95`	`84`
`96`		`- // const int remain = nkstot%this->kpar;`
`97`		`-`
`98`	`85`	`startk_pool[0] = 0;`
`99`	`86`	`for (int i = 1; i < this->kpar; i++)`
`100`		`-`
`101`	`87`	`{`
`102`	`88`	`startk_pool[i] = startk_pool[i - 1] + nks_pool[i - 1];`
`103`		`- // ofs_running << "\n startk_pool[i] = " << startk_pool[i];`
`104`	`89`	`}`
`105`	`90`	`return;`
`106`	`91`	`}`
`@@ -120,7 +105,6 @@ void Parallel_Kpoints::set_startpro_pool()`
`120`	`105`	`{`
`121`	`106`	`startpro_pool[i]++;`
`122`	`107`	`}`
`123`		`- // ofs_running << "\n startpro_pool[i] = " << startpro_pool[i];`
`124`	`108`	`}`
`125`	`109`	`return;`
`126`	`110`	`}`
`@@ -138,8 +122,6 @@ void Parallel_Kpoints::gatherkvec(const std::vector<ModuleBase::Vector3<double>>`
`138`	`122`	`{`
`139`	`123`	`vec_global[i + startk_pool[this->my_pool]] = vec_local[i];`
`140`	`124`	`}`
`141`		`- // vec_global[i + startk_pool[MY_POOL]] = vec_local[i] / double(NPROC_IN_POOL);`
`142`		`-`
`143`	`125`	`}`
`144`	`126`
`145`	`127`	`MPI_Allreduce(MPI_IN_PLACE, &vec_global[0], 3 * this->nkstot_np, MPI_DOUBLE, MPI_SUM, MPI_COMM_WORLD);`
`@@ -152,7 +134,6 @@ void Parallel_Kpoints::pool_collection(double& value, const double* wk, const in`
`152`	`134`	`#ifdef __MPI`
`153`	`135`
`154`	`136`	`const int ik_now = ik - this->startk_pool[this->my_pool];`
`155`		`- // ofs_running << "\n\n ik=" << ik << " ik_now=" << ik_now;`
`156`	`137`
`157`	`138`	`const int pool = this->whichpool[ik];`
`158`	`139`
`@@ -167,8 +148,6 @@ void Parallel_Kpoints::pool_collection(double& value, const double* wk, const in`
`167`	`148`	`}`
`168`	`149`	`else`
`169`	`150`	`{`
`170`		`-`
`171`		`- // ofs_running << " receive data.";`
`172`	`151`	`MPI_Status ierror;`
`173`	`152`	`MPI_Recv(&value, 1, MPI_DOUBLE, this->startpro_pool[pool], ik, MPI_COMM_WORLD, &ierror);`
`174`	`153`
`@@ -178,18 +157,12 @@ void Parallel_Kpoints::pool_collection(double& value, const double* wk, const in`
`178`	`157`	`{`
`179`	`158`	`if (this->my_pool == pool)`
`180`	`159`	`{`
`181`		`-`
`182`		`- // ofs_running << " send data.";`
`183`		`-`
`184`	`160`	`MPI_Send(&wk[ik_now], 1, MPI_DOUBLE, 0, ik, MPI_COMM_WORLD);`
`185`	`161`	`}`
`186`	`162`	`}`
`187`	`163`	`}`
`188`	`164`	`else`
`189`	`165`	`{`
`190`		`-`
`191`		`- // ofs_running << "\n do nothing.";`
`192`		`-`
`193`	`166`	`}`
`194`	`167`
`195`	`168`	`MPI_Barrier(MPI_COMM_WORLD);`
`@@ -238,13 +211,10 @@ void Parallel_Kpoints::pool_collection_aux(T* value, const V& w, const int& dim,`
`238`	`211`	`T* p = &w.ptr[begin];`
`239`	`212`	`// temprary restrict kpar=1 for NSPIN=2 case for generating_orbitals`
`240`	`213`	`int pool = 0;`
`241`		`- if (this->nspin != 2) {`
`242`		`- pool = this->whichpool[ik];`
`243`		`-}`
`244`		`-`
`245`		`-`
`246`		`- // ofs_running << "\n ik=" << ik;`
`247`		`-`
	`214`	`+ if (this->nspin != 2)`
	`215`	`+ {`
	`216`	`+ pool = this->whichpool[ik];`
	`217`	`+ }`
`248`	`218`
`249`	`219`	`if (this->rank_in_pool == 0)`
`250`	`220`	`{`
`@@ -261,7 +231,6 @@ void Parallel_Kpoints::pool_collection_aux(T* value, const V& w, const int& dim,`
`261`	`231`	`}`
`262`	`232`	`else`
`263`	`233`	`{`
`264`		`- // ofs_running << " receive data.";`
`265`	`234`	`MPI_Status ierror;`
`266`	`235`	`MPI_Recv(value, dim, MPI_DOUBLE, this->startpro_pool[pool], ik * 2 + 0, MPI_COMM_WORLD, &ierror);`
`267`	`236`	`}`
`@@ -270,14 +239,12 @@ void Parallel_Kpoints::pool_collection_aux(T* value, const V& w, const int& dim,`
`270`	`239`	`{`
`271`	`240`	`if (this->my_pool == pool)`
`272`	`241`	`{`
`273`		`- // ofs_running << " send data.";`
`274`	`242`	`MPI_Send(p, dim, MPI_DOUBLE, 0, ik * 2 + 0, MPI_COMM_WORLD);`
`275`	`243`	`}`
`276`	`244`	`}`
`277`	`245`	`}`
`278`	`246`	`else`
`279`	`247`	`{`
`280`		`- // ofs_running << "\n do nothing.";`
`281`	`248`	`}`
`282`	`249`	`MPI_Barrier(MPI_COMM_WORLD);`
`283`	`250`
`@@ -292,4 +259,4 @@ void Parallel_Kpoints::pool_collection_aux(T* value, const V& w, const int& dim,`
`292`	`259`	`}`
`293`	`260`	`// data transfer ends.`
`294`	`261`	`#endif`
`295`		`-}`
	`262`	`+}`