Code cleanup

Author: Growl1234

docs/advanced/input_files/input-main.md

@@ -856,11 +856,14 @@
 - **Type**: String
 - **Description**: 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.
 
+  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.
-  The k-point, ASPC, and GExt_PROJ paths will be enabled by later updates.
 - **Default**: none
 
 ### nb2d

docs/parameters.yaml

@@ -339,11 +339,14 @@ parameters:
     description: |
       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.
 
+      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.
-      The k-point, ASPC, and GExt_PROJ paths will be enabled by later updates.
     default_value: none
     unit: ""
     availability: ""

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)
@@ -182,19 +184,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_inp_sys.cpp

@@ -929,11 +929,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/CMakeLists.txt

@@ -22,43 +22,24 @@ if(ENABLE_LCAO)
         spar_hsr.cpp
         spar_st.cpp
         spar_u.cpp
-<<<<<<< HEAD
         lcao_set.cpp
-		lcao_set_fs.cpp
-		lcao_set_st.cpp
-		lcao_nl_mu.cpp	
+        lcao_set_fs.cpp
+        lcao_set_st.cpp
+        lcao_nl_mu.cpp
         lcao_set_zero.cpp
         lcao_allocate.cpp
         lcao_set_mat2d.cpp
         lcao_init_basis.cpp
-=======
-        LCAO_set.cpp
-        LCAO_set_fs.cpp
-        LCAO_set_st.cpp
-        LCAO_nl_mu.cpp
-        LCAO_set_zero.cpp
-        LCAO_allocate.cpp
-        LCAO_set_mat2d.cpp
-        LCAO_init_basis.cpp
->>>>>>> 2a092828f (CMake cleanup of source_lcao)
         setup_nonlocal.cpp
         setup_exx.cpp
         setup_deepks.cpp
         setup_dm.cpp
         rho_tau_lcao.cpp
-<<<<<<< HEAD
-		record_adj.cpp
+        record_adj.cpp
         center2orb.cpp
         center2orb_orb11.cpp
         center2orb_orb21.cpp
         center2orb_orb22.cpp
-=======
-        record_adj.cpp
-        center2_orb.cpp
-        center2_orb-orb11.cpp
-        center2_orb-orb21.cpp
-        center2_orb-orb22.cpp
->>>>>>> 2a092828f (CMake cleanup of source_lcao)
         wavefunc_in_pw.cpp
     )
 

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);
 }