Skip to content

Commit becf33a

Browse files
haozhihanhaozhihan
committed
Centralize all psi initializes code into WFInit Class
1 parent 10d4ef7 commit becf33a

File tree

3 files changed

+232
-72
lines changed

3 files changed

+232
-72
lines changed

‎examples/scf/pw_Si2/test.txt‎

Lines changed: 132 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,132 @@
1+
2+
ABACUS v3.8.2
3+
4+
Atomic-orbital Based Ab-initio Computation at UStc
5+
6+
Website: http://abacus.ustc.edu.cn/
7+
Documentation: https://abacus.deepmodeling.com/
8+
Repository: https://github.com/abacusmodeling/abacus-develop
9+
https://github.com/deepmodeling/abacus-develop
10+
Commit: a2fdb95dc (Wed Nov 13 14:01:56 2024 +0800)
11+
12+
Fri Nov 15 12:45:46 2024
13+
MAKE THE DIR : OUT.ABACUS/
14+
RUNNING WITH DEVICE : CPU / Intel(R) Xeon(R) Gold 6132 CPU @ 2.60GHz
15+
UNIFORM GRID DIM : 36 * 36 * 36
16+
UNIFORM GRID DIM(BIG) : 36 * 36 * 36
17+
DONE(0.0216239 SEC) : SETUP UNITCELL
18+
DONE(0.0497069 SEC) : SYMMETRY
19+
DONE(0.13171 SEC) : INIT K-POINTS
20+
---------------------------------------------------------
21+
Self-consistent calculations for electrons
22+
---------------------------------------------------------
23+
SPIN KPOINTS PROCESSORS THREADS
24+
1 8 1 1
25+
---------------------------------------------------------
26+
Use plane wave basis
27+
---------------------------------------------------------
28+
ELEMENT NATOM XC
29+
Si 2
30+
---------------------------------------------------------
31+
Initial plane wave basis and FFT box
32+
---------------------------------------------------------
33+
DONE(0.133689 SEC) : INIT PLANEWAVE
34+
DONE(0.137845 SEC) : LOCAL POTENTIAL
35+
DONE(0.148574 SEC) : NON-LOCAL POTENTIAL
36+
MEMORY FOR PSI (MB) : 3.62305
37+
DONE(0.159973 SEC) : INIT BASIS
38+
-------------------------------------------
39+
SELF-CONSISTENT :
40+
-------------------------------------------
41+
START CHARGE : atomic
42+
DONE(0.191484 SEC) : INIT SCF
43+
ITER ETOT/eV EDIFF/eV DRHO TIME/s
44+
DS1 -2.15454699e+02 0.00000000e+00 6.9965e-02 1.05
45+
DS2 -2.15504108e+02 -4.94089775e-02 1.7575e-03 0.36
46+
DS3 -2.15505651e+02 -1.54281648e-03 2.6651e-05 0.46
47+
DS4 -2.15505696e+02 -4.51483582e-05 3.6536e-07 0.36
48+
DS5 -2.15505698e+02 -2.73278040e-06 3.5488e-08 0.39
49+
TIME STATISTICS
50+
-----------------------------------------------------------------------------
51+
CLASS_NAME NAME TIME/s CALLS AVG/s PER/%
52+
-----------------------------------------------------------------------------
53+
total 2.81 17 0.17 100.00
54+
Driver reading 0.01 1 0.01 0.34
55+
Input_Conv Convert 0.00 1 0.00 0.02
56+
Driver driver_line 2.80 1 2.80 99.66
57+
UnitCell check_tau 0.00 1 0.00 0.00
58+
PW_Basis_Sup setuptransform 0.00 1 0.00 0.07
59+
PW_Basis_Sup distributeg 0.00 1 0.00 0.02
60+
mymath heapsort 0.00 174 0.00 0.07
61+
Symmetry analy_sys 0.03 1 0.03 1.00
62+
PW_Basis_K setuptransform 0.00 1 0.00 0.04
63+
PW_Basis_K distributeg 0.00 1 0.00 0.01
64+
PW_Basis setup_struc_factor 0.00 1 0.00 0.03
65+
ppcell_vnl init 0.00 1 0.00 0.01
66+
ppcell_vl init_vloc 0.00 1 0.00 0.03
67+
ppcell_vnl init_vnl 0.01 1 0.01 0.38
68+
WF_atomic init_at_1 0.01 1 0.01 0.40
69+
wavefunc wfcinit 0.00 1 0.00 0.00
70+
Ions opt_ions 2.65 1 2.65 94.27
71+
ESolver_KS_PW runner 2.65 1 2.65 94.27
72+
ESolver_KS_PW before_scf 0.03 1 0.03 1.12
73+
H_Ewald_pw compute_ewald 0.00 1 0.00 0.01
74+
Charge set_rho_core 0.00 1 0.00 0.00
75+
Charge atomic_rho 0.01 2 0.00 0.22
76+
PW_Basis_Sup recip2real 0.02 20 0.00 0.72
77+
PW_Basis_Sup gathers_scatterp 0.00 20 0.00 0.07
78+
Potential init_pot 0.01 1 0.01 0.20
79+
Potential update_from_charge 0.03 6 0.00 1.00
80+
Potential cal_fixed_v 0.00 1 0.00 0.04
81+
PotLocal cal_fixed_v 0.00 1 0.00 0.04
82+
Potential cal_v_eff 0.03 6 0.00 0.95
83+
H_Hartree_pw v_hartree 0.01 6 0.00 0.49
84+
PW_Basis_Sup real2recip 0.02 26 0.00 0.80
85+
PW_Basis_Sup gatherp_scatters 0.00 26 0.00 0.04
86+
PotXC cal_v_eff 0.01 6 0.00 0.44
87+
XC_Functional v_xc 0.01 6 0.00 0.43
88+
Potential interpolate_vrs 0.00 6 0.00 0.01
89+
Symmetry rhog_symmetry 0.02 7 0.00 0.67
90+
Symmetry group fft grids 0.00 7 0.00 0.11
91+
WFInit initialize_psi 0.02 1 0.02 0.60
92+
Nonlocal getvnl 0.03 64 0.00 1.15
93+
pp_cell_vnl getvnl 0.03 64 0.00 1.14
94+
Structure_Factor get_sk 0.01 96 0.00 0.42
95+
WF_atomic atomic_wfc 0.01 16 0.00 0.32
96+
Charge_Mixing init_mixing 0.00 1 0.00 0.00
97+
ESolver_KS_PW hamilt2density_single 2.57 6 0.43 91.39
98+
HSolverPW solve 2.54 6 0.42 90.33
99+
Diago_DavSubspace diag_once 2.11 48 0.04 75.21
100+
Diago_DavSubspace first 0.76 48 0.02 26.90
101+
DavSubspace hpsi_func 1.93 198 0.01 68.72
102+
Operator hPsi 1.93 198 0.01 68.71
103+
Operator EkineticPW 0.00 198 0.00 0.18
104+
Operator VeffPW 1.90 198 0.01 67.64
105+
PW_Basis_K recip2real 1.28 2447 0.00 45.64
106+
PW_Basis_K gathers_scatterp 0.08 2447 0.00 2.67
107+
PW_Basis_K real2recip 0.89 1775 0.00 31.71
108+
PW_Basis_K gatherp_scatters 0.02 1775 0.00 0.61
109+
Operator NonlocalPW 0.03 198 0.00 0.89
110+
Nonlocal add_nonlocal_pp 0.01 198 0.00 0.34
111+
Diago_DavSubspace cal_elem 0.05 198 0.00 1.92
112+
Diago_DavSubspace diag_zhegvx 0.05 198 0.00 1.95
113+
Diago_DavSubspace cal_grad 1.24 150 0.01 44.16
114+
Diago_DavSubspace check_update 0.00 150 0.00 0.00
115+
Diago_DavSubspace last 0.02 56 0.00 0.75
116+
Diago_DavSubspace refresh 0.00 8 0.00 0.11
117+
ElecStatePW psiToRho 0.36 6 0.06 12.92
118+
Charge_Mixing get_drho 0.01 6 0.00 0.37
119+
Charge_Mixing inner_product_recip_rho 0.00 6 0.00 0.01
120+
Charge mix_rho 0.01 4 0.00 0.24
121+
Charge Broyden_mixing 0.00 4 0.00 0.04
122+
Charge_Mixing inner_product_recip_hartree 0.00 12 0.00 0.02
123+
ESolver_KS_PW after_scf 0.01 1 0.01 0.26
124+
ModuleIO write_rhog 0.00 1 0.00 0.08
125+
ModuleIO write_istate_info 0.00 1 0.00 0.02
126+
-----------------------------------------------------------------------------
127+
128+
129+
START Time : Fri Nov 15 12:45:46 2024
130+
FINISH Time : Fri Nov 15 12:45:49 2024
131+
TOTAL Time : 3
132+
SEE INFORMATION IN : OUT.ABACUS/

‎source/module_hamilt_pw/hamilt_pwdft/wfinit.cpp‎

Lines changed: 98 additions & 72 deletions
Original file line numberDiff line numberDiff line change
@@ -38,9 +38,10 @@ void WFInit<T, Device>::prepare_init(Structure_Factor* p_sf,
3838
#endif
3939
pseudopot_cell_vnl* p_ppcell)
4040
{
41-
if (!this->use_psiinitializer) {
41+
if (!this->use_psiinitializer)
42+
{
4243
return;
43-
}
44+
}
4445
// under restriction of C++11, std::unique_ptr can not be allocate via std::make_unique
4546
// use new instead, but will cause asymmetric allocation and deallocation, in literal aspect
4647
ModuleBase::timer::tick("WFInit", "prepare_init");
@@ -136,7 +137,9 @@ void WFInit<T, Device>::allocate_psi(Psi<std::complex<double>>*& psi,
136137
template <typename T, typename Device>
137138
void WFInit<T, Device>::make_table(const int nks, Structure_Factor* p_sf)
138139
{
139-
if (this->use_psiinitializer) {} // do not need to do anything because the interpolate table is unchanged
140+
if (this->use_psiinitializer)
141+
{
142+
} // do not need to do anything because the interpolate table is unchanged
140143
else // old initialization method, used in EXX calculation
141144
{
142145
this->p_wf->init_after_vc(nks); // reallocate wanf2, the planewave expansion of lcao
@@ -150,95 +153,118 @@ void WFInit<T, Device>::initialize_psi(Psi<std::complex<double>>* psi,
150153
hamilt::Hamilt<T, Device>* p_hamilt,
151154
std::ofstream& ofs_running)
152155
{
153-
if (!this->use_psiinitializer) { return; }
154156
ModuleBase::timer::tick("WFInit", "initialize_psi");
155-
// if psig is not allocated before, allocate it
156-
if (!this->psi_init->psig_use_count()) { this->psi_init->allocate(/*psig_only=*/true); }
157-
158-
// loop over kpoints, make it possible to only allocate memory for psig at the only one kpt
159-
// like (1, nbands, npwx), in which npwx is the maximal npw of all kpoints
160-
for (int ik = 0; ik < this->pw_wfc->nks; ik++)
157+
158+
if (this->use_psiinitializer)
161159
{
162-
//! Fix the wavefunction to initialize at given kpoint
163-
psi->fix_k(ik);
160+
// if psig is not allocated before, allocate it
161+
if (!this->psi_init->psig_use_count())
162+
{
163+
this->psi_init->allocate(/*psig_only=*/true);
164+
}
164165

165-
//! Update Hamiltonian from other kpoint to the given one
166-
p_hamilt->updateHk(ik);
166+
// loop over kpoints, make it possible to only allocate memory for psig at the only one kpt
167+
// like (1, nbands, npwx), in which npwx is the maximal npw of all kpoints
168+
for (int ik = 0; ik < this->pw_wfc->nks; ik++)
169+
{
170+
//! Fix the wavefunction to initialize at given kpoint
171+
psi->fix_k(ik);
167172

168-
//! Project atomic orbitals on |k+G> planewave basis, where k is wavevector of kpoint
169-
//! and G is wavevector of the peroiodic part of the Bloch function
170-
this->psi_init->proj_ao_onkG(ik);
173+
//! Update Hamiltonian from other kpoint to the given one
174+
p_hamilt->updateHk(ik);
171175

172-
//! psi_initializer manages memory of psig with shared pointer,
173-
//! its access to use is shared here via weak pointer
174-
//! therefore once the psi_initializer is destructed, psig will be destructed, too
175-
//! this way, we can avoid memory leak and undefined behavior
176-
std::weak_ptr<psi::Psi<T, Device>> psig = this->psi_init->share_psig();
176+
//! Project atomic orbitals on |k+G> planewave basis, where k is wavevector of kpoint
177+
//! and G is wavevector of the peroiodic part of the Bloch function
178+
this->psi_init->proj_ao_onkG(ik);
177179

178-
if (psig.expired())
179-
{
180-
ModuleBase::WARNING_QUIT("WFInit::initialize_psi", "psig lifetime is expired");
181-
}
180+
//! psi_initializer manages memory of psig with shared pointer,
181+
//! its access to use is shared here via weak pointer
182+
//! therefore once the psi_initializer is destructed, psig will be destructed, too
183+
//! this way, we can avoid memory leak and undefined behavior
184+
std::weak_ptr<psi::Psi<T, Device>> psig = this->psi_init->share_psig();
182185

183-
//! to use psig, we need to lock it to get a shared pointer version,
184-
//! then switch kpoint of psig to the given one
185-
auto psig_ = psig.lock();
186-
// CHANGE LOG: if not lcaoinpw, the psig will only be used in psi-initialization
187-
// so we can only allocate memory for one kpoint with the maximal number of pw
188-
// over all kpoints, then the memory space will be always enough. Then for each
189-
// kpoint, the psig is calculated in an overwrite manner.
190-
const int ik_psig = (psig_->get_nk() == 1) ? 0 : ik;
191-
psig_->fix_k(ik_psig);
192-
193-
std::vector<typename GetTypeReal<T>::type> etatom(psig_->get_nbands(), 0.0);
194-
195-
// then adjust dimension from psig to psi
196-
// either by matrix-multiplication or by copying-discarding
197-
if (this->psi_init->method() != "random")
198-
{
199-
// lcaoinpw and pw share the same esolver. In the future, we will have different esolver
200-
if (((this->ks_solver == "cg") || (this->ks_solver == "lapack")) && (this->basis_type == "pw"))
186+
if (psig.expired())
201187
{
202-
// the following function is only run serially, to be improved
203-
hsolver::DiagoIterAssist<T, Device>::diagH_subspace_init(p_hamilt,
204-
psig_->get_pointer(),
205-
psig_->get_nbands(),
206-
psig_->get_nbasis(),
207-
*(kspw_psi),
208-
etatom.data());
209-
continue;
188+
ModuleBase::WARNING_QUIT("WFInit::initialize_psi", "psig lifetime is expired");
210189
}
211-
else if ((this->ks_solver == "lapack") && (this->basis_type == "lcao_in_pw"))
190+
191+
//! to use psig, we need to lock it to get a shared pointer version,
192+
//! then switch kpoint of psig to the given one
193+
auto psig_ = psig.lock();
194+
// CHANGE LOG: if not lcaoinpw, the psig will only be used in psi-initialization
195+
// so we can only allocate memory for one kpoint with the maximal number of pw
196+
// over all kpoints, then the memory space will be always enough. Then for each
197+
// kpoint, the psig is calculated in an overwrite manner.
198+
const int ik_psig = (psig_->get_nk() == 1) ? 0 : ik;
199+
psig_->fix_k(ik_psig);
200+
201+
std::vector<typename GetTypeReal<T>::type> etatom(psig_->get_nbands(), 0.0);
202+
203+
// then adjust dimension from psig to psi
204+
// either by matrix-multiplication or by copying-discarding
205+
if (this->psi_init->method() != "random")
212206
{
213-
if (ik == 0)
207+
// lcaoinpw and pw share the same esolver. In the future, we will have different esolver
208+
if (((this->ks_solver == "cg") || (this->ks_solver == "lapack")) && (this->basis_type == "pw"))
209+
{
210+
// the following function is only run serially, to be improved
211+
hsolver::DiagoIterAssist<T, Device>::diagH_subspace_init(p_hamilt,
212+
psig_->get_pointer(),
213+
psig_->get_nbands(),
214+
psig_->get_nbasis(),
215+
*(kspw_psi),
216+
etatom.data());
217+
continue;
218+
}
219+
else if ((this->ks_solver == "lapack") && (this->basis_type == "lcao_in_pw"))
214220
{
215-
ofs_running << " START WAVEFUNCTION: LCAO_IN_PW, psi initialization skipped " << std::endl;
221+
if (ik == 0)
222+
{
223+
ofs_running << " START WAVEFUNCTION: LCAO_IN_PW, psi initialization skipped " << std::endl;
224+
}
225+
continue;
216226
}
217-
continue;
227+
// else the case is davidson
218228
}
219-
// else the case is davidson
220-
}
221-
else
222-
{
223-
if (this->ks_solver == "cg")
229+
else
224230
{
225-
hsolver::DiagoIterAssist<T, Device>::diagH_subspace(p_hamilt, *(psig_), *(kspw_psi), etatom.data());
226-
continue;
231+
if (this->ks_solver == "cg")
232+
{
233+
hsolver::DiagoIterAssist<T, Device>::diagH_subspace(p_hamilt, *(psig_), *(kspw_psi), etatom.data());
234+
continue;
235+
}
236+
// else the case is davidson
227237
}
228-
// else the case is davidson
229-
}
230238

231-
// for the Davidson method, we just copy the wavefunction (partially)
232-
for (int iband = 0; iband < kspw_psi->get_nbands(); iband++)
233-
{
234-
for (int ibasis = 0; ibasis < kspw_psi->get_nbasis(); ibasis++)
239+
// for the Davidson method, we just copy the wavefunction (partially)
240+
for (int iband = 0; iband < kspw_psi->get_nbands(); iband++)
235241
{
236-
(*(kspw_psi))(iband, ibasis) = (*psig_)(iband, ibasis);
242+
for (int ibasis = 0; ibasis < kspw_psi->get_nbasis(); ibasis++)
243+
{
244+
(*(kspw_psi))(iband, ibasis) = (*psig_)(iband, ibasis);
245+
}
237246
}
247+
} // end k-point loop
248+
249+
if (this->basis_type != "lcao_in_pw")
250+
{
251+
this->psi_init->deallocate_psig();
238252
}
239-
} // end k-point loop
253+
}
254+
else
255+
{
256+
for (int ik = 0; ik < this->pw_wfc->nks; ++ik)
257+
{
258+
//! Update Hamiltonian from other kpoint to the given one
259+
p_hamilt->updateHk(ik);
260+
261+
//! Fix the wavefunction to initialize at given kpoint
262+
kspw_psi->fix_k(ik);
240263

241-
if (this->basis_type != "lcao_in_pw") { this->psi_init->deallocate_psig(); }
264+
/// for psi init guess!!!!
265+
hamilt::diago_PAO_in_pw_k2(this->ctx, ik, *kspw_psi, this->pw_wfc, this->p_wf, p_hamilt);
266+
}
267+
}
242268

243269
ModuleBase::timer::tick("WFInit", "initialize_psi");
244270
}

‎source/module_hamilt_pw/hamilt_pwdft/wfinit.h‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -84,6 +84,8 @@ class WFInit
8484
std::string basis_type = "none";
8585
// pw basis
8686
ModulePW::PW_Basis_K* pw_wfc = nullptr;
87+
88+
Device* ctx = {};
8789
};
8890

8991
} // namespace psi

0 commit comments

Comments
 (0)