Fix: bug in charge extrapolation (deepmodeling#5007)

* Fix: bug in charge extrapolation

* Test: update result.ref

source/module_elecstate/module_charge/charge_extra.cpp

@@ -20,11 +20,7 @@ Charge_Extra::~Charge_Extra()
     }
 }
 
-void Charge_Extra::Init_CE(const int& nspin,
-                           const int& natom,
-                           const double& volume,
-                           const int& nrxx,
-                           const std::string chg_extrap)
+void Charge_Extra::Init_CE(const int& nspin, const int& natom, const int& nrxx, const std::string chg_extrap)
 {
     if (chg_extrap == "none")
     {
@@ -47,13 +43,13 @@ void Charge_Extra::Init_CE(const int& nspin,
         ModuleBase::WARNING_QUIT("Charge_Extra","charge extrapolation method is not available !");
     }
 
-    this->omega_old = volume;
     this->nspin = nspin;
 
-    if (pot_order > 1)
+    if (pot_order > 0)
     {
         delta_rho1.resize(this->nspin, std::vector<double>(nrxx, 0.0));
         delta_rho2.resize(this->nspin, std::vector<double>(nrxx, 0.0));
+        delta_rho3.resize(this->nspin, std::vector<double>(nrxx, 0.0));
     }
 
     if(pot_order == 3)
@@ -110,39 +106,18 @@ void Charge_Extra::extrapolate_charge(
 
     // if(lsda || noncolin) rho2zeta();
 
-    double** rho_atom = new double*[this->nspin];
-    for (int is = 0; is < this->nspin; is++)
-    {
-        rho_atom[is] = new double[chr->rhopw->nrxx];
-    }
-    chr->atomic_rho(this->nspin, omega_old, rho_atom, sf->strucFac, ucell);
-#ifdef _OPENMP
-#pragma omp parallel for collapse(2) schedule(static, 512)
-#endif
-    for (int is = 0; is < this->nspin; is++)
-    {
-        for (int ir = 0; ir < chr->rhopw->nrxx; ir++)
-        {
-            chr->rho[is][ir] -= rho_atom[is][ir];
-            chr->rho[is][ir] *= omega_old;
-        }
-    }
-
     if(rho_extr == 1)
     {
         ofs_running << " NEW-OLD atomic charge density approx. for the potential !" << std::endl;
 
-        if (pot_order > 1)
-        {
 #ifdef _OPENMP
-#pragma omp parallel for collapse(2) schedule(static, 512)
+#pragma omp parallel for collapse(2) schedule(static, 128)
 #endif
-            for (int is = 0; is < this->nspin; is++)
+        for (int is = 0; is < this->nspin; is++)
+        {
+            for (int ir = 0; ir < chr->rhopw->nrxx; ir++)
             {
-                for (int ir = 0; ir < chr->rhopw->nrxx; ir++)
-                {
-                    delta_rho1[is][ir] = chr->rho[is][ir];
-                }
+                chr->rho[is][ir] = delta_rho1[is][ir];
             }
         }
     }
@@ -158,8 +133,6 @@ void Charge_Extra::extrapolate_charge(
         {
             for (int ir = 0; ir < chr->rhopw->nrxx; ir++)
             {
-                delta_rho2[is][ir] = delta_rho1[is][ir];
-                delta_rho1[is][ir] = chr->rho[is][ir];
                 chr->rho[is][ir] = 2 * delta_rho1[is][ir] - delta_rho2[is][ir];
             }
         }
@@ -174,24 +147,25 @@ void Charge_Extra::extrapolate_charge(
         const double one_add_alpha = 1 + alpha;
         const double beta_alpha = beta - alpha;
 
-        std::vector<std::vector<double>> delta_rho3(this->nspin, std::vector<double>(chr->rhopw->nrxx, 0.0));
 #ifdef _OPENMP
 #pragma omp parallel for collapse(2) schedule(static, 64)
 #endif
         for (int is = 0; is < this->nspin; is++)
         {
             for (int ir = 0; ir < chr->rhopw->nrxx; ir++)
             {
-                delta_rho3[is][ir] = delta_rho2[is][ir];
-                delta_rho2[is][ir] = delta_rho1[is][ir];
-                delta_rho1[is][ir] = chr->rho[is][ir];
                 chr->rho[is][ir]
                     = one_add_alpha * delta_rho1[is][ir] + beta_alpha * delta_rho2[is][ir] - beta * delta_rho3[is][ir];
             }
         }
     }
 
     sf->setup_structure_factor(&ucell, chr->rhopw);
+    double** rho_atom = new double*[this->nspin];
+    for (int is = 0; is < this->nspin; is++)
+    {
+        rho_atom[is] = new double[chr->rhopw->nrxx];
+    }
     chr->atomic_rho(this->nspin, ucell.omega, rho_atom, sf->strucFac, ucell);
 #ifdef _OPENMP
 #pragma omp parallel for collapse(2) schedule(static, 512)
@@ -205,8 +179,6 @@ void Charge_Extra::extrapolate_charge(
         }
     }
 
-    omega_old = ucell.omega;
-
     for (int is = 0; is < this->nspin; is++)
     {
         delete[] rho_atom[is];
@@ -296,4 +268,41 @@ void Charge_Extra::update_all_dis(const UnitCell& ucell)
         assert(iat == ucell.nat);
     }
     return;
+}
+
+void Charge_Extra::update_delta_rho(const UnitCell& ucell, const Charge* chr, const Structure_Factor* sf)
+{
+    if (pot_order == 0)
+    {
+        return;
+    }
+
+    // obtain the difference between chr->rho and atomic_rho
+    double** rho_atom = new double*[this->nspin];
+    for (int is = 0; is < this->nspin; is++)
+    {
+        rho_atom[is] = new double[chr->rhopw->nrxx];
+    }
+    chr->atomic_rho(this->nspin, ucell.omega, rho_atom, sf->strucFac, ucell);
+
+#ifdef _OPENMP
+#pragma omp parallel for collapse(2) schedule(static, 512)
+#endif
+    for (int is = 0; is < this->nspin; is++)
+    {
+        for (int ir = 0; ir < chr->rhopw->nrxx; ir++)
+        {
+            delta_rho3[is][ir] = delta_rho2[is][ir];
+            delta_rho2[is][ir] = delta_rho1[is][ir];
+            delta_rho1[is][ir] = chr->rho[is][ir] - rho_atom[is][ir];
+            delta_rho1[is][ir] *= ucell.omega;
+        }
+    }
+
+    for (int is = 0; is < this->nspin; is++)
+    {
+        delete[] rho_atom[is];
+    }
+    delete[] rho_atom;
+    return;
 }

source/module_elecstate/module_charge/charge_extra.h

@@ -47,15 +47,10 @@ class Charge_Extra
      *
      * @param nspin the number of spins
      * @param natom the number of atoms
-     * @param volume the volume of the cell
      * @param nrxx the number of grids
      * @param chg_extrap the charge extrapolation method
      */
-    void Init_CE(const int& nspin,
-                 const int& natom,
-                 const double& volume,
-                 const int& nrxx,
-                 const std::string chg_extrap);
+    void Init_CE(const int& nspin, const int& natom, const int& nrxx, const std::string chg_extrap);
 
     /**
      * @brief charge extrapolation method
@@ -87,11 +82,19 @@ class Charge_Extra
      */
     void update_all_dis(const UnitCell& ucell);
 
+    /**
+     * @brief update the difference of charge density
+     *
+     * @param ucell the cell information
+     * @param chr the charge density
+     * @param sf the structure factor
+     */
+    void update_delta_rho(const UnitCell& ucell, const Charge* chr, const Structure_Factor* sf);
+
   private:
     int istep = 0; ///< the current step
     int pot_order; ///< the specified charge extrapolation method
     int rho_extr;  ///< the actually used method
-    double omega_old; ///< the old volume of the last step
     int nspin;        ///< the number of spins
 
     ModuleBase::Vector3<double>* dis_old1 = nullptr; ///< dis_old2 = pos_old1 - pos_old2
@@ -100,6 +103,7 @@ class Charge_Extra
 
     std::vector<std::vector<double>> delta_rho1; ///< the last step difference of rho and atomic_rho
     std::vector<std::vector<double>> delta_rho2; ///< the second last step difference of rho and atomic_rho
+    std::vector<std::vector<double>> delta_rho3; ///< the third last step difference of rho and atomic_rho
 
     double alpha; ///< parameter used in the second order extrapolation
     double beta;  ///< parameter used in the second order extrapolation

source/module_elecstate/test/charge_extra_test.cpp

@@ -136,7 +136,7 @@ TEST_F(ChargeExtraTest, InitCEWarningQuit)
 {
     GlobalV::chg_extrap = "wwww";
     testing::internal::CaptureStdout();
-    EXPECT_EXIT(CE.Init_CE(GlobalV::NSPIN, ucell->nat, ucell->omega, charge.rhopw->nrxx, GlobalV::chg_extrap),
+    EXPECT_EXIT(CE.Init_CE(GlobalV::NSPIN, ucell->nat, charge.rhopw->nrxx, GlobalV::chg_extrap),
                 ::testing::ExitedWithCode(0),
                 "");
     std::string output = testing::internal::GetCapturedStdout();
@@ -146,21 +146,21 @@ TEST_F(ChargeExtraTest, InitCEWarningQuit)
 TEST_F(ChargeExtraTest, InitCECase1)
 {
     GlobalV::chg_extrap = "none";
-    CE.Init_CE(GlobalV::NSPIN, ucell->nat, ucell->omega, charge.rhopw->nrxx, GlobalV::chg_extrap);
+    CE.Init_CE(GlobalV::NSPIN, ucell->nat, charge.rhopw->nrxx, GlobalV::chg_extrap);
     EXPECT_EQ(CE.pot_order, 0);
 }
 
 TEST_F(ChargeExtraTest, InitCECase2)
 {
     GlobalV::chg_extrap = "atomic";
-    CE.Init_CE(GlobalV::NSPIN, ucell->nat, ucell->omega, charge.rhopw->nrxx, GlobalV::chg_extrap);
+    CE.Init_CE(GlobalV::NSPIN, ucell->nat, charge.rhopw->nrxx, GlobalV::chg_extrap);
     EXPECT_EQ(CE.pot_order, 1);
 }
 
 TEST_F(ChargeExtraTest, InitCECase3)
 {
     GlobalV::chg_extrap = "first-order";
-    CE.Init_CE(GlobalV::NSPIN, ucell->nat, ucell->omega, charge.rhopw->nrxx, GlobalV::chg_extrap);
+    CE.Init_CE(GlobalV::NSPIN, ucell->nat, charge.rhopw->nrxx, GlobalV::chg_extrap);
     EXPECT_EQ(CE.pot_order, 2);
     EXPECT_NE(CE.delta_rho1.size(), 0);
     EXPECT_NE(CE.delta_rho2.size(), 0);
@@ -169,7 +169,7 @@ TEST_F(ChargeExtraTest, InitCECase3)
 TEST_F(ChargeExtraTest, InitCECase4)
 {
     GlobalV::chg_extrap = "second-order";
-    CE.Init_CE(GlobalV::NSPIN, ucell->nat, ucell->omega, charge.rhopw->nrxx, GlobalV::chg_extrap);
+    CE.Init_CE(GlobalV::NSPIN, ucell->nat, charge.rhopw->nrxx, GlobalV::chg_extrap);
     EXPECT_EQ(CE.pot_order, 3);
     EXPECT_DOUBLE_EQ(CE.alpha, 1.0);
     EXPECT_DOUBLE_EQ(CE.beta, 0.0);
@@ -183,7 +183,7 @@ TEST_F(ChargeExtraTest, InitCECase4)
 TEST_F(ChargeExtraTest, ExtrapolateChargeCase1)
 {
     GlobalV::chg_extrap = "second-order";
-    CE.Init_CE(GlobalV::NSPIN, ucell->nat, ucell->omega, charge.rhopw->nrxx, GlobalV::chg_extrap);
+    CE.Init_CE(GlobalV::NSPIN, ucell->nat, charge.rhopw->nrxx, GlobalV::chg_extrap);
     CE.istep = 0;
     CE.pot_order = 3;
 
@@ -205,7 +205,7 @@ TEST_F(ChargeExtraTest, ExtrapolateChargeCase1)
 TEST_F(ChargeExtraTest, ExtrapolateChargeCase2)
 {
     GlobalV::chg_extrap = "second-order";
-    CE.Init_CE(GlobalV::NSPIN, ucell->nat, ucell->omega, charge.rhopw->nrxx, GlobalV::chg_extrap);
+    CE.Init_CE(GlobalV::NSPIN, ucell->nat, charge.rhopw->nrxx, GlobalV::chg_extrap);
     CE.istep = 1;
     CE.pot_order = 3;
 
@@ -227,7 +227,7 @@ TEST_F(ChargeExtraTest, ExtrapolateChargeCase2)
 TEST_F(ChargeExtraTest, ExtrapolateChargeCase3)
 {
     GlobalV::chg_extrap = "second-order";
-    CE.Init_CE(GlobalV::NSPIN, ucell->nat, ucell->omega, charge.rhopw->nrxx, GlobalV::chg_extrap);
+    CE.Init_CE(GlobalV::NSPIN, ucell->nat, charge.rhopw->nrxx, GlobalV::chg_extrap);
     CE.istep = 2;
     CE.pot_order = 3;
 
@@ -249,7 +249,7 @@ TEST_F(ChargeExtraTest, ExtrapolateChargeCase3)
 TEST_F(ChargeExtraTest, ExtrapolateChargeCase4)
 {
     GlobalV::chg_extrap = "second-order";
-    CE.Init_CE(GlobalV::NSPIN, ucell->nat, ucell->omega, charge.rhopw->nrxx, GlobalV::chg_extrap);
+    CE.Init_CE(GlobalV::NSPIN, ucell->nat, charge.rhopw->nrxx, GlobalV::chg_extrap);
     CE.istep = 3;
 
     GlobalV::ofs_running.open("log");
@@ -271,7 +271,7 @@ TEST_F(ChargeExtraTest, ExtrapolateChargeCase4)
 TEST_F(ChargeExtraTest, UpdateAllDis)
 {
     GlobalV::chg_extrap = "second-order";
-    CE.Init_CE(GlobalV::NSPIN, ucell->nat, ucell->omega, charge.rhopw->nrxx, GlobalV::chg_extrap);
+    CE.Init_CE(GlobalV::NSPIN, ucell->nat, charge.rhopw->nrxx, GlobalV::chg_extrap);
     CE.istep = 3;
     for (int i = 0; i < ucell->nat; ++i)
     {
@@ -293,7 +293,7 @@ TEST_F(ChargeExtraTest, UpdateAllDis)
 TEST_F(ChargeExtraTest, FindAlphaAndBeta)
 {
     GlobalV::chg_extrap = "second-order";
-    CE.Init_CE(GlobalV::NSPIN, ucell->nat, ucell->omega, charge.rhopw->nrxx, GlobalV::chg_extrap);
+    CE.Init_CE(GlobalV::NSPIN, ucell->nat, charge.rhopw->nrxx, GlobalV::chg_extrap);
     CE.istep = 3;
     for (int i = 0; i < ucell->nat; ++i)
     {

source/module_esolver/esolver_fp.cpp

@@ -106,7 +106,7 @@ void ESolver_FP::before_all_runners(const Input_para& inp, UnitCell& cell)
     this->print_rhofft(inp, GlobalV::ofs_running);
 
     //! 5) initialize the charge extrapolation method if necessary
-    this->CE.Init_CE(GlobalV::NSPIN, GlobalC::ucell.nat, GlobalC::ucell.omega, this->pw_rhod->nrxx, inp.chg_extrap);
+    this->CE.Init_CE(GlobalV::NSPIN, GlobalC::ucell.nat, this->pw_rhod->nrxx, inp.chg_extrap);
 
     return;
 }
@@ -120,6 +120,9 @@ void ESolver_FP::after_scf(const int istep)
     // 1) write fermi energy
     ModuleIO::output_efermi(this->conv_elec, this->pelec->eferm.ef);
 
+    // 2) update delta rho for charge extrapolation
+    CE.update_delta_rho(GlobalC::ucell, &(this->chr), &(this->sf));
+
     if (istep % PARAM.inp.out_interval == 0)
     {
         // 3) write charge density

tests/integrate/108_PW_RE_NEW/result.ref

@@ -1,5 +1,6 @@
-etotref -253.6810444502974633
-etotperatomref -63.4202611126
-totalforceref 0.094780
-totalstressref 2.603647
-totaltimeref 3.99
+etotref -253.6810827494790033
+etotperatomref -63.4202706874
+totalforceref 0.078384
+totalstressref 2.163967
+STRU_ION1_D_pass 0
+totaltimeref 4.77

tests/integrate/108_PW_RE_PINT_RKS/result.ref

@@ -1,5 +1,5 @@
-etotref -253.6810443355996085
-etotperatomref -63.4202610839
-totalforceref 0.094440
-totalstressref 2.593312
-totaltimeref 5.96
+etotref -253.6810827448143186
+etotperatomref -63.4202706862
+totalforceref 0.078204
+totalstressref 2.142930
+totaltimeref 6.65