Fix: optimize the IO of lcao-wfc and charge

source/module_elecstate/module_charge/charge_init.cpp

@@ -23,26 +23,8 @@ void Charge::init_rho(elecstate::efermi& eferm_iout, const ModuleBase::ComplexMa
     ModuleBase::GlobalFunc::OUT(GlobalV::ofs_running, "init_chg", GlobalV::init_chg);
 
     std::cout << " START CHARGE      : " << GlobalV::init_chg << std::endl;
-    if (GlobalV::init_chg == "atomic") // mohan add 2007-10-17
-    {
-        this->atomic_rho(GlobalV::NSPIN, GlobalC::ucell.omega, rho, strucFac, GlobalC::ucell);
-
-        // liuyu 2023-06-29 : move here from atomic_rho(), which will be called several times in charge extrapolation
-        // wenfei 2021-7-29 : initial tau = 3/5 rho^2/3, Thomas-Fermi
-        if (XC_Functional::get_func_type() == 3 || XC_Functional::get_func_type() == 5)
-        {
-            const double pi = 3.141592653589790;
-            const double fact = (3.0 / 5.0) * pow(3.0 * pi * pi, 2.0 / 3.0);
-            for (int is = 0; is < GlobalV::NSPIN; ++is)
-            {
-                for (int ir = 0; ir < this->rhopw->nrxx; ++ir)
-                {
-                    kin_r[is][ir] = fact * pow(std::abs(rho[is][ir]) * GlobalV::NSPIN, 5.0 / 3.0) / GlobalV::NSPIN;
-                }
-            }
-        }
-    }
-    else if (GlobalV::init_chg == "file")
+    bool read_error = false;
+    if (GlobalV::init_chg == "file")
     {
         GlobalV::ofs_running << " try to read charge from file : ";
         for (int is = 0; is < GlobalV::NSPIN; ++is)
@@ -84,7 +66,9 @@ void Charge::init_rho(elecstate::efermi& eferm_iout, const ModuleBase::ComplexMa
                 { // read up and down , then rearrange them.
                     if (is == 1)
                     {
-                        ModuleBase::WARNING_QUIT("Charge::init_rho", "Incomplete charge density file!");
+                        std::cout << "Incomplete charge density file!" << std::endl;
+                        read_error = true;
+                        break;
                     }
                     else if (is == 2)
                     {
@@ -106,10 +90,8 @@ void Charge::init_rho(elecstate::efermi& eferm_iout, const ModuleBase::ComplexMa
             }
             else
             {
-                ModuleBase::WARNING_QUIT(
-                    "init_rho",
-                    "!!! Couldn't find the charge file !!! The default directory \n of SPIN1_CHG.cube is OUT.suffix, "
-                    "or you must set read_file_dir \n to a specific directory. ");
+                read_error = true;
+                break;
             }
         }
 
@@ -141,9 +123,33 @@ void Charge::init_rho(elecstate::efermi& eferm_iout, const ModuleBase::ComplexMa
             }
         }
     }
-    else
+    if (read_error)
     {
-        ModuleBase::WARNING_QUIT("Charge::init_rho", "init_chg is wrong!");
+        std::cout << " WARNING: Failed to read charge density from file. Possible reasons: " << std::endl;
+        std::cout << " - not found: The default directory of SPIN1_CHG.cube is OUT.suffix, \n"
+            "or you must set read_file_dir to a specific directory. " << std::endl;
+        std::cout << " - parameter mismatch: check the previous warning." << std::endl;
+        std::cout << " Use atomic initialization instead." << std::endl;
+    }
+
+    if (GlobalV::init_chg != "file" || read_error) // mohan add 2007-10-17
+    {
+        this->atomic_rho(GlobalV::NSPIN, GlobalC::ucell.omega, rho, strucFac, GlobalC::ucell);
+
+        // liuyu 2023-06-29 : move here from atomic_rho(), which will be called several times in charge extrapolation
+        // wenfei 2021-7-29 : initial tau = 3/5 rho^2/3, Thomas-Fermi
+        if (XC_Functional::get_func_type() == 3 || XC_Functional::get_func_type() == 5)
+        {
+            const double pi = 3.141592653589790;
+            const double fact = (3.0 / 5.0) * pow(3.0 * pi * pi, 2.0 / 3.0);
+            for (int is = 0; is < GlobalV::NSPIN; ++is)
+            {
+                for (int ir = 0; ir < this->rhopw->nrxx; ++ir)
+                {
+                    kin_r[is][ir] = fact * pow(std::abs(rho[is][ir]) * GlobalV::NSPIN, 5.0 / 3.0) / GlobalV::NSPIN;
+                }
+            }
+        }
     }
 
     // Peize Lin add 2020.04.04

source/module_esolver/esolver_ks_lcao.cpp

@@ -486,7 +486,7 @@ namespace ModuleESolver
     // first need to calculate the weight according to
     // electrons number.
 
-    if (this->wf.init_wfc == "file")
+    if (istep == 0 && this->wf.init_wfc == "file" && this->LOWF.error == 0)
     {
         if (iter == 1)
         {

source/module_esolver/esolver_ks_lcao_elec.cpp

@@ -146,7 +146,7 @@ void ESolver_KS_LCAO<TK, TR>::beforesolver(const int istep)
                                                         DM);
     }
     // init density kernel and wave functions.
-    this->LOC.allocate_dm_wfc(this->GridT, this->pelec, this->LOWF, this->psi, this->kv);
+    this->LOC.allocate_dm_wfc(this->GridT, this->pelec, this->LOWF, this->psi, this->kv, istep);
 
     //======================================
     // do the charge extrapolation before the density matrix is regenerated.

source/module_hamilt_lcao/hamilt_lcaodft/DM_gamma.cpp

@@ -10,10 +10,11 @@
 // allocate density kernel may change once the ion
 // positions change
 void Local_Orbital_Charge::allocate_gamma(
-                const int& lgd, 
-                psi::Psi<double>* psid, 
-                elecstate::ElecState* pelec,
-                const int& nks)
+    const int& lgd,
+    psi::Psi<double>* psid,
+    elecstate::ElecState* pelec,
+    const int& nks,
+    const int& istep)
 {
      ModuleBase::TITLE("Local_Orbital_Charge","allocate_gamma");
 
@@ -77,71 +78,13 @@ void Local_Orbital_Charge::allocate_gamma(
 	// Peize Lin test 2019-01-16
     this->init_dm_2d(nks);
 
-    if (INPUT.init_wfc == "file")
+    if (istep == 0 && INPUT.init_wfc == "file")
     {
-        this->gamma_file(psid, this->LOWF[0], pelec);
+        this->LOWF->gamma_file(psid, pelec);
     }
     return;
 }
 
-void Local_Orbital_Charge::gamma_file(psi::Psi<double>* psid, Local_Orbital_wfc &lowf, elecstate::ElecState* pelec)
-{
-	ModuleBase::TITLE("Local_Orbital_Charge","gamma_file");
-
-	int error;
-	std::cout << " Read in gamma point wave function files " << std::endl;
-
-	double **ctot;
-
-    //allocate psi
-    int ncol = this->ParaV->ncol_bands;
-    if(GlobalV::KS_SOLVER=="genelpa" || GlobalV::KS_SOLVER=="lapack_gvx" || GlobalV::KS_SOLVER == "scalapack_gvx"
-#ifdef __CUSOLVER_LCAO
-    ||GlobalV::KS_SOLVER=="cusolver"
-#endif
-    )
-    {
-        ncol = this->ParaV->ncol;
-    }
-    if(psid == nullptr)
-    {
-        ModuleBase::WARNING_QUIT("gamma_file", "psid should be allocated first!");
-    }
-    else
-    {
-        psid->resize(GlobalV::NSPIN, ncol, this->ParaV->nrow);
-    }
-    ModuleBase::GlobalFunc::ZEROS( psid->get_pointer(), psid->size() );
-
-	for(int is=0; is<GlobalV::NSPIN; ++is)
-	{
-
-		GlobalV::ofs_running << " Read in wave functions " << is << std::endl;
-		error = ModuleIO::read_wfc_nao( ctot , is, this->ParaV, psid, pelec);
-#ifdef __MPI
-		Parallel_Common::bcast_int(error);
-#endif
-		GlobalV::ofs_running << " Error=" << error << std::endl;
-		if(error==1)
-		{
-			ModuleBase::WARNING_QUIT("Local_Orbital_wfc","Can't find the wave function file: LOWF.dat");
-		}
-		else if(error==2)
-		{
-			ModuleBase::WARNING_QUIT("Local_Orbital_wfc","In wave function file, band number doesn't match");
-		}
-		else if(error==3)
-		{
-			ModuleBase::WARNING_QUIT("Local_Orbital_wfc","In wave function file, nlocal doesn't match");
-		}
-		else if(error==4)
-		{
-			ModuleBase::WARNING_QUIT("Local_Orbital_wfc","In k-dependent wave function file, k point is not correct");
-		}
-
-	}//loop ispin
-}
-
 void Local_Orbital_Charge::cal_dk_gamma_from_2D_pub(void)
 {
     ModuleBase::TITLE("Local_Orbital_Charge","cal_dk_gamma_from_2D_pub");

source/module_hamilt_lcao/hamilt_lcaodft/local_orbital_charge.cpp

@@ -61,28 +61,30 @@ void Local_Orbital_Charge::allocate_dm_wfc(const Grid_Technique& gt,
     elecstate::ElecState* pelec,
     Local_Orbital_wfc& lowf,
     psi::Psi<double>* psi,
-    const K_Vectors& kv)
+    const K_Vectors& kv,
+    const int& istep)
 {
     ModuleBase::TITLE("Local_Orbital_Charge", "allocate_dm_wfc");
     this->LOWF = &lowf;
     this->LOWF->gridt = &gt;
     // here we reset the density matrix dimension.
-    this->allocate_gamma(gt.lgd, psi, pelec, kv.nks);
+    this->allocate_gamma(gt.lgd, psi, pelec, kv.nks, istep);
     return;
 }
 
 void Local_Orbital_Charge::allocate_dm_wfc(const Grid_Technique &gt,
     elecstate::ElecState* pelec,
     Local_Orbital_wfc& lowf,
     psi::Psi<std::complex<double>>* psi,
-    const K_Vectors& kv)
+    const K_Vectors& kv,
+    const int& istep)
 {
     ModuleBase::TITLE("Local_Orbital_Charge", "allocate_dm_wfc");
 
     this->LOWF = &lowf;
     this->LOWF->gridt = &gt;
     // here we reset the density matrix dimension.
-    lowf.allocate_k(gt.lgd, psi, pelec, kv.nks, kv.nkstot, kv.kvec_c);
+    lowf.allocate_k(gt.lgd, psi, pelec, kv.nks, kv.nkstot, kv.kvec_c, istep);
     this->allocate_DM_k(kv.nks, gt.nnrg);
     return;
 }

source/module_hamilt_lcao/hamilt_lcaodft/local_orbital_charge.h

@@ -25,18 +25,18 @@ class Local_Orbital_Charge
         elecstate::ElecState* pelec,
         Local_Orbital_wfc &lowf,
         psi::Psi<double>* psi,
-        const K_Vectors& kv);
+        const K_Vectors& kv,
+        const int& istep);
     void allocate_dm_wfc(const Grid_Technique& gt,
         elecstate::ElecState* pelec,
         Local_Orbital_wfc& lowf,
         psi::Psi<std::complex<double>>* psi,
-        const K_Vectors& kv);
-	//-----------------
+        const K_Vectors& kv,
+        const int& istep);
+    //-----------------
 	// in DM_gamma.cpp
 	//-----------------
-	void allocate_gamma(const int &lgd, psi::Psi<double>* psid, elecstate::ElecState* pelec, const int& nks);
-
-    void gamma_file(psi::Psi<double>* psid, Local_Orbital_wfc &lowf, elecstate::ElecState* pelec);
+    void allocate_gamma(const int& lgd, psi::Psi<double>* psid, elecstate::ElecState* pelec, const int& nks, const int& istep);
     void cal_dk_gamma_from_2D_pub(void);
 
 	//-----------------

source/module_hamilt_lcao/hamilt_lcaodft/local_orbital_wfc.cpp

@@ -31,12 +31,71 @@ Local_Orbital_wfc::~Local_Orbital_wfc()
 	}
 }
 
+void Local_Orbital_wfc::gamma_file(psi::Psi<double>* psid, elecstate::ElecState* pelec)
+{
+    ModuleBase::TITLE("Local_Orbital_Charge", "gamma_file");
+    std::cout << " Read in gamma point wave function files " << std::endl;
+
+    double** ctot;
+
+    //allocate psi
+    int ncol = this->ParaV->ncol_bands;
+    if (GlobalV::KS_SOLVER == "genelpa" || GlobalV::KS_SOLVER == "lapack_gvx" || GlobalV::KS_SOLVER == "scalapack_gvx"
+#ifdef __CUSOLVER_LCAO
+        || GlobalV::KS_SOLVER == "cusolver"
+#endif
+        )
+    {
+        ncol = this->ParaV->ncol;
+    }
+    if (psid == nullptr)
+    {
+        ModuleBase::WARNING_QUIT("gamma_file", "psid should be allocated first!");
+    }
+    else
+    {
+        psid->resize(GlobalV::NSPIN, ncol, this->ParaV->nrow);
+    }
+    ModuleBase::GlobalFunc::ZEROS(psid->get_pointer(), psid->size());
+
+    for (int is = 0; is < GlobalV::NSPIN; ++is)
+    {
+        this->error = ModuleIO::read_wfc_nao(ctot, is, this->ParaV, psid, pelec);
+#ifdef __MPI
+        Parallel_Common::bcast_int(this->error);
+#endif
+        switch (this->error)
+        {
+        case 1:
+            std::cout << "Can't find the wave function file: LOWF_GAMMA_S" << is + 1 << ".txt" << std::endl;
+            break;
+        case 2:
+            std::cout << "In wave function file, band number doesn't match" << std::endl;
+            break;
+        case 3:
+            std::cout << "In wave function file, nlocal doesn't match" << std::endl;
+            break;
+        case 4:
+            std::cout << "In k-dependent wave function file, k point is not correct" << std::endl;
+            break;
+        default:
+            std::cout << " Successfully read in wave functions " << is << std::endl;
+        }
+        if (this->error)
+        {
+            std::cout << "WARNING: Failed to read in wavefunction, use default initialization instead." << std::endl;
+            break;
+        }
+    }//loop ispin
+}
+
 void Local_Orbital_wfc::allocate_k(const int& lgd,
     psi::Psi<std::complex<double>>* psi,
     elecstate::ElecState* pelec,
     const int& nks,
     const int& nkstot,
-    const std::vector<ModuleBase::Vector3<double>>& kvec_c)
+    const std::vector<ModuleBase::Vector3<double>>& kvec_c,
+    const int& istep)
 {
     this->nks = nks;
 
@@ -92,40 +151,44 @@ void Local_Orbital_wfc::allocate_k(const int& lgd,
     }
     else if (INPUT.init_wfc == "file")
     {
-        int error;
+        if (istep > 0)return;
         std::cout << " Read in wave functions files: " << nkstot << std::endl;
-        if(psi == nullptr)
+        if (psi == nullptr)
         {
-            ModuleBase::WARNING_QUIT("allocate_k","psi should be allocated first!");
+            ModuleBase::WARNING_QUIT("allocate_k", "psi should be allocated first!");
         }
         else
         {
             psi->resize(nkstot, this->ParaV->ncol_bands, this->ParaV->nrow);
         }
-		for(int ik=0; ik<nkstot; ++ik)
-		{
-            GlobalV::ofs_running << " Read in wave functions " << ik + 1 << std::endl;
+        for (int ik = 0; ik < nkstot; ++ik)
+        {
             std::complex<double>** ctot;
-            error = ModuleIO::read_wfc_nao_complex(ctot, ik, kvec_c[ik], this->ParaV, psi, pelec);
+            this->error = ModuleIO::read_wfc_nao_complex(ctot, ik, kvec_c[ik], this->ParaV, psi, pelec);
 #ifdef __MPI
-            Parallel_Common::bcast_int(error);
+            Parallel_Common::bcast_int(this->error);
 #endif
-            GlobalV::ofs_running << " Error=" << error << std::endl;
-            if(error==1)
-            {
-                ModuleBase::WARNING_QUIT("Local_Orbital_wfc","Can't find the wave function file: LOWF.dat");
-            }
-            else if(error==2)
-            {
-                ModuleBase::WARNING_QUIT("Local_Orbital_wfc","In wave function file, band number doesn't match");
-            }
-            else if(error==3)
+            switch (this->error)
             {
-                ModuleBase::WARNING_QUIT("Local_Orbital_wfc","In wave function file, nlocal doesn't match");
+            case 1:
+                std::cout << "Can't find the wave function file: LOWF_K_" << ik + 1 << ".txt" << std::endl;
+                break;
+            case 2:
+                std::cout << "In wave function file, band number doesn't match" << std::endl;
+                break;
+            case 3:
+                std::cout << "In wave function file, nlocal doesn't match" << std::endl;
+                break;
+            case 4:
+                std::cout << "In k-dependent wave function file, k point is not correct" << std::endl;
+                break;
+            default:
+                std::cout << " Successfully read in wave functions " << ik + 1 << std::endl;
             }
-            else if(error==4)
+            if (this->error)
             {
-                ModuleBase::WARNING_QUIT("Local_Orbital_wfc","In k-dependent wave function file, k point is not correct");
+                std::cout << "WARNING: Failed to read in wavefunction, use default initialization instead." << std::endl;
+                break;
             }
         }
     }