refactor: separate the print_psi from wfc_2d_to_grid and remove the invalid call (#4268)

* refactor: seperate the print_psi from wfc_2d_to_grid and remove the invalid call

* correct fn

* transfer pointer to vector and fix compiling bug

* use Cpxgemr2d to do the gather of wfc

* fix the compiling error and ut error

* fix bug in diago_cusolver

* fix integrate test

* use orb_con.ParaV instead of LOWF.ParaV; add unit test

* move write_wfc_nao to write_wfc_lcao; and correct some description in doc.

* fix bug

* fxi bug in read_wfc_nao_test

* solve the conflict of CURRENT_SPIN

* fix bug in integrate test

* fix bug in integrate test

* fix bug in integrate

* rename wfc_lcao to wfc_nao

* fix bug in Makefile.Objiects; and add comments for Cpxgemr2d

---------

Co-authored-by: Mohan Chen <mohan.chen.chen.mohan@gmail.com>

Author: pxlxingliang

docs/advanced/elec_properties/wfc.md

@@ -5,7 +5,7 @@ ABACUS is able to output electron wave functions in both PW and LCAO basis calcu
 ## wave function in G space
 For the wave function in G space, one only needs to do a ground-state energy calculation with one additional keyword in the INPUT file: '***[out_wfc_pw](https://abacus-rtd.readthedocs.io/en/latest/advanced/input_files/input-main.html#out-wfc-pw)***' for PW basis calculation, and '***[out_wfc_lcao](https://abacus-rtd.readthedocs.io/en/latest/advanced/input_files/input-main.html#out-wfc-lcao)***' for LCAO basis calculation.
 In the PW basis case, the wave function is output in a file called `WAVEFUNC${k}.txt`, where `${k}` is the index of K point. \
-In the LCAO basis case, several `LOWF_K_${k}.dat` files will be output in multi-k calculation and `LOWF_GAMMA_S1.dat` in gamma-only calculation. 
+In the LCAO basis case, several `WFC_NAO_K${k}.dat` files will be output in multi-k calculation and `WFC_NAO_GAMMA1.dat` in gamma-only calculation. 
 
 ## wave function in real space
 

docs/advanced/input_files/input-main.md

@@ -1541,14 +1541,14 @@ These variables are used to control the output of properties.
 
 - **Type**: Integer
 - **Availability**: Numerical atomic orbital basis
-- **Description**: Whether to output the wavefunction coefficients into files in the folder `OUT.${suffix}`. The files are named as:
+- **Description**: Whether to output the wavefunction coefficients into files in the folder `OUT.${suffix}`. The files are named as `WFC_{GAMMA|K}{index of K point}`, and if [out_app_flag](#out_app_flag) is false, the file name will also contains `_ION{ION step}`, where `ION step` is the index of ionic step:
   - 0: no output
-  - 1: (txt format)
-    - gamma-only: `LOWF_GAMMA_S1.txt`;
-    - non-gamma-only: `LOWF_K_${k}.txt`, where `${k}` is the index of k points.
+  - 1: (txt format) 
+    - gamma-only: `WFC_NAO_GAMMA1_ION1.txt` or `WFC_NAO_GAMMA1.txt`, ...;
+    - non-gamma-only: `WFC_NAO_K1_ION1.txt` or `WFC_NAO_K1.txt`, ...;
   - 2: (binary format)
-    - gamma-only: `LOWF_GAMMA_S1.dat`;
-    - non-gamma-only: `LOWF_K_${k}.dat`, where `${k}` is the index of k points.
+    - gamma-only: `WFC_NAO_GAMMA1_ION1.dat` or `WFC_NAO_GAMMA1.dat`, ...;
+    - non-gamma-only: `WFC_NAO_K1_ION1.dat` or `WFC_NAO_K1.dat`, ....
 
   The corresponding sequence of the orbitals can be seen in [Basis Set](../pp_orb.md#basis-set).
 
@@ -2414,7 +2414,7 @@ These variables are used to control molecular dynamics calculations. For more in
 
 - **Type**: Boolean
 - **Description**: Control whether to restart molecular dynamics calculations and time-dependent density functional theory calculations.
-  - True: ABACUS will read in `${read_file_dir}/Restart_md.dat` to determine the current step `${md_step}`, then read in the corresponding `STRU_MD_${md_step}` in the folder `OUT.$suffix/STRU/` automatically. For tddft, ABACUS will also read in `LOWF_K_${kpoint}` of the last step (You need to set out_wfc_lcao=1 and out_app_flag=0 to obtain this file).
+  - True: ABACUS will read in `${read_file_dir}/Restart_md.dat` to determine the current step `${md_step}`, then read in the corresponding `STRU_MD_${md_step}` in the folder `OUT.$suffix/STRU/` automatically. For tddft, ABACUS will also read in `WFC_NAO_K${kpoint}` of the last step (You need to set out_wfc_lcao=1 and out_app_flag=0 to obtain this file).
   - False: ABACUS will start molecular dynamics calculations normally from the first step.
 - **Default**: False
 

examples/wfc/lcao_ienvelope_Si2/run.sh

@@ -15,6 +15,8 @@ if [[ ! -f scf.output ]] ||
    [[ ! -f get_wf.output ]] || 
    [[ ! -f OUT.ABACUS/running_scf.log ]] ||
    [[ ! -f OUT.ABACUS/running_get_wf.log ]] ||
+   [[ ! -f OUT.ABACUS/WFC_NAO_K1.txt ]] ||
+   [[ ! -f OUT.ABACUS/WFC_NAO_K36.txt ]] ||
    [[ ! ( "$(tail -1 OUT.ABACUS/running_scf.log)" == " Total  Time  :"* ) ]] ||
    [[ ! ( "$(tail -1 OUT.ABACUS/running_get_wf.log)" == " Total  Time  :"* ) ]] 
 then

examples/wfc/lcao_scf_Si2/run.sh

@@ -8,11 +8,13 @@ OMP_NUM_THREADS=${ABACUS_THREADS} mpirun -np ${ABACUS_NPROCS} ${ABACUS_PATH} | t
 
 if [[ ! -f output ]] ||
    [[ ! -f OUT.ABACUS/running_scf.log ]] ||
+   [[ ! -f OUT.ABACUS/WFC_NAO_K1.txt ]] ||
+   [[ ! -f OUT.ABACUS/WFC_NAO_K36.txt ]] ||
    [[ ! ( "$(tail -1 OUT.ABACUS/running_scf.log)" == " Total  Time  :"* ) ]] 
 then
 	echo "job is failed!"
 	exit 1
 else
 	echo "job is successed!"
 	exit 0
-fi
+fi

source/module_base/scalapack_connector.h

@@ -115,8 +115,56 @@ extern "C"
 	    std::complex<double> *A, const int *IA, const int *JA, const int *DESCA, 
 		std::complex<double> *B, const int *IB, const int *JB, const int *DESCB,
 		const int *ICTXT);
+
+	// Scalapack wrappers to copy 2D blocks of data
+	// more info: 
+	// https://netlib.org/scalapack/explore-html/da/db5/pigemr_8c.html
+	// https://netlib.org/scalapack/explore-html/dd/dcd/pdgemr_8c.html
+	// https://netlib.org/scalapack/explore-html/d5/dd4/pzgemr_8c.html
+	// https://netlib.org/scalapack/explore-html/d5/deb/psgemr_8c.html
+	// https://netlib.org/scalapack/explore-html/d4/dad/pcgemr_8c.html
+	void Cpigemr2d (int m, int n, int *ptrmyblock, int ia, int ja, int *ma, int *ptrmynewblock, int ib, int jb, int *mb, int globcontext);
+	void Cpdgemr2d (int m, int n, double *ptrmyblock, int ia, int ja, int *ma, double *ptrmynewblock, int ib, int jb, int *mb, int globcontext);
+	void Cpzgemr2d (int m, int n, std::complex<double> *ptrmyblock, int ia, int ja, int *ma, std::complex<double> *ptrmynewblock, int ib, int jb, int *mb, int globcontext);
+	void Cpsgemr2d (int m, int n, float *ptrmyblock, int ia, int ja, int *ma, float *ptrmynewblock, int ib, int jb, int *mb, int globcontext);
+	void Cpcgemr2d (int m, int n, std::complex<float> *ptrmyblock, int ia, int ja, int *ma, std::complex<float> *ptrmynewblock, int ib, int jb, int *mb, int globcontext);
 }
 
+	template <typename T>
+	struct block2d_data_type
+	{
+		constexpr static bool value = std::is_same<T, double>::value || std::is_same<T, std::complex<double>>::value || std::is_same<T, float>::value || std::is_same<T, std::complex<float>>::value || std::is_same<T, int>::value;
+	};
+
+
+	/**
+	 * Copies a 2D block of data from matrix A to matrix B using the Scalapack library.
+	 * This function supports different data types: double, std::complex<double>, float, std::complex<float>, and int.
+	 *
+	 * @tparam T The data type of the matrices A and B.
+	 * @param M The number of rows of matrix A.
+	 * @param N The number of columns of matrix A.
+	 * @param A Pointer to the source matrix A.
+	 * @param IA The starting row index of the block in matrix A.
+	 * @param JA The starting column index of the block in matrix A.
+	 * @param DESCA Descriptor array for matrix A.
+	 * @param B Pointer to the destination matrix B.
+	 * @param IB The starting row index of the block in matrix B.
+	 * @param JB The starting column index of the block in matrix B.
+	 * @param DESCB Descriptor array for matrix B.
+	 * @param ICTXT The context identifier.
+	 */
+	template <typename T>
+	typename std::enable_if<block2d_data_type<T>::value,void>::type Cpxgemr2d(int M, int N, T *A, int IA, int JA, int *DESCA, T *B, int IB, int JB, int *DESCB, int ICTXT)
+	{
+		if (std::is_same<T,double>::value) Cpdgemr2d(M, N, reinterpret_cast<double*>(A),IA, JA, DESCA,reinterpret_cast<double*>(B),IB,JB, DESCB,ICTXT);
+		if (std::is_same<T,std::complex<double>>::value) Cpzgemr2d(M, N, reinterpret_cast<std::complex<double>*>(A),IA, JA, DESCA,reinterpret_cast<std::complex<double>*>(B),IB,JB, DESCB,ICTXT);
+		if (std::is_same<T,float>::value) Cpsgemr2d(M, N, reinterpret_cast<float*>(A),IA, JA, DESCA,reinterpret_cast<float*>(B),IB,JB, DESCB,ICTXT);
+		if (std::is_same<T,std::complex<float>>::value) Cpcgemr2d(M, N, reinterpret_cast<std::complex<float>*>(A),IA, JA, DESCA,reinterpret_cast<std::complex<float>*>(B),IB,JB, DESCB,ICTXT);
+		if (std::is_same<T,int>::value) Cpigemr2d(M, N, reinterpret_cast<int*>(A),IA, JA, DESCA,reinterpret_cast<int*>(B),IB,JB, DESCB,ICTXT);
+	};
+	
+
 class ScalapackConnector
 {
 public:

source/module_elecstate/elecstate_lcao.cpp

@@ -12,81 +12,6 @@
 namespace elecstate
 {
 
-template <>
-void ElecStateLCAO<double>::print_psi(const psi::Psi<double>& psi_in, const int istep)
-{
-	if (!ElecStateLCAO<double>::out_wfc_lcao)
-	{
-		return;
-	}
-
-    // output but not do  "2d-to-grid" conversion
-    double** wfc_grid = nullptr;
-    const int ik = psi_in.get_current_k();
-#ifdef __MPI
-    this->lowf->wfc_2d_to_grid(istep, out_wfc_flag, psi_in.get_pointer(), wfc_grid, ik, this->ekb, this->wg);
-#endif
-    return;
-}
-
-template <>
-void ElecStateLCAO<std::complex<double>>::print_psi(const psi::Psi<std::complex<double>>& psi_in, const int istep)
-{
-	if (!ElecStateLCAO<std::complex<double>>::out_wfc_lcao 
-			&& !ElecStateLCAO<std::complex<double>>::need_psi_grid)
-	{
-		return;
-	}
-
-    // output but not do "2d-to-grid" conversion
-    std::complex<double>** wfc_grid = nullptr;
-    int ik = psi_in.get_current_k();
-    if (ElecStateLCAO<std::complex<double>>::need_psi_grid)
-    {
-        wfc_grid = this->lowf->wfc_k_grid[ik];
-    }
-
-#ifdef __MPI
-    this->lowf->wfc_2d_to_grid(istep,
-                               ElecStateLCAO<std::complex<double>>::out_wfc_flag,
-                               psi_in.get_pointer(),
-                               wfc_grid,
-                               ik,
-                               this->ekb,
-                               this->wg,
-                               this->klist->kvec_c);
-#else
-    for (int ib = 0; ib < GlobalV::NBANDS; ib++)
-    {
-        for (int iw = 0; iw < GlobalV::NLOCAL; iw++)
-        {
-            this->lowf->wfc_k_grid[ik][ib][iw] = psi_in(ib, iw);
-        }
-    }
-#endif
-
-    // added by zhengdy-soc, rearrange the wfc_k_grid from [up,down,up,down...] to [up,up...down,down...],
-    if (ElecStateLCAO<std::complex<double>>::need_psi_grid && GlobalV::NSPIN == 4)
-    {
-        int row = this->lowf->gridt->lgd;
-        std::vector<std::complex<double>> tmp(row);
-        for (int ib = 0; ib < GlobalV::NBANDS; ib++)
-        {
-            for (int iw = 0; iw < row / GlobalV::NPOL; iw++)
-            {
-                tmp[iw] = this->lowf->wfc_k_grid[ik][ib][iw * GlobalV::NPOL];
-                tmp[iw + row / GlobalV::NPOL] = this->lowf->wfc_k_grid[ik][ib][iw * GlobalV::NPOL + 1];
-            }
-            for (int iw = 0; iw < row; iw++)
-            {
-                this->lowf->wfc_k_grid[ik][ib][iw] = tmp[iw];
-            }
-        }
-    }
-
-    return;
-}
-
 // multi-k case
 template <>
 void ElecStateLCAO<std::complex<double>>::psiToRho(const psi::Psi<std::complex<double>>& psi)
@@ -129,15 +54,6 @@ if(!GlobalV::dm_to_rho)
 #endif
 
     }
-    if (GlobalV::KS_SOLVER == "genelpa" || GlobalV::KS_SOLVER == "scalapack_gvx" || GlobalV::KS_SOLVER == "lapack"
-        || GlobalV::KS_SOLVER == "cusolver" || GlobalV::KS_SOLVER == "cg_in_lcao"  ||  GlobalV::KS_SOLVER == "pexsi")
-    {
-        for (int ik = 0; ik < psi.get_nk(); ik++)
-        {
-            psi.fix_k(ik);
-            this->print_psi(psi);
-        }
-    }
 }
     // old 2D-to-Grid conversion has been replaced by new Gint Refactor 2023/09/25
     //this->loc->cal_dk_k(*this->lowf->gridt, this->wg, (*this->klist));
@@ -202,12 +118,6 @@ void ElecStateLCAO<double>::psiToRho(const psi::Psi<double>& psi)
         for (int ik = 0; ik < psi.get_nk(); ++ik)
         {
             // for gamma_only case, no convertion occured, just for print.
-            if (GlobalV::KS_SOLVER == "genelpa" || GlobalV::KS_SOLVER == "scalapack_gvx"
-                || GlobalV::KS_SOLVER == "cusolver" || GlobalV::KS_SOLVER == "cg_in_lcao")
-            {
-                psi.fix_k(ik);
-                this->print_psi(psi);
-            }
             // old 2D-to-Grid conversion has been replaced by new Gint Refactor 2023/09/25
             if (this->loc->out_dm) // keep interface for old Output_DM until new one is ready
             {

source/module_elecstate/elecstate_lcao.h

@@ -63,9 +63,6 @@ class ElecStateLCAO : public ElecState
     // update charge density for next scf step
     // void getNewRho() override;
 
-    virtual void print_psi(const psi::Psi<TK>& psi_in, const int istep = -1) override;
-    //virtual void print_psi(const psi::Psi<std::complex<double>>& psi_in, const int istep = -1) override;
-
     // initial density matrix
     void init_DM(const K_Vectors* kv, const Parallel_Orbitals* paraV, const int nspin);
     DensityMatrix<TK,double>* get_DM() const { return const_cast<DensityMatrix<TK,double>*>(this->DM); } 

source/module_elecstate/elecstate_lcao_tddft.cpp

@@ -42,14 +42,6 @@ void ElecStateLCAO_TDDFT::psiToRho_td(const psi::Psi<std::complex<double>>& psi)
 #endif
     }
 
-    if (GlobalV::KS_SOLVER == "genelpa" || GlobalV::KS_SOLVER == "scalapack_gvx" || GlobalV::KS_SOLVER == "lapack")
-    {
-        for (int ik = 0; ik < psi.get_nk(); ik++)
-        {
-            psi.fix_k(ik);
-            this->print_psi(psi);
-        }
-    }
 
     for (int is = 0; is < GlobalV::NSPIN; is++)
     {

source/module_esolver/esolver_ks_lcao.cpp

@@ -42,6 +42,8 @@
 
 #include "module_hamilt_lcao/module_deltaspin/spin_constrain.h"
 
+#include "module_io/write_wfc_nao.h"
+
 namespace ModuleESolver
 {
 
@@ -980,25 +982,17 @@ void ESolver_KS_LCAO<TK, TR>::update_pot(const int istep, const int iter)
     }
 
     // 2) print wavefunctions
-    if (this->conv_elec)
-    {
-        if (elecstate::ElecStateLCAO<TK>::out_wfc_lcao)
-        {
-            elecstate::ElecStateLCAO<TK>::out_wfc_flag = elecstate::ElecStateLCAO<TK>::out_wfc_lcao;
-        }
-
-        for (int ik = 0; ik < this->kv.get_nks(); ik++)
-        {
-            if (istep % GlobalV::out_interval == 0)
-            {
-                this->psi[0].fix_k(ik);
-                this->pelec->print_psi(this->psi[0], istep);
-            }
-        }
-        if (elecstate::ElecStateLCAO<TK>::out_wfc_lcao)
-        {
-            elecstate::ElecStateLCAO<TK>::out_wfc_flag = 0;
-        }
+    if (elecstate::ElecStateLCAO<TK>::out_wfc_lcao && 
+        (this->conv_elec || iter == GlobalV::SCF_NMAX) &&
+        (istep % GlobalV::out_interval == 0))
+    {
+            ModuleIO::write_wfc_nao(elecstate::ElecStateLCAO<TK>::out_wfc_lcao,
+                           this->psi[0],
+                           this->pelec->ekb,
+                           this->pelec->wg,
+                           this->pelec->klist->kvec_c,
+                           this->orb_con.ParaV,
+                           istep);
     }
 
     // 3) print potential

source/module_esolver/esolver_ks_lcao_tddft.cpp

@@ -7,6 +7,7 @@
 #include "module_io/td_current_io.h"
 #include "module_io/write_HS.h"
 #include "module_io/write_HS_R.h"
+#include "module_io/write_wfc_nao.h"
 
 //--------------temporary----------------------------
 #include "module_base/blas_connector.h"
@@ -284,22 +285,17 @@ void ESolver_KS_LCAO_TDDFT::update_pot(const int istep, const int iter)
         }
     }
 
-    if (this->conv_elec)
+    if (elecstate::ElecStateLCAO<std::complex<double>>::out_wfc_lcao &&
+        (this->conv_elec || iter == GlobalV::SCF_NMAX) &&
+        (istep % GlobalV::out_interval == 0) )
     {
-        if (elecstate::ElecStateLCAO<std::complex<double>>::out_wfc_lcao)
-        {
-            elecstate::ElecStateLCAO<std::complex<double>>::out_wfc_flag
-                = elecstate::ElecStateLCAO<std::complex<double>>::out_wfc_lcao;
-        }
-        for (int ik = 0; ik < kv.get_nks(); ik++)
-        {
-            if (istep % GlobalV::out_interval == 0)
-            {
-                this->psi[0].fix_k(ik);
-                this->pelec->print_psi(this->psi[0], istep);
-            }
-        }
-        elecstate::ElecStateLCAO<std::complex<double>>::out_wfc_flag = 0;
+            ModuleIO::write_wfc_nao(elecstate::ElecStateLCAO<std::complex<double>>::out_wfc_lcao,
+                           this->psi[0],
+                           this->pelec->ekb,
+                           this->pelec->wg,
+                           this->pelec->klist->kvec_c,
+                           this->orb_con.ParaV,
+                           istep);
     }
 
     // Calculate new potential according to new Charge Density