Skip to content

DeePKS: support collinear spin (nspin=2) #7433

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
ErjieWu wants to merge 14 commits into
base: develop
Choose a base branch
from
Open

DeePKS: support collinear spin (nspin=2) #7433

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
14 commits
Select commit Hold shift + click to select a range
ef5cc23
deepks(nspin2): add magnetization-channel data members
ErjieWu Jun 4, 2026
18ea114
deepks(nspin2): compute charge/magnetization channels, feed 2-channel…
ErjieWu Jun 4, 2026
eb0ccba
deepks(nspin2): apply per-spin V_delta and forces/stress
ErjieWu Jun 4, 2026
366f1d7
deepks(nspin2): write the dm_eig descriptor label with the spin channel
ErjieWu Jun 4, 2026
d949987
deepks(nspin2): spin channel for precalc labels and bandgap
ErjieWu Jun 4, 2026
b993569
test(deepks): add nspin=2 integration tests (29-31)
ErjieWu Jun 4, 2026
b8b46a3
deepks(nspin2): guard precalc/bandgap mag-channel stacks on rank 0
ErjieWu Jun 5, 2026
2272448
deepks(nspin2): guard undefined autograd grads and fix indentation
ErjieWu Jun 5, 2026
fac0998
deepks(nspin2): make _mag parameters adjacent and type-consistent
ErjieWu Jun 5, 2026
031caa6
deepks(nspin2): make _mag arguments a trailing optional group
ErjieWu Jun 5, 2026
89e6618
deepks(nspin2): apply clang-format to changed lines
ErjieWu Jun 5, 2026
e43ef68
deepks(nspin2): drop per-call gedm buffer, clarify contributeHR
ErjieWu Jun 5, 2026
ac36afc
deepks: bracket whole contributeHR with its timer
ErjieWu Jun 5, 2026
b2eef18
test(deepks): loosen tolerance for nspin=2 bandgap precalc case
ErjieWu Jun 5, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
40 changes: 32 additions & 8 deletions source/source_lcao/FORCE_STRESS.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -276,17 +276,41 @@ void Force_Stress_LCAO<T>::getForceStress(UnitCell& ucell,
const int nks = (PARAM.inp.nspin == 1 || PARAM.inp.nspin == 2) ? 1 : kv.get_nks();
if (PARAM.globalv.gamma_only_local)
{
DeePKS_domain::cal_f_delta<double>(deepks.ld.dm_r, ucell, orb, gd,
*flk.ParaV, nks, deepks.ld.deepks_param,
kv.kvec_d, deepks.ld.phialpha, deepks.ld.gedm,
fvnl_dalpha, isstress, svnl_dalpha);
DeePKS_domain::cal_f_delta<double>(
ucell,
orb,
gd,
*flk.ParaV,
nks,
deepks.ld.deepks_param,
kv.kvec_d,
deepks.ld.phialpha,
fvnl_dalpha,
isstress,
svnl_dalpha,
deepks.ld.dm_r,
deepks.ld.gedm,
(PARAM.inp.nspin == 2 && !PARAM.inp.deepks_equiv) ? deepks.ld.dm_r_mag : nullptr,
(PARAM.inp.nspin == 2 && !PARAM.inp.deepks_equiv) ? deepks.ld.gedm_mag : nullptr);
}
else
{
DeePKS_domain::cal_f_delta<std::complex<double>>(deepks.ld.dm_r, ucell, orb, gd,
*flk.ParaV, nks, deepks.ld.deepks_param,
kv.kvec_d, deepks.ld.phialpha, deepks.ld.gedm,
fvnl_dalpha, isstress, svnl_dalpha);
DeePKS_domain::cal_f_delta<std::complex<double>>(
ucell,
orb,
gd,
*flk.ParaV,
nks,
deepks.ld.deepks_param,
kv.kvec_d,
deepks.ld.phialpha,
fvnl_dalpha,
isstress,
svnl_dalpha,
deepks.ld.dm_r,
deepks.ld.gedm,
(PARAM.inp.nspin == 2 && !PARAM.inp.deepks_equiv) ? deepks.ld.dm_r_mag : nullptr,
(PARAM.inp.nspin == 2 && !PARAM.inp.deepks_equiv) ? deepks.ld.gedm_mag : nullptr);
}

if (isforce)
Expand Down
8 changes: 4 additions & 4 deletions source/source_lcao/FORCE_gamma.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -217,19 +217,19 @@ void Force_LCAO<double>::ftable(const bool isforce,
{
// No need to update E_delta here since it have been done in LCAO_Deepks_Interface in after_scf
const int nks = 1;
DeePKS_domain::cal_f_delta<double>(deepks.ld.dm_r,
ucell,
DeePKS_domain::cal_f_delta<double>(ucell,
orb,
gd,
*this->ParaV,
nks,
deepks.ld.deepks_param,
kv->kvec_d,
deepks.ld.phialpha,
deepks.ld.gedm,
fvnl_dalpha,
isstress,
svnl_dalpha);
svnl_dalpha,
deepks.ld.dm_r,
deepks.ld.gedm);
}
#endif

Expand Down
8 changes: 4 additions & 4 deletions source/source_lcao/FORCE_k.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -224,19 +224,19 @@ void Force_LCAO<std::complex<double>>::ftable(const bool isforce,
if (PARAM.inp.deepks_scf)
{
// No need to update E_delta since it have been done in LCAO_Deepks_Interface in after_scf
DeePKS_domain::cal_f_delta<std::complex<double>>(deepks.ld.dm_r,
ucell,
DeePKS_domain::cal_f_delta<std::complex<double>>(ucell,
orb,
gd,
pv,
kv->get_nks(),
deepks.ld.deepks_param,
kv->kvec_d,
deepks.ld.phialpha,
deepks.ld.gedm,
fvnl_dalpha,
isstress,
svnl_dalpha);
svnl_dalpha,
deepks.ld.dm_r,
deepks.ld.gedm);
}
#endif

Expand Down
35 changes: 35 additions & 0 deletions source/source_lcao/module_deepks/LCAO_deepks.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -11,6 +11,7 @@ LCAO_Deepks<T>::LCAO_Deepks()
{
deepks_param.inl_index = new ModuleBase::IntArray[1];
gedm = nullptr;
gedm_mag = nullptr;
this->phialpha.resize(1);
}

Expand All @@ -33,6 +34,16 @@ LCAO_Deepks<T>::~LCAO_Deepks()
}
delete[] gedm;
}
if (gedm_mag)
{
for (int inl = 0; inl < this->deepks_param.inlmax; inl++)
{
delete[] gedm_mag[inl];
}
delete[] gedm_mag;
}
pdm_mag.clear();
delete dm_r_mag;
}

template <typename T>
Expand Down Expand Up @@ -124,6 +135,15 @@ void LCAO_Deepks<T>::init(const LCAO_Orbitals& orb,
}
}

if (PARAM.inp.nspin == 2 && !PARAM.inp.deepks_equiv) // magnetization-channel PDM
{
this->pdm_mag.resize(this->deepks_param.inlmax);
for (int inl = 0; inl < this->deepks_param.inlmax; ++inl)
{
this->pdm_mag[inl] = torch::zeros_like(this->pdm[inl]);
}
}

this->pv = &pv_in;

ModuleBase::timer::end("LCAO_Deepks", "init");
Expand Down Expand Up @@ -202,6 +222,16 @@ void LCAO_Deepks<T>::allocate_V_delta(const int nat, const int nks)
ModuleBase::GlobalFunc::ZEROS(this->gedm[inl], pdm_size);
}

if (PARAM.inp.nspin == 2 && !PARAM.inp.deepks_equiv) // magnetization-channel gedm
{
this->gedm_mag = new double*[this->deepks_param.inlmax];
for (int inl = 0; inl < this->deepks_param.inlmax; inl++)
{
this->gedm_mag[inl] = new double[pdm_size];
ModuleBase::GlobalFunc::ZEROS(this->gedm_mag[inl], pdm_size);
}
}

ModuleBase::timer::end("LCAO_Deepks", "allocate_V_delta");
return;
}
Expand Down Expand Up @@ -249,6 +279,11 @@ void LCAO_Deepks<T>::init_DMR(const UnitCell& ucell,
this->dm_r->insert_pair(dm_pair);
});
this->dm_r->allocate(nullptr, true);

if (PARAM.inp.nspin == 2 && !PARAM.inp.deepks_equiv) // magnetization-channel real-space DM
{
this->dm_r_mag = new hamilt::HContainer<double>(*this->dm_r);
}
}

template <typename T>
Expand Down
5 changes: 5 additions & 0 deletions source/source_lcao/module_deepks/LCAO_deepks.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -79,6 +79,11 @@ class LCAO_Deepks
/// dE/dD, autograd from loaded model(E: Ry)
double** gedm = nullptr; //[tot_Inl][(2l+1)*(2l+1)]

// magnetization-channel (rho_up - rho_dn) counterparts, used only for nspin=2
hamilt::HContainer<double>* dm_r_mag = nullptr;
std::vector<torch::Tensor> pdm_mag;
double** gedm_mag = nullptr;

// functions for hr status: 1. get value; 2. set value;
int get_hr_cal()
{
Expand Down
Loading
Loading