Refactor: modify EXX realx/md calculation framework (#5403)

1041176461 · jiyuang · Qianruipku · web-flow · commit cc9d759da969 · 2024-11-06T22:37:39.000+08:00
* Use template to reconstruct parse_expression * Feature: output R matrix at each MD step * Modify'matrix_HS' to 'matrix' for R matrix output * Merge branches 'develop' and 'develop' of https://github.com/1041176461/abacus-develop into develop * Fix: modify index in parse_expression * Fix: add regfree for parse_expression * Doc: update phonopy doc * Doc: update phonopy doc * fix tdos plot for nspin=2 * optimize dosplot for nspin=2 * fix legend for dosplot * Refactor: modify exx loop for relax/md * add exx when istep>0 * Update dos.py * Update esolver_sdft_pw.cpp * Update lcao_before_scf.cpp * Update Exx_LRI_interface.h * Update Exx_LRI_interface.hpp * Fix: compile error * Fix: compile error * Fix: change HSE relax/md result.ref for new framework * Fix: compile error * compatible with exx_iter_finish --------- Co-authored-by: jiyuang <jiyuyang@mail.ustc.com> Co-authored-by: Qianrui <76200646+Qianruipku@users.noreply.github.com> Co-authored-by: maki49 <1579492865@qq.com>
diff --git a/source/module_esolver/esolver_ks.cpp b/source/module_esolver/esolver_ks.cpp
@@ -609,7 +609,7 @@ void ESolver_KS<T, Device>::runner(const int istep, UnitCell& ucell)
         this->update_pot(istep, iter);
 
         // 10) finish scf iterations
-        this->iter_finish(iter);
+        this->iter_finish(istep, iter);
 #ifdef __MPI
         double duration = (double)(MPI_Wtime() - iterstart);
 #else
@@ -663,7 +663,7 @@ void ESolver_KS<T, Device>::runner(const int istep, UnitCell& ucell)
 };
 
 template <typename T, typename Device>
-void ESolver_KS<T, Device>::iter_finish(int& iter)
+void ESolver_KS<T, Device>::iter_finish(const int istep, int& iter)
 {
     // 1 means Harris-Foulkes functional
     // 2 means Kohn-Sham functional
diff --git a/source/module_esolver/esolver_ks.h b/source/module_esolver/esolver_ks.h
@@ -58,7 +58,7 @@ class ESolver_KS : public ESolver_FP
 		virtual void iter_init(const int istep, const int iter) {};
 
 		//! Something to do after hamilt2density function in each iter loop.
-        virtual void iter_finish(int& iter);
+        virtual void iter_finish(const int istep, int& iter);
 
         //! Something to do after SCF iterations when SCF is converged or comes to the max iter step.
         virtual void after_scf(const int istep) override;
diff --git a/source/module_esolver/esolver_ks_lcao.cpp b/source/module_esolver/esolver_ks_lcao.cpp
@@ -625,13 +625,15 @@ void ESolver_KS_LCAO<TK, TR>::iter_init(const int istep, const int iter)
     // calculate exact-exchange
     if (GlobalC::exx_info.info_ri.real_number)
     {
-        this->exd->exx_eachiterinit(*dynamic_cast<const elecstate::ElecStateLCAO<TK>*>(this->pelec)->get_DM(),
+        this->exd->exx_eachiterinit(istep,
+                                    *dynamic_cast<const elecstate::ElecStateLCAO<TK>*>(this->pelec)->get_DM(),
                                     this->kv,
                                     iter);
     }
     else
     {
-        this->exc->exx_eachiterinit(*dynamic_cast<const elecstate::ElecStateLCAO<TK>*>(this->pelec)->get_DM(),
+        this->exc->exx_eachiterinit(istep,
+                                    *dynamic_cast<const elecstate::ElecStateLCAO<TK>*>(this->pelec)->get_DM(),
                                     this->kv,
                                     iter);
     }
@@ -914,12 +916,12 @@ void ESolver_KS_LCAO<TK, TR>::update_pot(const int istep, const int iter)
 //! 5) cal_MW? (why put it here?)
 //------------------------------------------------------------------------------
 template <typename TK, typename TR>
-void ESolver_KS_LCAO<TK, TR>::iter_finish(int& iter)
+void ESolver_KS_LCAO<TK, TR>::iter_finish(const int istep, int& iter)
 {
     ModuleBase::TITLE("ESolver_KS_LCAO", "iter_finish");
 
     // call iter_finish() of ESolver_KS
-    ESolver_KS<TK>::iter_finish(iter);
+    ESolver_KS<TK>::iter_finish(istep, iter);
 
     // 1) mix density matrix if mixing_restart + mixing_dmr + not first
     // mixing_restart at every iter
@@ -944,8 +946,8 @@ void ESolver_KS_LCAO<TK, TR>::iter_finish(int& iter)
     if (GlobalC::exx_info.info_global.cal_exx)
     {
         GlobalC::exx_info.info_ri.real_number ?
-            this->exd->exx_iter_finish(this->kv, GlobalC::ucell, *this->p_hamilt, *this->pelec, *this->p_chgmix, this->scf_ene_thr, iter, this->conv_esolver) :
-            this->exc->exx_iter_finish(this->kv, GlobalC::ucell, *this->p_hamilt, *this->pelec, *this->p_chgmix, this->scf_ene_thr, iter, this->conv_esolver);
+            this->exd->exx_iter_finish(this->kv, GlobalC::ucell, *this->p_hamilt, *this->pelec, *this->p_chgmix, this->scf_ene_thr, iter, istep, this->conv_esolver) :
+            this->exc->exx_iter_finish(this->kv, GlobalC::ucell, *this->p_hamilt, *this->pelec, *this->p_chgmix, this->scf_ene_thr, iter, istep, this->conv_esolver);
     }
 #endif
 
diff --git a/source/module_esolver/esolver_ks_lcao.h b/source/module_esolver/esolver_ks_lcao.h
@@ -56,7 +56,7 @@ class ESolver_KS_LCAO : public ESolver_KS<TK> {
 
     virtual void update_pot(const int istep, const int iter) override;
 
-    virtual void iter_finish(int& iter) override;
+    virtual void iter_finish(const int istep, int& iter) override;
 
     virtual void after_scf(const int istep) override;
 
diff --git a/source/module_esolver/esolver_ks_lcaopw.cpp b/source/module_esolver/esolver_ks_lcaopw.cpp
@@ -191,9 +191,9 @@ namespace ModuleESolver
     }
 
     template <typename T>
-    void ESolver_KS_LIP<T>::iter_finish(int& iter)
+    void ESolver_KS_LIP<T>::iter_finish(const int istep, int& iter)
     {
-        ESolver_KS_PW<T>::iter_finish(iter);
+        ESolver_KS_PW<T>::iter_finish(istep, iter);
 
 #ifdef __EXX
         if (GlobalC::exx_info.info_global.cal_exx && this->conv_esolver)
diff --git a/source/module_esolver/esolver_ks_lcaopw.h b/source/module_esolver/esolver_ks_lcaopw.h
@@ -28,7 +28,7 @@ namespace ModuleESolver
 
     protected:
       virtual void iter_init(const int istep, const int iter) override;
-      virtual void iter_finish(int& iter) override;
+      virtual void iter_finish(const int istep, int& iter) override;
 
       virtual void allocate_hamilt() override;
       virtual void deallocate_hamilt() override;
diff --git a/source/module_esolver/esolver_ks_pw.cpp b/source/module_esolver/esolver_ks_pw.cpp
@@ -434,10 +434,10 @@ void ESolver_KS_PW<T, Device>::update_pot(const int istep, const int iter)
 }
 
 template <typename T, typename Device>
-void ESolver_KS_PW<T, Device>::iter_finish(int& iter)
+void ESolver_KS_PW<T, Device>::iter_finish(const int istep, int& iter)
 {
     // call iter_finish() of ESolver_KS
-    ESolver_KS<T, Device>::iter_finish(iter);
+    ESolver_KS<T, Device>::iter_finish(istep, iter);
 
     // liuyu 2023-10-24
     // D in uspp need vloc, thus needs update when veff updated
diff --git a/source/module_esolver/esolver_ks_pw.h b/source/module_esolver/esolver_ks_pw.h
@@ -46,7 +46,7 @@ class ESolver_KS_PW : public ESolver_KS<T, Device>
 
     virtual void update_pot(const int istep, const int iter) override;
 
-    virtual void iter_finish(int& iter) override;
+    virtual void iter_finish(const int istep, int& iter) override;
 
     virtual void after_scf(const int istep) override;
 
diff --git a/source/module_esolver/esolver_sdft_pw.cpp b/source/module_esolver/esolver_sdft_pw.cpp
@@ -153,10 +153,10 @@ void ESolver_SDFT_PW<T, Device>::before_scf(const int istep)
 }
 
 template <typename T, typename Device>
-void ESolver_SDFT_PW<T, Device>::iter_finish(int& iter)
+void ESolver_SDFT_PW<T, Device>::iter_finish(const int istep, int& iter)
 {
     // call iter_finish() of ESolver_KS
-    ESolver_KS<T, Device>::iter_finish(iter);
+    ESolver_KS<T, Device>::iter_finish(istep, iter);
 }
 
 template <typename T, typename Device>
diff --git a/source/module_esolver/esolver_sdft_pw.h b/source/module_esolver/esolver_sdft_pw.h
@@ -39,7 +39,7 @@ class ESolver_SDFT_PW : public ESolver_KS_PW<T, Device>
 
     virtual void others(const int istep) override;
 
-    virtual void iter_finish(int& iter) override;
+    virtual void iter_finish(const int istep, int& iter) override;
 
     virtual void after_scf(const int istep) override;
 
diff --git a/source/module_esolver/lcao_before_scf.cpp b/source/module_esolver/lcao_before_scf.cpp
@@ -115,6 +115,7 @@ void ESolver_KS_LCAO<TK, TR>::beforesolver(const int istep)
             orb_,
             DM
 #ifdef __EXX
+            , istep
             , GlobalC::exx_info.info_ri.real_number ? &this->exd->two_level_step : &this->exc->two_level_step
             , GlobalC::exx_info.info_ri.real_number ? &exx_lri_double->Hexxs : nullptr
             , GlobalC::exx_info.info_ri.real_number ? nullptr : &exx_lri_complex->Hexxs
@@ -210,11 +211,11 @@ void ESolver_KS_LCAO<TK, TR>::before_scf(const int istep)
 #ifdef __EXX // set xc type before the first cal of xc in pelec->init_scf
     if (GlobalC::exx_info.info_ri.real_number)
     {
-        this->exd->exx_beforescf(this->kv, *this->p_chgmix, GlobalC::ucell, orb_);
+        this->exd->exx_beforescf(istep, this->kv, *this->p_chgmix, GlobalC::ucell, orb_);
     }
     else
     {
-        this->exc->exx_beforescf(this->kv, *this->p_chgmix, GlobalC::ucell, orb_);
+        this->exc->exx_beforescf(istep, this->kv, *this->p_chgmix, GlobalC::ucell, orb_);
     }
 #endif // __EXX
 
diff --git a/source/module_hamilt_lcao/hamilt_lcaodft/hamilt_lcao.cpp b/source/module_hamilt_lcao/hamilt_lcaodft/hamilt_lcao.cpp
@@ -73,6 +73,7 @@ HamiltLCAO<TK, TR>::HamiltLCAO(Gint_Gamma* GG_in,
     const LCAO_Orbitals& orb,
     elecstate::DensityMatrix<TK, double>* DM_in
 #ifdef __EXX
+    , const int istep
     , int* exx_two_level_step
     , std::vector<std::map<int, std::map<TAC, RI::Tensor<double>>>>* Hexxd
     , std::vector<std::map<int, std::map<TAC, RI::Tensor<std::complex<double>>>>>* Hexxc
@@ -397,6 +398,7 @@ HamiltLCAO<TK, TR>::HamiltLCAO(Gint_Gamma* GG_in,
             Hexxd,
             Hexxc,
             Add_Hexx_Type::R,
+            istep,
             exx_two_level_step,
             !GlobalC::restart.info_load.restart_exx
             && GlobalC::restart.info_load.load_H);
diff --git a/source/module_hamilt_lcao/hamilt_lcaodft/hamilt_lcao.h b/source/module_hamilt_lcao/hamilt_lcaodft/hamilt_lcao.h
@@ -37,6 +37,7 @@ class HamiltLCAO : public Hamilt<TK>
           const LCAO_Orbitals& orb,
           elecstate::DensityMatrix<TK, double>* DM_in
 #ifdef __EXX
+          , const int istep
           , int* exx_two_level_step = nullptr
           , std::vector<std::map<int, std::map<TAC, RI::Tensor<double>>>>* Hexxd = nullptr
           , std::vector<std::map<int, std::map<TAC, RI::Tensor<std::complex<double>>>>>* Hexxc = nullptr
@@ -105,6 +106,8 @@ class HamiltLCAO : public Hamilt<TK>
     /// current_spin for NSPIN=2, 0: hamiltonian for spin up, 1: hamiltonian for spin down
     int current_spin = 0;
 
+    const int istep = 0;
+
     // sk and hk will be refactored to HamiltLCAO later
     // std::vector<TK> sk;
     // std::vector<TK> hk;
diff --git a/source/module_hamilt_lcao/hamilt_lcaodft/operator_lcao/op_exx_lcao.h b/source/module_hamilt_lcao/hamilt_lcaodft/operator_lcao/op_exx_lcao.h
@@ -32,6 +32,7 @@ class OperatorEXX<OperatorLCAO<TK, TR>> : public OperatorLCAO<TK, TR>
         std::vector<std::map<int, std::map<TAC, RI::Tensor<double>>>>* Hexxd_in = nullptr,
         std::vector<std::map<int, std::map<TAC, RI::Tensor<std::complex<double>>>>>* Hexxc_in = nullptr,
         Add_Hexx_Type add_hexx_type_in = Add_Hexx_Type::R,
+        const int istep_in = 0,
         int* two_level_step_in = nullptr,
         const bool restart_in = false);
 
@@ -54,6 +55,8 @@ class OperatorEXX<OperatorLCAO<TK, TR>> : public OperatorLCAO<TK, TR>
       /// @brief if restart, read and save Hexx, and directly use it during the first outer loop.
       bool restart = false;
 
+      const int istep = 0; // the ion step
+
       void add_loaded_Hexx(const int ik);
 
       const K_Vectors& kv;
diff --git a/source/module_hamilt_lcao/hamilt_lcaodft/operator_lcao/op_exx_lcao.hpp b/source/module_hamilt_lcao/hamilt_lcaodft/operator_lcao/op_exx_lcao.hpp
@@ -85,13 +85,15 @@ OperatorEXX<OperatorLCAO<TK, TR>>::OperatorEXX(HS_Matrix_K<TK>* hsk_in,
 	std::vector<std::map<int, std::map<TAC, RI::Tensor<double>>>>* Hexxd_in,
 	std::vector<std::map<int, std::map<TAC, RI::Tensor<std::complex<double>>>>>* Hexxc_in,
     Add_Hexx_Type add_hexx_type_in,
-	int* two_level_step_in,
+    const int istep,
+    int* two_level_step_in,
 	const bool restart_in)
     : OperatorLCAO<TK, TR>(hsk_in, kv_in.kvec_d, hR_in),
     kv(kv_in),
     Hexxd(Hexxd_in),
     Hexxc(Hexxc_in),
     add_hexx_type(add_hexx_type_in),
+    istep(istep),
     two_level_step(two_level_step_in),
     restart(restart_in)
 {
@@ -201,7 +203,7 @@ void OperatorEXX<OperatorLCAO<TK, TR>>::contributeHR()
 {
     ModuleBase::TITLE("OperatorEXX", "contributeHR");
     // Peize Lin add 2016-12-03
-    if (PARAM.inp.calculation != "nscf" && this->two_level_step != nullptr && *this->two_level_step == 0 && !this->restart) { return; }  //in the non-exx loop, do nothing 
+    if (this->istep == 0 && PARAM.inp.calculation != "nscf" && this->two_level_step != nullptr && *this->two_level_step == 0 && !this->restart) { return; }  //in the non-exx loop, do nothing 
     if (this->add_hexx_type == Add_Hexx_Type::k) { return; }
 
     if (XC_Functional::get_func_type() == 4 || XC_Functional::get_func_type() == 5)
diff --git a/source/module_ri/Exx_LRI_interface.h b/source/module_ri/Exx_LRI_interface.h
@@ -54,10 +54,10 @@ class Exx_LRI_Interface
     void exx_before_all_runners(const K_Vectors& kv, const UnitCell& ucell, const Parallel_2D& pv);
 
     /// @brief in beforescf: set xc type, opt_orb, do DM mixing
-    void exx_beforescf(const K_Vectors& kv, const Charge_Mixing& chgmix, const UnitCell& ucell, const LCAO_Orbitals& orb);
+    void exx_beforescf(const int istep, const K_Vectors& kv, const Charge_Mixing& chgmix, const UnitCell& ucell, const LCAO_Orbitals& orb);
 
     /// @brief in eachiterinit:  do DM mixing and calculate Hexx when entering 2nd SCF
-    void exx_eachiterinit(const elecstate::DensityMatrix<T, double>& dm/**< double should be Tdata if complex-PBE-DM is supported*/,
+    void exx_eachiterinit(const int istep, const elecstate::DensityMatrix<T, double>& dm/**< double should be Tdata if complex-PBE-DM is supported*/,
         const K_Vectors& kv, const int& iter);
 
     /// @brief in hamilt2density: calculate Hexx and Eexx
@@ -66,18 +66,20 @@ class Exx_LRI_Interface
     /// @brief in iter_finish: write Hexx, do something according to whether SCF is converged
     void exx_iter_finish(const K_Vectors& kv, const UnitCell& ucell,
         hamilt::Hamilt<T>& hamilt, elecstate::ElecState& elec, Charge_Mixing& chgmix,
-        const double& scf_ene_thr, int& iter, bool& conv_esolver);
+        const double& scf_ene_thr, int& iter, const int istep, bool& conv_esolver);
     /// @brief: in do_after_converge: add exx operators; do DM mixing if seperate loop
     bool exx_after_converge(
         hamilt::Hamilt<T>& hamilt,
         const elecstate::DensityMatrix<T, double>& dm/**< double should be Tdata if complex-PBE-DM is supported*/,
         const K_Vectors& kv,
         const int& nspin,
         int& iter,
+        const int& istep,
         const double& etot,
         const double& scf_ene_thr);
     int two_level_step = 0;
     double etot_last_outer_loop = 0.0;
+    elecstate::DensityMatrix<T, double>* dm_last_step;
 private:
     std::shared_ptr<Exx_LRI<Tdata>> exx_ptr;
     Mix_DMk_2D mix_DMk_2D;
diff --git a/source/module_ri/Exx_LRI_interface.hpp b/source/module_ri/Exx_LRI_interface.hpp
@@ -53,13 +53,16 @@ void Exx_LRI_Interface<T, Tdata>::exx_before_all_runners(const K_Vectors& kv, co
 }
 
 template<typename T, typename Tdata>
-void Exx_LRI_Interface<T, Tdata>::exx_beforescf(const K_Vectors& kv, const Charge_Mixing& chgmix, const UnitCell& ucell, const LCAO_Orbitals& orb)
+void Exx_LRI_Interface<T, Tdata>::exx_beforescf(const int istep, const K_Vectors& kv, const Charge_Mixing& chgmix, const UnitCell& ucell, const LCAO_Orbitals& orb)
 {
 #ifdef __MPI
     if (GlobalC::exx_info.info_global.cal_exx)
     {
         if (GlobalC::restart.info_load.load_H_finish && !GlobalC::restart.info_load.restart_exx) { XC_Functional::set_xc_type(GlobalC::ucell.atoms[0].ncpp.xc_func);
-        } else
+        } 
+        else if (istep > 0) { XC_Functional::set_xc_type(GlobalC::ucell.atoms[0].ncpp.xc_func);
+        } 
+        else
         {
             if (ucell.atoms[0].ncpp.xc_func == "HF" || ucell.atoms[0].ncpp.xc_func == "PBE0" || ucell.atoms[0].ncpp.xc_func == "HSE")
             {
@@ -101,21 +104,28 @@ void Exx_LRI_Interface<T, Tdata>::exx_beforescf(const K_Vectors& kv, const Charg
 }
 
 template<typename T, typename Tdata>
-void Exx_LRI_Interface<T, Tdata>::exx_eachiterinit(const elecstate::DensityMatrix<T, double>& dm, const K_Vectors& kv, const int& iter)
+void Exx_LRI_Interface<T, Tdata>::exx_eachiterinit(const int istep, const elecstate::DensityMatrix<T, double>& dm, const K_Vectors& kv, const int& iter)
 {
     if (GlobalC::exx_info.info_global.cal_exx)
     {
-        if (!GlobalC::exx_info.info_global.separate_loop && this->two_level_step)
+        if (!GlobalC::exx_info.info_global.separate_loop && (this->two_level_step || istep > 0))
         {
             const bool flag_restart = (iter == 1) ? true : false;
-            if (this->exx_spacegroup_symmetry) { this->mix_DMk_2D.mix(symrot_.restore_dm(kv, dm.get_DMK_vector(), *dm.get_paraV_pointer()), flag_restart); }
-            else { this->mix_DMk_2D.mix(dm.get_DMK_vector(), flag_restart); }
-			const std::vector<std::map<int,std::map<std::pair<int, std::array<int, 3>>,RI::Tensor<Tdata>>>>
-				Ds = PARAM.globalv.gamma_only_local
-                ? RI_2D_Comm::split_m2D_ktoR<Tdata>(*this->exx_ptr->p_kv, this->mix_DMk_2D.get_DMk_gamma_out(), *dm.get_paraV_pointer(), PARAM.inp.nspin)
-                : RI_2D_Comm::split_m2D_ktoR<Tdata>(*this->exx_ptr->p_kv, this->mix_DMk_2D.get_DMk_k_out(), *dm.get_paraV_pointer(), PARAM.inp.nspin, this->exx_spacegroup_symmetry);
-            if (this->exx_spacegroup_symmetry && GlobalC::exx_info.info_global.exx_symmetry_realspace) { this->exx_ptr->cal_exx_elec(Ds, *dm.get_paraV_pointer(), &this->symrot_); }
-            else { this->exx_ptr->cal_exx_elec(Ds, *dm.get_paraV_pointer()); }
+            auto cal = [this, &kv, &flag_restart](const elecstate::DensityMatrix<T, double>& dm_in)
+            {
+                if (this->exx_spacegroup_symmetry) { this->mix_DMk_2D.mix(symrot_.restore_dm(kv,dm_in.get_DMK_vector(), *dm_in.get_paraV_pointer()), flag_restart); }
+                else { this->mix_DMk_2D.mix(dm_in.get_DMK_vector(), flag_restart); }
+			        const std::vector<std::map<int,std::map<std::pair<int, std::array<int, 3>>,RI::Tensor<Tdata>>>>
+				            Ds = PARAM.globalv.gamma_only_local
+                                ? RI_2D_Comm::split_m2D_ktoR<Tdata>(*this->exx_ptr->p_kv, this->mix_DMk_2D.get_DMk_gamma_out(), *dm_in.get_paraV_pointer(), PARAM.inp.nspin)
+                                : RI_2D_Comm::split_m2D_ktoR<Tdata>(*this->exx_ptr->p_kv, this->mix_DMk_2D.get_DMk_k_out(), *dm_in.get_paraV_pointer(), PARAM.inp.nspin, this->exx_spacegroup_symmetry);
+                if (this->exx_spacegroup_symmetry && GlobalC::exx_info.info_global.exx_symmetry_realspace) { this->exx_ptr->cal_exx_elec(Ds, *dm_in.get_paraV_pointer(), &this->symrot_); }
+                else { this->exx_ptr->cal_exx_elec(Ds, *dm_in.get_paraV_pointer()); }
+            };
+            if(istep > 0 && flag_restart)
+                cal(*dm_last_step);
+            else
+                cal(dm);
         }
     }
 }
@@ -147,9 +157,9 @@ void Exx_LRI_Interface<T, Tdata>::exx_hamilt2density(elecstate::ElecState& elec,
 template<typename T, typename Tdata>
 void Exx_LRI_Interface<T, Tdata>::exx_iter_finish(const K_Vectors& kv, const UnitCell& ucell,
     hamilt::Hamilt<T>& hamilt, elecstate::ElecState& elec, Charge_Mixing& chgmix,
-    const double& scf_ene_thr, int& iter, bool& conv_esolver)
+    const double& scf_ene_thr, int& iter, const int istep, bool& conv_esolver)
 {
-    if (GlobalC::restart.info_save.save_H && this->two_level_step > 0
+    if (GlobalC::restart.info_save.save_H && (this->two_level_step > 0 || istep > 0)
         && (!GlobalC::exx_info.info_global.separate_loop || iter == 1)) // to avoid saving the same value repeatedly
     {
         ////////// for Add_Hexx_Type::k
@@ -189,12 +199,14 @@ void Exx_LRI_Interface<T, Tdata>::exx_iter_finish(const K_Vectors& kv, const Uni
         {
             chgmix.close_kerker_gg0();
         }
+        this->dm_last_step = dynamic_cast<const elecstate::ElecStateLCAO<T>*>(&elec)->get_DM();
         conv_esolver = this->exx_after_converge(
             hamilt,
             *dynamic_cast<const elecstate::ElecStateLCAO<T>*>(&elec)->get_DM(),
             kv,
             PARAM.inp.nspin,
             iter,
+            istep,
             elec.f_en.etot,
             scf_ene_thr);
     }
@@ -208,6 +220,7 @@ bool Exx_LRI_Interface<T, Tdata>::exx_after_converge(
     const K_Vectors& kv,
     const int& nspin,
     int& iter,
+    const int& istep,
     const double& etot,
     const double& scf_ene_thr)
 {   // only called if (GlobalC::exx_info.info_global.cal_exx)
@@ -216,7 +229,7 @@ bool Exx_LRI_Interface<T, Tdata>::exx_after_converge(
             GlobalC::restart.info_load.restart_exx = true;
             this->exx_ptr->Eexx = 0;
         };
-
+        
         // no separate_loop case
         if (!GlobalC::exx_info.info_global.separate_loop)
         {
@@ -226,7 +239,7 @@ bool Exx_LRI_Interface<T, Tdata>::exx_after_converge(
             // in first scf loop, exx updated once in beginning,
             // in second scf loop, exx updated every iter
 
-            if (this->two_level_step)
+            if (this->two_level_step || istep > 0)
             {
                 restart_reset();
                 return true;
diff --git a/tests/integrate/285_NO_KP_RE_HSE/result.ref b/tests/integrate/285_NO_KP_RE_HSE/result.ref
@@ -1,3 +1,3 @@
-etotref -200.8969758707964
-etotperatomref -100.4484879354
+etotref -199.35070744
+etotperatomref -99.67535372
 totaltimeref 19.17
diff --git a/tests/integrate/286_NO_KP_CR_HSE/result.ref b/tests/integrate/286_NO_KP_CR_HSE/result.ref
diff --git a/tests/integrate/385_NO_GO_RE_S1_HSE/result.ref b/tests/integrate/385_NO_GO_RE_S1_HSE/result.ref
diff --git a/tests/integrate/386_NO_GO_MD_S1_HSE/result.ref b/tests/integrate/386_NO_GO_MD_S1_HSE/result.ref

Original file line number	Diff line number	Diff line change
`@@ -625,13 +625,15 @@ void ESolver_KS_LCAO<TK, TR>::iter_init(const int istep, const int iter)`
`625`	`625`	`// calculate exact-exchange`
`626`	`626`	`if (GlobalC::exx_info.info_ri.real_number)`
`627`	`627`	`{`
`628`		`- this->exd->exx_eachiterinit(dynamic_cast<const elecstate::ElecStateLCAO<TK>>(this->pelec)->get_DM(),`
	`628`	`+ this->exd->exx_eachiterinit(istep,`
	`629`	`+ dynamic_cast<const elecstate::ElecStateLCAO<TK>>(this->pelec)->get_DM(),`
`629`	`630`	`this->kv,`
`630`	`631`	`iter);`
`631`	`632`	`}`
`632`	`633`	`else`
`633`	`634`	`{`
`634`		`- this->exc->exx_eachiterinit(dynamic_cast<const elecstate::ElecStateLCAO<TK>>(this->pelec)->get_DM(),`
	`635`	`+ this->exc->exx_eachiterinit(istep,`
	`636`	`+ dynamic_cast<const elecstate::ElecStateLCAO<TK>>(this->pelec)->get_DM(),`
`635`	`637`	`this->kv,`
`636`	`638`	`iter);`
`637`	`639`	`}`
`@@ -914,12 +916,12 @@ void ESolver_KS_LCAO<TK, TR>::update_pot(const int istep, const int iter)`
`914`	`916`	`//! 5) cal_MW? (why put it here?)`
`915`	`917`	`//------------------------------------------------------------------------------`
`916`	`918`	`template <typename TK, typename TR>`
`917`		`-void ESolver_KS_LCAO<TK, TR>::iter_finish(int& iter)`
	`919`	`+void ESolver_KS_LCAO<TK, TR>::iter_finish(const int istep, int& iter)`
`918`	`920`	`{`
`919`	`921`	`ModuleBase::TITLE("ESolver_KS_LCAO", "iter_finish");`
`920`	`922`
`921`	`923`	`// call iter_finish() of ESolver_KS`
`922`		`- ESolver_KS<TK>::iter_finish(iter);`
	`924`	`+ ESolver_KS<TK>::iter_finish(istep, iter);`
`923`	`925`
`924`	`926`	`// 1) mix density matrix if mixing_restart + mixing_dmr + not first`
`925`	`927`	`// mixing_restart at every iter`
`@@ -944,8 +946,8 @@ void ESolver_KS_LCAO<TK, TR>::iter_finish(int& iter)`
`944`	`946`	`if (GlobalC::exx_info.info_global.cal_exx)`
`945`	`947`	`{`
`946`	`948`	`GlobalC::exx_info.info_ri.real_number ?`
`947`		`- this->exd->exx_iter_finish(this->kv, GlobalC::ucell, this->p_hamilt, this->pelec, *this->p_chgmix, this->scf_ene_thr, iter, this->conv_esolver) :`
`948`		`- this->exc->exx_iter_finish(this->kv, GlobalC::ucell, this->p_hamilt, this->pelec, *this->p_chgmix, this->scf_ene_thr, iter, this->conv_esolver);`
	`949`	`+ this->exd->exx_iter_finish(this->kv, GlobalC::ucell, this->p_hamilt, this->pelec, *this->p_chgmix, this->scf_ene_thr, iter, istep, this->conv_esolver) :`
	`950`	`+ this->exc->exx_iter_finish(this->kv, GlobalC::ucell, this->p_hamilt, this->pelec, *this->p_chgmix, this->scf_ene_thr, iter, istep, this->conv_esolver);`
`949`	`951`	`}`
`950`	`952`	`#endif`
`951`	`953`
Original file line number	Diff line number	Diff line change
`@@ -191,9 +191,9 @@ namespace ModuleESolver`
`191`	`191`	`}`
`192`	`192`
`193`	`193`	`template <typename T>`
`194`		`- void ESolver_KS_LIP<T>::iter_finish(int& iter)`
	`194`	`+ void ESolver_KS_LIP<T>::iter_finish(const int istep, int& iter)`
`195`	`195`	`{`
`196`		`- ESolver_KS_PW<T>::iter_finish(iter);`
	`196`	`+ ESolver_KS_PW<T>::iter_finish(istep, iter);`
`197`	`197`
`198`	`198`	`#ifdef __EXX`
`199`	`199`	`if (GlobalC::exx_info.info_global.cal_exx && this->conv_esolver)`
Original file line number	Diff line number	Diff line change
`@@ -434,10 +434,10 @@ void ESolver_KS_PW<T, Device>::update_pot(const int istep, const int iter)`
`434`	`434`	`}`
`435`	`435`
`436`	`436`	`template <typename T, typename Device>`
`437`		`-void ESolver_KS_PW<T, Device>::iter_finish(int& iter)`
	`437`	`+void ESolver_KS_PW<T, Device>::iter_finish(const int istep, int& iter)`
`438`	`438`	`{`
`439`	`439`	`// call iter_finish() of ESolver_KS`
`440`		`- ESolver_KS<T, Device>::iter_finish(iter);`
	`440`	`+ ESolver_KS<T, Device>::iter_finish(istep, iter);`
`441`	`441`
`442`	`442`	`// liuyu 2023-10-24`
`443`	`443`	`// D in uspp need vloc, thus needs update when veff updated`
Original file line number	Diff line number	Diff line change
`@@ -153,10 +153,10 @@ void ESolver_SDFT_PW<T, Device>::before_scf(const int istep)`
`153`	`153`	`}`
`154`	`154`
`155`	`155`	`template <typename T, typename Device>`
`156`		`-void ESolver_SDFT_PW<T, Device>::iter_finish(int& iter)`
	`156`	`+void ESolver_SDFT_PW<T, Device>::iter_finish(const int istep, int& iter)`
`157`	`157`	`{`
`158`	`158`	`// call iter_finish() of ESolver_KS`
`159`		`- ESolver_KS<T, Device>::iter_finish(iter);`
	`159`	`+ ESolver_KS<T, Device>::iter_finish(istep, iter);`
`160`	`160`	`}`
`161`	`161`
`162`	`162`	`template <typename T, typename Device>`