Merge branch 'pw' into pw-merge

Author: mohanchen

docs/advanced/input_files/input-main.md

@@ -1664,16 +1664,19 @@ These variables are used to control the output of properties.
 ### out_wfc_pw
 
 - **Type**: Integer
-- **Availability**: Plane wave basis or get_wf calculation in numerical atomic orbital basis
-- **Description**:
-  - 1: Output the coefficients of wave functions into text files named `OUT.${suffix}/WAVEFUNC${K}.txt`, where ${K} is the index of k points.
-  - 2: results are stored in binary files named `OUT.${suffix}/WAVEFUNC${K}.dat`.
+- **Availability**: Output electronic wave functions in plane wave basis, or transform the real-space electronic wave function into plane wave basis (see get_wf option in [calculation](#calculation) with NAO basis) 
+- **Description**: Whether to output the electronic wavefunction coefficients into files and store them in the folder `OUT.${suffix}`. The files are named as `wf{s}{spin index}{k}{k-point index}{_pw} + {".txt"/".dat"}`. Here, 's' refers to spin, where s1 means spin up channel while s2 means spin down channel, and 's12' refer to spinor wave functions that contains both spin channels with spin-orbital coupling or noncollinear calculations enabled. 
+  - 0: no output
+  - 1: (txt format) 
+    - non-gamma-only: `wfs1k1_pw.txt` or `wfs1k2_pw.txt`, ...;
+  - 2: (binary format)
+    - non-gamma-only: `wfs1k1_pw.dat` or `wfs1k2_Pw.dat`, ....
 - **Default**: 0
 
 ### out_wfc_r
 
 - **Type**: Boolean
-- **Availability**: Plane wave basis or get_wf calculation in numerical atomic orbital basis
+- **Availability**: Output electronic wave functions in plane wave basis, or transform the real-space electronic wave function into plane wave basis (see get_wf option in [calculation](#calculation) with NAO basis)
 - **Description**: Whether to output real-space wave functions into `OUT.suffix/wfc_realspace/wfc_realspace_${K}_${B}`, where `${K}` is the index of k points, `${B}` is the index of bands.
 - **Default**: False
 

source/module_cell/unitcell.cpp

@@ -247,7 +247,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_wf2rho_pw(pw_wfc, symm, kv->ik2iktot.data(), nkstot, isk, *this);
+        ModuleIO::read_wf2rho_pw(pw_wfc, symm, kv->ik2iktot, nkstot, isk, *this);
     }
 }
 

source/module_elecstate/module_charge/charge_init.cpp

@@ -650,9 +650,10 @@ void ESolver_KS_PW<T, Device>::iter_finish(UnitCell& ucell, const int istep, int
     {
         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);
+			ModuleIO::write_wfc_pw(
+					PARAM.inp.out_wfc_pw,  
+					PARAM.globalv.global_out_dir,
+					this->psi[0], this->kv, this->pw_wfc);
         }
     }
 
@@ -712,10 +713,10 @@ void ESolver_KS_PW<T, Device>::after_scf(UnitCell& ucell, const int istep, const
     //------------------------------------------------------------------
     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);
-    }
+		ModuleIO::write_wfc_pw(PARAM.inp.out_wfc_pw, 
+				PARAM.globalv.global_out_dir,
+				this->psi[0], this->kv, this->pw_wfc);
+	}
 
     //------------------------------------------------------------------
     // 5) calculate band-decomposed (partial) charge density in pw basis

source/module_esolver/esolver_ks_pw.cpp

@@ -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/CMakeLists.txt

@@ -0,0 +1,110 @@
+#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,
+			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>=0);
+	assert(ik2iktot.size() == nkstot);
+    assert(nspin>0);
+
+	// spin index
+	int is0 = -1;
+	// ik0 is the k-point index, starting from 1
+	int ik0 = ik2iktot[ik];
+
+    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.cpp

@@ -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: index of the k-point, 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,
+			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/filename.h

@@ -243,11 +243,12 @@ 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);
+		ModuleIO::write_wfc_pw(
+				out_wfc_pw,
+                PARAM.globalv.global_out_dir,
+				psi_g, 
+				kv, 
+				pw_wfc);
     }
 
     if (out_wfc_r)
@@ -275,7 +276,7 @@ void Get_wf_lcao::begin(const UnitCell& ucell,
                             const Parallel_Grid& pgrid,
                             const Parallel_Orbitals& para_orb,
                             Gint_k& gk,
-                            const int& out_wf,
+                            const int& out_wf_pw,
                             const int& out_wf_r,
                             const K_Vectors& kv,
                             const double nelec,
@@ -311,7 +312,7 @@ void Get_wf_lcao::begin(const UnitCell& ucell,
 
     // for pw_wfc in G space
     psi::Psi<std::complex<double>> psi_g;
-    if (out_wf || out_wf_r)
+    if (out_wf_pw || out_wf_r)
     {
         psi_g.resize(nks, nbands, pw_wfc->npwk_max);
     }
@@ -401,7 +402,7 @@ void Get_wf_lcao::begin(const UnitCell& ucell,
                                               3,
                                               1);
 
-                if (out_wf || out_wf_r)
+                if (out_wf_pw || out_wf_r)
                 {
                     psi_g.fix_k(ik);
                     this->set_pw_wfc(pw_wfc, ik, ib, nspin, pes_->charge->rho, psi_g);
@@ -410,15 +411,11 @@ void Get_wf_lcao::begin(const UnitCell& ucell,
         }
     }
 
-    if (out_wf || out_wf_r)
+    if (out_wf_pw || out_wf_r)
     {
-        if (out_wf)
+        if (out_wf_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);
+            ModuleIO::write_wfc_pw(out_wf_pw, global_out_dir, psi_g, kv, pw_wfc);
         }
         if (out_wf_r)
         {

source/module_io/get_wf_lcao.cpp

@@ -6,12 +6,13 @@
 #include "module_elecstate/module_charge/symmetry_rho.h"
 #include "module_parameter/parameter.h"
 #include "module_elecstate/kernels/elecstate_op.h"
+#include "module_io/filename.h"
 
 void ModuleIO::read_wf2rho_pw(const ModulePW::PW_Basis_K* pw_wfc,
                                ModuleSymmetry::Symmetry& symm,
-                               const int* ik2iktot,
+                               const std::vector<int> &ik2iktot,
                                const int nkstot,
-                               const std::vector<int>& isk,
+                               const std::vector<int> &isk,
                                Charge& chg)
 {
     ModuleBase::TITLE("ModuleIO", "read_wf2rho_pw");
@@ -78,11 +79,20 @@ void ModuleIO::read_wf2rho_pw(const ModulePW::PW_Basis_K* pw_wfc,
         {
             is = isk[ik];
         }
-        std::stringstream filename;
         const int ikstot = ik2iktot[ik];
-        filename << PARAM.globalv.global_readin_dir << "WAVEFUNC" << ikstot + 1 << ".dat";
-        ModuleIO::read_wfc_pw(filename.str(), pw_wfc, ik, ikstot, nkstot, wfc_tmp);
-        if (PARAM.inp.nspin == 4)
+
+        // mohan add 2025-05-17
+        // .dat file
+        const int out_type = 2; 
+        const bool out_app_flag = false;
+        const bool gamma_only = false;
+        const int istep = -1;
+        
+		std::string fn = filename_output(PARAM.globalv.global_readin_dir,"wf","pw",ik,ik2iktot,nspin,nkstot,
+				out_type,out_app_flag,gamma_only,istep);
+
+        ModuleIO::read_wfc_pw(fn, pw_wfc, ik, ikstot, nkstot, wfc_tmp);
+        if (nspin == 4)
         {
             std::vector<std::complex<double>> rho_tmp2(nrxx);
             for (int ib = 0; ib < nbands; ++ib)

source/module_io/read_wf2rho_pw.cpp

@@ -4,6 +4,7 @@
 #include "module_basis/module_pw/pw_basis_k.h"
 #include "module_elecstate/module_charge/charge.h"
 
+#include <vector>
 #include <string>
 
 namespace ModuleIO
@@ -19,9 +20,9 @@ namespace ModuleIO
  */
 void read_wf2rho_pw(const ModulePW::PW_Basis_K* pw_wfc,
                      ModuleSymmetry::Symmetry& symm,
-                     const int* ik2iktot,
+                     const std::vector<int> &ik2iktot,
                      const int nkstot,
-                     const std::vector<int>& isk,
+                     const std::vector<int> &isk,
                      Charge& chg);
 
 } // namespace ModuleIO