Skip to content

Commit 8da109e

Browse files
dyzhengdyzheng
committed
Refactor: save memory for kinetic and overlap force and stress
1 parent ed0b428 commit 8da109e

7 files changed

Lines changed: 696 additions & 31 deletions

File tree

‎source/source_lcao/FORCE_STRESS.cpp‎

Lines changed: 134 additions & 27 deletions
Original file line numberDiff line numberDiff line change
@@ -21,6 +21,9 @@
2121
#include "source_lcao/module_operator_lcao/dftu_lcao.h"
2222
#include "source_lcao/module_operator_lcao/dspin_lcao.h"
2323
#include "source_lcao/module_operator_lcao/nonlocal_new.h"
24+
#include "source_lcao/module_operator_lcao/ekinetic_new.h"
25+
#include "source_lcao/module_operator_lcao/overlap_new.h"
26+
#include "source_lcao/pulay_fs.h"
2427

2528

2629
// mohan add 2025-11-04
@@ -94,7 +97,8 @@ void Force_Stress_LCAO<T>::getForceStress(UnitCell& ucell,
9497

9598
const int nat = ucell.nat;
9699

97-
ForceStressArrays fsr; // mohan add 2024-06-15
100+
// NOTE: ForceStressArrays is no longer needed as we use operator-based force calculation
101+
// ForceStressArrays fsr; // removed - no longer needed
98102

99103
// total force : ModuleBase::matrix fcs;
100104

@@ -161,46 +165,149 @@ void Force_Stress_LCAO<T>::getForceStress(UnitCell& ucell,
161165
sigmaxc, pelec->f_en.etxc, pelec->charge, rhopw, locpp, sf);
162166
}
163167

164-
//! atomic forces from integration (4 terms)
165-
this->integral_part(PARAM.globalv.gamma_only_local, isforce, isstress,
166-
ucell, gd, fsr, pelec, dmat.dm, psi, foverlap, ftvnl_dphi, // add dmat.dm, mohan 20251104
167-
fvnl_dbeta, fvl_dphi, soverlap, stvnl_dphi, svnl_dbeta,
168-
svl_dphi, fvnl_dalpha, svnl_dalpha, deepks,
169-
two_center_bundle, orb, pv, kv);
168+
// Calculate forces and stresses using new operator-based methods
169+
// Step 1: Calculate Energy Density Matrix (EDM) for overlap force
170+
// EDM = Σ_k w_k * ε_k * |ψ_k><ψ_k|
171+
elecstate::DensityMatrix<T, double> edm = flk.cal_edm(pelec, *psi, *dmat.dm, kv, pv,
172+
PARAM.inp.nspin, PARAM.inp.nbands, ucell, *this->RA);
170173

171-
// calculate force and stress for Nonlocal part
174+
// Step 2: Handle different spin cases
172175
if (PARAM.inp.nspin == 1 || PARAM.inp.nspin == 2)
173176
{
174-
hamilt::NonlocalNew<hamilt::OperatorLCAO<T, double>> tmp_nonlocal(nullptr,
175-
kv.kvec_d, nullptr, &ucell, orb.cutoffs(), &gd, two_center_bundle.overlap_orb_beta.get());
176-
177+
// For nspin=1 or nspin=2, use double precision
178+
// Switch to spin channel 1 for DMR access
177179
if (PARAM.inp.nspin == 2)
178180
{
179181
dmat.dm->switch_dmr(1);
180182
}
181183

182-
const hamilt::HContainer<double>* dmr = dmat.dm->get_DMR_pointer(1);
183-
tmp_nonlocal.cal_force_stress(isforce, isstress, dmr, fvnl_dbeta, svnl_dbeta);
184+
const hamilt::HContainer<double>* dmR = dmat.dm->get_DMR_pointer(1);
185+
const hamilt::HContainer<double>* edmR = edm.get_DMR_pointer(1);
186+
187+
// Calculate kinetic force/stress (uses DM)
188+
if (PARAM.inp.t_in_h)
189+
{
190+
hamilt::EkineticNew<hamilt::OperatorLCAO<T, double>> tmp_ekinetic(
191+
nullptr, kv.kvec_d, nullptr, &ucell, orb.cutoffs(), &gd,
192+
two_center_bundle.kinetic_orb.get());
193+
tmp_ekinetic.cal_force_stress(isforce, isstress, dmR, ftvnl_dphi, stvnl_dphi);
194+
}
195+
196+
// Calculate overlap force/stress (uses EDM)
197+
hamilt::OverlapNew<hamilt::OperatorLCAO<T, double>> tmp_overlap(
198+
nullptr, kv.kvec_d, nullptr, nullptr, &ucell, orb.cutoffs(), &gd,
199+
two_center_bundle.overlap_orb.get());
200+
tmp_overlap.cal_force_stress(isforce, isstress, edmR, foverlap, soverlap);
201+
202+
// Calculate nonlocal force/stress (uses DM)
203+
hamilt::NonlocalNew<hamilt::OperatorLCAO<T, double>> tmp_nonlocal(
204+
nullptr, kv.kvec_d, nullptr, &ucell, orb.cutoffs(), &gd,
205+
two_center_bundle.overlap_orb_beta.get());
206+
tmp_nonlocal.cal_force_stress(isforce, isstress, dmR, fvnl_dbeta, svnl_dbeta);
207+
208+
// Calculate local potential force/stress (vl_dphi)
209+
// This uses grid integration, not operator-based method
210+
flk.ParaV = dmat.dm->get_paraV_pointer();
211+
PulayForceStress::cal_pulay_fs(fvl_dphi, svl_dphi, *dmat.dm, ucell, pelec->pot,
212+
isforce, isstress, false /*reset dm to gint*/);
213+
214+
// Switch back to spin channel 0
184215
if (PARAM.inp.nspin == 2)
185216
{
186217
dmat.dm->switch_dmr(0);
187218
}
188219
}
189220
else if (PARAM.inp.nspin == 4)
190221
{
191-
hamilt::NonlocalNew<hamilt::OperatorLCAO<std::complex<double>, std::complex<double>>> tmp_nonlocal(
192-
nullptr, kv.kvec_d, nullptr, &ucell, orb.cutoffs(), &gd,
193-
two_center_bundle.overlap_orb_beta.get());
194-
195-
// calculate temporary complex DMR for nonlocal force&stress
196-
// In fact, only SOC part need the imaginary part of DMR for correct force&stress
222+
// For nspin=4 (non-collinear), need complex DMR
223+
// Create temporary complex DMR for DM
197224
hamilt::HContainer<std::complex<double>> tmp_dmr(dmat.dm->get_DMR_pointer(1)->get_paraV());
198225
std::vector<int> ijrs = dmat.dm->get_DMR_pointer(1)->get_ijr_info();
199226
tmp_dmr.insert_ijrs(&ijrs);
200227
tmp_dmr.allocate();
201228
dmat.dm->cal_DMR_full(&tmp_dmr);
229+
230+
// Create temporary complex DMR for EDM
231+
hamilt::HContainer<std::complex<double>> tmp_edmr(edm.get_DMR_pointer(1)->get_paraV());
232+
tmp_edmr.insert_ijrs(&ijrs);
233+
tmp_edmr.allocate();
234+
edm.cal_DMR_full(&tmp_edmr);
235+
236+
// Calculate kinetic force/stress (uses DM)
237+
if (PARAM.inp.t_in_h)
238+
{
239+
hamilt::EkineticNew<hamilt::OperatorLCAO<std::complex<double>, std::complex<double>>> tmp_ekinetic(
240+
nullptr, kv.kvec_d, nullptr, &ucell, orb.cutoffs(), &gd,
241+
two_center_bundle.kinetic_orb.get());
242+
tmp_ekinetic.cal_force_stress(isforce, isstress, &tmp_dmr, ftvnl_dphi, stvnl_dphi);
243+
}
244+
245+
// Calculate overlap force/stress (uses EDM)
246+
hamilt::OverlapNew<hamilt::OperatorLCAO<std::complex<double>, std::complex<double>>> tmp_overlap(
247+
nullptr, kv.kvec_d, nullptr, nullptr, &ucell, orb.cutoffs(), &gd,
248+
two_center_bundle.overlap_orb.get());
249+
tmp_overlap.cal_force_stress(isforce, isstress, &tmp_edmr, foverlap, soverlap);
250+
251+
// Calculate nonlocal force/stress (uses DM)
252+
hamilt::NonlocalNew<hamilt::OperatorLCAO<std::complex<double>, std::complex<double>>> tmp_nonlocal(
253+
nullptr, kv.kvec_d, nullptr, &ucell, orb.cutoffs(), &gd,
254+
two_center_bundle.overlap_orb_beta.get());
202255
tmp_nonlocal.cal_force_stress(isforce, isstress, &tmp_dmr, fvnl_dbeta, svnl_dbeta);
256+
257+
// Calculate local potential force/stress (vl_dphi)
258+
flk.ParaV = dmat.dm->get_paraV_pointer();
259+
PulayForceStress::cal_pulay_fs(fvl_dphi, svl_dphi, *dmat.dm, ucell, pelec->pot,
260+
isforce, isstress, false /*reset dm to gint*/);
261+
}
262+
263+
// MPI reduction for forces
264+
if (isforce)
265+
{
266+
Parallel_Reduce::reduce_pool(foverlap.c, foverlap.nr * foverlap.nc);
267+
Parallel_Reduce::reduce_pool(ftvnl_dphi.c, ftvnl_dphi.nr * ftvnl_dphi.nc);
268+
Parallel_Reduce::reduce_pool(fvnl_dbeta.c, fvnl_dbeta.nr * fvnl_dbeta.nc);
269+
Parallel_Reduce::reduce_pool(fvl_dphi.c, fvl_dphi.nr * fvl_dphi.nc);
270+
}
271+
272+
// MPI reduction for stresses
273+
if (isstress)
274+
{
275+
Parallel_Reduce::reduce_pool(soverlap.c, soverlap.nr * soverlap.nc);
276+
Parallel_Reduce::reduce_pool(stvnl_dphi.c, stvnl_dphi.nr * stvnl_dphi.nc);
277+
Parallel_Reduce::reduce_pool(svnl_dbeta.c, svnl_dbeta.nr * svnl_dbeta.nc);
278+
Parallel_Reduce::reduce_pool(svl_dphi.c, svl_dphi.nr * svl_dphi.nc);
279+
}
280+
281+
// Handle DeePKS forces if enabled
282+
#ifdef __MLALGO
283+
if (PARAM.inp.deepks_scf)
284+
{
285+
const int nks = (PARAM.inp.nspin == 1 || PARAM.inp.nspin == 2) ? 1 : kv.get_nks();
286+
if (PARAM.globalv.gamma_only_local)
287+
{
288+
DeePKS_domain::cal_f_delta<double>(deepks.ld.dm_r, ucell, orb, gd,
289+
*flk.ParaV, nks, deepks.ld.deepks_param,
290+
kv.kvec_d, deepks.ld.phialpha, deepks.ld.gedm,
291+
fvnl_dalpha, isstress, svnl_dalpha);
292+
}
293+
else
294+
{
295+
DeePKS_domain::cal_f_delta<std::complex<double>>(deepks.ld.dm_r, ucell, orb, gd,
296+
*flk.ParaV, nks, deepks.ld.deepks_param,
297+
kv.kvec_d, deepks.ld.phialpha, deepks.ld.gedm,
298+
fvnl_dalpha, isstress, svnl_dalpha);
299+
}
300+
301+
if (isforce)
302+
{
303+
Parallel_Reduce::reduce_pool(fvnl_dalpha.c, fvnl_dalpha.nr * fvnl_dalpha.nc);
304+
}
305+
if (isstress)
306+
{
307+
Parallel_Reduce::reduce_pool(svnl_dalpha.c, svnl_dalpha.nr * svnl_dalpha.nc);
308+
}
203309
}
310+
#endif
204311

205312
//! forces and stress from vdw
206313
// Peize Lin add 2014-04-04, update 2021-03-09
@@ -275,13 +382,12 @@ void Force_Stress_LCAO<T>::getForceStress(UnitCell& ucell,
275382
}
276383
if (PARAM.inp.dft_plus_u == 2)
277384
{
278-
// dftu.force_stress(ucell, gd, pelec, pv, fsr, force_u, stress_u, kv);
279-
// mohan modify 2025-11-03
385+
// Old DFT+U implementation (dft_plus_u==2) still needs ForceStressArrays
386+
ForceStressArrays fsr_dftu;
280387
std::vector<std::vector<double>>* dmk_d = nullptr;
281388
std::vector<std::vector<std::complex<double>>>* dmk_c = nullptr;
282-
// add a new template function
283389
assign_dmk_ptr<T>(dmat.dm, dmk_d, dmk_c, PARAM.globalv.gamma_only_local);
284-
dftu.force_stress(ucell, gd, dmk_d, dmk_c, pv, fsr, force_u, stress_u, kv);
390+
dftu.force_stress(ucell, gd, dmk_d, dmk_c, pv, fsr_dftu, force_u, stress_u, kv);
285391
}
286392
else
287393
{
@@ -332,10 +438,11 @@ void Force_Stress_LCAO<T>::getForceStress(UnitCell& ucell,
332438
}
333439
}
334440

335-
if (!PARAM.globalv.gamma_only_local)
336-
{
337-
this->flk.finish_ftable(fsr);
338-
}
441+
// NOTE: finish_ftable is no longer needed as we don't use ForceStressArrays for overlap/kinetic
442+
// if (!PARAM.globalv.gamma_only_local)
443+
// {
444+
// this->flk.finish_ftable(fsr);
445+
// }
339446

340447
#ifdef __EXX
341448
// Force and Stress contribution from exx

0 commit comments

Comments
 (0)