Skip to content

Commit 6630a2b

Browse files
author
dyzheng
committed
Fix(dftu-pw): fix DFT+U for PW basis (nspin=1/2/4) and occupation matrix mixing
- Relax DFT+U nspin validation to support nspin=1/2/4 for PW basis - Encapsulate Plus_U with typed accessors (get/set_locale, get_orbital_corr, get_hubbard_u, is_locale_initialized, mark_locale_dirty, enable_mixing) - Rewrite cal_occ_pw for nspin=1/2/4 with correct becp indexing - Add DFT+U occupation matrix mixing via Broyden method in Charge_Mixing - Fix DFT+U locale double-counting when kpar>1 - Fix spin-channel selection: use isk[ik] instead of ik>=nk/2 in cal_occ_pw - Extract setup_pw_dftu_indices() from cal_ps_dftu - Propagate ld_psi for correct GEMM strides when ngk[ik] < npwx - Restructure eff_pot_pw layout: nspin=2 uses split spin_up|spin_down - Add get_eff_pot_pw_spin(isk) for nspin-aware access - Add becp_ready caching in OnsiteProjector - Update dftu_lcao.cpp for new Plus_U accessor interface - Add unit tests for PW DFT+U (dftu_pw_test.cpp) - Update dftu_io.cpp to use new accessors and onsite.dm format
1 parent bb2cb4b commit 6630a2b

9 files changed

Lines changed: 875 additions & 225 deletions

File tree

‎source/source_lcao/module_dftu/dftu.cpp‎

Lines changed: 4 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -131,6 +131,10 @@ void Plus_U::init(UnitCell& cell,
131131
// it:index of type of atom
132132
for (int it = 0; it < cell.ntype; ++it)
133133
{
134+
if(!has_correlated_orbital(it))
135+
{
136+
continue;
137+
}
134138
for (int ia = 0; ia < cell.atoms[it].na; ia++)
135139
{
136140
// ia:index of atoms of this type
@@ -140,14 +144,6 @@ void Plus_U::init(UnitCell& cell,
140144
locale[iat].resize(cell.atoms[it].nwl + 1);
141145
locale_save[iat].resize(cell.atoms[it].nwl + 1);
142146

143-
// initialize the arrry iatlnm2iwt[iat][l][n][m]
144-
this->iatlnmipol2iwt[iat].resize(cell.atoms[it].nwl + 1);
145-
146-
if(!has_correlated_orbital(it))
147-
{
148-
continue;
149-
}
150-
151147
const int tlp1_npol = (get_orbital_corr(it)*2+1)*npol;
152148
const int tlp1 = 2 * get_orbital_corr(it) + 1;
153149
const int elem_size = tlp1 * tlp1;

‎source/source_lcao/module_dftu/dftu.h‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -8,7 +8,7 @@
88
#ifdef __LCAO
99
#include "source_basis/module_ao/ORB_read.h"
1010
#include "source_hamilt/hamilt.h"
11-
#include "source_hamilt/module_hcontainer/hcontainer.h"
11+
#include "source_lcao/module_hcontainer/hcontainer.h"
1212
#include "source_estate/module_dm/density_matrix.h"
1313
#include "source_lcao/force_stress_arrays.h" // mohan add 2024-06-15
1414
#endif

‎source/source_lcao/module_dftu/dftu_folding.cpp‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -4,8 +4,8 @@
44
#include "source_io/module_parameter/parameter.h"
55
#include "source_cell/module_neighbor/sltk_grid_driver.h"
66
#include "source_lcao/hamilt_lcao.h"
7-
#include "source_hamilt/module_hcontainer/hcontainer.h"
8-
#include "source_hamilt/module_hcontainer/hcontainer_funcs.h"
7+
#include "source_lcao/module_hcontainer/hcontainer.h"
8+
#include "source_lcao/module_hcontainer/hcontainer_funcs.h"
99

1010
void Plus_U::fold_dSR_gamma(const UnitCell& ucell,
1111
const Parallel_Orbitals& pv,

‎source/source_lcao/module_dftu/dftu_io.cpp‎

Lines changed: 29 additions & 22 deletions
Original file line numberDiff line numberDiff line change
@@ -27,10 +27,10 @@ void Plus_U::output(const UnitCell& ucell,
2727

2828
if (L >= get_orbital_corr(T) && has_correlated_orbital(T))
2929
{
30-
if (L != get_orbital_corr(T))
31-
{
32-
continue;
33-
}
30+
if (L != get_orbital_corr(T))
31+
{
32+
continue;
33+
}
3434

3535
if (!Yukawa)
3636
{
@@ -97,11 +97,11 @@ void Plus_U::write_occup_m(const UnitCell& ucell,
9797

9898
for (int T = 0; T < ucell.ntype; T++)
9999
{
100-
if (!has_correlated_orbital(T))
101-
{
102-
continue;
103-
}
104-
const int NL = ucell.atoms[T].nwl + 1;
100+
if (!has_correlated_orbital(T))
101+
{
102+
continue;
103+
}
104+
const int NL = ucell.atoms[T].nwl + 1;
105105
const int LC = get_orbital_corr(T);
106106

107107
for (int I = 0; I < ucell.atoms[T].na; I++)
@@ -110,10 +110,10 @@ void Plus_U::write_occup_m(const UnitCell& ucell,
110110

111111
for (int l = 0; l < NL; l++)
112112
{
113-
if (l != get_orbital_corr(T))
114-
{
115-
continue;
116-
}
113+
if (l != get_orbital_corr(T))
114+
{
115+
continue;
116+
}
117117

118118
const int N = ucell.atoms[T].l_nchi[l];
119119

@@ -316,7 +316,14 @@ void Plus_U::read_occup_m(const UnitCell& ucell,
316316

317317
for (int l = 0; l < NL; l++)
318318
{
319-
if (l != get_orbital_corr(T))
319+
if (l != get_orbital_corr(T))
320+
{
321+
continue;
322+
}
323+
324+
ifdftu >> word;
325+
326+
if (strcmp("L", word) == 0)
320327
{
321328
continue;
322329
}
@@ -397,10 +404,10 @@ void Plus_U::local_occup_bcast(const UnitCell& ucell,
397404

398405
for (int T = 0; T < ucell.ntype; T++)
399406
{
400-
if (!has_correlated_orbital(T))
401-
{
402-
continue;
403-
}
407+
if (!has_correlated_orbital(T))
408+
{
409+
continue;
410+
}
404411

405412
for (int I = 0; I < ucell.atoms[T].na; I++)
406413
{
@@ -409,10 +416,10 @@ void Plus_U::local_occup_bcast(const UnitCell& ucell,
409416

410417
for (int l = 0; l <= ucell.atoms[T].nwl; l++)
411418
{
412-
if (l != get_orbital_corr(T))
413-
{
414-
continue;
415-
}
419+
if (l != get_orbital_corr(T))
420+
{
421+
continue;
422+
}
416423

417424
for (int n = 0; n < ucell.atoms[T].l_nchi[l]; n++)
418425
{

‎source/source_lcao/module_dftu/dftu_occup.cpp‎

Lines changed: 7 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -30,6 +30,7 @@ void Plus_U::copy_locale(const UnitCell& ucell)
3030
if (Plus_U::nspin == 4)
3131
{
3232
locale_save[iat][target_l][0][0] = locale[iat][target_l][0][0];
33+
// nspin=4 locale matrix already contains all spin components interleaved
3334
if(this->uom_save.size() != 0)
3435
{
3536
const int size = locale[iat][target_l][0][0].nr * locale[iat][target_l][0][0].nc;
@@ -43,6 +44,7 @@ void Plus_U::copy_locale(const UnitCell& ucell)
4344
{
4445
locale_save[iat][target_l][0][0] = locale[iat][target_l][0][0];
4546
locale_save[iat][target_l][0][1] = locale[iat][target_l][0][1];
47+
// save locale matrix for spin=0,1 to uom_save
4648
if(this->uom_save.size() != 0)
4749
{
4850
const int size = locale[iat][target_l][0][0].nr * locale[iat][target_l][0][0].nc;
@@ -107,13 +109,15 @@ void Plus_U::mix_locale(const UnitCell& ucell,
107109

108110
for (int T = 0; T < ucell.ntype; T++)
109111
{
110-
int target_l = get_orbital_corr(T);
111-
if (target_l == -1)
112-
continue;
112+
if (!has_correlated_orbital(T))
113+
{
114+
continue;
115+
}
113116

114117
for (int I = 0; I < ucell.atoms[T].na; I++)
115118
{
116119
const int iat = ucell.itia2iat(T, I);
120+
int target_l = get_orbital_corr(T);
117121

118122
if (Plus_U::nspin == 4)
119123
{

‎source/source_lcao/module_dftu/dftu_pw.cpp‎

Lines changed: 20 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -28,6 +28,8 @@ void Plus_U::cal_occ_pw(const int iter,
2828
ModuleBase::timer::start("Plus_U", "cal_occ_pw");
2929
this->copy_locale(cell);
3030
this->zero_locale(cell);
31+
const int nspin = PARAM.inp.nspin;
32+
const int kpar = PARAM.inp.kpar;
3133

3234
if(this->device == "cpu")
3335
{
@@ -37,7 +39,7 @@ void Plus_U::cal_occ_pw(const int iter,
3739
const int npol = psi_p->get_npol();
3840
for(int ik = 0; ik < psi_p->get_nk(); ik++)
3941
{
40-
int is = (Plus_U::nspin == 2) ? isk[ik] : 0;
42+
int is = (nspin == 2) ? isk[ik] : 0;
4143
psi_p->fix_k(ik);
4244
onsite_p->tabulate_atomic(ik);
4345

@@ -59,7 +61,7 @@ void Plus_U::cal_occ_pw(const int iter,
5961
const int m_begin = target_l * target_l;
6062
const int tlp1 = 2 * target_l + 1;
6163
const int tlp1_2 = tlp1 * tlp1;
62-
if(Plus_U::nspin == 4)
64+
if(nspin == 4)
6365
{
6466
for(int ib = 0;ib<nbands;ib++)
6567
{
@@ -117,7 +119,8 @@ void Plus_U::cal_occ_pw(const int iter,
117119
const int npol = psi_p->get_npol();
118120
for(int ik = 0; ik < psi_p->get_nk(); ik++)
119121
{
120-
int is = (Plus_U::nspin == 2) ? isk[ik] : 0;
122+
int is = (nspin == 2) ? isk[ik] : 0;
123+
const_cast<psi::Psi<std::complex<double>, base_device::DEVICE_GPU>*>(psi_p)->load_k_to_gpu(ik);
121124
psi_p->fix_k(ik);
122125
onsite_p->tabulate_atomic(ik);
123126

@@ -138,7 +141,7 @@ void Plus_U::cal_occ_pw(const int iter,
138141
const int m_begin = target_l * target_l;
139142
const int tlp1 = 2 * target_l + 1;
140143
const int tlp1_2 = tlp1 * tlp1;
141-
if(Plus_U::nspin == 4)
144+
if(nspin == 4)
142145
{
143146
for(int ib = 0;ib<nbands;ib++)
144147
{
@@ -199,23 +202,23 @@ void Plus_U::cal_occ_pw(const int iter,
199202
}
200203
const int size = (2 * target_l + 1) * (2 * target_l + 1);
201204

202-
if(Plus_U::nspin != 4)
205+
if(nspin != 4)
203206
{
204-
Parallel_Reduce::reduce_double_allpool(this->kpar,
207+
Parallel_Reduce::reduce_double_allpool(kpar,
205208
GlobalV::NPROC_IN_POOL,
206209
this->locale[iat][target_l][0][0].c,
207210
size);
208-
if(Plus_U::nspin == 2)
211+
if(nspin == 2)
209212
{
210-
Parallel_Reduce::reduce_double_allpool(this->kpar,
213+
Parallel_Reduce::reduce_double_allpool(kpar,
211214
GlobalV::NPROC_IN_POOL,
212215
this->locale[iat][target_l][0][1].c,
213216
size);
214217
}
215218
}
216219
else
217220
{
218-
Parallel_Reduce::reduce_double_allpool(this->kpar,
221+
Parallel_Reduce::reduce_double_allpool(kpar,
219222
GlobalV::NPROC_IN_POOL,
220223
this->locale[iat][target_l][0][0].c,
221224
size * 4);
@@ -228,7 +231,7 @@ void Plus_U::cal_occ_pw(const int iter,
228231
{
229232
this->uom_array[eff_pot_pw_index[iat]+mm] = this->locale[iat][target_l][0][0].c[mm];
230233
}
231-
if(Plus_U::nspin == 2)
234+
if(nspin == 2)
232235
{
233236
const int half_size = this->uom_array.size() / 2;
234237
for(int mm=0;mm<size;mm++)
@@ -247,8 +250,8 @@ void Plus_U::cal_occ_pw(const int iter,
247250
}
248251

249252
Plus_U::energy_u = 0.0;
250-
const double weight_eu = (Plus_U::nspin == 1) ? 0.25 : (Plus_U::nspin == 2) ? 0.5 : 0.25;
251-
const double diag_coeff = (Plus_U::nspin == 4) ? 1.0 : 0.5;
253+
const double weight_eu = (nspin == 1) ? 1.0 : (nspin == 2) ? 0.5 : 0.25;
254+
const double diag_coeff = (nspin == 4) ? 1.0 : 0.5;
252255
// calculate VU and energy (locale already reduced above)
253256
for(int iat = 0; iat < cell.nat; iat++)
254257
{
@@ -265,7 +268,7 @@ void Plus_U::cal_occ_pw(const int iter,
265268
std::complex<double>* vu_iat = &(this->eff_pot_pw[this->eff_pot_pw_index[iat]]);
266269
const int m_size = 2 * target_l + 1;
267270

268-
if(Plus_U::nspin == 4)
271+
if(nspin == 4)
269272
{
270273
for (int m1 = 0; m1 < m_size; m1++)
271274
{
@@ -309,8 +312,8 @@ void Plus_U::cal_occ_pw(const int iter,
309312
}
310313
vu_iat[index[0]] = 0.5 * (vu_tmp[0] + vu_tmp[3]);
311314
vu_iat[index[3]] = 0.5 * (vu_tmp[0] - vu_tmp[3]);
312-
vu_iat[index[1]] = 0.5 * (vu_tmp[1] - std::complex<double>(0.0, 1.0) * vu_tmp[2]);
313-
vu_iat[index[2]] = 0.5 * (vu_tmp[1] + std::complex<double>(0.0, 1.0) * vu_tmp[2]);
315+
vu_iat[index[1]] = 0.5 * (vu_tmp[1] + std::complex<double>(0.0, 1.0) * vu_tmp[2]);
316+
vu_iat[index[2]] = 0.5 * (vu_tmp[1] - std::complex<double>(0.0, 1.0) * vu_tmp[2]);
314317
}
315318
}
316319
}
@@ -328,7 +331,7 @@ void Plus_U::cal_occ_pw(const int iter,
328331
}
329332
}
330333
// spin-down channel for nspin=2
331-
if(Plus_U::nspin == 2)
334+
if(nspin == 2)
332335
{
333336
std::complex<double>* vu_iat1 = &(this->eff_pot_pw[this->eff_pot_pw.size()/2 + this->eff_pot_pw_index[iat]]);
334337
for (int m1 = 0; m1 < m_size; m1++)
@@ -348,3 +351,4 @@ void Plus_U::cal_occ_pw(const int iter,
348351
ModuleBase::timer::end("Plus_U", "cal_occ_pw");
349352
}
350353

354+

‎source/source_lcao/module_dftu/test/CMakeLists.txt‎

Lines changed: 1 addition & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -1,19 +1,5 @@
1-
abacus_disable_feature_definitions(__CUDA)
2-
31
AddTest(
42
TARGET dftu_pw_test
5-
LIBS base device parameter
3+
LIBS ${math_libs} base device parameter
64
SOURCES dftu_pw_test.cpp
75
)
8-
9-
AddTest(
10-
TARGET dftu_core_test
11-
LIBS base device
12-
SOURCES dftu_core_test.cpp
13-
)
14-
15-
AddTest(
16-
TARGET dftu_operator_test
17-
LIBS base device
18-
SOURCES dftu_operator_test.cpp
19-
)

0 commit comments

Comments
 (0)