Feature : printing initial charge density (#4254)

* Feature : loading equivariant deepks model

* Feature : printing initial charge density

* add for pw calculation

* add printing of initial potential

* update documentation

---------

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

docs/advanced/input_files/input-main.md

@@ -1470,16 +1470,21 @@ These variables are used to control the output of properties.
 
 ### out_chg
 
-- **Type**: Boolean
-- **Description**: Whether to output the charge density (in Bohr^-3) on real space grids into the density files in the folder `OUT.${suffix}`. The files are named as:
-  - npsin = 1: SPIN1_CHG.cube;
-  - npsin = 2: SPIN1_CHG.cube, and SPIN2_CHG.cube;
-  - npsin = 4: SPIN1_CHG.cube, SPIN2_CHG.cube, SPIN3_CHG.cube, and SPIN4_CHG.cube.
+- **Type**: Integer
+- **Description**: 
+  - 1. Output the charge density (in Bohr^-3) on real space grids into the density files in the folder `OUT.${suffix}`. The files are named as:
+    - npsin = 1: SPIN1_CHG.cube;
+    - npsin = 2: SPIN1_CHG.cube, and SPIN2_CHG.cube;
+    - npsin = 4: SPIN1_CHG.cube, SPIN2_CHG.cube, SPIN3_CHG.cube, and SPIN4_CHG.cube.
+  - 2. On top of 1, also output the initial charge density. The files are named as:
+    - nspin = 1: SPIN1_CHG_INI.cube
+    - npsin = 2: SPIN1_CHG_INI.cube, and SPIN2_CHG_INI.cube;
+    - npsin = 4: SPIN1_CHG_INI.cube, SPIN2_CHG_INI.cube, SPIN3_CHG_INI.cube, and SPIN4_CHG_INI.cube.
 
   The circle order of the charge density on real space grids is: x is the outer loop, then y and finally z (z is moving fastest).
 
   If EXX(exact exchange) is calculated, (i.e. *[dft_fuctional](#dft_functional)==hse/hf/pbe0/scan0/opt_orb* or *[rpa](#rpa)==True*), the Hexx(R) files will be output in the folder `OUT.${suffix}` too, which can be read in NSCF calculation.
-- **Default**: False
+- **Default**: 0
 
 ### out_pot
 
@@ -1492,6 +1497,10 @@ These variables are used to control the output of properties.
   - 2: Output the **electrostatic potential** on real space grids into `OUT.${suffix}/ElecStaticPot.cube`. The Python script named `tools/average_pot/aveElecStatPot.py` can be used to calculate the average electrostatic potential along the z-axis and outputs it into ElecStaticPot_AVE.
 
     Please note that the total local potential refers to the local component of the self-consistent potential, excluding the non-local pseudopotential. The distinction between the local potential and the electrostatic potential is as follows: local potential = electrostatic potential + XC potential.
+  - 3: Apart from 1, also output the **total local potential** of the initial charge density. The files are named as:
+    - npsin = 1: SPIN1_POT_INI.cube;
+    - npsin = 2: SPIN1_POT_INI.cube, and SPIN2_POT_INI.cube;
+    - npsin = 4: SPIN1_POT_INI.cube, SPIN2_POT_INI.cube, SPIN3_POT_INI.cube, and SPIN4_POT_INI.cube.
 - **Default**: 0
 
 ### out_dm

source/module_esolver/esolver_ks_lcao_elec.cpp

@@ -23,6 +23,8 @@
 #include "module_hamilt_lcao/hamilt_lcaodft/operator_lcao/op_exx_lcao.h"
 #include "module_hamilt_lcao/module_deltaspin/spin_constrain.h"
 #include "module_io/dm_io.h"
+#include "module_io/rho_io.h"
+#include "module_io/potential_io.h"
 
 namespace ModuleESolver
 {
@@ -251,6 +253,41 @@ void ESolver_KS_LCAO<TK, TR>::before_scf(int istep)
 #endif // __EXX
 
     this->pelec->init_scf(istep, this->sf.strucFac);
+    if(GlobalV::out_chg == 2)
+    {
+        for(int is = 0; is < GlobalV::NSPIN; is++)
+        {
+            std::stringstream ss;
+            ss << GlobalV::global_out_dir << "SPIN" << is+1 << "_CHG_INI.cube";
+            ModuleIO::write_rho(
+#ifdef __MPI
+                this->pw_big->nbz, this->pw_big->bz,
+                this->pw_rho->nplane, this->pw_rho->startz_current,
+#endif
+                this->pelec->charge->rho[is], is,
+                GlobalV::NSPIN, 0,
+                ss.str(),
+                this->pw_rho->nx, this->pw_rho->ny, this->pw_rho->nz,
+                this->pelec->eferm.ef, &(GlobalC::ucell) , 11);
+        }
+    }
+
+    if(GlobalV::out_pot == 3)
+    {
+        for(int is = 0; is < GlobalV::NSPIN; is++)
+        {
+            std::stringstream ss;
+            ss << GlobalV::global_out_dir << "SPIN" << is+1 << "_POT_INI.cube";
+            ModuleIO::write_potential(
+#ifdef __MPI
+                this->pw_big->nbz, this->pw_big->bz,
+                this->pw_rho->nplane, this->pw_rho->startz_current,
+#endif
+                is,0,ss.str(),
+                this->pw_rho->nx, this->pw_rho->ny, this->pw_rho->nz,
+                this->pelec->pot->get_effective_v(), 11);
+        }
+    }
     // initalize DMR
     // DMR should be same size with Hamiltonian(R)
     dynamic_cast<elecstate::ElecStateLCAO<TK>*>(this->pelec)

source/module_esolver/esolver_ks_pw.cpp

@@ -504,6 +504,41 @@ void ESolver_KS_PW<T, Device>::before_scf(int istep)
 
     //! calculate the total local pseudopotential in real space
     this->pelec->init_scf(istep, this->sf.strucFac);
+    if(GlobalV::out_chg == 2)
+    {
+        for(int is = 0; is < GlobalV::NSPIN; is++)
+        {
+            std::stringstream ss;
+            ss << GlobalV::global_out_dir << "SPIN" << is+1 << "_CHG_INI.cube";
+            ModuleIO::write_rho(
+#ifdef __MPI
+                this->pw_big->nbz, this->pw_big->bz,
+                this->pw_rho->nplane, this->pw_rho->startz_current,
+#endif
+                this->pelec->charge->rho[is], is,
+                GlobalV::NSPIN, 0,
+                ss.str(),
+                this->pw_rho->nx, this->pw_rho->ny, this->pw_rho->nz,
+                this->pelec->eferm.ef, &(GlobalC::ucell) , 11);
+        }
+    }
+
+    if(GlobalV::out_pot == 3)
+    {
+        for(int is = 0; is < GlobalV::NSPIN; is++)
+        {
+            std::stringstream ss;
+            ss << GlobalV::global_out_dir << "SPIN" << is+1 << "_POT_INI.cube";
+            ModuleIO::write_potential(
+#ifdef __MPI
+                this->pw_big->nbz, this->pw_big->bz,
+                this->pw_rho->nplane, this->pw_rho->startz_current,
+#endif
+                is,0,ss.str(),
+                this->pw_rho->nx, this->pw_rho->ny, this->pw_rho->nz,
+                this->pelec->pot->get_effective_v(), 11);
+        }
+    }
 
     //! Symmetry_rho should behind init_scf, because charge should be initialized first.
     //! liuyu comment: Symmetry_rho should be located between init_rho and v_of_rho?

source/module_esolver/esolver_of.cpp

@@ -544,7 +544,7 @@ void ESolver_OF::after_opt(const int istep, UnitCell& ucell)
                 3);
         }
 
-        if (GlobalV::out_pot == 1) // output the effective potential, sunliang 2023-03-16
+        if (GlobalV::out_pot == 1 || GlobalV::out_pot == 3) // output the effective potential, sunliang 2023-03-16
         {
             int precision = 3; // be consistent with esolver_ks_lcao.cpp
             std::stringstream ssp;

source/module_io/input.cpp

@@ -1393,7 +1393,7 @@ bool Input::Read(const std::string& fn)
         }
         else if (strcmp("out_chg", word) == 0)
         {
-            read_bool(ifs, out_chg);
+            read_value(ifs, out_chg);
         }
         else if (strcmp("out_dm", word) == 0)
         {
@@ -3583,7 +3583,7 @@ void Input::Bcast()
     Parallel_Common::bcast_string(chg_extrap); // xiaohui modify 2015-02-01
     Parallel_Common::bcast_int(out_freq_elec);
     Parallel_Common::bcast_int(out_freq_ion);
-    Parallel_Common::bcast_bool(out_chg);
+    Parallel_Common::bcast_int(out_chg);
     Parallel_Common::bcast_bool(out_dm);
     Parallel_Common::bcast_bool(out_dm1);
     Parallel_Common::bcast_bool(out_bandgap); // for bandgap printing

source/module_io/input.h

@@ -262,7 +262,7 @@ class Input
     int printe; // mohan add 2011-03-16
     int out_freq_elec;  // the frequency ( >= 0) of electronic iter to output charge density and wavefunction. 0: output only when converged
     int out_freq_ion;  // the frequency ( >= 0 ) of ionic step to output charge density and wavefunction. 0: output only when ion steps are finished
-    bool out_chg; // output charge density. 0: no; 1: yes
+    int out_chg; // output charge density. 0: no; 1: yes
     bool out_dm; // output density matrix.
     bool out_dm1;
     int out_pot; // yes or no

source/module_io/output_potential.cpp

@@ -36,7 +36,7 @@ Output_Potential::Output_Potential(const ModulePW::PW_Basis_Big* pw_big,
 
 void Output_Potential::write()
 {
-    if (_out_pot == 1)
+    if (_out_pot == 1 || _out_pot == 3)
     {
         for (int is = 0; is < _nspin; is++)
         {

source/module_io/parameter_pool.cpp

@@ -888,7 +888,7 @@ void input_parameters_set(std::map<std::string, InputParameter> input_parameters
     }
     else if (input_parameters.count("out_chg") != 0)
     {
-        INPUT.out_chg = *static_cast<bool*>(input_parameters["out_chg"].get());
+        INPUT.out_chg = *static_cast<int*>(input_parameters["out_chg"].get());
     }
     else if (input_parameters.count("out_dm") != 0)
     {

source/module_io/test/input_conv_test.cpp

@@ -141,7 +141,7 @@ TEST_F(InputConvTest, Conv)
 	EXPECT_EQ(GlobalV::OUT_FREQ_ION,0);
 	EXPECT_EQ(GlobalV::init_chg,"atomic");
 	EXPECT_EQ(GlobalV::chg_extrap,"atomic");
-	EXPECT_EQ(GlobalV::out_chg,false);
+	EXPECT_EQ(GlobalV::out_chg, 0);
     EXPECT_EQ(GlobalV::out_pot, 2);
     EXPECT_EQ(GlobalV::out_app_flag, false);
     EXPECT_EQ(GlobalV::out_bandgap, false);

source/module_io/test/support/INPUT

@@ -54,7 +54,7 @@ scf_thr_type                   2 #type of the criterion of scf_thr, 1: reci drho
 init_wfc                       atomic #start wave functions are from 'atomic', 'atomic+random', 'random' or 'file'
 init_chg                       atomic #start charge is from 'atomic' or file
 chg_extrap                     atomic #atomic; first-order; second-order; dm:coefficients of SIA
-out_chg                        FALSE #>0 output charge density for selected electron steps
+out_chg                        0 #>0 output charge density for selected electron steps
 out_pot                        2 #output realspace potential
 out_wfc_pw                     0 #output wave functions
 out_wfc_r                      0 #output wave functions in realspace

source/module_io/test/support/witestfile

@@ -50,7 +50,7 @@ scf_thr_type                   2 #type of the criterion of scf_thr, 1: reci drho
 init_wfc                       atomic #start wave functions are from 'atomic', 'atomic+random', 'random' or 'file'
 init_chg                       atomic #start charge is from 'atomic' or file
 chg_extrap                     atomic #atomic; first-order; second-order; dm:coefficients of SIA
-out_chg                        FALSE #>0 output charge density for selected electron steps
+out_chg                        0 #>0 output charge density for selected electron steps
 out_pot                        2 #output realspace potential
 out_wfc_pw                     0 #output wave functions
 out_wfc_r                      0 #output wave functions in realspace

source/module_io/winput.cpp

@@ -54,7 +54,7 @@ bool winput::plot_wanq;//add 2008-01-26
 std::string winput::plot_option;//(110),[110] etc.
 int winput::n_unitcell;//number of unitcell to plot
 bool winput::out_all;
-bool winput::out_chg;
+int winput::out_chg;
 std::string winput::charge_type;
 bool winput::cal_bands; //for wan  	   wan basis + wan charge
 bool winput::cal_bands2;//for semi-wan ;pw basis + wan charge add 2008-4-11
@@ -700,7 +700,7 @@ void winput::Bcast(void)
 	Parallel_Common::bcast_string( plot_option );
 	Parallel_Common::bcast_int( n_unitcell );
 	Parallel_Common::bcast_bool( out_all );
-	Parallel_Common::bcast_bool( out_chg );
+	Parallel_Common::bcast_int( out_chg );
 
 	Parallel_Common::bcast_string( charge_type );
 	Parallel_Common::bcast_bool( cal_bands );

source/module_io/winput.h

@@ -115,7 +115,7 @@ class winput
 	static std::string plot_option;//(110),[110] etc.
 	static int n_unitcell;//number of unitcell to plot
 	static bool out_all;
-	static bool out_chg;
+	static int out_chg;
 	static std::string charge_type;
 	static bool cal_bands; //for wan  	   wan basis + wan charge
 	static bool cal_bands2;//for semi-wan ;pw basis + wan charge add 2008-4-11

source/module_io/write_potential.cpp

@@ -28,7 +28,7 @@ void write_potential(
     const int& hartree)
 {
     ModuleBase::TITLE("potential", "write_potential");
-    if (GlobalV::out_pot != 1)
+    if (GlobalV::out_pot != 1 && GlobalV::out_pot != 3)
     {
         return;
     }