Code cleanup

Author: Growl1234

source/source_esolver/esolver_fp.cpp

@@ -73,7 +73,9 @@ void ESolver_FP::before_all_runners(BaseCell& basecell, const Input_para& inp)
     this->sf.set(this->pw_rhod, inp.nbspline);
 
     //! 4) init charge extrapolation
-    this->CE.Init_CE(inp.nspin, ucell.nat, this->pw_rhod->nrxx, inp.chg_extrap);
+    this->use_wfc_extrapolation_ = inp.wfc_extrap != "none";
+    this->CE.Init_CE(inp.nspin, ucell.nat, this->pw_rhod->nrxx,
+                     this->use_wfc_extrapolation_ ? "none" : inp.chg_extrap);
 
     //! 5) symmetry analysis should be performed every time the cell is changed
     if (ModuleSymmetry::Symmetry::symm_flag == 1)
@@ -180,19 +182,13 @@ void ESolver_FP::before_scf(UnitCell& ucell, const int istep)
     // charge extrapolation
     if (ucell.ionic_position_updated)
     {
-        this->CE.update_all_dis(ucell);
-
-        const bool skip_charge_extrap_for_wfc = PARAM.inp.basis_type == "lcao"
-                                                && PARAM.inp.wfc_extrap != "none"
-                                                && istep > 0;
-        if (skip_charge_extrap_for_wfc)
+        if (this->use_wfc_extrapolation_)
         {
             this->sf.setup(&ucell, this->Pgrid, this->pw_rhod);
-            GlobalV::ofs_running << " charge density extrapolation is skipped because wfc_extrap = "
-                                 << PARAM.inp.wfc_extrap << "." << std::endl;
         }
         else
         {
+            this->CE.update_all_dis(ucell);
             this->CE.extrapolate_charge(&this->Pgrid, ucell, &this->chr, &this->sf,
                                         GlobalV::ofs_running, GlobalV::ofs_warning);
         }

source/source_esolver/esolver_fp.h

@@ -82,6 +82,7 @@ class ESolver_FP : public ESolver
 
     //! charge extrapolation method
     Charge_Extra CE;
+    bool use_wfc_extrapolation_ = false;
 
     //! solvent model
     surchem solvent;

source/source_esolver/esolver_ks_lcao.cpp

@@ -90,8 +90,6 @@ void ESolver_KS_LCAO<TK, TR>::before_all_runners(BaseCell& basecell, const Input
     LCAO_domain::set_psi_occ_dm_chg<TK>(this->kv, this->psi, this->pv, this->pelec,
       this->dmat, this->chr, inp);
 
-    this->wf_history_lcao_.set_method(ModuleExtrap::wfc_extrap_method_from_string(inp.wfc_extrap));
-
     LCAO_domain::set_pot<TK>(ucell, this->kv, this->sf, *this->pw_rho, *this->pw_rhod,
       this->pelec, this->orb_, this->pv, this->locpp, this->dftu,
       this->solvent, this->exx_nao, this->deepks, inp);
@@ -205,13 +203,13 @@ 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 (!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 (this->use_wfc_extrapolation_)
+        {
+            this->wf_history_lcao_.initialize_gamma_density(
+                *hamilt_lcao, this->pv, *(this->psi), this->pelec->wg, *(this->dmat.dm), this->chr);
+        }
+        else
         {
-            // 13.1.2) two cases are considered:
-            // 1. DMK in DensityMatrix is not empty (istep > 0), then DMR is initialized by DMK
-            // 2. DMK in DensityMatrix is empty (istep == 0), then DMR is initialized by zeros
             this->dmat.dm->cal_DMR();
         }
     }
@@ -579,7 +577,7 @@ void ESolver_KS_LCAO<TK, TR>::after_scf(UnitCell& ucell, const int istep, const
             this->rdmft_solver, this->deepks, this->exx_nao,
             this->conv_esolver, this->scf_nmax_flag, istep);
 
-    if (conv_esolver && this->psi != nullptr)
+    if (conv_esolver && this->psi != nullptr && this->use_wfc_extrapolation_)
     {
         this->wf_history_lcao_.update_after_scf(istep, *(this->psi), this->pelec->wg);
     }

source/source_io/module_parameter/read_input_item_system.cpp

@@ -927,11 +927,14 @@ Available options are:
         item.type = "String";
         item.description = R"(Wavefunction extrapolation method for LCAO calculations.
 
-* none: Disable wavefunction-based extrapolation.
-* use_prev_wf: Use the previous ionic step wavefunctions as the initial guess.
+* none: Disable wavefunction-based extrapolation and use chg_extrap instead.
+* use_prev_wf: Restore the previous converged ionic-step wavefunctions,
+  reorthonormalize their occupied subspace in the current AO metric, and rebuild rho.
 
-This option is currently limited to Gamma-only LCAO calculations.
-The k-point, ASPC, and GExt_PROJ paths will be enabled by later updates.)";
+The first ionic step uses the normal init_wfc/init_chg initialization because no
+history exists yet. From the second ionic step onward, the selected WFN method
+must succeed; ABACUS does not silently fall back to charge-density extrapolation.
+This option is currently limited to Gamma-only LCAO calculations.)";
         item.default_value = "none";
         read_sync_string(input.wfc_extrap);
         item.check_value = [](const Input_Item& item, const Parameter& para) {

source/source_io/test/read_input_ptest.cpp

@@ -182,6 +182,7 @@ TEST_F(InputParaTest, ParaRead)
     EXPECT_EQ(param.inp.mem_saver, 0);
     EXPECT_EQ(param.inp.init_chg, "atomic");
     EXPECT_EQ(param.inp.chg_extrap, "atomic");
+    EXPECT_EQ(param.inp.wfc_extrap, "use_prev_wf");
     EXPECT_EQ(param.inp.out_freq_elec, 50);
     EXPECT_EQ(param.inp.out_freq_ion, 0);
     EXPECT_EQ(param.inp.out_freq_td, 0);

source/source_io/test/support/INPUT

@@ -57,6 +57,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
+wfc_extrap                     use_prev_wf #Gamma-only LCAO wavefunction initialization
 out_chg                        0 #>0 output charge density for selected electron steps
 out_pot                        2 #output realspace potential
 out_wfc_pw                     0 #output wave functions

source/source_lcao/module_extrap/test/wf_extrap_test.cpp

@@ -14,6 +14,9 @@ namespace
 constexpr int nstate = 1;
 constexpr int nbands = 3;
 constexpr int nbasis = 3;
+constexpr double occupation_threshold = 1.0e-12;
+constexpr double pivot_threshold = 1.0e-14;
+constexpr double check_tolerance = 1.0e-8;
 
 void initialize_serial_orbitals(Parallel_Orbitals& pv)
 {
@@ -72,6 +75,14 @@ const std::vector<double> coeff = {
     7.0, 8.0, 9.0,
 };
 
+ModuleExtrap::WfOrthonormalizeResult reorthonormalize(const Parallel_Orbitals& pv,
+                                                       psi::Psi<double>& wfc,
+                                                       const ModuleBase::matrix& occupations)
+{
+    return ModuleExtrap::reorthonormalize_gamma_lcao(overlap.data(), pv, wfc, occupations,
+                                                      occupation_threshold, pivot_threshold, check_tolerance);
+}
+
 } // namespace
 
 TEST(WfSnapshotLCAO, OwnsAndRestoresData)
@@ -102,8 +113,7 @@ TEST(WfOrthonormalizeLCAO, ReorthonormalizesOccupiedSubspace)
     psi::Psi<double> wfc = make_psi(coeff);
     const ModuleBase::matrix occupations = make_occupations();
 
-    const ModuleExtrap::WfOrthonormalizeResult result
-        = ModuleExtrap::reorthonormalize_gamma_lcao(overlap.data(), pv, wfc, occupations);
+    const ModuleExtrap::WfOrthonormalizeResult result = reorthonormalize(pv, wfc, occupations);
 
     ASSERT_TRUE(result.ok());
     EXPECT_EQ(result.nstate, nstate);
@@ -133,8 +143,7 @@ TEST(WfOrthonormalizeLCAO, DoesNotModifyPsiOnFailure)
     psi::Psi<double> wfc = make_psi(rank_deficient_coeff);
     const ModuleBase::matrix occupations = make_occupations();
 
-    const ModuleExtrap::WfOrthonormalizeResult result
-        = ModuleExtrap::reorthonormalize_gamma_lcao(overlap.data(), pv, wfc, occupations);
+    const ModuleExtrap::WfOrthonormalizeResult result = reorthonormalize(pv, wfc, occupations);
 
     EXPECT_EQ(result.status, ModuleExtrap::WfcExtrapStatus::OrthogonalizationFailed);
     for (std::size_t i = 0; i < rank_deficient_coeff.size(); ++i)
@@ -143,12 +152,45 @@ TEST(WfOrthonormalizeLCAO, DoesNotModifyPsiOnFailure)
     }
 }
 
+TEST(WfOrthonormalizeLCAO, HandlesBandDistributedCoefficients)
+{
+    Parallel_Orbitals pv;
+    initialize_serial_orbitals(pv);
+    pv.ncol_bands = 2;
+    pv.nbands = 2;
+#ifdef __MPI
+    pv.desc_wfc[3] = 2;
+#endif
+    psi::Psi<double> wfc(nstate, 2, nbasis, nbasis, true);
+    wfc.set_all_psi(coeff.data(), 2 * nbasis);
+
+    const ModuleExtrap::WfOrthonormalizeResult result = reorthonormalize(pv, wfc, make_occupations());
+    EXPECT_TRUE(result.ok());
+}
+
+TEST(WfHistoryLCAO, ReportsMissingHistory)
+{
+    Parallel_Orbitals pv;
+    initialize_serial_orbitals(pv);
+    ModuleExtrap::WfHistoryLCAO<double> history;
+    psi::Psi<double> predicted = make_psi(coeff);
+
+    const ModuleExtrap::WfExtrapApplyResult result
+        = history.try_use_prev_wf_gamma(overlap.data(),
+                                        pv,
+                                        predicted,
+                                        make_occupations(),
+                                        pivot_threshold,
+                                        check_tolerance);
+    EXPECT_EQ(result.status, ModuleExtrap::WfcExtrapStatus::EmptyHistory);
+}
+
 TEST(WfHistoryLCAO, UsesLatestOwnedSnapshot)
 {
     Parallel_Orbitals pv;
     initialize_serial_orbitals(pv);
     const ModuleBase::matrix occupations = make_occupations();
-    ModuleExtrap::WfHistoryLCAO<double> history(ModuleExtrap::WfcExtrapMethod::UsePrevWf, 2);
+    ModuleExtrap::WfHistoryLCAO<double> history;
 
     psi::Psi<double> first = make_psi(coeff);
     history.update_after_scf(3, first, occupations);
@@ -164,15 +206,16 @@ TEST(WfHistoryLCAO, UsesLatestOwnedSnapshot)
 
     psi::Psi<double> predicted(nstate, nbands, nbasis, nbasis, true);
     const ModuleExtrap::WfExtrapApplyResult result
-        = history.try_use_prev_wf_gamma(overlap.data(), pv, predicted, occupations);
+        = history.try_use_prev_wf_gamma(overlap.data(),
+                                        pv,
+                                        predicted,
+                                        occupations,
+                                        pivot_threshold,
+                                        check_tolerance);
 
     ASSERT_TRUE(result.ok());
     EXPECT_EQ(result.snapshot_istep, 4);
-    EXPECT_EQ(history.size(), 2U);
     EXPECT_NEAR(metric_element(predicted, overlap, 0, 0), 1.0, 1.0e-10);
     EXPECT_NEAR(metric_element(predicted, overlap, 1, 1), 1.0, 1.0e-10);
     EXPECT_NEAR(metric_element(predicted, overlap, 0, 1), 0.0, 1.0e-10);
-
-    history.set_max_depth(1);
-    EXPECT_EQ(history.size(), 1U);
 }

source/source_lcao/module_extrap/wf_density_init_lcao.cpp

@@ -1,7 +1,5 @@
 #include "source_lcao/module_extrap/wf_history_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"
@@ -35,44 +33,45 @@ std::string wfc_extrapolation_failure_message(const WfExtrapApplyResult& result)
         << " 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";
+        << " max_orthonormality_deviation=" << result.max_orthonormality_deviation << "\n"
+        << "No fallback was attempted; set wfc_extrap to none to use chg_extrap.\n";
     return oss.str();
 }
 
 } // namespace
 
 template <typename TK>
-bool WfHistoryLCAO<TK>::initialize_gamma_density(hamilt::Hamilt<TK>&,
+void 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&)
+                                                  Charge&)
 {
-    if (!this->enabled() || this->empty())
+    if (!this->has_snapshot_)
     {
-        return false;
+        ModuleBase::WARNING_QUIT("WfHistoryLCAO::initialize_gamma_density",
+                                 "wfc_extrap was selected, but no converged wavefunction history is available. "
+                                 "No fallback was attempted; set wfc_extrap to none to use chg_extrap.");
     }
     ModuleBase::WARNING_QUIT("WfHistoryLCAO::initialize_gamma_density",
-                             "WFN extrapolation is currently supported only for the real Gamma-only NAO path.");
-    return false;
+                             "WFN extrapolation is supported only for the real Gamma-only NAO path. "
+                             "No fallback was attempted; set wfc_extrap to none to use chg_extrap.");
 }
 
 template <>
-bool WfHistoryLCAO<double>::initialize_gamma_density(hamilt::Hamilt<double>& hamiltonian,
+void 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)
+                                                       Charge& charge)
 {
-    if (!this->enabled() || this->empty())
+    if (!this->has_snapshot_)
     {
-        return false;
+        ModuleBase::WARNING_QUIT("WfHistoryLCAO::initialize_gamma_density",
+                                 "wfc_extrap was selected, but no converged wavefunction history is available. "
+                                 "No fallback was attempted; set wfc_extrap to none to use chg_extrap.");
     }
 
     auto* hamilt_lcao = dynamic_cast<hamilt::HamiltLCAO<double, double>*>(&hamiltonian);
@@ -91,12 +90,12 @@ bool WfHistoryLCAO<double>::initialize_gamma_density(hamilt::Hamilt<double>& ham
 
     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);
+    hamilt_lcao->updateSk(0, 1);
     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);
+    const WfExtrapApplyResult result
+        = this->try_use_prev_wf_gamma(hamilt_lcao->getSk(), pv, psi, wg_now, 1.0e-14, 1.0e-8);
     ModuleBase::timer::end("WFN_Extrap", "apply");
     if (!result.ok())
     {
@@ -107,18 +106,12 @@ bool WfHistoryLCAO<double>::initialize_gamma_density(hamilt::Hamilt<double>& ham
     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);
+    LCAO_domain::dm2rho(dmat.get_DMR_vector(), charge.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 bool WfHistoryLCAO<std::complex<double>>::initialize_gamma_density(
+template void WfHistoryLCAO<std::complex<double>>::initialize_gamma_density(
     hamilt::Hamilt<std::complex<double>>&, const Parallel_Orbitals&, psi::Psi<std::complex<double>>&,
-    const ModuleBase::matrix&, elecstate::DensityMatrix<std::complex<double>, double>&, Charge&, int,
-    const std::string&);
+    const ModuleBase::matrix&, elecstate::DensityMatrix<std::complex<double>, double>&, Charge&);
 
 } // namespace ModuleExtrap