Refactor WFN extrapolation into module_extrap

Author: Growl1234

source/source_esolver/esolver_ks_lcao.cpp

@@ -19,10 +19,8 @@
 #endif
 #include "source_lcao/module_rdmft/rdmft.h"
 #include "source_lcao/module_extrap/wf_history_lcao.h"
-#include "source_lcao/module_operator_lcao/operator_lcao.h" // build overlap SR before WFN reorthonormalization
 #include "source_estate/module_charge/chgmixing.h" // use charge mixing, mohan add 20251006
 #include "source_estate/module_dm/init_dm.h" // init dm from electronic wave functions
-#include "source_estate/module_dm/cal_dm_psi.h" // rebuild DMK from extrapolated WFN
 #include "source_io/module_ctrl/ctrl_runner_lcao.h" // use ctrl_runner_lcao() 
 #include "source_io/module_ctrl/ctrl_iter_lcao.h" // use ctrl_iter_lcao() 
 #include "source_io/module_ctrl/ctrl_scf_lcao.h" // use ctrl_scf_lcao()
@@ -31,10 +29,6 @@
 #include "source_lcao/LCAO_set.h" // mohan add 20251111
 #include "source_psi/setup_psi.h" // use Setup_Psi for deallocate_psi
 
-#include <iomanip>
-#include <sstream>
-#include <string>
-#include <type_traits>
 
 namespace ModuleESolver
 {
@@ -215,86 +209,17 @@ void ESolver_KS_LCAO<TK, TR>::before_scf(UnitCell& ucell, const int istep)
     }
     else if(PARAM.inp.esolver_type!="tddft")//initialize DMR from WFN history if required
     {
-        if (std::is_same<TK, double>::value)
+        if (this->wf_history_lcao_ != nullptr)
         {
-            if (this->wf_history_lcao_ != nullptr)
-            {
-                // The newly-created HamiltLCAO has allocated S(R), but the actual
-                // overlap values are normally filled later by updateHk() inside the
-                // eigensolver.  WFN extrapolation needs S before entering SCF, so build
-                // only the overlap real-space matrix here and then fold it to S(Gamma).
-                auto* lcao_op = dynamic_cast<hamilt::OperatorLCAO<TK, TR>*>(hamilt_lcao->getOperator());
-                if (lcao_op == nullptr)
-                {
-                    ModuleBase::WARNING_QUIT("ESolver_KS_LCAO::before_scf",
-                                             "Failed to access the LCAO operator chain for WFN extrapolation.");
-                }
-                ModuleBase::timer::start("WFN_Extrap", "prepare_overlap");
-                lcao_op->contributeHR();
-                // Refresh the current Gamma-only overlap matrix before reusing the previous WFN.
-                // The layout must follow the active LCAO solver because the WFN orthonormalizer
-                // assembles distributed local blocks using the same column/row-major convention.
-                const int sk_layout =
-                    ModuleBase::GlobalFunc::IS_COLUMN_MAJOR_KS_SOLVER(PARAM.inp.ks_solver) ? 1 : 0;
-                hamilt_lcao->updateSk(0, sk_layout);
-                ModuleBase::timer::end("WFN_Extrap", "prepare_overlap");
-
-                ModuleBase::timer::start("WFN_Extrap", "apply");
-                const ModuleExtrap::WfExtrapApplyResult wfc_result
-                    = this->wf_history_lcao_->try_use_prev_wf_gamma(hamilt_lcao->getSk(),
-                                                                     this->pv,
-                                                                     *(this->psi),
-                                                                     this->pelec->wg);
-                ModuleBase::timer::end("WFN_Extrap", "apply");
-                if (wfc_result.ok())
-                {
-                    ModuleBase::timer::start("WFN_Extrap", "rebuild_density");
-                    elecstate::cal_dm_psi(this->dmat.dm->get_paraV_pointer(),
-                                          this->pelec->wg,
-                                          *(this->psi),
-                                          *(this->dmat.dm));
-                    this->dmat.dm->cal_DMR();
-                    LCAO_domain::dm2rho(this->dmat.dm->get_DMR_vector(), PARAM.inp.nspin, &this->chr);
-                    ModuleBase::timer::end("WFN_Extrap", "rebuild_density");
-                    initialized_by_wfc_extrap = true;
-
-                    GlobalV::ofs_running << " WFN extrapolation: use_prev_wf from ionic step "
-                                         << wfc_result.snapshot_istep
-                                         << ", active bands = " << wfc_result.nactive_bands
-                                         << ", max |C^T S C - I| = "
-                                         << wfc_result.max_orthonormality_deviation << std::endl;
-                }
-                else
-                {
-                    std::ostringstream oss;
-                    oss << std::scientific << std::setprecision(16);
-                    oss << "WFN extrapolation failed with status: "
-                        << ModuleExtrap::to_string(wfc_result.status) << "\n\n"
-                        << " snapshot_istep=" << wfc_result.snapshot_istep << "\n"
-                        << " failed_state=" << wfc_result.failed_state << "\n"
-                        << " failed_pivot_index=" << wfc_result.failed_pivot_index << "\n"
-                        << " failed_pivot=" << wfc_result.failed_pivot << "\n"
-                        << " nstate=" << wfc_result.nstate << "\n"
-                        << " nbands=" << wfc_result.nbands << "\n"
-                        << " nbasis=" << wfc_result.nbasis << "\n"
-                        << " min_metric_diag=" << wfc_result.min_metric_diag << "\n"
-                        << " max_metric_diag=" << wfc_result.max_metric_diag << "\n"
-                        << " max_metric_abs=" << wfc_result.max_metric_abs << "\n"
-                        << " max_metric_asymmetry=" << wfc_result.max_metric_asymmetry << "\n"
-                        << " max_orthonormality_deviation="
-                        << wfc_result.max_orthonormality_deviation << "\n";
-                    const std::string message = oss.str();
-                    ModuleBase::WARNING_QUIT("ESolver_KS_LCAO::before_scf", message);
-                }
-            }
-        }
-        else
-        {
-            if (this->wf_history_lcao_ != nullptr)
-            {
-                ModuleBase::WARNING_QUIT("ESolver_KS_LCAO::before_scf",
-                                         "WFN extrapolation is currently supported only for the real Gamma-only NAO path.");
-            }
+            initialized_by_wfc_extrap = this->wf_history_lcao_->initialize_gamma_density(
+                *hamilt_lcao,
+                this->pv,
+                *(this->psi),
+                this->pelec->wg,
+                *(this->dmat.dm),
+                this->chr,
+                PARAM.inp.nspin,
+                PARAM.inp.ks_solver);
         }
 
         if (!initialized_by_wfc_extrap)

source/source_lcao/module_extrap/wf_history_lcao.cpp

@@ -1,15 +1,48 @@
 #include "source_lcao/module_extrap/wf_history_lcao.h"
 
 #include "source_lcao/module_extrap/wf_orthonormalize_lcao.h"
+#include "source_base/global_function.h"
+#include "source_base/global_variable.h"
 #include "source_base/timer.h"
+#include "source_base/tool_quit.h"
+#include "source_estate/module_dm/cal_dm_psi.h"
+#include "source_lcao/hamilt_lcao.h"
+#include "source_lcao/module_operator_lcao/operator_lcao.h"
+#include "source_lcao/rho_tau_lcao.h"
 
 #include <algorithm>
 #include <complex>
-#include <type_traits>
+#include <iomanip>
+#include <sstream>
 
 namespace ModuleExtrap
 {
 
+namespace
+{
+
+std::string wfc_extrapolation_failure_message(const WfExtrapApplyResult& result)
+{
+    std::ostringstream oss;
+    oss << std::scientific << std::setprecision(16);
+    oss << "WFN extrapolation failed with status: " << to_string(result.status) << "\n\n"
+        << " snapshot_istep=" << result.snapshot_istep << "\n"
+        << " failed_state=" << result.failed_state << "\n"
+        << " failed_pivot_index=" << result.failed_pivot_index << "\n"
+        << " failed_pivot=" << result.failed_pivot << "\n"
+        << " nstate=" << result.nstate << "\n"
+        << " nbands=" << result.nbands << "\n"
+        << " nbasis=" << result.nbasis << "\n"
+        << " min_metric_diag=" << result.min_metric_diag << "\n"
+        << " max_metric_diag=" << result.max_metric_diag << "\n"
+        << " max_metric_abs=" << result.max_metric_abs << "\n"
+        << " max_metric_asymmetry=" << result.max_metric_asymmetry << "\n"
+        << " max_orthonormality_deviation=" << result.max_orthonormality_deviation << "\n";
+    return oss.str();
+}
+
+} // namespace
+
 template <typename TK>
 WfHistoryLCAO<TK>::WfHistoryLCAO(const WfcExtrapMethod method, const std::size_t max_depth)
     : method_(method), max_depth_(std::max<std::size_t>(max_depth, 1))
@@ -196,6 +229,71 @@ WfExtrapApplyResult WfHistoryLCAO<double>::try_use_prev_wf_gamma(const double* c
     return result;
 }
 
+template <typename TK>
+bool WfHistoryLCAO<TK>::initialize_gamma_density(hamilt::Hamilt<TK>&,
+                                                  const Parallel_Orbitals&,
+                                                  psi::Psi<TK>&,
+                                                  const ModuleBase::matrix&,
+                                                  elecstate::DensityMatrix<TK, double>&,
+                                                  Charge&,
+                                                  const int,
+                                                  const std::string&)
+{
+    ModuleBase::WARNING_QUIT("WfHistoryLCAO::initialize_gamma_density",
+                             "WFN extrapolation is currently supported only for the real Gamma-only NAO path.");
+}
+
+template <>
+bool WfHistoryLCAO<double>::initialize_gamma_density(hamilt::Hamilt<double>& hamiltonian,
+                                                       const Parallel_Orbitals& pv,
+                                                       psi::Psi<double>& psi,
+                                                       const ModuleBase::matrix& wg_now,
+                                                       elecstate::DensityMatrix<double, double>& dmat,
+                                                       Charge& charge,
+                                                       const int nspin,
+                                                       const std::string& ks_solver)
+{
+    auto* hamilt_lcao = dynamic_cast<hamilt::HamiltLCAO<double, double>*>(&hamiltonian);
+    if (hamilt_lcao == nullptr)
+    {
+        ModuleBase::WARNING_QUIT("WfHistoryLCAO::initialize_gamma_density",
+                                 "Failed to access the Gamma-only LCAO Hamiltonian for WFN extrapolation.");
+    }
+
+    auto* lcao_op = dynamic_cast<hamilt::OperatorLCAO<double, double>*>(hamilt_lcao->getOperator());
+    if (lcao_op == nullptr)
+    {
+        ModuleBase::WARNING_QUIT("WfHistoryLCAO::initialize_gamma_density",
+                                 "Failed to access the LCAO operator chain for WFN extrapolation.");
+    }
+
+    ModuleBase::timer::start("WFN_Extrap", "prepare_overlap");
+    lcao_op->contributeHR();
+    const int sk_layout = ModuleBase::GlobalFunc::IS_COLUMN_MAJOR_KS_SOLVER(ks_solver) ? 1 : 0;
+    hamilt_lcao->updateSk(0, sk_layout);
+    ModuleBase::timer::end("WFN_Extrap", "prepare_overlap");
+
+    ModuleBase::timer::start("WFN_Extrap", "apply");
+    const WfExtrapApplyResult result = this->try_use_prev_wf_gamma(hamilt_lcao->getSk(), pv, psi, wg_now);
+    ModuleBase::timer::end("WFN_Extrap", "apply");
+    if (!result.ok())
+    {
+        ModuleBase::WARNING_QUIT("WfHistoryLCAO::initialize_gamma_density",
+                                 wfc_extrapolation_failure_message(result));
+    }
+
+    ModuleBase::timer::start("WFN_Extrap", "rebuild_density");
+    elecstate::cal_dm_psi(dmat.get_paraV_pointer(), wg_now, psi, dmat);
+    dmat.cal_DMR();
+    LCAO_domain::dm2rho(dmat.get_DMR_vector(), nspin, &charge);
+    ModuleBase::timer::end("WFN_Extrap", "rebuild_density");
+
+    GlobalV::ofs_running << " WFN extrapolation: use_prev_wf from ionic step " << result.snapshot_istep
+                         << ", active bands = " << result.nactive_bands
+                         << ", max |C^T S C - I| = " << result.max_orthonormality_deviation << std::endl;
+    return true;
+}
+
 template <typename TK>
 void WfHistoryLCAO<TK>::prune_history()
 {

source/source_lcao/module_extrap/wf_history_lcao.h

@@ -4,9 +4,19 @@
 #include "source_lcao/module_extrap/wf_extrap_method.h"
 #include "source_lcao/module_extrap/wf_snapshot_lcao.h"
 #include "source_basis/module_ao/parallel_orbitals.h"
+#include "source_hamilt/hamilt.h"
 
 #include <cstddef>
 #include <deque>
+#include <string>
+
+namespace elecstate
+{
+template <typename TK, typename TR>
+class DensityMatrix;
+}
+
+class Charge;
 
 namespace ModuleExtrap
 {
@@ -79,6 +89,23 @@ class WfHistoryLCAO
                                               double pivot_threshold = 1.0e-14,
                                               double check_tolerance = 1.0e-8);
 
+    /**
+     * Initialize the Gamma-only charge density from the latest WFN snapshot.
+     *
+     * The real Gamma-only implementation prepares the current overlap, restores
+     * and reorthonormalizes the previous WFN, and rebuilds DMK, DMR, and rho.
+     * Unsupported paths terminate with a clear diagnostic rather than silently
+     * falling back to charge-density extrapolation.
+     */
+    bool initialize_gamma_density(hamilt::Hamilt<TK>& hamiltonian,
+                                  const Parallel_Orbitals& pv,
+                                  psi::Psi<TK>& psi,
+                                  const ModuleBase::matrix& wg_now,
+                                  elecstate::DensityMatrix<TK, double>& dmat,
+                                  Charge& charge,
+                                  int nspin,
+                                  const std::string& ks_solver);
+
   private:
     void prune_history();