Skip to content

Refactor: Remove some redundant variables and global dependence in DeePKS. #5791

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

Merged
merged 3 commits into from
Jan 1, 2025
Merged

Refactor: Remove some redundant variables and global dependence in DeePKS. #5791

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
42066f8
Remove cal_orbital_precalc() and corresponding temporary variables fr…
ErjieWu Dec 31, 2024
e99ff8c
Merge branch 'refactor' of https://github.com/ErjieWu/abacus-develop …
ErjieWu Dec 31, 2024
f85c16d
clang-format adjustment.
ErjieWu Dec 31, 2024
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
2 changes: 1 addition & 1 deletion source/Makefile.Objects
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -200,7 +200,7 @@ OBJS_DEEPKS=LCAO_deepks.o\
LCAO_deepks_vdelta.o\
deepks_hmat.o\
LCAO_deepks_interface.o\
orbital_precalc.o\
deepks_orbpre.o\
cal_gdmx.o\
cal_gedm.o\
cal_gvx.o\
Expand Down
10 changes: 7 additions & 3 deletions source/module_hamilt_lcao/hamilt_lcaodft/FORCE_STRESS.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -7,12 +7,12 @@
// new
#include "module_base/timer.h"
#include "module_cell/module_neighbor/sltk_grid_driver.h"
#include "module_elecstate/elecstate_lcao.h"
#include "module_elecstate/potentials/efield.h" // liuyu add 2022-05-18
#include "module_elecstate/potentials/gatefield.h" // liuyu add 2022-09-13
#include "module_hamilt_general/module_surchem/surchem.h" //sunml add 2022-08-10
#include "module_hamilt_general/module_vdw/vdw.h"
#include "module_parameter/parameter.h"
#include "module_elecstate/elecstate_lcao.h"
#ifdef __DEEPKS
#include "module_hamilt_lcao/module_deepks/LCAO_deepks.h" //caoyu add for deepks 2021-06-03
#include "module_hamilt_lcao/module_deepks/LCAO_deepks_io.h" // mohan add 2024-07-22
Expand Down Expand Up @@ -540,7 +540,9 @@ void Force_Stress_LCAO<T>::getForceStress(UnitCell& ucell,
{
GlobalC::ld.check_gdmx(ucell.nat);
}
GlobalC::ld.cal_gvx(ucell.nat);
std::vector<torch::Tensor> gevdm;
GlobalC::ld.cal_gevdm(ucell.nat, gevdm);
GlobalC::ld.cal_gvx(ucell.nat, gevdm);

if (PARAM.inp.deepks_out_unittest)
{
Expand Down Expand Up @@ -758,7 +760,9 @@ void Force_Stress_LCAO<T>::getForceStress(UnitCell& ucell,

if (!PARAM.inp.deepks_equiv) // training with stress label not supported by equivariant version now
{
GlobalC::ld.cal_gvepsl(ucell.nat);
std::vector<torch::Tensor> gevdm;
GlobalC::ld.cal_gevdm(ucell.nat, gevdm);
GlobalC::ld.cal_gvepsl(ucell.nat, gevdm);

LCAO_deepks_io::save_npy_gvepsl(ucell.nat,
GlobalC::ld.des_per_atom,
Expand Down
2 changes: 1 addition & 1 deletion source/module_hamilt_lcao/module_deepks/CMakeLists.txt
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -11,7 +11,7 @@ if(ENABLE_DEEPKS)
LCAO_deepks_vdelta.cpp
deepks_hmat.cpp
LCAO_deepks_interface.cpp
orbital_precalc.cpp
deepks_orbpre.cpp
cal_gdmx.cpp
cal_gedm.cpp
cal_gvx.cpp
Expand Down
87 changes: 26 additions & 61 deletions source/module_hamilt_lcao/module_deepks/LCAO_deepks.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -45,12 +45,7 @@ LCAO_Deepks::~LCAO_Deepks()
delete[] inl_l;

//=======1. to use deepks, pdm is required==========
// delete pdm**
for (int inl = 0; inl < this->inlmax; inl++)
{
delete[] pdm[inl];
}
delete[] pdm;
pdm.clear();
//=======2. "deepks_scf" part==========
// if (PARAM.inp.deepks_scf)
if (gedm)
Expand Down Expand Up @@ -100,12 +95,33 @@ void LCAO_Deepks::init(const LCAO_Orbitals& orb,

int pdm_size = 0;
this->inlmax = tot_inl;
this->pdm.resize(this->inlmax);

// cal n(descriptor) per atom , related to Lmax, nchi(L) and m. (not total_nchi!)
if (!PARAM.inp.deepks_equiv)
{
this->des_per_atom = 0; // mohan add 2021-04-21
for (int l = 0; l <= this->lmaxd; l++)
{
this->des_per_atom += orb.Alpha[0].getNchi(l) * (2 * l + 1);
}
this->n_descriptor = nat * this->des_per_atom;

this->init_index(ntype, nat, na, tot_inl, orb);
}

if (!PARAM.inp.deepks_equiv)
{
GlobalV::ofs_running << " total basis (all atoms) for descriptor = " << std::endl;

// init pdm**
pdm_size = (this->lmaxd * 2 + 1) * (this->lmaxd * 2 + 1);
// init pdm
for (int inl = 0; inl < this->inlmax; inl++)
{
int nm = 2 * inl_l[inl] + 1;
pdm_size += nm * nm;
this->pdm[inl] = torch::zeros({nm, nm}, torch::kFloat64);
// this->pdm[inl].requires_grad_(true);
}
}
else
{
Expand All @@ -116,26 +132,10 @@ void LCAO_Deepks::init(const LCAO_Orbitals& orb,
pdm_size = pdm_size * pdm_size;
this->des_per_atom = pdm_size;
GlobalV::ofs_running << " Equivariant version, size of pdm matrices : " << pdm_size << std::endl;
}

this->pdm = new double*[this->inlmax];
for (int inl = 0; inl < this->inlmax; inl++)
{
this->pdm[inl] = new double[pdm_size];
ModuleBase::GlobalFunc::ZEROS(this->pdm[inl], pdm_size);
}

// cal n(descriptor) per atom , related to Lmax, nchi(L) and m. (not total_nchi!)
if (!PARAM.inp.deepks_equiv)
{
this->des_per_atom = 0; // mohan add 2021-04-21
for (int l = 0; l <= this->lmaxd; l++)
for (int inl = 0; inl < this->inlmax; inl++)
{
this->des_per_atom += orb.Alpha[0].getNchi(l) * (2 * l + 1);
this->pdm[inl] = torch::zeros({pdm_size}, torch::kFloat64);
}
this->n_descriptor = nat * this->des_per_atom;

this->init_index(ntype, nat, na, tot_inl, orb);
}

this->pv = &pv_in;
Expand Down Expand Up @@ -343,41 +343,6 @@ void LCAO_Deepks::allocate_V_delta(const int nat, const int nks)
return;
}

void LCAO_Deepks::init_orbital_pdm_shell(const int nks)
{

this->orbital_pdm_shell = new double**[nks];

for (int iks = 0; iks < nks; iks++)
{
this->orbital_pdm_shell[iks] = new double*[this->inlmax];
for (int inl = 0; inl < this->inlmax; inl++)
{
this->orbital_pdm_shell[iks][inl] = new double[(2 * this->lmaxd + 1) * (2 * this->lmaxd + 1)];
ModuleBase::GlobalFunc::ZEROS(orbital_pdm_shell[iks][inl],
(2 * this->lmaxd + 1) * (2 * this->lmaxd + 1));
}

}

return;
}

void LCAO_Deepks::del_orbital_pdm_shell(const int nks)
{
for (int iks = 0; iks < nks; iks++)
{
for (int inl = 0; inl < this->inlmax; inl++)
{
delete[] this->orbital_pdm_shell[iks][inl];
}
delete[] this->orbital_pdm_shell[iks];
}
delete[] this->orbital_pdm_shell;

return;
}

void LCAO_Deepks::init_v_delta_pdm_shell(const int nks, const int nlocal)
{
const int mn_size = (2 * this->lmaxd + 1) * (2 * this->lmaxd + 1);
Expand Down
68 changes: 26 additions & 42 deletions source/module_hamilt_lcao/module_deepks/LCAO_deepks.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -6,6 +6,7 @@
#include "deepks_force.h"
#include "deepks_hmat.h"
#include "deepks_orbital.h"
#include "deepks_orbpre.h"
#include "module_base/complexmatrix.h"
#include "module_base/intarray.h"
#include "module_base/matrix.h"
Expand Down Expand Up @@ -57,12 +58,6 @@ class LCAO_Deepks
/// Correction term to Hamiltonian, for multi-k
std::vector<std::vector<std::complex<double>>> H_V_delta_k;

// k index of HOMO for multi-k bandgap label. QO added 2022-01-24
int h_ind = 0;

// k index of LUMO for multi-k bandgap label. QO added 2022-01-24
int l_ind = 0;

// functions for hr status: 1. get value; 2. set value;
int get_hr_cal()
{
Expand Down Expand Up @@ -109,8 +104,9 @@ class LCAO_Deepks
std::vector<hamilt::HContainer<double>*> phialpha;

// projected density matrix
double** pdm; //[tot_Inl][2l+1][2l+1] caoyu modified 2021-05-07; if equivariant version: [nat][nlm*nlm]
std::vector<torch::Tensor> pdm_tensor;
// [tot_Inl][2l+1][2l+1], here l is corresponding to inl;
// [nat][nlm*nlm] for equivariant version
std::vector<torch::Tensor> pdm;

// descriptors
std::vector<torch::Tensor> d_tensor;
Expand Down Expand Up @@ -138,17 +134,9 @@ class LCAO_Deepks
// gvx:d(d)/dX, [natom][3][natom][des_per_atom]
torch::Tensor gvx_tensor;

// d(d)/dD, autograd from torch::linalg::eigh
std::vector<torch::Tensor> gevdm_vector;

// dD/dX, tensor form of gdmx
std::vector<torch::Tensor> gdmr_vector;

// orbital_pdm_shell:[Inl,nm*nm]; \langle \phi_\mu|\alpha\rangle\langle\alpha|\phi_\nu\ranlge
double*** orbital_pdm_shell;
// orbital_precalc:[1,NAt,NDscrpt]; gvdm*orbital_pdm_shell
torch::Tensor orbital_precalc_tensor;

// v_delta_pdm_shell[nks,nlocal,nlocal,Inl,nm*nm] = overlap * overlap
double***** v_delta_pdm_shell;
std::complex<double>***** v_delta_pdm_shell_complex; // for multi-k
Expand Down Expand Up @@ -223,12 +211,6 @@ class LCAO_Deepks
private:
// arrange index of descriptor in all atoms
void init_index(const int ntype, const int nat, std::vector<int> na, const int tot_inl, const LCAO_Orbitals& orb);
// data structure that saves <phi|alpha>
void allocate_nlm(const int nat);

// for bandgap label calculation; QO added on 2022-1-7
void init_orbital_pdm_shell(const int nks);
void del_orbital_pdm_shell(const int nks);

// for v_delta label calculation; xinyuan added on 2023-2-22
void init_v_delta_pdm_shell(const int nks, const int nlocal);
Expand Down Expand Up @@ -373,16 +355,16 @@ class LCAO_Deepks
// descriptors wrt strain tensor, calculated by
// d(des)/d\epsilon_{ab} = d(pdm)/d\epsilon_{ab} * d(des)/d(pdm) = gdm_epsl * gvdm
// using einsum
// 6. cal_gvdm : d(des)/d(pdm)
// 6. cal_gevdm : d(des)/d(pdm)
// calculated using torch::autograd::grad
// 7. load_model : loads model for applying V_delta
// 8. cal_gedm : calculates d(E_delta)/d(pdm)
// this is the term V(D) that enters the expression H_V_delta = |alpha>V(D)<alpha|
// caculated using torch::autograd::grad
// 9. check_gedm : prints gedm for checking
// 10. cal_orbital_precalc : orbital_precalc is usted for training with orbital label,
// which equals gvdm * orbital_pdm_shell,
// orbital_pdm_shell[Inl,nm*nm] = dm_hl * overlap * overlap
// which equals gvdm * orbital_pdm,
// orbital_pdm[nks,Inl,nm,nm] = dm_hl * overlap * overlap
// 11. cal_v_delta_precalc : v_delta_precalc is used for training with v_delta label,
// which equals gvdm * v_delta_pdm_shell,
// v_delta_pdm_shell = overlap * overlap
Expand All @@ -408,11 +390,11 @@ class LCAO_Deepks
/// - b: the atoms whose force being calculated)
/// gvdm*gdmx->gvx
///----------------------------------------------------
void cal_gvx(const int nat);
void cal_gvx(const int nat, const std::vector<torch::Tensor>& gevdm);
void check_gvx(const int nat);

// for stress
void cal_gvepsl(const int nat);
void cal_gvepsl(const int nat, const std::vector<torch::Tensor>& gevdm);

// load the trained neural network model
void load_model(const std::string& model_file);
Expand All @@ -423,20 +405,22 @@ class LCAO_Deepks
void cal_gedm_equiv(const int nat);

// calculates orbital_precalc
template <typename TK, typename TH>
void cal_orbital_precalc(const std::vector<TH>& dm_hl,
const int lmaxd,
const int inlmax,
const int nat,
const int nks,
const int* inl_l,
const std::vector<ModuleBase::Vector3<double>>& kvec_d,
const std::vector<hamilt::HContainer<double>*> phialpha,
const ModuleBase::IntArray* inl_index,
const UnitCell& ucell,
const LCAO_Orbitals& orb,
const Parallel_Orbitals& pv,
const Grid_Driver& GridD);
// template <typename TK, typename TH>
// void cal_orbital_precalc(const std::vector<TH>& dm_hl,
// const int lmaxd,
// const int inlmax,
// const int nat,
// const int nks,
// const int* inl_l,
// const std::vector<ModuleBase::Vector3<double>>& kvec_d,
// const std::vector<hamilt::HContainer<double>*> phialpha,
// const std::vector<torch::Tensor> gevdm,
// const ModuleBase::IntArray* inl_index,
// const UnitCell& ucell,
// const LCAO_Orbitals& orb,
// const Parallel_Orbitals& pv,
// const Grid_Driver& GridD,
// torch::Tensor& orbital_precalc);

// calculates v_delta_precalc
template <typename TK>
Expand Down Expand Up @@ -466,11 +450,11 @@ class LCAO_Deepks

// prepare gevdm for outputting npy file
void prepare_gevdm(const int nat, const LCAO_Orbitals& orb);
void cal_gevdm(const int nat, std::vector<torch::Tensor>& gevdm);
void check_vdp_gevdm(const int nat);

private:
const Parallel_Orbitals* pv;
void cal_gvdm(const int nat);

#ifdef __MPI

Expand Down
26 changes: 22 additions & 4 deletions source/module_hamilt_lcao/module_deepks/LCAO_deepks_interface.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -113,14 +113,32 @@ void LCAO_Deepks_Interface<TK, TR>::out_deepks_labels(const double& etot,

std::vector<double> o_delta(nks, 0.0);

ld->cal_orbital_precalc<TK, TH>(dm_bandgap, ld->lmaxd, ld->inlmax, nat, nks, ld->inl_l, kvec_d, ld->phialpha, ld->inl_index, ucell, orb, *ParaV, GridD);
// calculate and save orbital_precalc: [nks,NAt,NDscrpt]
torch::Tensor orbital_precalc;
std::vector<torch::Tensor> gevdm;
ld->cal_gevdm(nat, gevdm);
DeePKS_domain::cal_orbital_precalc<TK, TH>(dm_bandgap,
ld->lmaxd,
ld->inlmax,
nat,
nks,
ld->inl_l,
kvec_d,
ld->phialpha,
gevdm,
ld->inl_index,
ucell,
orb,
*ParaV,
GridD,
orbital_precalc);
DeePKS_domain::cal_o_delta<TK, TH>(dm_bandgap, *h_delta, o_delta, *ParaV, nks);

// save obase and orbital_precalc
LCAO_deepks_io::save_npy_orbital_precalc(nat,
nks,
ld->des_per_atom,
ld->orbital_precalc_tensor,
orbital_precalc,
PARAM.globalv.global_out_dir,
my_rank);
const std::string file_obase = PARAM.globalv.global_out_dir + "deepks_obase.npy";
Expand All @@ -135,8 +153,8 @@ void LCAO_Deepks_Interface<TK, TR>::out_deepks_labels(const double& etot,
{
const std::string file_obase = PARAM.globalv.global_out_dir + "deepks_obase.npy";
LCAO_deepks_io::save_npy_o(o_tot, file_obase, nks, my_rank); // no scf, o_tot=o_base
} // end deepks_scf == 0
} // end bandgap label
} // end deepks_scf == 0
} // end bandgap label

// save H(R) matrix
if (true) // should be modified later!
Expand Down
Loading
Loading