Skip to content

Commit 78e5645

Browse files
AsTonyshmentAsTonyshment
committed
Update DeePKS RT-TDDFT code
1 parent 4b1fc7d commit 78e5645

12 files changed

Lines changed: 1646 additions & 10 deletions

File tree

‎source/source_esolver/esolver_ks_lcao_tddft.cpp‎

Lines changed: 42 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -20,6 +20,9 @@
2020
#ifdef __EXX
2121
#include "source_lcao/module_ri/Exx_LRI_interface.h"
2222
#endif
23+
#ifdef __MLALGO
24+
#include "source_lcao/module_deepks/deepks_vcomm_r.h"
25+
#endif
2326

2427
namespace ModuleESolver
2528
{
@@ -138,6 +141,45 @@ void ESolver_KS_LCAO_TDDFT<TR, Device>::runner(UnitCell& ucell, const int istep)
138141
// calculate velocity operator
139142
velocity_mat->calculate_grad_term();
140143
velocity_mat->calculate_vcomm_r();
144+
145+
#ifdef __MLALGO
146+
if (PARAM.inp.deepks_scf)
147+
{
148+
std::vector<hamilt::HContainer<std::complex<double>>*> current_term(3);
149+
for (int i = 0; i < 3; i++)
150+
{
151+
current_term[i] = velocity_mat->get_current_term_pointer(i);
152+
}
153+
154+
if (PARAM.inp.td_stype == 1)
155+
{
156+
// Velocity gauge: use grid integration with Peierls phase
157+
DeePKS_domain::cal_deepks_vcomm_r_vel<std::complex<double>>(ucell,
158+
this->orb_,
159+
this->gd,
160+
&this->pv,
161+
this->deepks.ld,
162+
current_term);
163+
}
164+
else
165+
{
166+
// Length gauge: use two-center integrals
167+
static bool alpha_init_done = false;
168+
if (!alpha_init_done)
169+
{
170+
TD_info::td_vel_op->r_calculator.init_alpha(ucell, this->pv, this->orb_);
171+
alpha_init_done = true;
172+
}
173+
DeePKS_domain::cal_deepks_vcomm_r<std::complex<double>>(ucell,
174+
this->orb_,
175+
this->gd,
176+
&this->pv,
177+
this->deepks.ld,
178+
TD_info::td_vel_op->r_calculator,
179+
current_term);
180+
}
181+
}
182+
#endif
141183
}
142184
int estep_max = (istep == 0 && !PARAM.inp.mdp.md_restart) ? 1 : PARAM.inp.estep_per_md;
143185
// mohan change md_nstep from 0 to 1, 2026-01-04

‎source/source_io/module_hs/cal_r_overlap_R.cpp‎

Lines changed: 217 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -16,22 +16,26 @@ cal_r_overlap_R::~cal_r_overlap_R()
1616
}
1717

1818
void cal_r_overlap_R::initialize_orb_table(const UnitCell& ucell,
19-
const LCAO_Orbitals& orb)
19+
const LCAO_Orbitals& orb,
20+
const int lmax_extra)
2021
{
2122
const int ntype = orb.get_ntype();
2223
int lmax_orb = -1;
2324
for (int it = 0; it < ntype; it++)
2425
{
2526
lmax_orb = std::max(lmax_orb, orb.Phi[it].getLmax());
2627
}
28+
// lmax_extra (e.g. Alpha[0].getLmax()) ensures the spherical Bessel and Gaunt
29+
// tables also cover integrals against projector sets with higher angular momentum.
30+
const int lmax_eff = std::max(lmax_orb, lmax_extra);
2731
const double dr = orb.get_dR();
2832
const double dk = orb.get_dk();
2933
const int kmesh = orb.get_kmesh() * 4 + 1;
3034
int Rmesh = static_cast<int>(orb.get_Rmax() / dr) + 4;
3135
Rmesh += 1 - Rmesh % 2;
3236

33-
const int Lmax = lmax_orb + 1;
34-
const int Lmax_used = 2 * lmax_orb + 1;
37+
const int Lmax = lmax_eff + 1;
38+
const int Lmax_used = lmax_orb + lmax_eff + 1;
3539
Center2_Orb::init_Table_Spherical_Bessel(Lmax_used,
3640
dr,
3741
dk,
@@ -635,6 +639,216 @@ void cal_r_overlap_R::get_psi_r_beta(const UnitCell& ucell,
635639
}
636640

637641

642+
#ifdef __MLALGO
643+
644+
void cal_r_overlap_R::init_alpha(const UnitCell& ucell,const Parallel_Orbitals& pv, const LCAO_Orbitals& orb)
645+
{
646+
ModuleBase::TITLE("cal_r_overlap_R", "init_alpha");
647+
ModuleBase::timer::start("cal_r_overlap_R", "init_alpha");
648+
this->ParaV = &pv;
649+
650+
// Initialize (or re-initialize) the spherical Bessel and Gaunt tables with
651+
// alpha_lmax to cover the orbital↔alpha integrals. construct_orbs_and_orb_r
652+
// is only needed once and skipped if already built (e.g. by init_nonlocal).
653+
const int alpha_lmax = orb.Alpha[0].getLmax();
654+
initialize_orb_table(ucell, orb, alpha_lmax);
655+
if (orbs.empty())
656+
{
657+
construct_orbs_and_orb_r(ucell, orb);
658+
}
659+
construct_orbs_and_alpha_and_orb_r(ucell, orb);
660+
661+
ModuleBase::timer::end("cal_r_overlap_R", "init_alpha");
662+
return;
663+
}
664+
665+
void cal_r_overlap_R::construct_orbs_and_alpha_and_orb_r(const UnitCell& ucell,const LCAO_Orbitals& orb)
666+
{
667+
// Build orbs_alpha from orb.Alpha[0]
668+
// Alpha orbitals are Numerical_Orbital_Lm (same type as Phi), so copy them directly.
669+
int lmax_alpha = orb.Alpha[0].getLmax();
670+
nproj_alpha = 0;
671+
alpha_ip2ln.clear();
672+
for (int L = 0; L <= lmax_alpha; ++L)
673+
{
674+
int nchi_L = orb.Alpha[0].getNchi(L);
675+
for (int N = 0; N < nchi_L; ++N)
676+
{
677+
alpha_ip2ln.push_back({L, N});
678+
nproj_alpha++;
679+
}
680+
}
681+
682+
orbs_alpha.resize(1); // single "type" for all alpha orbitals
683+
orbs_alpha[0].resize(nproj_alpha);
684+
685+
int ip = 0;
686+
for (int L = 0; L <= lmax_alpha; ++L)
687+
{
688+
int nchi_L = orb.Alpha[0].getNchi(L);
689+
for (int N = 0; N < nchi_L; ++N)
690+
{
691+
const auto& alpha_ln = orb.Alpha[0].PhiLN(L, N);
692+
orbs_alpha[0][ip].set_orbital_info(alpha_ln.getLabel(),
693+
0, // type 0 for descriptor
694+
alpha_ln.getL(),
695+
alpha_ln.getChi(),
696+
alpha_ln.getNr(),
697+
alpha_ln.getRab(),
698+
alpha_ln.getRadial(),
699+
Numerical_Orbital_Lm::Psi_Type::Psi,
700+
alpha_ln.getPsi(),
701+
static_cast<int>(alpha_ln.getNk() * kmesh_times) | 1,
702+
alpha_ln.getDk(),
703+
alpha_ln.getDruniform(),
704+
false,
705+
true,
706+
PARAM.inp.cal_force);
707+
ip++;
708+
}
709+
}
710+
711+
// Build two-center tables: Orb11 (overlap) and Orb21 (r-operator)
712+
int ntype = orb.get_ntype();
713+
for (int TA = 0; TA < ntype; ++TA)
714+
{
715+
for (int LA = 0; LA <= orb.Phi[TA].getLmax(); ++LA)
716+
{
717+
for (int NA = 0; NA < orb.Phi[TA].getNchi(LA); ++NA)
718+
{
719+
for (int aip = 0; aip < nproj_alpha; aip++)
720+
{
721+
center2_orb11_alpha[TA][0][LA][NA].insert(
722+
std::make_pair(aip, Center2_Orb::Orb11(orbs[TA][LA][NA], orbs_alpha[0][aip], psb_, MGT)));
723+
}
724+
}
725+
}
726+
}
727+
728+
for (int TA = 0; TA < ntype; ++TA)
729+
{
730+
for (int LA = 0; LA <= orb.Phi[TA].getLmax(); ++LA)
731+
{
732+
for (int NA = 0; NA < orb.Phi[TA].getNchi(LA); ++NA)
733+
{
734+
for (int aip = 0; aip < nproj_alpha; aip++)
735+
{
736+
center2_orb21_r_alpha[TA][0][LA][NA].insert(
737+
std::make_pair(aip, Center2_Orb::Orb21(orbs[TA][LA][NA], orb_r, orbs_alpha[0][aip], psb_, MGT)));
738+
}
739+
}
740+
}
741+
}
742+
743+
// Initialize radial tables
744+
for (auto& co1: center2_orb11_alpha)
745+
{
746+
for (auto& co2: co1.second)
747+
{
748+
for (auto& co3: co2.second)
749+
{
750+
for (auto& co4: co3.second)
751+
{
752+
for (auto& co5: co4.second)
753+
{
754+
co5.second.init_radial_table();
755+
}
756+
}
757+
}
758+
}
759+
}
760+
761+
for (auto& co1: center2_orb21_r_alpha)
762+
{
763+
for (auto& co2: co1.second)
764+
{
765+
for (auto& co3: co2.second)
766+
{
767+
for (auto& co4: co3.second)
768+
{
769+
for (auto& co5: co4.second)
770+
{
771+
co5.second.init_radial_table();
772+
}
773+
}
774+
}
775+
}
776+
}
777+
}
778+
779+
void cal_r_overlap_R::get_psi_r_alpha(
780+
std::vector<std::vector<double>>& nlm,
781+
const ModuleBase::Vector3<double>& R1,
782+
const int& T1,
783+
const int& L1,
784+
const int& m1,
785+
const int& N1,
786+
const ModuleBase::Vector3<double>& R0)
787+
{
788+
ModuleBase::Vector3<double> origin_point(0.0, 0.0, 0.0);
789+
double factor = sqrt(ModuleBase::FOUR_PI / 3.0);
790+
const ModuleBase::Vector3<double>& distance = R0 - R1;
791+
792+
// Count total m-components across all alpha channels
793+
int natomwfc = 0;
794+
for (int ip = 0; ip < nproj_alpha; ip++)
795+
{
796+
int L = alpha_ip2ln[ip].first;
797+
natomwfc += 2 * L + 1;
798+
}
799+
800+
nlm.resize(4);
801+
for (int i = 0; i < 4; i++)
802+
{
803+
nlm[i].resize(natomwfc);
804+
}
805+
806+
int index = 0;
807+
for (int ip = 0; ip < nproj_alpha; ip++)
808+
{
809+
int L2 = alpha_ip2ln[ip].first;
810+
int nm2 = 2 * L2 + 1;
811+
812+
for (int m2 = 0; m2 < nm2; m2++)
813+
{
814+
// <phi|alpha> overlap
815+
double overlap_o
816+
= center2_orb11_alpha[T1][0][L1][N1].at(ip).cal_overlap(origin_point, distance, m1, m2);
817+
818+
// <phi|(r-R1)|alpha> : local dipole term using orb_r
819+
// m_middle = 1 → x, 2 → y, 0 → z (with sign conventions matching get_psi_r_beta)
820+
double overlap_x = -1 * factor
821+
* center2_orb21_r_alpha[T1][0][L1][N1].at(ip).cal_overlap(origin_point,
822+
distance,
823+
m1,
824+
1, // m_middle for x
825+
m2);
826+
double overlap_y = -1 * factor
827+
* center2_orb21_r_alpha[T1][0][L1][N1].at(ip).cal_overlap(origin_point,
828+
distance,
829+
m1,
830+
2, // m_middle for y
831+
m2);
832+
double overlap_z = factor
833+
* center2_orb21_r_alpha[T1][0][L1][N1].at(ip).cal_overlap(origin_point,
834+
distance,
835+
m1,
836+
0, // m_middle for z
837+
m2);
838+
839+
// <phi|r|alpha> = <phi|(r-R1)|alpha> + R1 * <phi|alpha>
840+
nlm[0][index] = overlap_o;
841+
nlm[1][index] = overlap_x + (R1 * overlap_o).x;
842+
nlm[2][index] = overlap_y + (R1 * overlap_o).y;
843+
nlm[3][index] = overlap_z + (R1 * overlap_o).z;
844+
index++;
845+
}
846+
}
847+
}
848+
849+
#endif
850+
851+
638852
void cal_r_overlap_R::out_rR(const UnitCell& ucell, const Grid_Driver& gd, const int& istep)
639853
{
640854
ModuleBase::TITLE("cal_r_overlap_R", "out_rR");

‎source/source_io/module_hs/cal_r_overlap_R.h‎

Lines changed: 39 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -70,13 +70,32 @@ class cal_r_overlap_R
7070
const ModuleBase::Vector3<double>& R2,
7171
const int& T2
7272
);
73+
#ifdef __MLALGO
74+
void init_alpha(const UnitCell& ucell,const Parallel_Orbitals& pv, const LCAO_Orbitals& orb);
75+
void get_psi_r_alpha(
76+
std::vector<std::vector<double>>& nlm,
77+
const ModuleBase::Vector3<double>& R1,
78+
const int& T1,
79+
const int& L1,
80+
const int& m1,
81+
const int& N1,
82+
const ModuleBase::Vector3<double>& R0
83+
);
84+
#endif
7385
void out_rR(const UnitCell& ucell, const Grid_Driver& gd, const int& istep);
7486
void out_rR_other(const UnitCell& ucell, const int& istep, const std::set<Abfs::Vector3_Order<int>>& output_R_coor);
7587

7688
private:
77-
void initialize_orb_table(const UnitCell& ucell, const LCAO_Orbitals& orb);
89+
// lmax_extra (default 0) enlarges the spherical Bessel / Gaunt tables to also
90+
// cover integrals against a higher-L projector set (e.g. DeePKS alpha orbitals).
91+
// With the default it reduces exactly to the orbital-only sizing, so existing
92+
// callers (init / init_nonlocal) are unaffected.
93+
void initialize_orb_table(const UnitCell& ucell, const LCAO_Orbitals& orb, const int lmax_extra = 0);
7894
void construct_orbs_and_orb_r(const UnitCell& ucell,const LCAO_Orbitals& orb);
7995
void construct_orbs_and_nonlocal_and_orb_r(const UnitCell& ucell,const LCAO_Orbitals& orb);
96+
#ifdef __MLALGO
97+
void construct_orbs_and_alpha_and_orb_r(const UnitCell& ucell,const LCAO_Orbitals& orb);
98+
#endif
8099

81100
std::vector<int> iw2ia;
82101
std::vector<int> iw2iL;
@@ -111,6 +130,25 @@ class cal_r_overlap_R
111130
std::map<size_t, std::map<size_t, std::map<size_t, std::map<size_t, Center2_Orb::Orb21>>>>>
112131
center2_orb21_r_nonlocal;
113132

133+
#ifdef __MLALGO
134+
/// @brief alpha orbital channels from orb.Alpha[0] (descriptor projectors)
135+
std::vector<std::vector<Numerical_Orbital_Lm>> orbs_alpha;
136+
/// @brief total number of alpha channels (flat index across all (L,N) pairs)
137+
int nproj_alpha = 0;
138+
/// @brief map from flat alpha channel index to {L, N}
139+
std::vector<std::pair<int, int>> alpha_ip2ln;
140+
/// @brief Orb11 table: <phi|alpha>
141+
std::map<
142+
size_t,
143+
std::map<size_t, std::map<size_t, std::map<size_t, std::map<size_t, Center2_Orb::Orb11>>>>>
144+
center2_orb11_alpha;
145+
/// @brief Orb21 table: <phi|r|alpha> using orb_r (L=1) on the phi center
146+
std::map<
147+
size_t,
148+
std::map<size_t, std::map<size_t, std::map<size_t, std::map<size_t, Center2_Orb::Orb21>>>>>
149+
center2_orb21_r_alpha;
150+
#endif
151+
114152
const Parallel_Orbitals* ParaV = nullptr;
115153
};
116154
#endif

‎source/source_lcao/hamilt_lcao.cpp‎

Lines changed: 11 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -202,6 +202,11 @@ HamiltLCAO<TK, TR>::HamiltLCAO(const UnitCell& ucell,
202202
#ifdef __MLALGO
203203
if (PARAM.inp.deepks_scf)
204204
{
205+
if (PARAM.inp.esolver_type == "tddft" && PARAM.inp.td_stype == 2)
206+
{
207+
ModuleBase::WARNING_QUIT("hamilt_lcao",
208+
"DeePKS with hybrid gauge (td_stype=2) is not supported yet.");
209+
}
205210
Operator<TK>* deepks_op = new DeePKS<OperatorLCAO<TK, TR>>(this->hsk,
206211
this->kv->kvec_d,
207212
this->hR, // no explicit call yet
@@ -327,6 +332,11 @@ HamiltLCAO<TK, TR>::HamiltLCAO(const UnitCell& ucell,
327332
#ifdef __MLALGO
328333
if (PARAM.inp.deepks_scf)
329334
{
335+
if (PARAM.inp.esolver_type == "tddft" && PARAM.inp.td_stype == 2)
336+
{
337+
ModuleBase::WARNING_QUIT("hamilt_lcao",
338+
"DeePKS with hybrid gauge (td_stype=2) is not supported yet.");
339+
}
330340
Operator<TK>* deepks_op = new DeePKS<OperatorLCAO<TK, TR>>(this->hsk,
331341
this->kv->kvec_d,
332342
hR,
@@ -374,6 +384,7 @@ HamiltLCAO<TK, TR>::HamiltLCAO(const UnitCell& ucell,
374384
two_center_bundle.kinetic_orb.get());
375385
this->getOperator()->add(td_pot_hybrid);
376386
}
387+
377388
if (PARAM.inp.dft_plus_u)
378389
{
379390
Operator<TK>* plus_u = nullptr;

0 commit comments

Comments
 (0)