Update the output and readin of electronic wave functions in plane wave basis set

.github/workflows/test.yml

@@ -49,13 +49,41 @@ jobs:
           cmake --build build -j8
           cmake --install build
 
-      - name: Integrated Test
+      - name: Integrated Tests Preparation
         env:
           GTEST_COLOR: 'yes'
           OMP_NUM_THREADS: '2'
         run: |
           cmake --build build --target test ARGS="-V --timeout 1700 -R integrated_test"
 
+      - name: Module_Base Unittests
+        env:
+          GTEST_COLOR: 'yes'
+          OMP_NUM_THREADS: '2'
+        run: |
+          cmake --build build --target test ARGS="-V --timeout 1700 -R MODULE_BASE"
+
+      - name: Module_IO Unittests
+        env:
+          GTEST_COLOR: 'yes'
+          OMP_NUM_THREADS: '2'
+        run: |
+          cmake --build build --target test ARGS="-V --timeout 1700 -R MODULE_IO"
+
+      - name: Module_HSolver Unittests
+        env:
+          GTEST_COLOR: 'yes'
+          OMP_NUM_THREADS: '2'
+        run: |
+          cmake --build build --target test ARGS="-V --timeout 1700 -R MODULE_HSOLVER"
+
+      - name: Module_Cell Unittests
+        env:
+          GTEST_COLOR: 'yes'
+          OMP_NUM_THREADS: '2'
+        run: |
+          cmake --build build --target test ARGS="-V --timeout 1700 -R MODULE_CELL"
+
       - name: 01_PW Test
         env:
           GTEST_COLOR: 'yes'
@@ -126,35 +154,7 @@ jobs:
         run: |
           cmake --build build --target test ARGS="-V --timeout 1700 -R 10_others"
 
-      - name: Module Base Unit Test
-        env:
-          GTEST_COLOR: 'yes'
-          OMP_NUM_THREADS: '2'
-        run: |
-          cmake --build build --target test ARGS="-V --timeout 1700 -R MODULE_BASE"
-
-      - name: Module IO Unit Test
-        env:
-          GTEST_COLOR: 'yes'
-          OMP_NUM_THREADS: '2'
-        run: |
-          cmake --build build --target test ARGS="-V --timeout 1700 -R MODULE_IO"
-
-      - name: Module HSolver Unit Test
-        env:
-          GTEST_COLOR: 'yes'
-          OMP_NUM_THREADS: '2'
-        run: |
-          cmake --build build --target test ARGS="-V --timeout 1700 -R MODULE_HSOLVER"
-
-      - name: Module Cell Unit Test
-        env:
-          GTEST_COLOR: 'yes'
-          OMP_NUM_THREADS: '2'
-        run: |
-          cmake --build build --target test ARGS="-V --timeout 1700 -R MODULE_CELL"
-
-      - name: Other Unit Tests
+      - name: Other Unittests
         env:
           GTEST_COLOR: 'yes'
           OMP_NUM_THREADS: '2'

source/Makefile.Objects

@@ -561,6 +561,7 @@ OBJS_IO=input_conv.o\
     read_set_globalv.o\
     orb_io.o\
     cal_pLpR.o\
+    filename.o\
 
 OBJS_IO_LCAO=cal_r_overlap_R.o\
       write_orb_info.o\

source/module_elecstate/module_charge/charge_init.cpp

@@ -243,11 +243,16 @@ void Charge::init_rho(elecstate::efermi& eferm_iout,
         {
             ModuleBase::WARNING_QUIT("Charge::init_rho", "wfc is only supported for PW-KSDFT.");
         }
+
         const ModulePW::PW_Basis_K* pw_wfc = reinterpret_cast<ModulePW::PW_Basis_K*>(const_cast<void*>(wfcpw));
         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_wf2rho_pw(pw_wfc, symm, kv->ik2iktot.data(), nkstot, isk, *this);
+
+		ModuleIO::read_wf2rho_pw(pw_wfc, symm, *this,
+                PARAM.globalv.global_readin_dir,
+				GlobalV::KPAR, GlobalV::MY_POOL, GlobalV::MY_RANK, 
+                GlobalV::NPROC_IN_POOL, GlobalV::RANK_IN_POOL,
+				PARAM.inp.nbands, PARAM.inp.nspin, PARAM.globalv.npol,
+				kv->get_nkstot(),kv->ik2iktot,kv->isk,GlobalV::ofs_running);
     }
 }
 

source/module_esolver/esolver.cpp

@@ -96,7 +96,7 @@ std::string determine_type()
         ModuleBase::WARNING_QUIT("ESolver", "No such esolver_type combined with basis_type");
     }
 
-    GlobalV::ofs_running << "\n The esolver type: " << esolver_type << std::endl;
+    GlobalV::ofs_running << "\n Energy Solver:" << esolver_type << std::endl;
 
     auto device_info = PARAM.inp.device;
 

source/module_esolver/esolver_ks.cpp

@@ -58,7 +58,7 @@ void ESolver_KS<T, Device>::before_all_runners(UnitCell& ucell, const Input_para
     ESolver_FP::before_all_runners(ucell, inp);
 
     classname = "ESolver_KS";
-    basisname = "PLEASE ADD BASISNAME FOR CURRENT ESOLVER.";
+    basisname = "";
 
     scf_thr = PARAM.inp.scf_thr;
     scf_ene_thr = PARAM.inp.scf_ene_thr;

source/module_esolver/esolver_ks_lcaopw.cpp

@@ -5,7 +5,6 @@
 #include "module_io/nscf_band.h"
 #include "module_io/output_log.h"
 #include "module_io/write_istate_info.h"
-#include "module_io/write_wfc_pw.h"
 
 #include <iostream>
 

source/module_esolver/esolver_ks_pw.cpp

@@ -647,15 +647,14 @@ void ESolver_KS_PW<T, Device>::iter_finish(UnitCell& ucell, const int istep, int
     //----------------------------------------------------------
     // 3) Print out electronic wavefunctions in pw basis
     //----------------------------------------------------------
-    if (PARAM.inp.out_wfc_pw == 1 || PARAM.inp.out_wfc_pw == 2)
-    {
-        if (iter % PARAM.inp.out_freq_elec == 0 || iter == PARAM.inp.scf_nmax || conv_esolver)
-        {
-            std::stringstream ssw;
-            ssw << PARAM.globalv.global_out_dir << "WAVEFUNC";
-            ModuleIO::write_wfc_pw(ssw.str(), this->psi[0], this->kv, this->pw_wfc);
-        }
-    }
+	if (iter % PARAM.inp.out_freq_elec == 0 || iter == PARAM.inp.scf_nmax || conv_esolver)
+	{
+		ModuleIO::write_wfc_pw(GlobalV::KPAR, GlobalV::MY_POOL, GlobalV::MY_RANK, 
+				PARAM.inp.nbands, PARAM.inp.nspin, PARAM.globalv.npol,
+				GlobalV::RANK_IN_POOL, GlobalV::NPROC_IN_POOL, 
+				PARAM.inp.out_wfc_pw, PARAM.inp.ecutwfc, PARAM.globalv.global_out_dir,
+				this->psi[0], this->kv, this->pw_wfc, GlobalV::ofs_running);
+	}
 
     //----------------------------------------------------------
     // 4) check if oscillate for delta_spin method
@@ -709,17 +708,7 @@ void ESolver_KS_PW<T, Device>::after_scf(UnitCell& ucell, const int istep, const
     }
 
     //------------------------------------------------------------------
-    // 4) output wavefunctions in pw basis
-    //------------------------------------------------------------------
-    if (PARAM.inp.out_wfc_pw == 1 || PARAM.inp.out_wfc_pw == 2)
-    {
-        std::stringstream ssw;
-        ssw << PARAM.globalv.global_out_dir << "WAVEFUNC";
-        ModuleIO::write_wfc_pw(ssw.str(), this->psi[0], this->kv, this->pw_wfc);
-    }
-
-    //------------------------------------------------------------------
-    // 5) calculate band-decomposed (partial) charge density in pw basis
+    // 4) calculate band-decomposed (partial) charge density in pw basis
     //------------------------------------------------------------------
     const std::vector<int> out_pchg = PARAM.inp.out_pchg;
     if (out_pchg.size() > 0)
@@ -747,8 +736,16 @@ void ESolver_KS_PW<T, Device>::after_scf(UnitCell& ucell, const int istep, const
                               PARAM.inp.if_separate_k);
     }
 
+
+	// tmp 2025-05-17, mohan note
+	ModuleIO::write_wfc_pw(GlobalV::KPAR, GlobalV::MY_POOL, GlobalV::MY_RANK, 
+			PARAM.inp.nbands, PARAM.inp.nspin, PARAM.globalv.npol,
+			GlobalV::RANK_IN_POOL, GlobalV::NPROC_IN_POOL, 
+			PARAM.inp.out_wfc_pw, PARAM.inp.ecutwfc, PARAM.globalv.global_out_dir,
+			this->psi[0], this->kv, this->pw_wfc, GlobalV::ofs_running);
+
     //------------------------------------------------------------------
-    //! 6) calculate Wannier functions in pw basis
+    //! 5) calculate Wannier functions in pw basis
     //------------------------------------------------------------------
     if (PARAM.inp.calculation == "nscf" && PARAM.inp.towannier90)
     {
@@ -766,7 +763,7 @@ void ESolver_KS_PW<T, Device>::after_scf(UnitCell& ucell, const int istep, const
     }
 
     //------------------------------------------------------------------
-    //! 7) calculate Berry phase polarization in pw basis
+    //! 6) calculate Berry phase polarization in pw basis
     //------------------------------------------------------------------
     if (PARAM.inp.calculation == "nscf" && berryphase::berry_phase_flag && ModuleSymmetry::Symmetry::symm_flag != 1)
     {
@@ -777,7 +774,7 @@ void ESolver_KS_PW<T, Device>::after_scf(UnitCell& ucell, const int istep, const
     }
 
     //------------------------------------------------------------------
-    // 8) write spin constrian results in pw basis
+    // 7) write spin constrian results in pw basis
     // spin constrain calculations, write atomic magnetization and magnetic force.
     //------------------------------------------------------------------
     if (PARAM.inp.sc_mag_switch)
@@ -789,7 +786,7 @@ void ESolver_KS_PW<T, Device>::after_scf(UnitCell& ucell, const int istep, const
     }
 
     //------------------------------------------------------------------
-    // 9) write onsite occupations for charge and magnetizations
+    // 8) write onsite occupations for charge and magnetizations
     //------------------------------------------------------------------
     if (PARAM.inp.onsite_radius > 0)
     { // float type has not been implemented

source/module_esolver/pw_others.cpp

@@ -6,7 +6,6 @@
 #include "module_io/nscf_band.h"
 #include "module_io/output_log.h"
 #include "module_io/write_istate_info.h"
-#include "module_io/write_wfc_pw.h"
 
 #include <iostream>
 

source/module_io/CMakeLists.txt

@@ -33,6 +33,7 @@ list(APPEND objects
     para_json.cpp
     parse_args.cpp
     orb_io.cpp
+    filename.cpp
 )
 
 list(APPEND objects_advanced

source/module_io/cal_dos.cpp

@@ -17,7 +17,6 @@ void ModuleIO::prepare_dos(std::ofstream& ofs_running,
 		double &emin)
 {
 	ofs_running << " DOS CALCULATIONS BEGINS" << std::endl;
-
 	ofs_running << " >>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>"
 		">>>>>>>>>>>>>>>>>>>>>>>>>" << std::endl;
 	ofs_running << " |                                            "

source/module_io/filename.cpp

@@ -0,0 +1,111 @@
+#include <set>
+#include "filename.h"
+#include "module_base/tool_quit.h"
+
+namespace ModuleIO
+{
+
+std::string filename_output(
+			const std::string &directory,
+			const std::string &property,
+			const std::string &basis,
+			const int ik_local, // the ik index within each pool
+			const std::vector<int> &ik2iktot,
+			const int nspin,
+			const int nkstot,
+			const int out_type,
+			const bool out_app_flag,
+			const bool gamma_only,
+			const int istep)
+{
+    // output filename = "{PARAM.globalv.global_out_dir}/property{s}{spin index}
+    // {k(optinal)}{k-point index}{g(optional)}{geometry index1}{_basis(nao|pw)} 
+    // + {".txt"/".dat"}"
+
+	std::set<std::string> valid_properties = {"wf", "chg", "h", "s"};
+	if (valid_properties.find(property) == valid_properties.end()) 
+	{
+		ModuleBase::WARNING_QUIT("ModuleIO::filename_output", "unknown property");
+	}
+
+	std::set<std::string> valid_basis = {"pw", "nao"};
+	if (valid_basis.find(basis) == valid_basis.end()) 
+	{
+		ModuleBase::WARNING_QUIT("ModuleIO::filename_output", "unknown basis");
+	}
+
+    assert(ik_local>=0);
+    // mohan update 2025.05.07, if KPAR>1, "<" works
+	assert(ik2iktot.size() <= nkstot);
+    assert(nspin>0);
+
+	// spin index
+	int is0 = -1;
+	// ik0 is the k-point index, starting from 0
+	int ik0 = ik2iktot[ik_local];
+
+	if(nspin == 1)
+	{
+		is0 = 1;
+	}
+	else if(nspin == 2)
+	{
+		const int half_k = nkstot/2;
+		if(ik0 >= half_k)
+		{
+			is0 = 2;
+			ik0 -= half_k;
+		}
+		else
+		{
+			is0 = 1;
+		}
+	}
+	else if(nspin==4)
+	{
+		is0 = 12;
+	}
+
+	// spin part
+	std::string spin_block;
+    spin_block = "s" + std::to_string(is0);
+
+    // k-point part
+    std::string kpoint_block;
+    if(gamma_only)
+    {
+        // do nothing;
+    }
+    else
+    {
+        kpoint_block = "k" + std::to_string(ik0+1);
+    }
+
+    std::string istep_block
+        = (istep >= 0 && (!out_app_flag))
+              ? "g" + std::to_string(istep + 1)
+              : ""; // only when istep >= 0 and out_app_flag is false will write each wfc to a separate file
+
+    std::string suffix_block;
+    if (out_type == 1)
+    {
+        suffix_block = ".txt";
+    }
+    else if (out_type == 2)
+    {
+        suffix_block = ".dat";
+    }
+    else
+    {
+        std::cout << "WARNING: the type of output wave function is not 1 or 2, so 1 is chosen." << std::endl;
+        suffix_block = ".txt";
+    }
+
+    std::string fn_out
+        = directory + property + spin_block + kpoint_block 
+          + istep_block + "_" + basis + suffix_block;
+
+    return fn_out;
+}
+
+}

source/module_io/filename.h

@@ -0,0 +1,38 @@
+#ifndef FILENAME_H
+#define FILENAME_H
+#include <vector>
+#include <string>
+
+namespace ModuleIO
+{
+
+/**
+ * Generates the filename for the output files
+ * @param directory: directory of the file 
+ * @param property: wave function (wf), charge density (chg) or matrix (mat)
+ * @param basis: nao or pw 
+ * @param ik_local: index of the k-points within each pool, and starting from 0.
+ * @param ik2iktot: map from ik to iktot
+ * @param nspin: number of spin channels, 1,2 or 4 
+ * @param nkstot: number of total k-points 
+ * @param out_type: two types of output file format, 1 for .txt and 2 for .dat (binary)
+ * @param out_app_flag: whether to append to existing file.
+ * @param gamma_only: gamma_only algorithm or not.
+ * @param istep: index of the ion step starting from 0. If < 0, the step number is not included in the file name.
+ * @return The generated filename.
+ */
+std::string filename_output(
+            const std::string &directory,
+            const std::string &property,
+            const std::string &basis,
+			const int ik_local,
+			const std::vector<int> &ik2iktot,
+			const int nspin,
+			const int nkstot,
+            const int out_type,
+			const bool out_app_flag,
+			const bool gamma_only,
+			const int istep=-1);
+
+}
+#endif