refactor esolver_ks_pw (deepmodeling#7008)

* refactor(esolver): extract update_cell_pw function from esolver_ks_pw

- Create new files update_cell_pw.h and update_cell_pw.cpp in source_pw/module_pwdft
- Extract cell parameter update logic from ESolver_KS_PW::before_scf()
- The new function handles:
  1. Rescaling non-local pseudopotential (ppcell.rescale_vnl)
  2. Reinitializing plane wave basis grids (pw_wfc->initgrids/initparameters/collect_local_pw)
- Keep psi initialization (p_psi_init->prepare_init) in esolver to avoid template dependency
- Update CMakeLists.txt and Makefile.Objects for new source files

This refactoring improves code organization by moving PW-specific cell update
logic out of the esolver, making the esolver code cleaner and more focused
on high-level workflow control.

* refactor(esolver): extract EXX initialization into Exx_Helper::init

- Add init() function to Exx_Helper class for EXX initialization
- The init function handles:
  1. Check if calculation type is scf/relax/cell-relax/md
  2. Check if cal_exx is enabled
  3. Set XC first loop if separate_loop is true
  4. Set wg pointer for EXX calculation
- Simplify ESolver_KS_PW::before_all_runners() by calling exx_helper.init()
- Move EXX-specific logic out of esolver, improving code organization

This refactoring makes the esolver code cleaner and more focused
on high-level workflow control.

* refactor(esolver): extract DFT+U initialization into pw::iter_init_dftu_pw

- Create new files dftu_pw.h and dftu_pw.cpp in source_pw/module_pwdft
- Extract DFT+U occupation update logic from ESolver_KS_PW::iter_init()
- The new function handles:
  1. Check if DFT+U is enabled
  2. Check iteration and step conditions
  3. Call cal_occ_pw for occupation calculation
  4. Output DFT+U results
- Use void* for psi parameter to avoid template dependency
- Update CMakeLists.txt and Makefile.Objects for new source files

This refactoring improves code organization by moving DFT+U specific
logic out of the esolver, making the esolver code cleaner and more
focused on high-level workflow control.

* refactor(esolver): extract DeltaSpin lambda loop into pw::run_deltaspin_lambda_loop

- Create new files deltaspin_pw.h and deltaspin_pw.cpp in source_pw/module_pwdft
- Extract DeltaSpin lambda loop logic from ESolver_KS_PW::hamilt2rho_single()
- The new function handles:
  1. Check if DeltaSpin (sc_mag_switch) is enabled
  2. Get SpinConstrain singleton instance
  3. Run lambda loop to constrain atomic magnetic moments
  4. Return skip_solve flag to control solver execution
- Add Doxygen-style comments in English
- Update CMakeLists.txt and Makefile.Objects for new source files

This refactoring improves code organization by moving DeltaSpin-specific
logic out of the esolver, making the esolver code cleaner and more
focused on high-level workflow control.

* refactor(esolver): extract DeltaSpin oscillation check into pw::check_deltaspin_oscillation

- Add check_deltaspin_oscillation() function to deltaspin_pw.h/cpp
- Extract DeltaSpin SCF oscillation check logic from ESolver_KS_PW::iter_finish()
- The new function handles:
  1. Check if DeltaSpin (sc_mag_switch) is enabled
  2. Get SpinConstrain singleton instance
  3. Detect SCF oscillation using if_scf_oscillate()
  4. Set mixing_restart_step if oscillation detected
- Add Doxygen-style comments in English

This refactoring consolidates all DeltaSpin-related functions in one place,
making the code more modular and easier to maintain.

* refactor(esolver): extract EXX before_scf setup into Exx_Helper::before_scf

- Add before_scf() function to Exx_Helper class
- Extract EXX setup logic from ESolver_KS_PW::before_scf()
- The new function handles:
  1. Check if calculation type is valid (scf/relax/cell-relax/md)
  2. Check if EXX is enabled and basis type is PW
  3. Set EXX helper to Hamiltonian
  4. Set psi for EXX calculation
- Use void* for p_hamilt parameter to avoid circular dependency
- Add Doxygen-style comments in English

This refactoring consolidates EXX-related setup logic in the Exx_Helper class,
making the code more modular and easier to maintain.

* refactor(esolver): extract EXX iter_finish logic into Exx_Helper::iter_finish

- Add iter_finish() function to Exx_Helper class
- Extract EXX convergence handling logic from ESolver_KS_PW::iter_finish()
- The new function handles:
  1. Check if EXX is enabled
  2. Handle separate_loop mode for EXX convergence
  3. Calculate EXX energy difference for energy threshold
  4. Update potential if SCF not converged
  5. Increment EXX iteration counter
- Use Charge* and void* parameters to avoid circular dependency
- Add Doxygen-style comments in English

This refactoring consolidates all EXX-related functions in the Exx_Helper class,
making the code more modular and easier to maintain.

---------

Co-authored-by: abacus_fixer <mohanchen@pku.eud.cn>

Author: mohanchen

source/Makefile.Objects

@@ -705,6 +705,9 @@ OBJS_SRCPW=H_Ewald_pw.o\
     setup_pot.o\
     setup_pwrho.o\
     setup_pwwfc.o\
+    update_cell_pw.o\
+    dftu_pw.o\
+    deltaspin_pw.o\
     forces.o\
     forces_us.o\
     forces_nl.o\

source/source_esolver/esolver_ks_pw.cpp

@@ -28,6 +28,9 @@
 #include "source_io/module_ctrl/ctrl_output_pw.h" // mohan add 20250927
 #include "source_estate/module_charge/chgmixing.h" // use charge mixing, mohan add 20251006 
 #include "source_estate/update_pot.h" // mohan add 20251016
+#include "source_pw/module_pwdft/update_cell_pw.h" // mohan add 20250309
+#include "source_pw/module_pwdft/dftu_pw.h" // mohan add 20250309
+#include "source_pw/module_pwdft/deltaspin_pw.h" // mohan add 20250309
 
 #include "source_hamilt/module_xc/exx_info.h" // use GlobalC::exx_info
 
@@ -93,20 +96,7 @@ void ESolver_KS_PW<T, Device>::before_all_runners(UnitCell& ucell, const Input_p
     ModuleBase::GlobalFunc::DONE(GlobalV::ofs_running, "INIT BASIS");
 
     //! Initialize exx pw
-    if (inp.calculation == "scf" || inp.calculation == "relax" || inp.calculation == "cell-relax"
-        || inp.calculation == "md")
-    {
-        if (GlobalC::exx_info.info_global.cal_exx && GlobalC::exx_info.info_global.separate_loop == true)
-        {
-            XC_Functional::set_xc_first_loop(ucell);
-            exx_helper.set_firstiter();
-        }
-
-        if (GlobalC::exx_info.info_global.cal_exx)
-        {
-            exx_helper.set_wg(&this->pelec->wg);
-        }
-    }
+    this->exx_helper.init(ucell, inp, this->pelec->wg);
 }
 
 template <typename T, typename Device>
@@ -119,17 +109,10 @@ void ESolver_KS_PW<T, Device>::before_scf(UnitCell& ucell, const int istep)
     ESolver_KS<T, Device>::before_scf(ucell, istep);
 
     //! Init variables (once the cell has changed)
+    pw::update_cell_pw(ucell, this->ppcell, this->kv, this->pw_wfc, PARAM.inp);
+
     if (ucell.cell_parameter_updated)
     {
-        this->ppcell.rescale_vnl(ucell.omega);
-        ModuleBase::GlobalFunc::DONE(GlobalV::ofs_running, "NON-LOCAL POTENTIAL");
-
-        this->pw_wfc->initgrids(ucell.lat0, ucell.latvec, this->pw_wfc->nx, this->pw_wfc->ny, this->pw_wfc->nz);
-
-        this->pw_wfc->initparameters(false, PARAM.inp.ecutwfc, this->kv.get_nks(), this->kv.kvec_d.data());
-
-        this->pw_wfc->collect_local_pw(PARAM.inp.erf_ecut, PARAM.inp.erf_height, PARAM.inp.erf_sigma);
-
         this->stp.p_psi_init->prepare_init(PARAM.inp.pw_seed);
     }
 
@@ -151,17 +134,8 @@ void ESolver_KS_PW<T, Device>::before_scf(UnitCell& ucell, const int istep)
     // setup psi (electronic wave functions)
     this->stp.init(this->p_hamilt);
 
-    //! Exx calculations
-    if (PARAM.inp.calculation == "scf" || PARAM.inp.calculation == "relax" 
-        || PARAM.inp.calculation == "cell-relax" || PARAM.inp.calculation == "md")
-    {
-        if (GlobalC::exx_info.info_global.cal_exx && PARAM.inp.basis_type == "pw")
-        {
-            auto hamilt_pw = reinterpret_cast<hamilt::HamiltPW<T, Device>*>(this->p_hamilt);
-            hamilt_pw->set_exx_helper(exx_helper);
-            exx_helper.set_psi(this->stp.psi_t);
-        }
-    }
+    //! Setup EXX helper for Hamiltonian and psi
+    exx_helper.before_scf(this->p_hamilt, this->stp.psi_t, PARAM.inp);
 
     ModuleBase::timer::tick("ESolver_KS_PW", "before_scf");
 }
@@ -181,16 +155,7 @@ void ESolver_KS_PW<T, Device>::iter_init(UnitCell& ucell, const int istep, const
 
     // 4) update local occupations for DFT+U
     // should before lambda loop in DeltaSpin
-    if (PARAM.inp.dft_plus_u && (iter != 1 || istep != 0))
-    {
-        // only old DFT+U method should calculate energy correction in esolver,
-        // new DFT+U method will calculate energy when evaluating the Hamiltonian
-        if (this->dftu.omc != 2)
-        {
-            this->dftu.cal_occ_pw(iter, this->stp.psi_t, this->pelec->wg, ucell, PARAM.inp.mixing_beta);
-        }
-        this->dftu.output(ucell);
-    }
+    pw::iter_init_dftu_pw(iter, istep, this->dftu, this->stp.psi_t, this->pelec->wg, ucell, PARAM.inp);
 }
 
 // Temporary, it should be replaced by hsolver later.
@@ -218,26 +183,7 @@ void ESolver_KS_PW<T, Device>::hamilt2rho_single(UnitCell& ucell, const int iste
     bool skip_charge = PARAM.inp.calculation == "nscf" ? true : false;
 
     // run the inner lambda loop to contrain atomic moments with the DeltaSpin method
-    bool skip_solve = false;
-
-    if (PARAM.inp.sc_mag_switch)
-    {
-        spinconstrain::SpinConstrain<std::complex<double>>& sc
-            = spinconstrain::SpinConstrain<std::complex<double>>::getScInstance();
-        if (!sc.mag_converged() && this->drho > 0 && this->drho < PARAM.inp.sc_scf_thr)
-        {
-            // optimize lambda to get target magnetic moments, but the lambda is not near target
-            sc.run_lambda_loop(iter - 1);
-            sc.set_mag_converged(true);
-            skip_solve = true;
-        }
-        else if (sc.mag_converged())
-        {
-            // optimize lambda to get target magnetic moments, but the lambda is not near target
-            sc.run_lambda_loop(iter - 1);
-            skip_solve = true;
-        }
-    }
+    bool skip_solve = pw::run_deltaspin_lambda_loop(iter - 1, this->drho, PARAM.inp);
 
     if (!skip_solve)
     {
@@ -293,65 +239,11 @@ void ESolver_KS_PW<T, Device>::iter_finish(UnitCell& ucell, const int istep, int
         this->ppcell.cal_effective_D(veff, this->pw_rhod, ucell);
     }
 
-    // Related to EXX
-    if (GlobalC::exx_info.info_global.cal_exx)
-    {
-        if (GlobalC::exx_info.info_global.separate_loop)
-        {
-            if (conv_esolver)
-            {
-                auto start = std::chrono::high_resolution_clock::now();
-                exx_helper.set_firstiter(false);
-                exx_helper.op_exx->first_iter = false;
-                double dexx = 0.0;
-                if (PARAM.inp.exx_thr_type == "energy")
-                {
-                    dexx = exx_helper.cal_exx_energy(this->stp.psi_t);
-                    exx_helper.set_psi(this->stp.psi_t);
-                    dexx -= exx_helper.cal_exx_energy(this->stp.psi_t);
-                    // std::cout << "dexx = " << dexx << std::endl;
-                }
-                bool conv_ene = std::abs(dexx) < PARAM.inp.exx_ene_thr;
-
-                conv_esolver = exx_helper.exx_after_converge(iter, conv_ene);
-                if (!conv_esolver)
-                {
-                    if (PARAM.inp.exx_thr_type != "energy")
-                    {
-                        exx_helper.set_psi(this->stp.psi_t);
-                    }
-                    auto duration = std::chrono::high_resolution_clock::now() - start;
-                    std::cout << " Setting Psi for EXX PW Inner Loop took "
-                              << std::chrono::duration_cast<std::chrono::milliseconds>(duration).count() / 1000.0 << "s"
-                              << std::endl;
-                    exx_helper.op_exx->first_iter = false;
-                    XC_Functional::set_xc_type(ucell.atoms[0].ncpp.xc_func);
-					elecstate::update_pot(ucell, this->pelec, this->chr, conv_esolver);
-					exx_helper.iter_inc();
-                }
-            }
-        }
-        else
-        {
-            exx_helper.set_psi(this->stp.psi_t);
-        }
-    }
+    // Handle EXX-related operations after SCF iteration
+    exx_helper.iter_finish(this->pelec, &this->chr, this->stp.psi_t, ucell, PARAM.inp, conv_esolver, iter);
 
     // check if oscillate for delta_spin method
-    if (PARAM.inp.sc_mag_switch)
-    {
-        spinconstrain::SpinConstrain<std::complex<double>>& sc
-            = spinconstrain::SpinConstrain<std::complex<double>>::getScInstance();
-        if (!sc.higher_mag_prec)
-        {
-            sc.higher_mag_prec = this->p_chgmix->if_scf_oscillate(iter, 
-              this->drho, PARAM.inp.sc_os_ndim, PARAM.inp.scf_os_thr);
-            if (sc.higher_mag_prec)
-            { // if oscillate, increase the precision of magnetization and do mixing_restart in next iteration
-                this->p_chgmix->mixing_restart_step = iter + 1;
-            }
-        }
-    }
+    pw::check_deltaspin_oscillation(iter, this->drho, this->p_chgmix, PARAM.inp);
 
     // the output quantities
     ModuleIO::ctrl_iter_pw(istep, iter, conv_esolver, this->stp.psi_cpu, 

source/source_pw/module_pwdft/CMakeLists.txt

@@ -15,6 +15,9 @@ list(APPEND objects
     setup_pot.cpp
     setup_pwrho.cpp
     setup_pwwfc.cpp
+    update_cell_pw.cpp
+    dftu_pw.cpp
+    deltaspin_pw.cpp
     forces_nl.cpp
     forces_cc.cpp
     forces_scc.cpp

source/source_pw/module_pwdft/deltaspin_pw.cpp

@@ -0,0 +1,72 @@
+#include "source_pw/module_pwdft/deltaspin_pw.h"
+#include "source_lcao/module_deltaspin/spin_constrain.h"
+#include "source_estate/module_charge/charge_mixing.h"
+
+namespace pw
+{
+
+bool run_deltaspin_lambda_loop(const int iter,
+                               const double drho,
+                               const Input_para& inp)
+{
+    /// Return false if DeltaSpin is not enabled
+    if (!inp.sc_mag_switch)
+    {
+        return false;
+    }
+
+    /// Get the singleton instance of SpinConstrain
+    spinconstrain::SpinConstrain<std::complex<double>>& sc
+        = spinconstrain::SpinConstrain<std::complex<double>>::getScInstance();
+
+    /// Case 1: Magnetic moments not yet converged and SCF is close to convergence.
+    /// This is the first time we enter the lambda loop after SCF is nearly converged.
+    if (!sc.mag_converged() && drho > 0 && drho < inp.sc_scf_thr)
+    {
+        /// Optimize lambda to get target magnetic moments
+        sc.run_lambda_loop(iter);
+        sc.set_mag_converged(true);
+        return true;
+    }
+    /// Case 2: Magnetic moments already converged in previous iteration.
+    /// Continue to refine lambda in subsequent SCF iterations.
+    else if (sc.mag_converged())
+    {
+        sc.run_lambda_loop(iter);
+        return true;
+    }
+
+    /// Default: run the normal solver
+    return false;
+}
+
+void check_deltaspin_oscillation(const int iter,
+                                 const double drho,
+                                 Charge_Mixing* p_chgmix,
+                                 const Input_para& inp)
+{
+    /// Return if DeltaSpin is not enabled
+    if (!inp.sc_mag_switch)
+    {
+        return;
+    }
+
+    /// Get the singleton instance of SpinConstrain
+    spinconstrain::SpinConstrain<std::complex<double>>& sc
+        = spinconstrain::SpinConstrain<std::complex<double>>::getScInstance();
+
+    /// Check if higher magnetization precision is needed
+    if (!sc.higher_mag_prec)
+    {
+        /// Detect SCF oscillation
+        sc.higher_mag_prec = p_chgmix->if_scf_oscillate(iter, drho, inp.sc_os_ndim, inp.scf_os_thr);
+
+        /// If oscillation detected, set mixing restart step for next iteration
+        if (sc.higher_mag_prec)
+        {
+            p_chgmix->mixing_restart_step = iter + 1;
+        }
+    }
+}
+
+}

source/source_pw/module_pwdft/deltaspin_pw.h

@@ -0,0 +1,46 @@
+#ifndef DELTASPIN_PW_H
+#define DELTASPIN_PW_H
+
+#include "source_io/module_parameter/parameter.h"
+
+class Charge_Mixing;
+
+namespace pw
+{
+
+/**
+ * @brief Run the inner lambda loop for DeltaSpin method to constrain atomic magnetic moments.
+ *
+ * This function is used in the PW basis SCF iteration to optimize lambda parameters
+ * for constraining atomic magnetic moments to target values using the DeltaSpin method.
+ *
+ * @param iter The current iteration number (0-indexed).
+ * @param drho The current charge density difference.
+ * @param inp The input parameters.
+ * @return true if the solver should be skipped (lambda loop was executed),
+ *         false otherwise.
+ */
+bool run_deltaspin_lambda_loop(const int iter,
+                               const double drho,
+                               const Input_para& inp);
+
+/**
+ * @brief Check if SCF oscillation occurs for DeltaSpin method.
+ *
+ * This function checks if the SCF iteration is oscillating and sets the
+ * mixing restart step if oscillation is detected. This is used to increase
+ * the precision of magnetization calculation.
+ *
+ * @param iter The current iteration number (1-indexed).
+ * @param drho The current charge density difference.
+ * @param p_chgmix Pointer to the Charge_Mixing object.
+ * @param inp The input parameters.
+ */
+void check_deltaspin_oscillation(const int iter,
+                                 const double drho,
+                                 Charge_Mixing* p_chgmix,
+                                 const Input_para& inp);
+
+}
+
+#endif

source/source_pw/module_pwdft/dftu_pw.cpp

@@ -0,0 +1,32 @@
+#include "source_pw/module_pwdft/dftu_pw.h"
+#include "source_lcao/module_dftu/dftu.h"
+
+namespace pw
+{
+
+void iter_init_dftu_pw(const int iter,
+                       const int istep,
+                       Plus_U& dftu,
+                       const void* psi,
+                       const ModuleBase::matrix& wg,
+                       const UnitCell& ucell,
+                       const Input_para& inp)
+{
+    if (!inp.dft_plus_u)
+    {
+        return;
+    }
+
+    if (iter == 1 && istep == 0)
+    {
+        return;
+    }
+
+    if (dftu.omc != 2)
+    {
+        dftu.cal_occ_pw(iter, psi, wg, ucell, inp.mixing_beta);
+    }
+    dftu.output(ucell);
+}
+
+}

source/source_pw/module_pwdft/dftu_pw.h

@@ -0,0 +1,23 @@
+#ifndef DFTU_PW_H
+#define DFTU_PW_H
+
+#include "source_io/module_parameter/parameter.h"
+#include "source_cell/unitcell.h"
+#include "source_base/matrix.h"
+
+class Plus_U;
+
+namespace pw
+{
+
+void iter_init_dftu_pw(const int iter,
+                       const int istep,
+                       Plus_U& dftu,
+                       const void* psi,
+                       const ModuleBase::matrix& wg,
+                       const UnitCell& ucell,
+                       const Input_para& inp);
+
+}
+
+#endif