Skip to content

Commit 8045f36

Browse files
author
abacus_fixer
committed
refactor(dftu): collapse cal_occ_mat_k/gamma signatures onto Plus_U&
cal_occ_mat_k: 13 params -> 9 (add Plus_U& dftu, drop nspin/npol/nlocal/ ks_solver/iatlnmipol2iwt/orbital_corr/occ_mat/occ_mat_save/occ_mat_initialized). cal_occ_mat_gamma: 13 params -> 7 (add Plus_U& dftu). Both now read all occupation-matrix state from dftu.occmat() (mat/mat_save/data/data_save/iatlnmipol2iwt/nspin/npol) and the Plus_U_Base accessors (get_orbital_corr_vec / is_occ_mat_initialized / mark_occ_mat_initialized), and use occmat().copy_to_save()/zero() instead of the duplicated hand-written copy/zero loops. The intermediate cal_occ_mat template forwarders shrink to a single direct call. MPI_Allreduce(MPI_COMM_WORLD) inside the accumulation is intentionally left unchanged: it has no exact Parallel_Reduce counterpart and was kept per the "only replace what maps cleanly" rule. Verified: cmake --build build -j4 (exit 0). This build has no test targets.
1 parent 816f302 commit 8045f36

2 files changed

Lines changed: 76 additions & 181 deletions

File tree

‎source/source_lcao/module_dftu/dftu_lcao_occ.cpp‎

Lines changed: 74 additions & 164 deletions
Original file line numberDiff line numberDiff line change
@@ -3,15 +3,15 @@
33
#include "dftu_folding.h"
44
#include "source_base/timer.h"
55
#include "source_base/module_external/scalapack_connector.h"
6-
#include "source_io/module_parameter/parameter.h"
76
#include "source_estate/occ_matrix.h"
7+
#include "source_io/module_parameter/parameter.h"
88
#ifdef __LCAO
99
#include "source_lcao/hamilt_lcao.h"
1010
#endif
1111

12-
// copy_occ_mat(), zero_occ_mat(), set_occ_mat(ucell),
13-
// get_occ_mat_flat(), set_occ_mat_flat()
14-
// are now implemented in dftu_base.cpp as Plus_U_Base methods (inherited by Plus_U).
12+
// cal_occ_mat_k / cal_occ_mat_gamma take Plus_U& dftu directly and read all
13+
// occupation-matrix state (occ/save arrays, lookup table, nspin/npol, and the
14+
// occ_mat_initialized flag) from dftu.occmat() and the Plus_U_Base accessors.
1515

1616
#ifdef __LCAO
1717

@@ -23,63 +23,21 @@ void DFTU_LCAO::cal_occ_mat_k(const Parallel_Orbitals* pv,
2323
const double& mixing_beta,
2424
hamilt::Hamilt<std::complex<double>>* p_ham,
2525
const bool gamma_only_local,
26-
const int nspin,
27-
const int npol,
28-
const int nlocal,
29-
const std::string& ks_solver,
30-
const std::vector<std::vector<std::vector<std::vector<std::vector<int>>>>>& iatlnmipol2iwt,
31-
const std::vector<int>& orbital_corr,
32-
std::vector<std::vector<std::vector<std::vector<ModuleBase::matrix>>>>& occ_mat,
33-
std::vector<std::vector<std::vector<std::vector<ModuleBase::matrix>>>>& occ_mat_save,
34-
bool& occ_mat_initialized)
26+
Plus_U& dftu)
3527
{
3628
ModuleBase::TITLE("DFTU_LCAO", "cal_occ_mat_k");
3729
ModuleBase::timer::start("DFTU_LCAO", "cal_occ_mat_k");
3830

39-
// copy occ_mat to occ_mat_save
40-
for (int T = 0; T < ucell.ntype; T++)
41-
{
42-
int target_l = orbital_corr[T];
43-
if (target_l == -1) continue;
44-
for (int I = 0; I < ucell.atoms[T].na; I++)
45-
{
46-
const int iat = ucell.itia2iat(T, I);
47-
if (nspin == 4)
48-
{
49-
occ_mat_save[iat][target_l][0][0] = occ_mat[iat][target_l][0][0];
50-
}
51-
else if (nspin == 1 || nspin == 2)
52-
{
53-
occ_mat_save[iat][target_l][0][0] = occ_mat[iat][target_l][0][0];
54-
occ_mat_save[iat][target_l][0][1] = occ_mat[iat][target_l][0][1];
55-
}
56-
}
57-
}
58-
// zero occ_mat
59-
for (int T = 0; T < ucell.ntype; T++)
60-
{
61-
if (orbital_corr[T] == -1) continue;
62-
for (int I = 0; I < ucell.atoms[T].na; I++)
63-
{
64-
const int iat = ucell.itia2iat(T, I);
65-
for (int l = 0; l < ucell.atoms[T].nwl + 1; l++)
66-
{
67-
const int N = ucell.atoms[T].l_nchi[l];
68-
for (int n = 0; n < N; n++)
69-
{
70-
if (nspin == 4)
71-
{
72-
occ_mat[iat][l][n][0].zero_out();
73-
}
74-
else if (nspin == 1 || nspin == 2)
75-
{
76-
occ_mat[iat][l][n][0].zero_out();
77-
occ_mat[iat][l][n][1].zero_out();
78-
}
79-
}
80-
}
81-
}
82-
}
31+
const int nspin = dftu.occmat().nspin();
32+
const int npol = dftu.occmat().npol();
33+
const int nlocal = pv->get_global_row_size();
34+
const std::string& ks_solver = PARAM.inp.ks_solver;
35+
const auto& iatlnmipol2iwt = dftu.occmat().iatlnmipol2iwt();
36+
const std::vector<int>& orbital_corr = dftu.get_orbital_corr_vec();
37+
38+
// copy occ_mat to occ_mat_save, then zero occ_mat
39+
dftu.occmat().copy_to_save(ucell, orbital_corr);
40+
dftu.occmat().zero(ucell, orbital_corr);
8341

8442
//=================Part 1======================
8543
// call SCALAPACK routine to calculate the product of the S and density matrix
@@ -161,6 +119,7 @@ void DFTU_LCAO::cal_occ_mat_k(const Parallel_Orbitals* pv,
161119
}
162120

163121
// Calculate the local occupation number matrix
122+
ModuleBase::matrix& occ = dftu.occmat().mat(iat, l, n, spin);
164123
for (int m0 = 0; m0 < 2 * l + 1; m0++)
165124
{
166125
for (int ipol0 = 0; ipol0 < npol; ipol0++)
@@ -185,12 +144,12 @@ void DFTU_LCAO::cal_occ_mat_k(const Parallel_Orbitals* pv,
185144

186145
if ((nu >= 0) && (mu >= 0))
187146
{
188-
occ_mat[iat][l][n][spin](m0_all, m1_all) += (srho[irc]).real() / 4.0;
147+
occ(m0_all, m1_all) += (srho[irc]).real() / 4.0;
189148
}
190149

191150
if ((nu_prime >= 0) && (mu_prime >= 0))
192151
{
193-
occ_mat[iat][l][n][spin](m0_all, m1_all)
152+
occ(m0_all, m1_all)
194153
+= (std::conj(srho[irc_prime])).real() / 4.0;
195154
}
196155
} // ipol1
@@ -238,27 +197,30 @@ void DFTU_LCAO::cal_occ_mat_k(const Parallel_Orbitals* pv,
238197
#ifdef __MPI
239198
if (nspin == 1 || nspin == 4)
240199
{
241-
ModuleBase::matrix temp(occ_mat[iat][l][n][0]);
200+
ModuleBase::matrix& occ0 = dftu.occmat().mat(iat, l, n, 0);
201+
ModuleBase::matrix temp(occ0);
242202
MPI_Allreduce(&temp(0, 0),
243-
&occ_mat[iat][l][n][0](0, 0),
203+
&occ0(0, 0),
244204
(2 * l + 1) * npol * (2 * l + 1) * npol,
245205
MPI_DOUBLE,
246206
MPI_SUM,
247207
MPI_COMM_WORLD);
248208
}
249209
else if (nspin == 2)
250210
{
251-
ModuleBase::matrix temp0(occ_mat[iat][l][n][0]);
211+
ModuleBase::matrix& occ0 = dftu.occmat().mat(iat, l, n, 0);
212+
ModuleBase::matrix temp0(occ0);
252213
MPI_Allreduce(&temp0(0, 0),
253-
&occ_mat[iat][l][n][0](0, 0),
214+
&occ0(0, 0),
254215
(2 * l + 1) * (2 * l + 1),
255216
MPI_DOUBLE,
256217
MPI_SUM,
257218
MPI_COMM_WORLD);
258219

259-
ModuleBase::matrix temp1(occ_mat[iat][l][n][1]);
220+
ModuleBase::matrix& occ1 = dftu.occmat().mat(iat, l, n, 1);
221+
ModuleBase::matrix temp1(occ1);
260222
MPI_Allreduce(&temp1(0, 0),
261-
&occ_mat[iat][l][n][1](0, 0),
223+
&occ1(0, 0),
262224
(2 * l + 1) * (2 * l + 1),
263225
MPI_DOUBLE,
264226
MPI_SUM,
@@ -269,19 +231,28 @@ void DFTU_LCAO::cal_occ_mat_k(const Parallel_Orbitals* pv,
269231
switch (nspin)
270232
{
271233
case 1:
272-
occ_mat[iat][l][n][0] += transpose(occ_mat[iat][l][n][0]);
273-
occ_mat[iat][l][n][0] *= 0.5;
274-
occ_mat[iat][l][n][1] += occ_mat[iat][l][n][0];
234+
{
235+
ModuleBase::matrix& occ0 = dftu.occmat().mat(iat, l, n, 0);
236+
occ0 += transpose(occ0);
237+
occ0 *= 0.5;
238+
dftu.occmat().mat(iat, l, n, 1) += occ0;
275239
break;
240+
}
276241

277242
case 2:
278243
for (int is = 0; is < nspin; is++)
279-
occ_mat[iat][l][n][is] += transpose(occ_mat[iat][l][n][is]);
244+
{
245+
ModuleBase::matrix& occ_is = dftu.occmat().mat(iat, l, n, is);
246+
occ_is += transpose(occ_is);
247+
}
280248
break;
281249

282250
case 4:
283-
occ_mat[iat][l][n][0] += transpose(occ_mat[iat][l][n][0]);
251+
{
252+
ModuleBase::matrix& occ0 = dftu.occmat().mat(iat, l, n, 0);
253+
occ0 += transpose(occ0);
284254
break;
255+
}
285256

286257
default:
287258
std::cout << "Not supported NSPIN parameter" << std::endl;
@@ -292,12 +263,13 @@ void DFTU_LCAO::cal_occ_mat_k(const Parallel_Orbitals* pv,
292263
} // end ia
293264
} // end it
294265

295-
if(PARAM.inp.mixing_dftu && occ_mat_initialized)
266+
if(PARAM.inp.mixing_dftu && dftu.is_occ_mat_initialized())
296267
{
297-
elecstate::mix_occ_with_save(occ_mat, occ_mat_save, ucell, orbital_corr, nspin, mixing_beta);
268+
elecstate::mix_occ_with_save(dftu.occmat().data(), dftu.occmat().data_save(),
269+
ucell, orbital_corr, nspin, mixing_beta);
298270
}
299271

300-
occ_mat_initialized = true;
272+
dftu.mark_occ_mat_initialized();
301273
ModuleBase::timer::end("DFTU_LCAO", "cal_occ_mat_k");
302274
return;
303275
}
@@ -308,61 +280,20 @@ void DFTU_LCAO::cal_occ_mat_gamma(const Parallel_Orbitals* pv,
308280
const std::vector<std::vector<double>> &dm_gamma,
309281
const double& mixing_beta,
310282
hamilt::Hamilt<double>* p_ham,
311-
const int nspin,
312-
const int npol,
313-
const int nlocal,
314-
const std::vector<std::vector<std::vector<std::vector<std::vector<int>>>>>& iatlnmipol2iwt,
315-
const std::vector<int>& orbital_corr,
316-
std::vector<std::vector<std::vector<std::vector<ModuleBase::matrix>>>>& occ_mat,
317-
std::vector<std::vector<std::vector<std::vector<ModuleBase::matrix>>>>& occ_mat_save,
318-
bool& occ_mat_initialized)
283+
Plus_U& dftu)
319284
{
320285
ModuleBase::TITLE("DFTU_LCAO", "cal_occ_mat_gamma");
321286
ModuleBase::timer::start("DFTU_LCAO", "cal_occ_mat_gamma");
322-
// copy occ_mat to occ_mat_save
323-
for (int T = 0; T < ucell.ntype; T++)
324-
{
325-
int target_l = orbital_corr[T];
326-
if (target_l == -1) continue;
327-
for (int I = 0; I < ucell.atoms[T].na; I++)
328-
{
329-
const int iat = ucell.itia2iat(T, I);
330-
if (nspin == 4)
331-
{
332-
occ_mat_save[iat][target_l][0][0] = occ_mat[iat][target_l][0][0];
333-
}
334-
else if (nspin == 1 || nspin == 2)
335-
{
336-
occ_mat_save[iat][target_l][0][0] = occ_mat[iat][target_l][0][0];
337-
occ_mat_save[iat][target_l][0][1] = occ_mat[iat][target_l][0][1];
338-
}
339-
}
340-
}
341-
// zero occ_mat
342-
for (int T = 0; T < ucell.ntype; T++)
343-
{
344-
if (orbital_corr[T] == -1) continue;
345-
for (int I = 0; I < ucell.atoms[T].na; I++)
346-
{
347-
const int iat = ucell.itia2iat(T, I);
348-
for (int l = 0; l < ucell.atoms[T].nwl + 1; l++)
349-
{
350-
const int N = ucell.atoms[T].l_nchi[l];
351-
for (int n = 0; n < N; n++)
352-
{
353-
if (nspin == 4)
354-
{
355-
occ_mat[iat][l][n][0].zero_out();
356-
}
357-
else if (nspin == 1 || nspin == 2)
358-
{
359-
occ_mat[iat][l][n][0].zero_out();
360-
occ_mat[iat][l][n][1].zero_out();
361-
}
362-
}
363-
}
364-
}
365-
}
287+
288+
const int nspin = dftu.occmat().nspin();
289+
const int npol = dftu.occmat().npol();
290+
const int nlocal = pv->get_global_row_size();
291+
const auto& iatlnmipol2iwt = dftu.occmat().iatlnmipol2iwt();
292+
const std::vector<int>& orbital_corr = dftu.get_orbital_corr_vec();
293+
294+
// copy occ_mat to occ_mat_save, then zero occ_mat
295+
dftu.occmat().copy_to_save(ucell, orbital_corr);
296+
dftu.occmat().zero(ucell, orbital_corr);
366297

367298
//=================Part 1======================
368299
// call PBLAS routine to calculate the product of the S and density matrix
@@ -428,6 +359,7 @@ void DFTU_LCAO::cal_occ_mat_gamma(const Parallel_Orbitals* pv,
428359
}
429360

430361
// Calculate the local occupation number matrix
362+
ModuleBase::matrix& occ_is = dftu.occmat().mat(iat, l, n, is);
431363
for (int m0 = 0; m0 < 2 * l + 1; m0++)
432364
{
433365
for (int ipol0 = 0; ipol0 < npol; ipol0++)
@@ -452,26 +384,26 @@ void DFTU_LCAO::cal_occ_mat_gamma(const Parallel_Orbitals* pv,
452384
int m0_all = m0 + (2 * l + 1) * ipol0;
453385
int m1_all = m0 + (2 * l + 1) * ipol1;
454386

455-
occ_mat[iat][l][n][is](m0, m1) += srho[irc] / 4.0;
387+
occ_is(m0, m1) += srho[irc] / 4.0;
456388
}
457389

458390
if ((nu_prime >= 0) && (mu_prime >= 0))
459391
{
460392
int m0_all = m0 + (2 * l + 1) * ipol0;
461393
int m1_all = m0 + (2 * l + 1) * ipol1;
462394

463-
occ_mat[iat][l][n][is](m0, m1) += srho[irc_prime] / 4.0;
395+
occ_is(m0, m1) += srho[irc_prime] / 4.0;
464396
}
465397
}
466398
}
467399
}
468400
}
469401

470-
ModuleBase::matrix temp(occ_mat[iat][l][n][is]);
402+
ModuleBase::matrix temp(occ_is);
471403

472404
#ifdef __MPI
473405
MPI_Allreduce(&temp(0, 0),
474-
&occ_mat[iat][l][n][is](0, 0),
406+
&occ_is(0, 0),
475407
(2 * l + 1) * npol * (2 * l + 1) * npol,
476408
MPI_DOUBLE,
477409
MPI_SUM,
@@ -482,13 +414,16 @@ void DFTU_LCAO::cal_occ_mat_gamma(const Parallel_Orbitals* pv,
482414
switch (nspin)
483415
{
484416
case 1:
485-
occ_mat[iat][l][n][0] += transpose(occ_mat[iat][l][n][0]);
486-
occ_mat[iat][l][n][0] *= 0.5;
487-
occ_mat[iat][l][n][1] += occ_mat[iat][l][n][0];
417+
{
418+
ModuleBase::matrix& occ0 = dftu.occmat().mat(iat, l, n, 0);
419+
occ0 += transpose(occ0);
420+
occ0 *= 0.5;
421+
dftu.occmat().mat(iat, l, n, 1) += occ0;
488422
break;
423+
}
489424

490425
case 2:
491-
occ_mat[iat][l][n][is] += transpose(occ_mat[iat][l][n][is]);
426+
occ_is += transpose(occ_is);
492427
break;
493428

494429
default:
@@ -502,12 +437,13 @@ void DFTU_LCAO::cal_occ_mat_gamma(const Parallel_Orbitals* pv,
502437
} // it
503438
} // is
504439

505-
if(PARAM.inp.mixing_dftu && occ_mat_initialized)
440+
if(PARAM.inp.mixing_dftu && dftu.is_occ_mat_initialized())
506441
{
507-
elecstate::mix_occ_with_save(occ_mat, occ_mat_save, ucell, orbital_corr, nspin, mixing_beta);
442+
elecstate::mix_occ_with_save(dftu.occmat().data(), dftu.occmat().data_save(),
443+
ucell, orbital_corr, nspin, mixing_beta);
508444
}
509445

510-
occ_mat_initialized = true;
446+
dftu.mark_occ_mat_initialized();
511447
ModuleBase::timer::end("DFTU_LCAO", "cal_occ_mat_gamma");
512448
return;
513449
}
@@ -527,20 +463,7 @@ void cal_occ_mat(const Parallel_Orbitals* pv,
527463
const bool gamma_only_local,
528464
const int nspin)
529465
{
530-
bool occ_mat_initialized = dftu.is_occ_mat_initialized();
531-
DFTU_LCAO::cal_occ_mat_gamma(pv, iter, ucell, dm, mixing_beta, p_ham, nspin,
532-
ucell.get_npol(), pv->get_global_row_size(), dftu.occmat().iatlnmipol2iwt(),
533-
dftu.get_orbital_corr_vec(),
534-
dftu.occmat().data(), dftu.occmat().data_save(),
535-
occ_mat_initialized);
536-
if (occ_mat_initialized)
537-
{
538-
dftu.mark_occ_mat_initialized();
539-
}
540-
else
541-
{
542-
dftu.mark_occ_mat_dirty();
543-
}
466+
DFTU_LCAO::cal_occ_mat_gamma(pv, iter, ucell, dm, mixing_beta, p_ham, dftu);
544467
}
545468

546469
//! dftu occupation matrix for multiple k-points using dm(complex)
@@ -556,20 +479,7 @@ void cal_occ_mat(const Parallel_Orbitals* pv,
556479
const bool gamma_only_local,
557480
const int nspin)
558481
{
559-
bool occ_mat_initialized = dftu.is_occ_mat_initialized();
560-
DFTU_LCAO::cal_occ_mat_k(pv, iter, ucell, dm, kv, mixing_beta, p_ham, gamma_only_local, nspin,
561-
ucell.get_npol(), pv->get_global_row_size(), PARAM.inp.ks_solver, dftu.occmat().iatlnmipol2iwt(),
562-
dftu.get_orbital_corr_vec(),
563-
dftu.occmat().data(), dftu.occmat().data_save(),
564-
occ_mat_initialized);
565-
if (occ_mat_initialized)
566-
{
567-
dftu.mark_occ_mat_initialized();
568-
}
569-
else
570-
{
571-
dftu.mark_occ_mat_dirty();
572-
}
482+
DFTU_LCAO::cal_occ_mat_k(pv, iter, ucell, dm, kv, mixing_beta, p_ham, gamma_only_local, dftu);
573483
}
574484

575485
} // namespace DFTU_LCAO

0 commit comments

Comments
 (0)