Feature: Support cube format real-space wave function output for Plane Wave basis set

Author: AsTonyshment

docs/advanced/input_files/input-main.md

@@ -1670,7 +1670,7 @@ These variables are used to control the output of properties.
   - 2: results are stored in binary files named `OUT.${suffix}/WAVEFUNC${K}.dat`.
 - **Default**: 0
 
-### out_wfc_r
+### out_wfc_r (deprecated)
 
 - **Type**: Boolean
 - **Availability**: Plane wave basis or get_wf calculation in numerical atomic orbital basis
@@ -1911,14 +1911,14 @@ The band (KS orbital) energy for each (k-point, spin, band) will be printed in t
 ### out_wfc_norm
 
 - **Type**: String
-- **Availability**: For LCAO, used when `calculation = get_wf`.
+- **Availability**: For both PW and LCAO. When `basis_type = lcao`, used when `calculation = get_wf`.
 - **Description**: Specifies the bands to calculate the real-space wave function modulus (norm, or known as the envelope function) $|\psi(r)|$. The syntax and band selection rules are identical to [`out_pchg`](#out_pchg), but the output is the norm of the wave function.
 - **Default**: none
 
 ### out_wfc_re_im
 
 - **Type**: String
-- **Availability**: For LCAO, used when `calculation = get_wf`.
+- **Availability**: For both PW and LCAO. When `basis_type = lcao`, used when `calculation = get_wf`.
 - **Description**: Specifies the bands to calculate the real and imaginary parts of the wave function $\text{Re}(\psi(r))$ and $\text{Im}(\psi(r))$. The syntax and band selection rules are identical to [`out_pchg`](#out_pchg), but the output contains both the real and imaginary components of the wave function.
 - **Default**: none
 

source/module_esolver/esolver_ks_pw.cpp

@@ -24,6 +24,7 @@
 #include "module_io/berryphase.h"
 #include "module_io/cal_ldos.h"
 #include "module_io/get_pchg_pw.h"
+#include "module_io/get_wf_pw.h"
 #include "module_io/numerical_basis.h"
 #include "module_io/numerical_descriptor.h"
 #include "module_io/to_wannier90_pw.h"
@@ -933,9 +934,41 @@ void ESolver_KS_PW<T, Device>::after_all_runners(UnitCell& ucell)
     //----------------------------------------------------------
     //! 5) Print out electronic wave functions in real space
     //----------------------------------------------------------
-    if (PARAM.inp.out_wfc_r == 1) // Peize Lin add 2021.11.21
+
+    //----------------------------------------------------------
+    //! The write_psi_r_1 interface will be removed in the very 
+    //! near future. Don't use it!
+    //----------------------------------------------------------
+    // if (PARAM.inp.out_wfc_r == 1) // Peize Lin add 2021.11.21
+    // {
+    //     ModuleIO::write_psi_r_1(ucell, this->psi[0], this->pw_wfc, "wfc_realspace", true, this->kv);
+    // }
+
+    const std::vector<int> out_wfc_norm = PARAM.inp.out_wfc_norm;
+    const std::vector<int> out_wfc_re_im = PARAM.inp.out_wfc_re_im;
+    if (out_wfc_norm.size() > 0)
     {
-        ModuleIO::write_psi_r_1(ucell, this->psi[0], this->pw_wfc, "wfc_realspace", true, this->kv);
+        ModuleIO::get_wf_pw(out_wfc_norm,
+                            out_wfc_re_im,
+                            this->kspw_psi->get_nbands(),
+                            PARAM.inp.nspin,
+                            this->pw_rhod->nx,
+                            this->pw_rhod->ny,
+                            this->pw_rhod->nz,
+                            this->pw_rhod->nxyz,
+                            this->kv.get_nks(),
+                            this->kv.isk,
+                            this->kv.wk,
+                            this->pw_big->bz,
+                            this->pw_big->nbz,
+                            this->chr.ngmc,
+                            &ucell,
+                            this->psi,
+                            this->pw_rhod,
+                            this->pw_wfc,
+                            this->ctx,
+                            this->Pgrid,
+                            PARAM.globalv.global_out_dir);
     }
 
     //----------------------------------------------------------