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

source/module_elecstate/module_charge/charge_init.cpp

@@ -247,7 +247,8 @@ 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_wf2rho_pw(pw_wfc, symm, kv->ik2iktot.data(), nkstot, isk, *this);
+        ModuleIO::read_wf2rho_pw(pw_wfc, symm, kv->ik2iktot, 
+          nkstot, isk, *this, 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,11 @@ 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(PARAM.inp.out_wfc_pw, PARAM.globalv.global_out_dir, 
+				this->psi[0], this->kv, this->pw_wfc);
+	}
 
     //----------------------------------------------------------
     // 4) check if oscillate for delta_spin method
@@ -709,17 +705,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 +733,13 @@ 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(PARAM.inp.out_wfc_pw, PARAM.globalv.global_out_dir, 
+			this->psi[0], this->kv, this->pw_wfc);
+
     //------------------------------------------------------------------
-    //! 6) calculate Wannier functions in pw basis
+    //! 5) calculate Wannier functions in pw basis
     //------------------------------------------------------------------
     if (PARAM.inp.calculation == "nscf" && PARAM.inp.towannier90)
     {
@@ -766,7 +757,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 +768,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 +780,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

source/module_io/get_wf_lcao.cpp

@@ -221,19 +221,7 @@ void Get_wf_lcao::begin(const UnitCell& ucell,
         }
     }
 
-    if (out_wfc_pw)
-    {
-        std::stringstream ssw;
-        ssw << global_out_dir << "WAVEFUNC";
-        std::cout << " Write G-space wave functions into \"" << global_out_dir << "/" << ssw.str() << "\" files."
-                  << std::endl;
-        ModuleIO::write_wfc_pw(ssw.str(), psi_g, kv, pw_wfc);
-    }
-
-    // if (out_wfc_r)
-    // {
-    //     ModuleIO::write_psi_r_1(ucell, psi_g, pw_wfc, "wfc_realspace", false, kv);
-    // }
+    ModuleIO::write_wfc_pw(out_wfc_pw,global_out_dir,psi_g, kv, pw_wfc);
 
     for (int is = 0; is < nspin; ++is)
     {
@@ -381,18 +369,7 @@ void Get_wf_lcao::begin(const UnitCell& ucell,
         }
     }
 
-    if (out_wf)
-    {
-        std::stringstream ssw;
-        ssw << global_out_dir << "WAVEFUNC";
-        std::cout << " write G-space wave functions into \"" << global_out_dir << "/" << ssw.str() << "\" files."
-                  << std::endl;
-        ModuleIO::write_wfc_pw(ssw.str(), psi_g, kv, pw_wfc);
-    }
-    // if (out_wf_r)
-    // {
-    //     ModuleIO::write_psi_r_1(ucell, psi_g, pw_wfc, "wfc_realspace", false, kv);
-    // }
+    ModuleIO::write_wfc_pw(out_wf,global_out_dir,psi_g, kv, pw_wfc);
 
     std::cout << " Outputting real-space wave functions in cube format..." << std::endl;