Feature: binary format of backup charge density (#5782)

* Feature: binary format of backup charge density * Tests: update tests * move backup charge density output to esolver_fp
deepmodeling · Jan 1, 2025 · cd48c62 · cd48c62
1 parent 6fadfb9
commit cd48c62
Show file tree

Hide file tree

Showing 10 changed files with 67 additions and 74 deletions.
diff --git a/docs/advanced/input_files/input-main.md b/docs/advanced/input_files/input-main.md
@@ -1561,10 +1561,8 @@ These variables are used to control the output of properties.
 ### out_freq_elec
 
 - **Type**: Integer
-- **Description**: The output frequency of the charge density (controlled by [out_chg](#out_chg)), wavefunction (controlled by [out_wfc_pw](#out_wfc_pw) or [out_wfc_r](#out_wfc_r)), and density matrix of localized orbitals (controlled by [out_dm](#out_dm)).
-  - \>0: Output them every `out_freq_elec` iteration numbers in electronic iterations.
-  - 0: Output them when the electronic iteration is converged or reaches the maximal iteration number.
-- **Default**: 0
+- **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).
+- **Default**: [scf_nmax](#scf_nmax)
 
 ### out_chg
 

diff --git a/source/module_esolver/esolver_fp.cpp b/source/module_esolver/esolver_fp.cpp
@@ -183,24 +183,6 @@ void ESolver_FP::after_scf(UnitCell& ucell, const int istep)
                 }
             }
         }
-        if (PARAM.inp.out_chg[0] != -1)
-        {
-            std::complex<double>** rhog_tot = (PARAM.inp.dm_to_rho)? this->pelec->charge->rhog : this->pelec->charge->rhog_save;
-            double** rhor_tot = (PARAM.inp.dm_to_rho)? this->pelec->charge->rho : this->pelec->charge->rho_save;
-            for (int is = 0; is < PARAM.inp.nspin; is++)
-            {
-                this->pw_rhod->real2recip(rhor_tot[is], rhog_tot[is]);
-            }
-            ModuleIO::write_rhog(PARAM.globalv.global_out_dir + PARAM.inp.suffix + "-CHARGE-DENSITY.restart",
-                                 PARAM.globalv.gamma_only_pw || PARAM.globalv.gamma_only_local,
-                                 this->pw_rhod,
-                                 PARAM.inp.nspin,
-                                 ucell.GT,
-                                 rhog_tot,
-                                 GlobalV::MY_POOL,
-                                 GlobalV::RANK_IN_POOL,
-                                 GlobalV::NPROC_IN_POOL);
-        }
 
         // 4) write potential
         if (PARAM.inp.out_pot == 1 || PARAM.inp.out_pot == 3)
@@ -304,4 +286,51 @@ void ESolver_FP::before_scf(UnitCell& ucell, const int istep)
     return;
 }
 
+void ESolver_FP::iter_finish(UnitCell& ucell, const int istep, int& iter)
+{
+    //! output charge density
+    if (PARAM.inp.out_chg[0] != -1)
+    {
+        if (iter % PARAM.inp.out_freq_elec == 0 || iter == PARAM.inp.scf_nmax || this->conv_esolver)
+        {
+            std::complex<double>** rhog_tot
+                = (PARAM.inp.dm_to_rho) ? this->pelec->charge->rhog : this->pelec->charge->rhog_save;
+            double** rhor_tot = (PARAM.inp.dm_to_rho) ? this->pelec->charge->rho : this->pelec->charge->rho_save;
+            for (int is = 0; is < PARAM.inp.nspin; is++)
+            {
+                this->pw_rhod->real2recip(rhor_tot[is], rhog_tot[is]);
+            }
+            ModuleIO::write_rhog(PARAM.globalv.global_out_dir + PARAM.inp.suffix + "-CHARGE-DENSITY.restart",
+                                 PARAM.globalv.gamma_only_pw || PARAM.globalv.gamma_only_local,
+                                 this->pw_rhod,
+                                 PARAM.inp.nspin,
+                                 ucell.GT,
+                                 rhog_tot,
+                                 GlobalV::MY_POOL,
+                                 GlobalV::RANK_IN_POOL,
+                                 GlobalV::NPROC_IN_POOL);
+
+            if (XC_Functional::get_func_type() == 3 || XC_Functional::get_func_type() == 5)
+            {
+                std::vector<std::complex<double>> kin_g_space(PARAM.inp.nspin * this->pelec->charge->ngmc, {0.0, 0.0});
+                std::vector<std::complex<double>*> kin_g;
+                for (int is = 0; is < PARAM.inp.nspin; is++)
+                {
+                    kin_g.push_back(kin_g_space.data() + is * this->pelec->charge->ngmc);
+                    this->pw_rhod->real2recip(this->pelec->charge->kin_r_save[is], kin_g[is]);
+                }
+                ModuleIO::write_rhog(PARAM.globalv.global_out_dir + PARAM.inp.suffix + "-TAU-DENSITY.restart",
+                                     PARAM.globalv.gamma_only_pw || PARAM.globalv.gamma_only_local,
+                                     this->pw_rhod,
+                                     PARAM.inp.nspin,
+                                     ucell.GT,
+                                     kin_g.data(),
+                                     GlobalV::MY_POOL,
+                                     GlobalV::RANK_IN_POOL,
+                                     GlobalV::NPROC_IN_POOL);
+            }
+        }
+    }
+}
+
 } // namespace ModuleESolver
diff --git a/source/module_esolver/esolver_fp.h b/source/module_esolver/esolver_fp.h
@@ -41,6 +41,9 @@ class ESolver_FP : public ESolver
     //! Something to do after SCF iterations when SCF is converged or comes to the max iter step.
     virtual void after_scf(UnitCell& ucell, const int istep);
 
+    //! Something to do after hamilt2density function in each iter loop.
+    virtual void iter_finish(UnitCell& ucell, const int istep, int& iter);
+
     //! ------------------------------------------------------------------------------
     //! These pointers will be deleted in the free_pointers() function every ion step.
     //! ------------------------------------------------------------------------------

diff --git a/source/module_esolver/esolver_ks.cpp b/source/module_esolver/esolver_ks.cpp
@@ -46,9 +46,6 @@ ESolver_KS<T, Device>::ESolver_KS()
     maxniter = PARAM.inp.scf_nmax;
     niter = maxniter;
 
-    // should not use GlobalV here, mohan 2024-05-12
-    out_freq_elec = PARAM.inp.out_freq_elec;
-
     // pw_rho = new ModuleBase::PW_Basis();
     // temporary, it will be removed
     std::string fft_device = PARAM.inp.device;
@@ -693,45 +690,7 @@ void ESolver_KS<T, Device>::iter_finish(UnitCell& ucell, const int istep, int& i
         std::cout << " SCF restart after this step!" << std::endl;
     }
 
-    //! output charge density and density matrix
-    if (this->out_freq_elec && iter % this->out_freq_elec == 0)
-    {
-        for (int is = 0; is < PARAM.inp.nspin; is++)
-        {
-            double* data = nullptr;
-            if (PARAM.inp.dm_to_rho)
-            {
-                data = this->pelec->charge->rho[is];
-            }
-            else
-            {
-                data = this->pelec->charge->rho_save[is];
-            }
-            std::string fn = PARAM.globalv.global_out_dir + "/tmp_SPIN" + std::to_string(is + 1) + "_CHG.cube";
-            ModuleIO::write_vdata_palgrid(Pgrid,
-                                          data,
-                                          is,
-                                          PARAM.inp.nspin,
-                                          0,
-                                          fn,
-                                          this->pelec->eferm.get_efval(is),
-                                          &(ucell),
-                                          3,
-                                          1);
-            if (XC_Functional::get_func_type() == 3 || XC_Functional::get_func_type() == 5)
-            {
-                fn = PARAM.globalv.global_out_dir + "/tmp_SPIN" + std::to_string(is + 1) + "_TAU.cube";
-                ModuleIO::write_vdata_palgrid(Pgrid,
-                                              this->pelec->charge->kin_r_save[is],
-                                              is,
-                                              PARAM.inp.nspin,
-                                              0,
-                                              fn,
-                                              this->pelec->eferm.get_efval(is),
-                                              &(ucell));
-            }
-        }
-    }
+    ESolver_FP::iter_finish(ucell, istep, iter);
 }
 
 //! Something to do after SCF iterations when SCF is converged or comes to the max iter step.

diff --git a/source/module_esolver/esolver_ks.h b/source/module_esolver/esolver_ks.h
@@ -95,7 +95,6 @@ class ESolver_KS : public ESolver_FP
     double hsolver_error;           //! the error of HSolver
     int maxniter;                   //! maximum iter steps for scf
     int niter;                      //! iter steps actually used in scf
-    int out_freq_elec;              //! frequency for output
 };
 } // namespace ModuleESolver
 #endif
diff --git a/source/module_esolver/esolver_ks_pw.cpp b/source/module_esolver/esolver_ks_pw.cpp
@@ -619,10 +619,10 @@ void ESolver_KS_PW<T, Device>::iter_finish(UnitCell& ucell, const int istep, int
         this->ppcell.cal_effective_D(veff, this->pw_rhod, ucell);
     }
 
-    if (this->out_freq_elec && iter % this->out_freq_elec == 0)
+    // 4) Print out electronic wavefunctions
+    if (PARAM.inp.out_wfc_pw == 1 || PARAM.inp.out_wfc_pw == 2)
     {
-        // 4) Print out electronic wavefunctions
-        if (PARAM.inp.out_wfc_pw == 1 || PARAM.inp.out_wfc_pw == 2)
+        if (iter % PARAM.inp.out_freq_elec == 0 || iter == PARAM.inp.scf_nmax || this->conv_esolver)
         {
             std::stringstream ssw;
             ssw << PARAM.globalv.global_out_dir << "WAVEFUNC";

diff --git a/source/module_esolver/esolver_of.cpp b/source/module_esolver/esolver_of.cpp
@@ -182,6 +182,8 @@ void ESolver_OF::runner(UnitCell& ucell, const int istep)
         this->update_rho();
 
         this->iter_++;
+
+        ESolver_FP::iter_finish(ucell, istep, this->iter_);
     }
 
     this->after_opt(istep, ucell);

diff --git a/source/module_io/read_input_item_output.cpp b/source/module_io/read_input_item_output.cpp
@@ -21,9 +21,13 @@ void ReadInput::item_output()
     }
     {
         Input_Item item("out_freq_elec");
-        item.annotation = "the frequency ( >= 0) of electronic iter to output "
-                          "charge density and wavefunction. 0: "
-                          "output only when converged";
+        item.annotation = "the frequency of electronic iter to output charge density and wavefunction ";
+        item.reset_value = [](const Input_Item& item, Parameter& para) {
+            if (para.input.out_freq_elec <= 0)
+            {
+                para.input.out_freq_elec = para.input.scf_nmax;
+            }
+        };
         read_sync_int(input.out_freq_elec);
         this->add_item(item);
     }

diff --git a/source/module_io/test/read_input_ptest.cpp b/source/module_io/test/read_input_ptest.cpp
@@ -184,7 +184,7 @@ TEST_F(InputParaTest, ParaRead)
     EXPECT_EQ(param.inp.printe, 100);
     EXPECT_EQ(param.inp.init_chg, "atomic");
     EXPECT_EQ(param.inp.chg_extrap, "atomic");
-    EXPECT_EQ(param.inp.out_freq_elec, 0);
+    EXPECT_EQ(param.inp.out_freq_elec, 50);
     EXPECT_EQ(param.inp.out_freq_ion, 0);
     EXPECT_EQ(param.inp.out_chg[0], 0);
     EXPECT_EQ(param.inp.out_chg[1], 3);

diff --git a/source/module_parameter/input_parameter.h b/source/module_parameter/input_parameter.h
@@ -316,8 +316,7 @@ struct Input_para
     std::vector<int> aims_nbasis = {};  ///< the number of basis functions for each atom type used in FHI-aims (for benchmark)
     // ==============   #Parameters (11.Output) ===========================
     bool out_stru = false;                ///< outut stru file each ion step
-    int out_freq_elec = 0;                ///< the frequency ( >= 0) of electronic iter to output charge
-                                          ///< 0: output only when converged
+    int out_freq_elec = 0;                ///< the frequency of electronic iter to output charge and wavefunction
     int out_freq_ion = 0;                 ///< the frequency ( >= 0 ) of ionic step to output charge density;
                                           ///< 0: output only when ion steps are finished
     std::vector<int> out_chg = {0, 3};    ///< output charge density. 0: no; 1: yes