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

docs/advanced/elec_properties/wfc.md

@@ -9,8 +9,6 @@ In the LCAO basis case, several `WFC_NAO_K${k}.dat` files will be output in mult
 
 ## wave function in real space
 
-One can also choose to output real-space wave function in PW basis calculation with the key word ***[out_wfc_r](https://abacus-rtd.readthedocs.io/en/latest/advanced/input_files/input-main.html#out_wfc_r)***.
-
-After calculation, an additional directory named `wfc_realspace` will appear in the `OUT.${system}` directory.
+One can also choose to output real-space wave function in PW basis calculation with the key word ***[out_wfc_norm](https://abacus-rtd.readthedocs.io/en/latest/advanced/input_files/input-main.html#out-wfc-norm)***.
 
 Notice: when the ***[basis_type](https://abacus-rtd.readthedocs.io/en/latest/advanced/input_files/input-main.html#basis_type)*** is `lcao`, only `get_wf` ***[calculation](https://abacus-rtd.readthedocs.io/en/latest/advanced/input_files/input-main.html#calculation)*** is effective. An example is [examples/wfc/lcao_ienvelope_Si2](https://github.com/deepmodeling/abacus-develop/tree/develop/examples/wfc/lcao_ienvelope_Si2). 

docs/advanced/input_files/input-main.md

@@ -136,7 +136,6 @@
     - [out\_dm](#out_dm)
     - [out\_dm1](#out_dm1)
     - [out\_wfc\_pw](#out_wfc_pw)
-    - [out\_wfc\_r](#out_wfc_r)
     - [out\_wfc\_lcao](#out_wfc_lcao)
     - [out\_dos](#out_dos)
     - [out\_ldos](#out_ldos)
@@ -1593,7 +1592,7 @@ These variables are used to control the output of properties.
 ### out_freq_elec
 
 - **Type**: Integer
-- **Description**: Output the charge density (only binary format, controlled by [out_chg](#out_chg)), wavefunction (controlled by [out_wfc_pw](#out_wfc_pw) or [out_wfc_r](#out_wfc_r)) per `out_freq_elec` electronic iterations. Note that they are always output when converged or reach the maximum iterations [scf_nmax](#scf_nmax).
+- **Description**: Output the charge density (only binary format, controlled by [out_chg](#out_chg)), wavefunction (controlled by [out_wfc_pw](#out_wfc_pw)) per `out_freq_elec` electronic iterations. Note that they are always output when converged or reach the maximum iterations [scf_nmax](#scf_nmax).
 - **Default**: [scf_nmax](#scf_nmax)
 
 ### out_chg
@@ -1670,13 +1669,6 @@ 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
-
-- **Type**: Boolean
-- **Availability**: Plane wave basis or get_wf calculation in numerical atomic orbital 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
-
 ### out_wfc_lcao
 
 - **Type**: Integer
@@ -1911,14 +1903,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 || out_wfc_re_im.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);
     }
 
     //----------------------------------------------------------

source/module_esolver/lcao_others.cpp

@@ -362,7 +362,6 @@ void ESolver_KS_LCAO<TK, TR>::others(UnitCell& ucell, const int istep)
                       this->pv,
                       this->GG,
                       PARAM.inp.out_wfc_pw,
-                      PARAM.inp.out_wfc_r,
                       this->kv,
                       PARAM.inp.nelec,
                       PARAM.inp.nbands_istate,
@@ -384,7 +383,6 @@ void ESolver_KS_LCAO<TK, TR>::others(UnitCell& ucell, const int istep)
                       this->pv,
                       this->GK,
                       PARAM.inp.out_wfc_pw,
-                      PARAM.inp.out_wfc_r,
                       this->kv,
                       PARAM.inp.nelec,
                       PARAM.inp.nbands_istate,