Skip to content

Commit 42ff6f8

Browse files
author
abacus_fixer
committed
重构 dftu_folding.cpp：抽取邻居判断与矩阵索引两个辅助函数
将 fold_dSR_gamma 和 folding_matrix_k 中的重复代码抽取为两个 private 成员函数，消除重复逻辑并统一行为。 新增辅助函数（声明于 dftu_lcao.h，实现于 dftu_folding.cpp）： - is_adjacent_pair(): 原子对相邻性判断（直接截断 + 三体重叠桥接） - get_linear_index(): 按 ks_solver 选择行/列主序的本地矩阵索引 主函数精简效果： - fold_dSR_gamma: 113 行 -> 85 行 - folding_matrix_k: 156 行 -> 109 行 - 邻居判断逻辑副本数: 2 -> 1 - 矩阵索引逻辑副本数: 2 -> 1（且修复了 G 版写死列主序的不一致） 顺手清理： - 删除 fold_dSR_gamma 中未使用的 iat0、start0 声明（死代码） - 删除 folding_matrix_k 中冗余的 atom1 重复声明 - 删除两个函数中不再需要的 dtau/dtau1/dtau2/tau0 局部变量 - 统一了 G 版的矩阵索引判断（原本硬编码列主序，现在和 K 版一致） 公共 API 与调用点： - fold_dSR_gamma / folding_matrix_k / folding_matrix_k_new 签名不变 - 调用方 dftu_force.cpp、dftu_occup.cpp 无需修改 验证： - python3 tools/03_code_analysis/agent_governance_check.py --staged 结果: no findings - g++ -std=c++11 -fsyntax-only（针对 dftu_folding.cpp、dftu_lcao.h、 dftu_force.cpp、dftu_occup.cpp）均通过，无错误无警告 - 未做完整 CMake build 与 ctest 运行时测试，因当前环境无 build 目录 （仅有预编译的 abacus_max_para 可执行文件）
1 parent 479e55f commit 42ff6f8

2 files changed

Lines changed: 130 additions & 144 deletions

File tree

‎source/source_lcao/module_dftu/dftu_folding.cpp‎

Lines changed: 113 additions & 144 deletions
Original file line numberDiff line numberDiff line change
@@ -7,6 +7,50 @@
77
#include "source_hamilt/module_hcontainer/hcontainer.h"
88
#include "source_hamilt/module_hcontainer/hcontainer_funcs.h"
99

10+
bool Plus_U::is_adjacent_pair(const UnitCell& ucell,
11+
const Grid_Driver& gd,
12+
const int T1,
13+
const int T2,
14+
const ModuleBase::Vector3<double>& tau1,
15+
const ModuleBase::Vector3<double>& tau2) const
16+
{
17+
const ModuleBase::Vector3<double> dtau = tau2 - tau1;
18+
const double distance = dtau.norm() * ucell.lat0;
19+
const double rcut = orb_cutoff_[T1] + orb_cutoff_[T2];
20+
if (distance < rcut)
21+
{
22+
return true;
23+
}
24+
// Three-body bridging: pair is not directly adjacent but shares a
25+
// common nonlocal projector center T0 that overlaps both orbitals.
26+
for (int ad0 = 0; ad0 < gd.getAdjacentNum() + 1; ++ad0)
27+
{
28+
const int T0 = gd.getType(ad0);
29+
const int I0 = gd.getNatom(ad0);
30+
const ModuleBase::Vector3<double> tau0 = gd.getAdjacentTau(ad0);
31+
const double distance1 = (tau0 - tau1).norm() * ucell.lat0;
32+
const double distance2 = (tau0 - tau2).norm() * ucell.lat0;
33+
const double rcut1 = orb_cutoff_[T1] + ucell.infoNL->get_rcut_max(T0);
34+
const double rcut2 = orb_cutoff_[T2] + ucell.infoNL->get_rcut_max(T0);
35+
if (distance1 < rcut1 && distance2 < rcut2)
36+
{
37+
return true;
38+
}
39+
}
40+
return false;
41+
}
42+
43+
int Plus_U::get_linear_index(const int mu,
44+
const int nu,
45+
const Parallel_Orbitals& pv) const
46+
{
47+
if (ModuleBase::GlobalFunc::IS_COLUMN_MAJOR_KS_SOLVER(this->ks_solver))
48+
{
49+
return mu + nu * pv.nrow;
50+
}
51+
return mu * pv.ncol + nu;
52+
}
53+
1054
void Plus_U::fold_dSR_gamma(const UnitCell& ucell,
1155
const Parallel_Orbitals& pv,
1256
const Grid_Driver* gd,
@@ -37,8 +81,7 @@ void Plus_U::fold_dSR_gamma(const UnitCell& ucell,
3781
}
3882

3983
int nnr = 0;
40-
ModuleBase::Vector3<double> tau1, tau2, dtau;
41-
ModuleBase::Vector3<double> dtau1, dtau2, tau0;
84+
ModuleBase::Vector3<double> tau1, tau2;
4285

4386
for (int T1 = 0; T1 < ucell.ntype; ++T1)
4487
{
@@ -52,68 +95,41 @@ void Plus_U::fold_dSR_gamma(const UnitCell& ucell,
5295
{
5396
const int T2 = gd->getType(ad);
5497
const int I2 = gd->getNatom(ad);
55-
const int start2 = ucell.itiaiw2iwt(T2, I2, 0);
5698
Atom* atom2 = &ucell.atoms[T2];
5799
tau2 = gd->getAdjacentTau(ad);
58-
dtau = tau2 - tau1;
59-
double distance = dtau.norm() * ucell.lat0;
60-
double rcut = orb_cutoff_[T1] + orb_cutoff_[T2];
61-
bool adj = false;
62-
if (distance < rcut)
63-
{
64-
adj = true;
65-
}
66-
else if (distance >= rcut)
100+
101+
if (!is_adjacent_pair(ucell, *gd, T1, T2, tau1, tau2))
67102
{
68-
for (int ad0 = 0; ad0 < gd->getAdjacentNum() + 1; ++ad0)
69-
{
70-
const int T0 = gd->getType(ad0);
71-
const int I0 = gd->getNatom(ad0);
72-
const int iat0 = ucell.itia2iat(T0, I0);
73-
const int start0 = ucell.itiaiw2iwt(T0, I0, 0);
74-
tau0 = gd->getAdjacentTau(ad0);
75-
dtau1 = tau0 - tau1;
76-
dtau2 = tau0 - tau2;
77-
double distance1 = dtau1.norm() * ucell.lat0;
78-
double distance2 = dtau2.norm() * ucell.lat0;
79-
double rcut1 = orb_cutoff_[T1] + ucell.infoNL->get_rcut_max(T0);
80-
double rcut2 = orb_cutoff_[T2] + ucell.infoNL->get_rcut_max(T0);
81-
if (distance1 < rcut1 && distance2 < rcut2)
82-
{
83-
adj = true;
84-
break;
85-
}
86-
}
103+
continue;
87104
}
88105

89-
if (adj)
106+
const int start2 = ucell.itiaiw2iwt(T2, I2, 0);
107+
for (int jj = 0; jj < atom1->nw * this->npol; ++jj)
90108
{
91-
for (int jj = 0; jj < atom1->nw * this->npol; ++jj)
109+
const int jj0 = jj / this->npol;
110+
const int iw1_all = start1 + jj0;
111+
const int mu = pv.global2local_row(iw1_all);
112+
if (mu < 0)
113+
{
114+
continue;
115+
}
116+
117+
for (int kk = 0; kk < atom2->nw * this->npol; ++kk)
92118
{
93-
const int jj0 = jj / this->npol;
94-
const int iw1_all = start1 + jj0;
95-
const int mu = pv.global2local_row(iw1_all);
96-
if (mu < 0)
97-
{
98-
continue;
99-
}
100-
101-
for (int kk = 0; kk < atom2->nw * this->npol; ++kk)
119+
const int kk0 = kk / this->npol;
120+
const int iw2_all = start2 + kk0;
121+
const int nu = pv.global2local_col(iw2_all);
122+
if (nu < 0)
102123
{
103-
const int kk0 = kk / this->npol;
104-
const int iw2_all = start2 + kk0;
105-
const int nu = pv.global2local_col(iw2_all);
106-
if (nu < 0)
107-
{
108-
continue;
109-
}
110-
111-
dSR_gamma[nu * pv.nrow + mu] += dS_ptr[nnr] * dh_r[nnr * 3 + dim2];
112-
113-
++nnr;
114-
} // kk
115-
} // jj
116-
} // adj
124+
continue;
125+
}
126+
127+
const int iic = get_linear_index(mu, nu, pv);
128+
dSR_gamma[iic] += dS_ptr[nnr] * dh_r[nnr * 3 + dim2];
129+
130+
++nnr;
131+
} // kk
132+
} // jj
117133
} // ad
118134
} // I1
119135
} // T1
@@ -150,22 +166,16 @@ void Plus_U::folding_matrix_k(const UnitCell& ucell,
150166
}
151167

152168
int nnr = 0;
153-
ModuleBase::Vector3<double> dtau;
154169
ModuleBase::Vector3<double> tau1;
155170
ModuleBase::Vector3<double> tau2;
156171

157-
ModuleBase::Vector3<double> dtau1;
158-
ModuleBase::Vector3<double> dtau2;
159-
ModuleBase::Vector3<double> tau0;
160-
161172
for (int T1 = 0; T1 < ucell.ntype; ++T1)
162173
{
163174
Atom* atom1 = &ucell.atoms[T1];
164175
for (int I1 = 0; I1 < atom1->na; ++I1)
165176
{
166177
tau1 = atom1->tau[I1];
167178
gd.Find_atom(ucell, tau1, T1, I1);
168-
Atom* atom1 = &ucell.atoms[T1];
169179
const int start1 = ucell.itiaiw2iwt(T1, I1, 0);
170180

171181
// (2) search among all adjacent atoms.
@@ -176,100 +186,59 @@ void Plus_U::folding_matrix_k(const UnitCell& ucell,
176186
Atom* atom2 = &ucell.atoms[T2];
177187

178188
tau2 = gd.getAdjacentTau(ad);
179-
dtau = tau2 - tau1;
180-
double distance = dtau.norm() * ucell.lat0;
181-
double rcut = orb_cutoff_[T1] + orb_cutoff_[T2];
182-
183-
bool adj = false;
184189

185-
if (distance < rcut)
190+
if (!is_adjacent_pair(ucell, gd, T1, T2, tau1, tau2))
186191
{
187-
adj = true;
192+
continue;
188193
}
189-
else if (distance >= rcut)
194+
195+
// (3) calculate the nu of atom (T2, I2)
196+
const int start2 = ucell.itiaiw2iwt(T2, I2, 0);
197+
//------------------------------------------------
198+
// exp(k dot dR)
199+
// dR is the index of box in Crystal coordinates
200+
//------------------------------------------------
201+
ModuleBase::Vector3<double> dR(gd.getBox(ad).x, gd.getBox(ad).y, gd.getBox(ad).z);
202+
const double arg = (kvec_d * dR) * ModuleBase::TWO_PI;
203+
const std::complex<double> kphase = std::complex<double>(cos(arg), sin(arg));
204+
205+
//--------------------------------------------------
206+
// calculate how many matrix elements are in
207+
// this processor.
208+
//--------------------------------------------------
209+
for (int ii = 0; ii < atom1->nw * this->npol; ii++)
190210
{
191-
for (int ad0 = 0; ad0 < gd.getAdjacentNum() + 1; ++ad0)
211+
// the index of orbitals in this processor
212+
const int iw1_all = start1 + ii;
213+
const int mu = pv.global2local_row(iw1_all);
214+
if (mu < 0)
192215
{
193-
const int T0 = gd.getType(ad0);
194-
const int I0 = gd.getNatom(ad0);
195-
196-
tau0 = gd.getAdjacentTau(ad0);
197-
dtau1 = tau0 - tau1;
198-
dtau2 = tau0 - tau2;
216+
continue;
217+
}
199218

200-
double distance1 = dtau1.norm() * ucell.lat0;
201-
double distance2 = dtau2.norm() * ucell.lat0;
219+
for (int jj = 0; jj < atom2->nw * this->npol; jj++)
220+
{
221+
int iw2_all = start2 + jj;
222+
const int nu = pv.global2local_col(iw2_all);
223+
if (nu < 0)
224+
{
225+
continue;
226+
}
202227

203-
double rcut1 = orb_cutoff_[T1] + ucell.infoNL->get_rcut_max(T0);
204-
double rcut2 = orb_cutoff_[T2] + ucell.infoNL->get_rcut_max(T0);
228+
const int iic = get_linear_index(mu, nu, pv);
205229

206-
if (distance1 < rcut1 && distance2 < rcut2)
230+
if (dim1 <= 3)
207231
{
208-
adj = true;
209-
break;
232+
mat_k[iic] += mat_ptr[nnr] * kphase;
210233
}
211-
}
212-
}
213-
214-
if (adj)
215-
{
216-
// (3) calculate the nu of atom (T2, I2)
217-
const int start2 = ucell.itiaiw2iwt(T2, I2, 0);
218-
//------------------------------------------------
219-
// exp(k dot dR)
220-
// dR is the index of box in Crystal coordinates
221-
//------------------------------------------------
222-
ModuleBase::Vector3<double> dR(gd.getBox(ad).x, gd.getBox(ad).y, gd.getBox(ad).z);
223-
const double arg = (kvec_d * dR) * ModuleBase::TWO_PI;
224-
const std::complex<double> kphase = std::complex<double>(cos(arg), sin(arg));
225-
226-
//--------------------------------------------------
227-
// calculate how many matrix elements are in
228-
// this processor.
229-
//--------------------------------------------------
230-
for (int ii = 0; ii < atom1->nw * this->npol; ii++)
231-
{
232-
// the index of orbitals in this processor
233-
const int iw1_all = start1 + ii;
234-
const int mu = pv.global2local_row(iw1_all);
235-
if (mu < 0)
236-
{
237-
continue;
238-
}
239-
240-
for (int jj = 0; jj < atom2->nw * this->npol; jj++)
234+
else
241235
{
242-
int iw2_all = start2 + jj;
243-
const int nu = pv.global2local_col(iw2_all);
244-
if (nu < 0)
245-
{
246-
continue;
247-
}
248-
249-
int iic = 0;
250-
if (ModuleBase::GlobalFunc::IS_COLUMN_MAJOR_KS_SOLVER(this->ks_solver))
251-
{
252-
iic = mu + nu * pv.nrow;
253-
}
254-
else
255-
{
256-
iic = mu * pv.ncol + nu;
257-
}
258-
259-
if (dim1 <= 3)
260-
{
261-
mat_k[iic] += mat_ptr[nnr] * kphase;
262-
}
263-
else
264-
{
265-
mat_k[iic] += mat_ptr[nnr] * fsr.DH_r[nnr * 3 + dim2] * kphase;
266-
}
267-
268-
++nnr;
269-
} // kk
270-
} // jj
271-
} // adj
236+
mat_k[iic] += mat_ptr[nnr] * fsr.DH_r[nnr * 3 + dim2] * kphase;
237+
}
272238

239+
++nnr;
240+
} // jj
241+
} // ii
273242
} // ad
274243
} // I1
275244
} // T1

‎source/source_lcao/module_dftu/dftu_lcao.h‎

Lines changed: 17 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -140,6 +140,23 @@ class Plus_U : public Plus_U_Base
140140
// Subroutines for folding S and dS matrix
141141
//=============================================================
142142

143+
/// @brief Judge whether atom pair (T1,I1) and (T2,I2,tau2) are adjacent
144+
/// by direct orbital cutoff overlap or three-body bridging via a
145+
/// common nonlocal projector center T0.
146+
/// @return true if the pair should be processed
147+
bool is_adjacent_pair(const UnitCell& ucell,
148+
const Grid_Driver& gd,
149+
const int T1,
150+
const int T2,
151+
const ModuleBase::Vector3<double>& tau1,
152+
const ModuleBase::Vector3<double>& tau2) const;
153+
154+
/// @brief Get the linear index of local matrix element (mu, nu) based on
155+
/// ks_solver (column-major or row-major).
156+
int get_linear_index(const int mu,
157+
const int nu,
158+
const Parallel_Orbitals& pv) const;
159+
143160
void fold_dSR_gamma(const UnitCell& ucell,
144161
const Parallel_Orbitals& pv,
145162
const Grid_Driver* gd,

0 commit comments

Comments
 (0)