Feature: support user defined precision of elf

sunliang98 · sunliang98 · commit 8cc9014b3fa8 · 2024-09-20T17:40:21.000+08:00
diff --git a/source/module_elecstate/elecstate_pw.cpp b/source/module_elecstate/elecstate_pw.cpp
@@ -34,7 +34,7 @@ ElecStatePW<T, Device>::~ElecStatePW()
     if (base_device::get_device_type<Device>(this->ctx) == base_device::GpuDevice)
     {
         delmem_var_op()(this->ctx, this->rho_data);
-        if (get_xc_func_type() == 3 || PARAM.inp.out_elf)
+        if (get_xc_func_type() == 3 || PARAM.inp.out_elf[0] > 0)
         {
             delmem_var_op()(this->ctx, this->kin_r_data);
         }
@@ -53,7 +53,7 @@ void ElecStatePW<T, Device>::init_rho_data()
         for (int ii = 0; ii < this->charge->nspin; ii++) {
             this->rho[ii] = this->rho_data + ii * this->charge->nrxx;
         }
-        if (get_xc_func_type() == 3 || PARAM.inp.out_elf)
+        if (get_xc_func_type() == 3 || PARAM.inp.out_elf[0] > 0)
         {
             this->kin_r = new Real*[this->charge->nspin];
             resmem_var_op()(this->ctx, this->kin_r_data, this->charge->nspin * this->charge->nrxx);
@@ -64,7 +64,7 @@ void ElecStatePW<T, Device>::init_rho_data()
     }
     else {
         this->rho = reinterpret_cast<Real **>(this->charge->rho);
-        if (get_xc_func_type() == 3 || PARAM.inp.out_elf)
+        if (get_xc_func_type() == 3 || PARAM.inp.out_elf[0] > 0)
         {
             this->kin_r = reinterpret_cast<Real **>(this->charge->kin_r);
         }
diff --git a/source/module_elecstate/module_charge/charge.cpp b/source/module_elecstate/module_charge/charge.cpp
@@ -80,7 +80,7 @@ void Charge::destroy()
         delete[] _space_rhog_save;
         delete[] _space_kin_r;
         delete[] _space_kin_r_save;
-        if (elecstate::get_xc_func_type() == 3 || elecstate::get_xc_func_type() == 5 || PARAM.inp.out_elf)
+        if (elecstate::get_xc_func_type() == 3 || elecstate::get_xc_func_type() == 5 || PARAM.inp.out_elf[0] > 0)
         {
             delete[] kin_r;
             delete[] kin_r_save;
@@ -121,7 +121,7 @@ void Charge::allocate(const int& nspin_in)
     _space_rho_save = new double[nspin * nrxx];
     _space_rhog = new std::complex<double>[nspin * ngmc];
     _space_rhog_save = new std::complex<double>[nspin * ngmc];
-    if (elecstate::get_xc_func_type() == 3 || elecstate::get_xc_func_type() == 5 || PARAM.inp.out_elf)
+    if (elecstate::get_xc_func_type() == 3 || elecstate::get_xc_func_type() == 5 || PARAM.inp.out_elf[0] > 0)
     {
         _space_kin_r = new double[nspin * nrxx];
         _space_kin_r_save = new double[nspin * nrxx];
@@ -130,7 +130,7 @@ void Charge::allocate(const int& nspin_in)
     rhog = new std::complex<double>*[nspin];
     rho_save = new double*[nspin];
     rhog_save = new std::complex<double>*[nspin];
-    if (elecstate::get_xc_func_type() == 3 || elecstate::get_xc_func_type() == 5 || PARAM.inp.out_elf)
+    if (elecstate::get_xc_func_type() == 3 || elecstate::get_xc_func_type() == 5 || PARAM.inp.out_elf[0] > 0)
     {
         kin_r = new double*[nspin];
         kin_r_save = new double*[nspin];
@@ -151,7 +151,7 @@ void Charge::allocate(const int& nspin_in)
         ModuleBase::GlobalFunc::ZEROS(rhog[is], ngmc);
         ModuleBase::GlobalFunc::ZEROS(rho_save[is], nrxx);
         ModuleBase::GlobalFunc::ZEROS(rhog_save[is], ngmc);
-        if (elecstate::get_xc_func_type() == 3 || elecstate::get_xc_func_type() == 5 || PARAM.inp.out_elf)
+        if (elecstate::get_xc_func_type() == 3 || elecstate::get_xc_func_type() == 5 || PARAM.inp.out_elf[0] > 0)
         {
             kin_r[is] = _space_kin_r + is * nrxx;
             ModuleBase::GlobalFunc::ZEROS(kin_r[is], nrxx);
@@ -171,7 +171,7 @@ void Charge::allocate(const int& nspin_in)
     ModuleBase::Memory::record("Chg::rho_save", sizeof(double) * nspin * nrxx);
     ModuleBase::Memory::record("Chg::rhog", sizeof(double) * nspin * ngmc);
     ModuleBase::Memory::record("Chg::rhog_save", sizeof(double) * nspin * ngmc);
-    if (elecstate::get_xc_func_type() == 3 || elecstate::get_xc_func_type() == 5 || PARAM.inp.out_elf)
+    if (elecstate::get_xc_func_type() == 3 || elecstate::get_xc_func_type() == 5 || PARAM.inp.out_elf[0] > 0)
     {
         ModuleBase::Memory::record("Chg::kin_r", sizeof(double) * nspin * ngmc);
         ModuleBase::Memory::record("Chg::kin_r_save", sizeof(double) * nspin * ngmc);
diff --git a/source/module_elecstate/module_charge/charge_mpi.cpp b/source/module_elecstate/module_charge/charge_mpi.cpp
@@ -132,7 +132,7 @@ void Charge::rho_mpi()
     for (int is = 0; is < PARAM.inp.nspin; ++is)
     {
         reduce_diff_pools(this->rho[is]);
-        if (elecstate::get_xc_func_type() == 3 || elecstate::get_xc_func_type() == 5 || PARAM.inp.out_elf)
+        if (elecstate::get_xc_func_type() == 3 || elecstate::get_xc_func_type() == 5 || PARAM.inp.out_elf[0] > 0)
         {
             reduce_diff_pools(this->kin_r[is]);
         }
diff --git a/source/module_esolver/esolver_fp.cpp b/source/module_esolver/esolver_fp.cpp
@@ -256,7 +256,7 @@ void ESolver_FP::after_scf(const int istep)
     }
 
     // 5) write ELF
-    if (PARAM.inp.out_elf)
+    if (PARAM.inp.out_elf[0] > 0)
     {
         this->pelec->charge->cal_elf = true;
         Symmetry_rho srho;
@@ -277,7 +277,8 @@ void ESolver_FP::after_scf(const int istep)
             this->pelec->charge->rho,
             this->pelec->charge->kin_r,
             this->pw_rhod,
-            &(GlobalC::ucell));
+            &(GlobalC::ucell),
+            PARAM.inp.out_elf[1]);
     }
 }
 
diff --git a/source/module_esolver/esolver_ks_lcao.cpp b/source/module_esolver/esolver_ks_lcao.cpp
@@ -1097,7 +1097,7 @@ void ESolver_KS_LCAO<TK, TR>::after_scf(const int istep)
 {
     ModuleBase::TITLE("ESolver_KS_LCAO", "after_scf");
     // 1) calculate the kinetic energy density tau, sunliang 2024-09-18
-    if (PARAM.inp.out_elf)
+    if (PARAM.inp.out_elf[0] > 0)
     {
         this->pelec->cal_tau(*(this->psi));
     }
diff --git a/source/module_esolver/esolver_ks_pw.cpp b/source/module_esolver/esolver_ks_pw.cpp
@@ -539,7 +539,7 @@ template <typename T, typename Device>
 void ESolver_KS_PW<T, Device>::after_scf(const int istep)
 {
     // 1) calculate the kinetic energy density tau, sunliang 2024-09-18
-    if (PARAM.inp.out_elf)
+    if (PARAM.inp.out_elf[0] > 0)
     {
         this->pelec->cal_tau(*(this->psi));
     }
diff --git a/source/module_esolver/esolver_of.cpp b/source/module_esolver/esolver_of.cpp
@@ -511,7 +511,7 @@ bool ESolver_OF::check_exit()
 void ESolver_OF::after_opt(const int istep, UnitCell& ucell)
 {
     // 1) calculate the kinetic energy density
-    if (PARAM.inp.out_elf)
+    if (PARAM.inp.out_elf[0] > 0)
     {
         this->kinetic_energy_density(this->pelec->charge->rho, this->pphi_, this->pelec->charge->kin_r);
     }
diff --git a/source/module_io/read_input_item_output.cpp b/source/module_io/read_input_item_output.cpp
@@ -494,14 +494,22 @@ void ReadInput::item_output()
     }
     {
         Input_Item item("out_elf");
-        item.annotation = "output electron localization function (ELF)";
-        read_sync_bool(input.out_elf);
+        item.annotation = "> 0 output electron localization function (ELF) for selected electron steps"
+                          ", second parameter controls the precision, default is 3.";
+        item.read_value = [](const Input_Item& item, Parameter& para) {
+            size_t count = item.get_size();
+            std::vector<int> out_elf(count); // create a placeholder vector
+            std::transform(item.str_values.begin(), item.str_values.end(), out_elf.begin(), [](std::string s) { return std::stoi(s); });
+            // assign non-negative values to para.input.out_elf
+            std::copy(out_elf.begin(), out_elf.end(), para.input.out_elf.begin());
+        };
         item.check_value = [](const Input_Item& item, const Parameter& para) {
-            if (para.input.out_elf && para.input.esolver_type != "ksdft" && para.input.esolver_type != "ofdft")
+            if (para.input.out_elf[0] > 0 && para.input.esolver_type != "ksdft" && para.input.esolver_type != "ofdft")
             {
                 ModuleBase::WARNING_QUIT("ReadInput", "ELF is only aviailable for ksdft and ofdft");
             }
         };
+        sync_intvec(input.out_elf, 2, 0);
         this->add_item(item);
     }
 }
diff --git a/source/module_io/write_elf.cpp b/source/module_io/write_elf.cpp
@@ -14,7 +14,8 @@ void write_elf(
     const double* const* rho,
     const double* const* tau,
     ModulePW::PW_Basis* rho_basis,
-    const UnitCell* ucell_)
+    const UnitCell* ucell_,
+    const int& precision)
 {
     std::vector<std::vector<double>> elf(nspin, std::vector<double>(rho_basis->nrxx, 0.));
     // 1) calculate the kinetic energy density of vW KEDF
@@ -78,7 +79,6 @@ void write_elf(
     }
 
     // 4) output the ELF = 1 / (1 + F^2) to cube file
-    int precision = 9;
     double ef_tmp = 0.0;
     int out_fermi = 0;
 
diff --git a/source/module_io/write_elf.h b/source/module_io/write_elf.h
@@ -18,7 +18,8 @@ void write_elf(
     const double* const* rho,
     const double* const* tau,
     ModulePW::PW_Basis* rho_basis,
-    const UnitCell* ucell_);
+    const UnitCell* ucell_,
+    const int& precision);
 }
 
 #endif
diff --git a/source/module_parameter/input_parameter.h b/source/module_parameter/input_parameter.h
@@ -356,7 +356,7 @@ struct Input_para
     int nbands_istate = 5;                ///< number of bands around fermi level for get_pchg calculation.
     std::vector<int> bands_to_print = {}; ///< specify the bands to be calculated in the get_pchg
     bool if_separate_k = false; ///< whether to write partial charge for all k-points to individual files or merge them
-    bool out_elf = false;       ///< output the electron localization function (ELF)
+    std::vector<int> out_elf = {0, 3};    ///< output the electron localization function (ELF). 0: no; 1: yes
 
     // ==============   #Parameters (12.Postprocess) ===========================
     double dos_emin_ev = -15.0;

Original file line number	Diff line number	Diff line change
`@@ -34,7 +34,7 @@ ElecStatePW<T, Device>::~ElecStatePW()`
`34`	`34`	`if (base_device::get_device_type<Device>(this->ctx) == base_device::GpuDevice)`
`35`	`35`	`{`
`36`	`36`	`delmem_var_op()(this->ctx, this->rho_data);`
`37`		`- if (get_xc_func_type() == 3 \|\| PARAM.inp.out_elf)`
	`37`	`+ if (get_xc_func_type() == 3 \|\| PARAM.inp.out_elf[0] > 0)`
`38`	`38`	`{`
`39`	`39`	`delmem_var_op()(this->ctx, this->kin_r_data);`
`40`	`40`	`}`
`@@ -53,7 +53,7 @@ void ElecStatePW<T, Device>::init_rho_data()`
`53`	`53`	`for (int ii = 0; ii < this->charge->nspin; ii++) {`
`54`	`54`	`this->rho[ii] = this->rho_data + ii * this->charge->nrxx;`
`55`	`55`	`}`
`56`		`- if (get_xc_func_type() == 3 \|\| PARAM.inp.out_elf)`
	`56`	`+ if (get_xc_func_type() == 3 \|\| PARAM.inp.out_elf[0] > 0)`
`57`	`57`	`{`
`58`	`58`	`this->kin_r = new Real*[this->charge->nspin];`
`59`	`59`	`resmem_var_op()(this->ctx, this->kin_r_data, this->charge->nspin * this->charge->nrxx);`
`@@ -64,7 +64,7 @@ void ElecStatePW<T, Device>::init_rho_data()`
`64`	`64`	`}`
`65`	`65`	`else {`
`66`	`66`	`this->rho = reinterpret_cast<Real **>(this->charge->rho);`
`67`		`- if (get_xc_func_type() == 3 \|\| PARAM.inp.out_elf)`
	`67`	`+ if (get_xc_func_type() == 3 \|\| PARAM.inp.out_elf[0] > 0)`
`68`	`68`	`{`
`69`	`69`	`this->kin_r = reinterpret_cast<Real **>(this->charge->kin_r);`
`70`	`70`	`}`
Original file line number	Diff line number	Diff line change
`@@ -132,7 +132,7 @@ void Charge::rho_mpi()`
`132`	`132`	`for (int is = 0; is < PARAM.inp.nspin; ++is)`
`133`	`133`	`{`
`134`	`134`	`reduce_diff_pools(this->rho[is]);`
`135`		`- if (elecstate::get_xc_func_type() == 3 \|\| elecstate::get_xc_func_type() == 5 \|\| PARAM.inp.out_elf)`
	`135`	`+ if (elecstate::get_xc_func_type() == 3 \|\| elecstate::get_xc_func_type() == 5 \|\| PARAM.inp.out_elf[0] > 0)`
`136`	`136`	`{`
`137`	`137`	`reduce_diff_pools(this->kin_r[is]);`
`138`	`138`	`}`
Original file line number	Diff line number	Diff line change
`@@ -256,7 +256,7 @@ void ESolver_FP::after_scf(const int istep)`
`256`	`256`	`}`
`257`	`257`
`258`	`258`	`// 5) write ELF`
`259`		`- if (PARAM.inp.out_elf)`
	`259`	`+ if (PARAM.inp.out_elf[0] > 0)`
`260`	`260`	`{`
`261`	`261`	`this->pelec->charge->cal_elf = true;`
`262`	`262`	`Symmetry_rho srho;`
`@@ -277,7 +277,8 @@ void ESolver_FP::after_scf(const int istep)`
`277`	`277`	`this->pelec->charge->rho,`
`278`	`278`	`this->pelec->charge->kin_r,`
`279`	`279`	`this->pw_rhod,`
`280`		`- &(GlobalC::ucell));`
	`280`	`+ &(GlobalC::ucell),`
	`281`	`+ PARAM.inp.out_elf[1]);`
`281`	`282`	`}`
`282`	`283`	`}`
`283`	`284`
Original file line number	Diff line number	Diff line change
`@@ -1097,7 +1097,7 @@ void ESolver_KS_LCAO<TK, TR>::after_scf(const int istep)`
`1097`	`1097`	`{`
`1098`	`1098`	`ModuleBase::TITLE("ESolver_KS_LCAO", "after_scf");`
`1099`	`1099`	`// 1) calculate the kinetic energy density tau, sunliang 2024-09-18`
`1100`		`- if (PARAM.inp.out_elf)`
	`1100`	`+ if (PARAM.inp.out_elf[0] > 0)`
`1101`	`1101`	`{`
`1102`	`1102`	`this->pelec->cal_tau(*(this->psi));`
`1103`	`1103`	`}`
Original file line number	Diff line number	Diff line change
`@@ -539,7 +539,7 @@ template <typename T, typename Device>`
`539`	`539`	`void ESolver_KS_PW<T, Device>::after_scf(const int istep)`
`540`	`540`	`{`
`541`	`541`	`// 1) calculate the kinetic energy density tau, sunliang 2024-09-18`
`542`		`- if (PARAM.inp.out_elf)`
	`542`	`+ if (PARAM.inp.out_elf[0] > 0)`
`543`	`543`	`{`
`544`	`544`	`this->pelec->cal_tau(*(this->psi));`
`545`	`545`	`}`
Original file line number	Diff line number	Diff line change
`@@ -511,7 +511,7 @@ bool ESolver_OF::check_exit()`
`511`	`511`	`void ESolver_OF::after_opt(const int istep, UnitCell& ucell)`
`512`	`512`	`{`
`513`	`513`	`// 1) calculate the kinetic energy density`
`514`		`- if (PARAM.inp.out_elf)`
	`514`	`+ if (PARAM.inp.out_elf[0] > 0)`
`515`	`515`	`{`
`516`	`516`	`this->kinetic_energy_density(this->pelec->charge->rho, this->pphi_, this->pelec->charge->kin_r);`
`517`	`517`	`}`
Original file line number	Diff line number	Diff line change
`@@ -494,14 +494,22 @@ void ReadInput::item_output()`
`494`	`494`	`}`
`495`	`495`	`{`
`496`	`496`	`Input_Item item("out_elf");`
`497`		`- item.annotation = "output electron localization function (ELF)";`
`498`		`- read_sync_bool(input.out_elf);`
	`497`	`+ item.annotation = "> 0 output electron localization function (ELF) for selected electron steps"`
	`498`	`+ ", second parameter controls the precision, default is 3.";`
	`499`	`+ item.read_value = [](const Input_Item& item, Parameter& para) {`
	`500`	`+ size_t count = item.get_size();`
	`501`	`+ std::vector<int> out_elf(count); // create a placeholder vector`
	`502`	`+ std::transform(item.str_values.begin(), item.str_values.end(), out_elf.begin(), [](std::string s) { return std::stoi(s); });`
	`503`	`+ // assign non-negative values to para.input.out_elf`
	`504`	`+ std::copy(out_elf.begin(), out_elf.end(), para.input.out_elf.begin());`
	`505`	`+ };`
`499`	`506`	`item.check_value = [](const Input_Item& item, const Parameter& para) {`
`500`		`- if (para.input.out_elf && para.input.esolver_type != "ksdft" && para.input.esolver_type != "ofdft")`
	`507`	`+ if (para.input.out_elf[0] > 0 && para.input.esolver_type != "ksdft" && para.input.esolver_type != "ofdft")`
`501`	`508`	`{`
`502`	`509`	`ModuleBase::WARNING_QUIT("ReadInput", "ELF is only aviailable for ksdft and ofdft");`
`503`	`510`	`}`
`504`	`511`	`};`
	`512`	`+ sync_intvec(input.out_elf, 2, 0);`
`505`	`513`	`this->add_item(item);`
`506`	`514`	`}`
`507`	`515`	`}`
Original file line number	Diff line number	Diff line change
`@@ -18,7 +18,8 @@ void write_elf(`
`18`	`18`	`const double* const* rho,`
`19`	`19`	`const double* const* tau,`
`20`	`20`	`ModulePW::PW_Basis* rho_basis,`
`21`		`- const UnitCell* ucell_);`
	`21`	`+ const UnitCell* ucell_,`
	`22`	`+ const int& precision);`
`22`	`23`	`}`
`23`	`24`
`24`	`25`	`#endif`