Skip to content

Commit 3c0fcbb

Browse files
author
dyzheng
committed
Fix: apply the S matrix in dav_subspace for USPP (backport deepmodeling#6428)
For ultrasoft pseudopotentials the eigenvalue problem is the generalized one H psi = eps S psi, but Diago_DavSubspace only formed the reduced overlap scc from the raw wavefunctions (psi^T psi) and used the standard residual, so USPP + ks_solver dav_subspace produced wrong eigenvalues. - Thread an spsi_func (S x block-vector) through diag/diag_once and use it to build sphi = S * psi in the reduced basis: scc is now psi^T (S psi), and the subspace residual uses S psi, so the preconditioned directions are (H - eps S) psi. - Update sphi together with hphi in refresh when the basis is rotated. - Pass the sPsi operator from hsolver_pw.cpp; update the module_lr call site to the new diag signature (identity S, unchanged behavior). Verified on the USPP NaCl test: ks_solver dav_subspace now converges to the same total energy as dav (-1675.06210 eV in 9 SCF iterations), while before the fix the eigenvalues were unphysical (~-6e4 eV). The identity-S (non-USPP) path is unchanged: 101_PW_15_pseudopots with dav_subspace matches its reference total energy, and the official 101_PW_upf201_uspp_Fe test still passes.
1 parent 9c3faeb commit 3c0fcbb

4 files changed

Lines changed: 93 additions & 15 deletions

File tree

‎source/module_hsolver/diago_dav_subspace.cpp‎

Lines changed: 67 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -43,6 +43,10 @@ Diago_DavSubspace<T, Device>::Diago_DavSubspace(const std::vector<Real>& precond
4343
resmem_complex_op()(this->ctx, this->hphi, this->nbase_x * this->dim, "DAV::hphi");
4444
setmem_complex_op()(this->ctx, this->hphi, 0, this->nbase_x * this->dim);
4545

46+
// the product of S and psi in the reduced psi set
47+
resmem_complex_op()(this->ctx, this->sphi, this->nbase_x * this->dim, "DAV::sphi");
48+
setmem_complex_op()(this->ctx, this->sphi, 0, this->nbase_x * this->dim);
49+
4650
// Hamiltonian on the reduced psi set
4751
resmem_complex_op()(this->ctx, this->hcc, this->nbase_x * this->nbase_x, "DAV::hcc");
4852
setmem_complex_op()(this->ctx, this->hcc, 0, this->nbase_x * this->nbase_x);
@@ -71,6 +75,7 @@ Diago_DavSubspace<T, Device>::~Diago_DavSubspace()
7175
delmem_complex_op()(this->ctx, this->psi_in_iter);
7276

7377
delmem_complex_op()(this->ctx, this->hphi);
78+
delmem_complex_op()(this->ctx, this->sphi);
7479
delmem_complex_op()(this->ctx, this->hcc);
7580
delmem_complex_op()(this->ctx, this->scc);
7681
delmem_complex_op()(this->ctx, this->vcc);
@@ -85,6 +90,7 @@ Diago_DavSubspace<T, Device>::~Diago_DavSubspace()
8590

8691
template <typename T, typename Device>
8792
int Diago_DavSubspace<T, Device>::diag_once(const HPsiFunc& hpsi_func,
93+
const HPsiFunc& spsi_func,
8894
T* psi_in,
8995
const int psi_in_dmax,
9096
Real* eigenvalue_in_hsolver,
@@ -125,7 +131,18 @@ int Diago_DavSubspace<T, Device>::diag_once(const HPsiFunc& hpsi_func,
125131
// hphi[:, 0:nbase_x] = H * psi_in_iter[:, 0:nbase_x]
126132
hpsi_func(this->psi_in_iter, this->hphi, this->dim, this->notconv);
127133

128-
this->cal_elem(this->dim, nbase, this->notconv, this->psi_in_iter, this->hphi, this->hcc, this->scc);
134+
// compute s*psi_in_iter
135+
// sphi[:, 0:nbase_x] = S * psi_in_iter[:, 0:nbase_x]
136+
spsi_func(this->psi_in_iter, this->sphi, this->dim, this->notconv);
137+
138+
this->cal_elem(this->dim,
139+
nbase,
140+
this->notconv,
141+
this->psi_in_iter,
142+
this->sphi,
143+
this->hphi,
144+
this->hcc,
145+
this->scc);
129146

130147
this->diag_zhegvx(nbase, this->notconv, this->hcc, this->scc, this->nbase_x, &eigenvalue_iter, this->vcc);
131148

@@ -143,16 +160,25 @@ int Diago_DavSubspace<T, Device>::diag_once(const HPsiFunc& hpsi_func,
143160
dav_iter++;
144161

145162
this->cal_grad(hpsi_func,
163+
spsi_func,
146164
this->dim,
147165
nbase,
148166
this->notconv,
149167
this->psi_in_iter,
150168
this->hphi,
169+
this->sphi,
151170
this->vcc,
152171
unconv.data(),
153172
&eigenvalue_iter);
154173

155-
this->cal_elem(this->dim, nbase, this->notconv, this->psi_in_iter, this->hphi, this->hcc, this->scc);
174+
this->cal_elem(this->dim,
175+
nbase,
176+
this->notconv,
177+
this->psi_in_iter,
178+
this->sphi,
179+
this->hphi,
180+
this->hcc,
181+
this->scc);
156182

157183
this->diag_zhegvx(nbase, this->n_band, this->hcc, this->scc, this->nbase_x, &eigenvalue_iter, this->vcc);
158184

@@ -231,6 +257,7 @@ int Diago_DavSubspace<T, Device>::diag_once(const HPsiFunc& hpsi_func,
231257
eigenvalue_in_hsolver,
232258
this->psi_in_iter,
233259
this->hphi,
260+
this->sphi,
234261
this->hcc,
235262
this->scc,
236263
this->vcc);
@@ -248,11 +275,13 @@ int Diago_DavSubspace<T, Device>::diag_once(const HPsiFunc& hpsi_func,
248275

249276
template <typename T, typename Device>
250277
void Diago_DavSubspace<T, Device>::cal_grad(const HPsiFunc& hpsi_func,
278+
const HPsiFunc& spsi_func,
251279
const int& dim,
252280
const int& nbase,
253281
const int& notconv,
254282
T* psi_iter,
255283
T* hphi,
284+
T* spsi,
256285
T* vcc,
257286
const int* unconv,
258287
std::vector<Real>* eigenvalue_iter)
@@ -326,7 +355,7 @@ void Diago_DavSubspace<T, Device>::cal_grad(const HPsiFunc& hpsi_func,
326355
notconv,
327356
nbase,
328357
this->one,
329-
psi_iter,
358+
sphi,
330359
this->dim,
331360
vcc,
332361
this->nbase_x,
@@ -392,6 +421,9 @@ void Diago_DavSubspace<T, Device>::cal_grad(const HPsiFunc& hpsi_func,
392421
// hpsi[:, nbase:nbase+notconv] = H * psi_iter[:, nbase:nbase+notconv]
393422
hpsi_func(psi_iter + nbase * dim, hphi + nbase * this->dim, this->dim, notconv);
394423

424+
// spsi[:, nbase:nbase+notconv] = S * psi_iter[:, nbase:nbase+notconv]
425+
spsi_func(psi_iter + nbase * dim, spsi + nbase * this->dim, this->dim, notconv);
426+
395427
ModuleBase::timer::tick("Diago_DavSubspace", "cal_grad");
396428
return;
397429
}
@@ -401,6 +433,7 @@ void Diago_DavSubspace<T, Device>::cal_elem(const int& dim,
401433
int& nbase,
402434
const int& notconv,
403435
const T* psi_iter,
436+
const T* spsi,
404437
const T* hphi,
405438
T* hcc,
406439
T* scc)
@@ -441,7 +474,7 @@ void Diago_DavSubspace<T, Device>::cal_elem(const int& dim,
441474
this->one,
442475
psi_iter,
443476
this->dim,
444-
psi_iter + nbase * this->dim,
477+
spsi + nbase * this->dim,
445478
this->dim,
446479
this->zero,
447480
&scc[nbase * this->nbase_x],
@@ -633,10 +666,11 @@ void Diago_DavSubspace<T, Device>::refresh(const int& dim,
633666
const Real* eigenvalue_in_hsolver,
634667
// const psi::Psi<T, Device>& psi,
635668
T* psi_iter,
636-
T* hp,
637-
T* sp,
638-
T* hc,
639-
T* vc)
669+
T* hphi,
670+
T* sphi,
671+
T* hcc,
672+
T* scc,
673+
T* vcc)
640674
{
641675
ModuleBase::timer::tick("Diago_DavSubspace", "refresh");
642676

@@ -663,6 +697,29 @@ void Diago_DavSubspace<T, Device>::refresh(const int& dim,
663697
// update hphi
664698
syncmem_complex_op()(this->ctx, this->ctx, hphi, psi_iter + nband * this->dim, this->dim * nband);
665699

700+
#ifdef __DSP
701+
gemm_op_mt<T, Device>()
702+
#else
703+
gemm_op<T, Device>()
704+
#endif
705+
(this->ctx,
706+
'N',
707+
'N',
708+
this->dim,
709+
nband,
710+
nbase,
711+
this->one,
712+
this->sphi,
713+
this->dim,
714+
this->vcc,
715+
this->nbase_x,
716+
this->zero,
717+
psi_iter + nband * this->dim,
718+
this->dim);
719+
720+
// update sphi
721+
syncmem_complex_op()(this->ctx, this->ctx, sphi, psi_iter + nband * this->dim, this->dim * nband);
722+
666723
nbase = nband;
667724

668725
// set hcc/scc/vcc to 0
@@ -728,6 +785,7 @@ void Diago_DavSubspace<T, Device>::refresh(const int& dim,
728785

729786
template <typename T, typename Device>
730787
int Diago_DavSubspace<T, Device>::diag(const HPsiFunc& hpsi_func,
788+
const HPsiFunc& spsi_func,
731789
T* psi_in,
732790
const int psi_in_dmax,
733791
Real* eigenvalue_in_hsolver,
@@ -743,7 +801,7 @@ int Diago_DavSubspace<T, Device>::diag(const HPsiFunc& hpsi_func,
743801
do
744802
{
745803

746-
sum_iter += this->diag_once(hpsi_func, psi_in, psi_in_dmax, eigenvalue_in_hsolver, ethr_band);
804+
sum_iter += this->diag_once(hpsi_func, spsi_func, psi_in, psi_in_dmax, eigenvalue_in_hsolver, ethr_band);
747805

748806
++ntry;
749807

‎source/module_hsolver/diago_dav_subspace.h‎

Lines changed: 17 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -39,6 +39,7 @@ class Diago_DavSubspace
3939
using HPsiFunc = std::function<void(T*, T*, const int, const int)>;
4040

4141
int diag(const HPsiFunc& hpsi_func,
42+
const HPsiFunc& spsi_func,
4243
T* psi_in,
4344
const int psi_in_dmax,
4445
Real* eigenvalue_in,
@@ -79,6 +80,9 @@ class Diago_DavSubspace
7980
/// the product of H and psi in the reduced basis set
8081
T* hphi = nullptr;
8182

83+
/// the product of S and psi in the reduced basis set
84+
T* sphi = nullptr;
85+
8286
/// Hamiltonian on the reduced basis
8387
T* hcc = nullptr;
8488

@@ -94,23 +98,33 @@ class Diago_DavSubspace
9498
base_device::AbacusDevice_t device = {};
9599

96100
void cal_grad(const HPsiFunc& hpsi_func,
101+
const HPsiFunc& spsi_func,
97102
const int& dim,
98103
const int& nbase,
99104
const int& notconv,
100105
T* psi_iter,
101106
T* hphi,
107+
T* spsi,
102108
T* vcc,
103109
const int* unconv,
104110
std::vector<Real>* eigenvalue_iter);
105111

106-
void cal_elem(const int& dim, int& nbase, const int& notconv, const T* psi_iter, const T* hphi, T* hcc, T* scc);
112+
void cal_elem(const int& dim,
113+
int& nbase,
114+
const int& notconv,
115+
const T* psi_iter,
116+
const T* spsi,
117+
const T* hphi,
118+
T* hcc,
119+
T* scc);
107120

108121
void refresh(const int& dim,
109122
const int& nband,
110123
int& nbase,
111124
const Real* eigenvalue,
112125
T* psi_iter,
113126
T* hphi,
127+
T* sphi,
114128
T* hcc,
115129
T* scc,
116130
T* vcc);
@@ -132,6 +146,7 @@ class Diago_DavSubspace
132146
T* vcc);
133147

134148
int diag_once(const HPsiFunc& hpsi_func,
149+
const HPsiFunc& spsi_func,
135150
T* psi_in,
136151
const int psi_in_dmax,
137152
Real* eigenvalue_in,
@@ -174,4 +189,4 @@ class Diago_DavSubspace
174189

175190
} // namespace hsolver
176191

177-
#endif
192+
#endif

‎source/module_hsolver/hsolver_pw.cpp‎

Lines changed: 7 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -583,6 +583,10 @@ void HSolverPW<T, Device>::hamiltSolvePsiK(hamilt::Hamilt<T, Device>* hm,
583583
};
584584
bool scf = this->calculation_type == "nscf" ? false : true;
585585

586+
auto spsi_func = [hm](T* psi_in, T* spsi_out, const int ld_psi, const int nvec) {
587+
hm->sPsi(psi_in, spsi_out, ld_psi, ld_psi, nvec);
588+
};
589+
586590
Diago_DavSubspace<T, Device> dav_subspace(pre_condition,
587591
psi.get_nbands(),
588592
psi.get_k_first() ? psi.get_current_nbas()
@@ -594,7 +598,8 @@ void HSolverPW<T, Device>::hamiltSolvePsiK(hamilt::Hamilt<T, Device>* hm,
594598
comm_info);
595599

596600
DiagoIterAssist<T, Device>::avg_iter += static_cast<double>(
597-
dav_subspace.diag(hpsi_func, psi.get_pointer(), psi.get_nbasis(), eigenvalue, this->ethr_band, scf));
601+
dav_subspace
602+
.diag(hpsi_func, spsi_func, psi.get_pointer(), psi.get_nbasis(), eigenvalue, this->ethr_band, scf));
598603
}
599604
else if (this->method == "dav")
600605
{
@@ -831,4 +836,4 @@ template class HSolverPW<std::complex<float>, base_device::DEVICE_GPU>;
831836
template class HSolverPW<std::complex<double>, base_device::DEVICE_GPU>;
832837
#endif
833838

834-
} // namespace hsolver
839+
} // namespace hsolver

‎source/module_lr/hsolver_lrtd.hpp‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -105,7 +105,7 @@ namespace LR
105105
std::vector<double> ethr_band(nband, diag_ethr);
106106
hsolver::DiagoIterAssist<T>::avg_iter
107107
+= static_cast<double>(dav_subspace.diag(
108-
hpsi_func, psi,
108+
hpsi_func, spsi_func, psi,
109109
dim,
110110
eigenvalue.data(),
111111
ethr_band,
@@ -184,4 +184,4 @@ namespace LR
184184
<< " ; where current threshold is: " << hsolver::DiagoIterAssist<T>::PW_DIAG_THR << " . " << std::endl;
185185
}
186186
}
187-
}
187+
}

0 commit comments

Comments
 (0)