Skip to content

Commit 9bea38b

Browse files
committed
Refactor WFN extrapolation into module_extrap
1 parent d4fd81f commit 9bea38b

5 files changed

Lines changed: 277 additions & 104 deletions

File tree

source/source_esolver/esolver_ks_lcao.cpp

Lines changed: 14 additions & 100 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,4 @@
11
#include "esolver_ks_lcao.h"
2-
#include "source_base/global_function.h"
32
#include "source_base/module_external/blacs_connector.h"
43
#include "source_cell/module_neighbor/sltk_atom_arrange.h"
54
#include "source_estate/elecstate_tools.h"
@@ -18,11 +17,9 @@
1817
#include "../source_lcao/module_ri/exx_opt_orb.h"
1918
#endif
2019
#include "source_lcao/module_rdmft/rdmft.h"
21-
#include "source_lcao/module_extrap/wf_history_lcao.h"
22-
#include "source_lcao/module_operator_lcao/operator_lcao.h" // build overlap SR before WFN reorthonormalization
20+
#include "source_lcao/module_extrap/wf_extrap_lcao.h"
2321
#include "source_estate/module_charge/chgmixing.h" // use charge mixing, mohan add 20251006
2422
#include "source_estate/module_dm/init_dm.h" // init dm from electronic wave functions
25-
#include "source_estate/module_dm/cal_dm_psi.h" // rebuild DMK from extrapolated WFN
2623
#include "source_io/module_ctrl/ctrl_runner_lcao.h" // use ctrl_runner_lcao()
2724
#include "source_io/module_ctrl/ctrl_iter_lcao.h" // use ctrl_iter_lcao()
2825
#include "source_io/module_ctrl/ctrl_scf_lcao.h" // use ctrl_scf_lcao()
@@ -31,10 +28,6 @@
3128
#include "source_lcao/LCAO_set.h" // mohan add 20251111
3229
#include "source_psi/setup_psi.h" // use Setup_Psi for deallocate_psi
3330

34-
#include <iomanip>
35-
#include <sstream>
36-
#include <string>
37-
#include <type_traits>
3831

3932
namespace ModuleESolver
4033
{
@@ -92,13 +85,7 @@ void ESolver_KS_LCAO<TK, TR>::before_all_runners(UnitCell& ucell, const Input_pa
9285
LCAO_domain::set_psi_occ_dm_chg<TK>(this->kv, this->psi, this->pv, this->pelec,
9386
this->dmat, this->chr, inp);
9487

95-
const auto wfc_extrap_method = ModuleExtrap::wfc_extrap_method_from_string(inp.wfc_extrap);
96-
if (wfc_extrap_method != ModuleExtrap::WfcExtrapMethod::None)
97-
{
98-
this->wf_history_lcao_.reset(new ModuleExtrap::WfHistoryLCAO<TK>(wfc_extrap_method));
99-
GlobalV::ofs_running << " WFN extrapolation method: "
100-
<< ModuleExtrap::to_string(wfc_extrap_method) << std::endl;
101-
}
88+
this->wf_history_lcao_ = ModuleExtrap::make_wf_history_lcao<TK>(inp.wfc_extrap);
10289

10390
LCAO_domain::set_pot<TK>(ucell, this->kv, this->sf, *this->pw_rho, *this->pw_rhod,
10491
this->pelec, this->orb_, this->pv, this->locpp, this->dftu,
@@ -215,87 +202,16 @@ void ESolver_KS_LCAO<TK, TR>::before_scf(UnitCell& ucell, const int istep)
215202
}
216203
else if(PARAM.inp.esolver_type!="tddft")//initialize DMR from WFN history if required
217204
{
218-
if (std::is_same<TK, double>::value)
219-
{
220-
if (this->wf_history_lcao_ != nullptr)
221-
{
222-
// The newly-created HamiltLCAO has allocated S(R), but the actual
223-
// overlap values are normally filled later by updateHk() inside the
224-
// eigensolver. WFN extrapolation needs S before entering SCF, so build
225-
// only the overlap real-space matrix here and then fold it to S(Gamma).
226-
auto* lcao_op = dynamic_cast<hamilt::OperatorLCAO<TK, TR>*>(hamilt_lcao->getOperator());
227-
if (lcao_op == nullptr)
228-
{
229-
ModuleBase::WARNING_QUIT("ESolver_KS_LCAO::before_scf",
230-
"Failed to access the LCAO operator chain for WFN extrapolation.");
231-
}
232-
ModuleBase::timer::start("WFN_Extrap", "prepare_overlap");
233-
lcao_op->contributeHR();
234-
// Refresh the current Gamma-only overlap matrix before reusing the previous WFN.
235-
// The layout must follow the active LCAO solver because the WFN orthonormalizer
236-
// assembles distributed local blocks using the same column/row-major convention.
237-
const int sk_layout =
238-
ModuleBase::GlobalFunc::IS_COLUMN_MAJOR_KS_SOLVER(PARAM.inp.ks_solver) ? 1 : 0;
239-
hamilt_lcao->updateSk(0, sk_layout);
240-
ModuleBase::timer::end("WFN_Extrap", "prepare_overlap");
241-
242-
ModuleBase::timer::start("WFN_Extrap", "apply");
243-
const ModuleExtrap::WfExtrapApplyResult wfc_result
244-
= this->wf_history_lcao_->try_use_prev_wf_gamma(hamilt_lcao->getSk(),
245-
this->pv,
246-
*(this->psi),
247-
this->pelec->wg);
248-
ModuleBase::timer::end("WFN_Extrap", "apply");
249-
if (wfc_result.ok())
250-
{
251-
ModuleBase::timer::start("WFN_Extrap", "rebuild_density");
252-
elecstate::cal_dm_psi(this->dmat.dm->get_paraV_pointer(),
253-
this->pelec->wg,
254-
*(this->psi),
255-
*(this->dmat.dm));
256-
this->dmat.dm->cal_DMR();
257-
LCAO_domain::dm2rho(this->dmat.dm->get_DMR_vector(), PARAM.inp.nspin, &this->chr);
258-
ModuleBase::timer::end("WFN_Extrap", "rebuild_density");
259-
initialized_by_wfc_extrap = true;
260-
261-
GlobalV::ofs_running << " WFN extrapolation: use_prev_wf from ionic step "
262-
<< wfc_result.snapshot_istep
263-
<< ", active bands = " << wfc_result.nactive_bands
264-
<< ", max |C^T S C - I| = "
265-
<< wfc_result.max_orthonormality_deviation << std::endl;
266-
}
267-
else
268-
{
269-
std::ostringstream oss;
270-
oss << std::scientific << std::setprecision(16);
271-
oss << "WFN extrapolation failed with status: "
272-
<< ModuleExtrap::to_string(wfc_result.status) << "\n\n"
273-
<< " snapshot_istep=" << wfc_result.snapshot_istep << "\n"
274-
<< " failed_state=" << wfc_result.failed_state << "\n"
275-
<< " failed_pivot_index=" << wfc_result.failed_pivot_index << "\n"
276-
<< " failed_pivot=" << wfc_result.failed_pivot << "\n"
277-
<< " nstate=" << wfc_result.nstate << "\n"
278-
<< " nbands=" << wfc_result.nbands << "\n"
279-
<< " nbasis=" << wfc_result.nbasis << "\n"
280-
<< " min_metric_diag=" << wfc_result.min_metric_diag << "\n"
281-
<< " max_metric_diag=" << wfc_result.max_metric_diag << "\n"
282-
<< " max_metric_abs=" << wfc_result.max_metric_abs << "\n"
283-
<< " max_metric_asymmetry=" << wfc_result.max_metric_asymmetry << "\n"
284-
<< " max_orthonormality_deviation="
285-
<< wfc_result.max_orthonormality_deviation << "\n";
286-
const std::string message = oss.str();
287-
ModuleBase::WARNING_QUIT("ESolver_KS_LCAO::before_scf", message);
288-
}
289-
}
290-
}
291-
else
292-
{
293-
if (this->wf_history_lcao_ != nullptr)
294-
{
295-
ModuleBase::WARNING_QUIT("ESolver_KS_LCAO::before_scf",
296-
"WFN extrapolation is currently supported only for the real Gamma-only NAO path.");
297-
}
298-
}
205+
initialized_by_wfc_extrap = ModuleExtrap::initialize_gamma_density_from_history(
206+
this->wf_history_lcao_.get(),
207+
*hamilt_lcao,
208+
this->pv,
209+
*(this->psi),
210+
this->pelec->wg,
211+
*(this->dmat.dm),
212+
this->chr,
213+
PARAM.inp.nspin,
214+
PARAM.inp.ks_solver);
299215

300216
if (!initialized_by_wfc_extrap)
301217
{
@@ -654,10 +570,8 @@ void ESolver_KS_LCAO<TK, TR>::after_scf(UnitCell& ucell, const int istep, const
654570
this->rdmft_solver, this->deepks, this->exx_nao,
655571
this->conv_esolver, this->scf_nmax_flag, istep);
656572

657-
if (conv_esolver && this->wf_history_lcao_ != nullptr && this->psi != nullptr)
658-
{
659-
this->wf_history_lcao_->update_after_scf(istep, *(this->psi), this->pelec->wg);
660-
}
573+
ModuleExtrap::update_wf_history_after_scf(
574+
this->wf_history_lcao_.get(), conv_esolver, istep, this->psi, this->pelec->wg);
661575

662576
//! 3) Clean up RA, which is used to serach for adjacent atoms
663577
if (!PARAM.inp.cal_force && !PARAM.inp.cal_stress)

source/source_lcao/module_extrap/CMakeLists.txt

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,7 @@ add_library(
22
lcao_extrap
33
OBJECT
44
wf_history_lcao.cpp
5+
wf_extrap_lcao.cpp
56
wf_orthonormalize_lcao.cpp
67
)
78

Lines changed: 185 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,185 @@
1+
#include "source_lcao/module_extrap/wf_extrap_lcao.h"
2+
3+
#include "source_base/global_function.h"
4+
#include "source_base/global_variable.h"
5+
#include "source_base/timer.h"
6+
#include "source_base/tool_quit.h"
7+
#include "source_estate/module_charge/charge.h"
8+
#include "source_estate/module_dm/cal_dm_psi.h"
9+
#include "source_estate/module_dm/density_matrix.h"
10+
#include "source_lcao/hamilt_lcao.h"
11+
#include "source_lcao/module_operator_lcao/operator_lcao.h"
12+
#include "source_lcao/rho_tau_lcao.h"
13+
14+
#include <complex>
15+
#include <iomanip>
16+
#include <sstream>
17+
18+
namespace
19+
{
20+
21+
std::string wfc_extrapolation_failure_message(const ModuleExtrap::WfcExtrapApplyResult& result)
22+
{
23+
std::ostringstream oss;
24+
oss << std::scientific << std::setprecision(16);
25+
oss << "WFN extrapolation failed with status: " << ModuleExtrap::to_string(result.status) << "\n\n"
26+
<< " snapshot_istep=" << result.snapshot_istep << "\n"
27+
<< " failed_state=" << result.failed_state << "\n"
28+
<< " failed_pivot_index=" << result.failed_pivot_index << "\n"
29+
<< " failed_pivot=" << result.failed_pivot << "\n"
30+
<< " nstate=" << result.nstate << "\n"
31+
<< " nbands=" << result.nbands << "\n"
32+
<< " nbasis=" << result.nbasis << "\n"
33+
<< " min_metric_diag=" << result.min_metric_diag << "\n"
34+
<< " max_metric_diag=" << result.max_metric_diag << "\n"
35+
<< " max_metric_abs=" << result.max_metric_abs << "\n"
36+
<< " max_metric_asymmetry=" << result.max_metric_asymmetry << "\n"
37+
<< " max_orthonormality_deviation=" << result.max_orthonormality_deviation << "\n";
38+
return oss.str();
39+
}
40+
41+
} // namespace
42+
43+
namespace ModuleExtrap
44+
{
45+
46+
template <typename TK>
47+
std::unique_ptr<WfHistoryLCAO<TK>> make_wf_history_lcao(const std::string& method)
48+
{
49+
const WfcExtrapMethod wfc_extrap_method = wfc_extrap_method_from_string(method);
50+
if (wfc_extrap_method == WfcExtrapMethod::None)
51+
{
52+
return std::unique_ptr<WfHistoryLCAO<TK>>();
53+
}
54+
55+
GlobalV::ofs_running << " WFN extrapolation method: " << to_string(wfc_extrap_method) << std::endl;
56+
return std::unique_ptr<WfHistoryLCAO<TK>>(new WfHistoryLCAO<TK>(wfc_extrap_method));
57+
}
58+
59+
bool initialize_gamma_density_from_history(WfHistoryLCAO<double>* const history,
60+
hamilt::HamiltLCAO<double, double>& hamilt_lcao,
61+
const Parallel_Orbitals& pv,
62+
psi::Psi<double>& psi,
63+
const ModuleBase::matrix& wg,
64+
elecstate::DensityMatrix<double, double>& density_matrix,
65+
Charge& charge,
66+
const int nspin,
67+
const std::string& ks_solver)
68+
{
69+
if (history == nullptr)
70+
{
71+
return false;
72+
}
73+
74+
// S(R) is allocated with HamiltLCAO, but its values are normally filled later
75+
// by updateHk(). Build and fold the current overlap before restoring the WFN.
76+
auto* const lcao_op = dynamic_cast<hamilt::OperatorLCAO<double, double>*>(hamilt_lcao.getOperator());
77+
if (lcao_op == nullptr)
78+
{
79+
ModuleBase::WARNING_QUIT("ModuleExtrap::initialize_gamma_density_from_history",
80+
"Failed to access the LCAO operator chain for WFN extrapolation.");
81+
return false;
82+
}
83+
84+
ModuleBase::timer::start("WFN_Extrap", "prepare_overlap");
85+
lcao_op->contributeHR();
86+
const int sk_layout = ModuleBase::GlobalFunc::IS_COLUMN_MAJOR_KS_SOLVER(ks_solver) ? 1 : 0;
87+
hamilt_lcao.updateSk(0, sk_layout);
88+
ModuleBase::timer::end("WFN_Extrap", "prepare_overlap");
89+
90+
ModuleBase::timer::start("WFN_Extrap", "apply");
91+
const WfcExtrapApplyResult result
92+
= history->try_use_prev_wf_gamma(hamilt_lcao.getSk(), pv, psi, wg);
93+
ModuleBase::timer::end("WFN_Extrap", "apply");
94+
if (!result.ok())
95+
{
96+
ModuleBase::WARNING_QUIT("ModuleExtrap::initialize_gamma_density_from_history",
97+
wfc_extrapolation_failure_message(result));
98+
return false;
99+
}
100+
101+
ModuleBase::timer::start("WFN_Extrap", "rebuild_density");
102+
elecstate::cal_dm_psi(density_matrix.get_paraV_pointer(), wg, psi, density_matrix);
103+
density_matrix.cal_DMR();
104+
LCAO_domain::dm2rho(density_matrix.get_DMR_vector(), nspin, &charge);
105+
ModuleBase::timer::end("WFN_Extrap", "rebuild_density");
106+
107+
GlobalV::ofs_running << " WFN extrapolation: use_prev_wf from ionic step " << result.snapshot_istep
108+
<< ", active bands = " << result.nactive_bands
109+
<< ", max |C^T S C - I| = " << result.max_orthonormality_deviation << std::endl;
110+
return true;
111+
}
112+
113+
template <typename TK, typename TR>
114+
bool initialize_gamma_density_from_history(WfHistoryLCAO<TK>* const history,
115+
hamilt::HamiltLCAO<TK, TR>&,
116+
const Parallel_Orbitals&,
117+
psi::Psi<TK>&,
118+
const ModuleBase::matrix&,
119+
elecstate::DensityMatrix<TK, TR>&,
120+
Charge&,
121+
const int,
122+
const std::string&)
123+
{
124+
if (history == nullptr)
125+
{
126+
return false;
127+
}
128+
129+
ModuleBase::WARNING_QUIT("ModuleExtrap::initialize_gamma_density_from_history",
130+
"WFN extrapolation is currently supported only for the real Gamma-only NAO path.");
131+
return false;
132+
}
133+
134+
template <typename TK>
135+
void update_wf_history_after_scf(WfHistoryLCAO<TK>* const history,
136+
const bool converged,
137+
const int istep,
138+
const psi::Psi<TK>* const psi,
139+
const ModuleBase::matrix& wg)
140+
{
141+
if (converged && history != nullptr && psi != nullptr)
142+
{
143+
history->update_after_scf(istep, *psi, wg);
144+
}
145+
}
146+
147+
template std::unique_ptr<WfHistoryLCAO<double>> make_wf_history_lcao<double>(const std::string& method);
148+
template std::unique_ptr<WfHistoryLCAO<std::complex<double>>>
149+
make_wf_history_lcao<std::complex<double>>(const std::string& method);
150+
151+
template bool initialize_gamma_density_from_history<std::complex<double>, double>(
152+
WfHistoryLCAO<std::complex<double>>* history,
153+
hamilt::HamiltLCAO<std::complex<double>, double>& hamilt_lcao,
154+
const Parallel_Orbitals& pv,
155+
psi::Psi<std::complex<double>>& psi,
156+
const ModuleBase::matrix& wg,
157+
elecstate::DensityMatrix<std::complex<double>, double>& density_matrix,
158+
Charge& charge,
159+
int nspin,
160+
const std::string& ks_solver);
161+
162+
template bool initialize_gamma_density_from_history<std::complex<double>, std::complex<double>>(
163+
WfHistoryLCAO<std::complex<double>>* history,
164+
hamilt::HamiltLCAO<std::complex<double>, std::complex<double>>& hamilt_lcao,
165+
const Parallel_Orbitals& pv,
166+
psi::Psi<std::complex<double>>& psi,
167+
const ModuleBase::matrix& wg,
168+
elecstate::DensityMatrix<std::complex<double>, std::complex<double>>& density_matrix,
169+
Charge& charge,
170+
int nspin,
171+
const std::string& ks_solver);
172+
173+
template void update_wf_history_after_scf<double>(WfHistoryLCAO<double>* history,
174+
bool converged,
175+
int istep,
176+
const psi::Psi<double>* psi,
177+
const ModuleBase::matrix& wg);
178+
template void update_wf_history_after_scf<std::complex<double>>(
179+
WfHistoryLCAO<std::complex<double>>* history,
180+
bool converged,
181+
int istep,
182+
const psi::Psi<std::complex<double>>* psi,
183+
const ModuleBase::matrix& wg);
184+
185+
} // namespace ModuleExtrap

0 commit comments

Comments
 (0)