From 42066f8b6d67d20240f0d6d146b7b3013a4a44a0 Mon Sep 17 00:00:00 2001 From: ErjieWu Date: Tue, 31 Dec 2024 17:47:06 +0800 Subject: [PATCH 1/2] Remove cal_orbital_precalc() and corresponding temporary variables from LCAO_Deepks; Partially change cal_gevdm() for future refactor; Remove the double pointer pdm and use tensor vector for replacement. --- source/Makefile.Objects | 2 +- .../hamilt_lcaodft/FORCE_STRESS.cpp | 8 +- .../module_deepks/CMakeLists.txt | 2 +- .../module_deepks/LCAO_deepks.cpp | 87 +++++---------- .../module_deepks/LCAO_deepks.h | 68 +++++------- .../module_deepks/LCAO_deepks_interface.cpp | 8 +- .../module_deepks/LCAO_deepks_io.cpp | 4 +- .../module_deepks/LCAO_deepks_io.h | 2 +- .../module_deepks/LCAO_deepks_pdm.cpp | 105 +++++++++++------- .../module_deepks/LCAO_deepks_torch.cpp | 32 +++--- .../module_deepks/cal_descriptor.cpp | 37 ++---- .../module_deepks/cal_gedm.cpp | 3 +- .../module_deepks/cal_gvx.cpp | 9 +- ...{orbital_precalc.cpp => deepks_orbpre.cpp} | 84 +++++++------- .../module_deepks/deepks_orbpre.h | 46 ++++++++ .../module_deepks/test/LCAO_deepks_test.cpp | 6 +- .../module_deepks/v_delta_precalc.cpp | 9 +- 17 files changed, 260 insertions(+), 252 deletions(-) rename source/module_hamilt_lcao/module_deepks/{orbital_precalc.cpp => deepks_orbpre.cpp} (82%) create mode 100644 source/module_hamilt_lcao/module_deepks/deepks_orbpre.h diff --git a/source/Makefile.Objects b/source/Makefile.Objects index b366ae524f..b04a96624b 100644 --- a/source/Makefile.Objects +++ b/source/Makefile.Objects @@ -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\ diff --git a/source/module_hamilt_lcao/hamilt_lcaodft/FORCE_STRESS.cpp b/source/module_hamilt_lcao/hamilt_lcaodft/FORCE_STRESS.cpp index f19d980d85..42954019f7 100644 --- a/source/module_hamilt_lcao/hamilt_lcaodft/FORCE_STRESS.cpp +++ b/source/module_hamilt_lcao/hamilt_lcaodft/FORCE_STRESS.cpp @@ -540,7 +540,9 @@ void Force_Stress_LCAO::getForceStress(UnitCell& ucell, { GlobalC::ld.check_gdmx(ucell.nat); } - GlobalC::ld.cal_gvx(ucell.nat); + std::vector gevdm; + GlobalC::ld.cal_gevdm(ucell.nat, gevdm); + GlobalC::ld.cal_gvx(ucell.nat, gevdm); if (PARAM.inp.deepks_out_unittest) { @@ -758,7 +760,9 @@ void Force_Stress_LCAO::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 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, diff --git a/source/module_hamilt_lcao/module_deepks/CMakeLists.txt b/source/module_hamilt_lcao/module_deepks/CMakeLists.txt index 1a2d618cbf..180b52a955 100644 --- a/source/module_hamilt_lcao/module_deepks/CMakeLists.txt +++ b/source/module_hamilt_lcao/module_deepks/CMakeLists.txt @@ -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 diff --git a/source/module_hamilt_lcao/module_deepks/LCAO_deepks.cpp b/source/module_hamilt_lcao/module_deepks/LCAO_deepks.cpp index ad64939097..bffeb97efd 100644 --- a/source/module_hamilt_lcao/module_deepks/LCAO_deepks.cpp +++ b/source/module_hamilt_lcao/module_deepks/LCAO_deepks.cpp @@ -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) @@ -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 { @@ -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; @@ -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); diff --git a/source/module_hamilt_lcao/module_deepks/LCAO_deepks.h b/source/module_hamilt_lcao/module_deepks/LCAO_deepks.h index 8737515004..9238ed73fd 100644 --- a/source/module_hamilt_lcao/module_deepks/LCAO_deepks.h +++ b/source/module_hamilt_lcao/module_deepks/LCAO_deepks.h @@ -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" @@ -57,12 +58,6 @@ class LCAO_Deepks /// Correction term to Hamiltonian, for multi-k std::vector>> 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() { @@ -109,8 +104,9 @@ class LCAO_Deepks std::vector*> 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 pdm_tensor; + // [tot_Inl][2l+1][2l+1], here l is corresponding to inl; + // [nat][nlm*nlm] for equivariant version + std::vector pdm; // descriptors std::vector d_tensor; @@ -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 gevdm_vector; - // dD/dX, tensor form of gdmx std::vector 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***** v_delta_pdm_shell_complex; // for multi-k @@ -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 na, const int tot_inl, const LCAO_Orbitals& orb); - // data structure that saves - 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); @@ -373,7 +355,7 @@ 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) @@ -381,8 +363,8 @@ class LCAO_Deepks // 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 @@ -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& gevdm); void check_gvx(const int nat); // for stress - void cal_gvepsl(const int nat); + void cal_gvepsl(const int nat, const std::vector& gevdm); // load the trained neural network model void load_model(const std::string& model_file); @@ -423,20 +405,22 @@ class LCAO_Deepks void cal_gedm_equiv(const int nat); // calculates orbital_precalc - template - void cal_orbital_precalc(const std::vector& dm_hl, - const int lmaxd, - const int inlmax, - const int nat, - const int nks, - const int* inl_l, - const std::vector>& kvec_d, - const std::vector*> phialpha, - const ModuleBase::IntArray* inl_index, - const UnitCell& ucell, - const LCAO_Orbitals& orb, - const Parallel_Orbitals& pv, - const Grid_Driver& GridD); + // template + // void cal_orbital_precalc(const std::vector& dm_hl, + // const int lmaxd, + // const int inlmax, + // const int nat, + // const int nks, + // const int* inl_l, + // const std::vector>& kvec_d, + // const std::vector*> phialpha, + // const std::vector 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 @@ -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& gevdm); void check_vdp_gevdm(const int nat); private: const Parallel_Orbitals* pv; - void cal_gvdm(const int nat); #ifdef __MPI diff --git a/source/module_hamilt_lcao/module_deepks/LCAO_deepks_interface.cpp b/source/module_hamilt_lcao/module_deepks/LCAO_deepks_interface.cpp index 37a5a2cc1b..4f8e5010d3 100644 --- a/source/module_hamilt_lcao/module_deepks/LCAO_deepks_interface.cpp +++ b/source/module_hamilt_lcao/module_deepks/LCAO_deepks_interface.cpp @@ -113,14 +113,18 @@ void LCAO_Deepks_Interface::out_deepks_labels(const double& etot, std::vector o_delta(nks, 0.0); - ld->cal_orbital_precalc(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 gevdm; + ld->cal_gevdm(nat, gevdm); + DeePKS_domain::cal_orbital_precalc(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(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"; diff --git a/source/module_hamilt_lcao/module_deepks/LCAO_deepks_io.cpp b/source/module_hamilt_lcao/module_deepks/LCAO_deepks_io.cpp index 5e2eff5860..8eb065b1bf 100644 --- a/source/module_hamilt_lcao/module_deepks/LCAO_deepks_io.cpp +++ b/source/module_hamilt_lcao/module_deepks/LCAO_deepks_io.cpp @@ -346,7 +346,7 @@ void LCAO_deepks_io::save_npy_o(const std::vector& bandgap, void LCAO_deepks_io::save_npy_orbital_precalc(const int nat, const int nks, const int des_per_atom, - const torch::Tensor& orbital_precalc_tensor, + const torch::Tensor& orbital_precalc, const std::string& out_dir, const int rank) { @@ -369,7 +369,7 @@ void LCAO_deepks_io::save_npy_orbital_precalc(const int nat, { for(int p=0; p& bandgap, /**<[in] \f$E_{base}\f$ or \ void save_npy_orbital_precalc(const int nat, const int nks, const int des_per_atom, - const torch::Tensor& orbital_precalc_tensor, + const torch::Tensor& orbital_precalc, const std::string& out_dir, const int rank); diff --git a/source/module_hamilt_lcao/module_deepks/LCAO_deepks_pdm.cpp b/source/module_hamilt_lcao/module_deepks/LCAO_deepks_pdm.cpp index c1c30989d1..1d961d84e6 100644 --- a/source/module_hamilt_lcao/module_deepks/LCAO_deepks_pdm.cpp +++ b/source/module_hamilt_lcao/module_deepks/LCAO_deepks_pdm.cpp @@ -21,42 +21,57 @@ #include "module_base/timer.h" #include "module_base/vector3.h" #include "module_hamilt_lcao/module_hcontainer/atom_pair.h" +#ifdef __MPI +#include "module_base/parallel_reduce.h" +#endif void LCAO_Deepks::read_projected_DM(bool read_pdm_file, bool is_equiv, const Numerical_Orbital& alpha) { if (read_pdm_file && !this->init_pdm) // for DeePKS NSCF calculation { - int pdm_size = 0; + const std::string file_projdm = PARAM.globalv.global_out_dir + "deepks_projdm.dat"; + std::ifstream ifs(file_projdm.c_str()); + + if (!ifs) + { + ModuleBase::WARNING_QUIT("LCAO_Deepks::read_projected_DM", "Cannot find the file deepks_projdm.dat"); + } if (!is_equiv) { - pdm_size = (this->lmaxd * 2 + 1) * (this->lmaxd * 2 + 1); + for (int inl = 0; inl < this->inlmax; inl++) + { + int nm = this->inl_l[inl] * 2 + 1; + for (int m1 = 0; m1 < nm; m1++) + { + for (int m2 = 0; m2 < nm; m2++) + { + double c; + ifs >> c; + this->pdm[inl][m1][m2] = c; + } + } + } } else { + int pdm_size = 0; int nproj = 0; for (int il = 0; il < this->lmaxd + 1; il++) { nproj += (2 * il + 1) * alpha.getNchi(il); } pdm_size = nproj * nproj; - } - - const std::string file_projdm = PARAM.globalv.global_out_dir + "deepks_projdm.dat"; - std::ifstream ifs(file_projdm.c_str()); - - if (!ifs) - { - ModuleBase::WARNING_QUIT("LCAO_Deepks::read_projected_DM", "Cannot find the file deepks_projdm.dat"); - } - for (int inl = 0; inl < this->inlmax; inl++) - { - for (int ind = 0; ind < pdm_size; ind++) + for (int inl = 0; inl < this->inlmax; inl++) { - double c; - ifs >> c; - pdm[inl][ind] = c; + for (int ind = 0; ind < pdm_size; ind++) + { + double c; + ifs >> c; + this->pdm[inl][ind] = c; + } } } + this->init_pdm = true; } } @@ -78,32 +93,32 @@ void LCAO_Deepks::cal_projected_DM(const elecstate::DensityMatrix* d this->init_pdm = false; return; } - int pdm_size = 0; + if (!PARAM.inp.deepks_equiv) { - pdm_size = (this->lmaxd * 2 + 1) * (this->lmaxd * 2 + 1); + for (int inl = 0; inl < this->inlmax; inl++) + { + int nm = this->inl_l[inl] * 2 + 1; + this->pdm[inl] = torch::zeros({nm, nm}, torch::kFloat64); + } } else { + int pdm_size = 0; int nproj = 0; for (int il = 0; il < this->lmaxd + 1; il++) { nproj += (2 * il + 1) * orb.Alpha[0].getNchi(il); } pdm_size = nproj * nproj; + for (int inl = 0; inl < inlmax; inl++) + { + this->pdm[inl] = torch::zeros({pdm_size}, torch::kFloat64); + } } - // if(dm.size() == 0 || dm[0].size() == 0) - //{ - // return; - // } ModuleBase::timer::tick("LCAO_Deepks", "cal_projected_DM"); - for (int inl = 0; inl < inlmax; inl++) - { - ModuleBase::GlobalFunc::ZEROS(pdm[inl], pdm_size); - } - const double Rcut_Alpha = orb.Alpha[0].getRcut(); for (int T0 = 0; T0 < ucell.ntype; T0++) { @@ -299,7 +314,6 @@ void LCAO_Deepks::cal_projected_DM(const elecstate::DensityMatrix* d g_1dmt.data(), &row_size); } // ad2 - // do dot of g_1dmt and s_1t to get orbital_pdm_shell if (!PARAM.inp.deepks_equiv) { int ib = 0, index = 0, inc = 1; @@ -315,11 +329,11 @@ void LCAO_Deepks::cal_projected_DM(const elecstate::DensityMatrix* d for (int m2 = 0; m2 < nm; ++m2) // m1 = 1 for s, 3 for p, 5 for d { int ind = m1 * nm + m2; - pdm[inl][ind] += ddot_(&row_size, - g_1dmt.data() + index * row_size, - &inc, - s_1t.data() + index * row_size, - &inc); + pdm[inl][m1][m2] += ddot_(&row_size, + g_1dmt.data() + index * row_size, + &inc, + s_1t.data() + index * row_size, + &inc); index++; } } @@ -340,10 +354,10 @@ void LCAO_Deepks::cal_projected_DM(const elecstate::DensityMatrix* d for (int jproj = 0; jproj < nproj; jproj++) { pdm[iat][iproj * nproj + jproj] += ddot_(&row_size, - g_1dmt.data() + index * row_size, - &inc, - s_1t.data() + index * row_size, - &inc); + g_1dmt.data() + index * row_size, + &inc, + s_1t.data() + index * row_size, + &inc); index++; } } @@ -353,7 +367,11 @@ void LCAO_Deepks::cal_projected_DM(const elecstate::DensityMatrix* d } // T0 #ifdef __MPI - allsum_deepks(this->inlmax, pdm_size, this->pdm); + for(int inl=0; inl(), pdm_size); + } #endif ModuleBase::timer::tick("LCAO_Deepks", "cal_projected_DM"); return; @@ -364,13 +382,16 @@ void LCAO_Deepks::check_projected_dm() const std::string file_projdm = PARAM.globalv.global_out_dir + "deepks_projdm.dat"; std::ofstream ofs(file_projdm.c_str()); - const int pdm_size = (this->lmaxd * 2 + 1) * (this->lmaxd * 2 + 1); ofs << std::setprecision(10); for (int inl = 0; inl < inlmax; inl++) { - for (int ind = 0; ind < pdm_size; ind++) + const int nm = 2 * this->inl_l[inl] + 1; + for (int m1 = 0; m1 < nm; m1++) { - ofs << pdm[inl][ind] << " "; + for (int m2 = 0; m2 < nm; m2++) + { + ofs << pdm[inl][m1][m2].item() << " "; + } } ofs << std::endl; } diff --git a/source/module_hamilt_lcao/module_deepks/LCAO_deepks_torch.cpp b/source/module_hamilt_lcao/module_deepks/LCAO_deepks_torch.cpp index 350c3db0e7..5d881f82f9 100644 --- a/source/module_hamilt_lcao/module_deepks/LCAO_deepks_torch.cpp +++ b/source/module_hamilt_lcao/module_deepks/LCAO_deepks_torch.cpp @@ -7,7 +7,7 @@ // 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 -// cal_gvdm : d(des)/d(pdm) +// cal_gevdm : d(des)/d(pdm) // calculated using torch::autograd::grad // load_model : loads model for applying V_delta // prepare_phialpha : prepare phialpha for outputting npy file @@ -25,11 +25,9 @@ #include "module_parameter/parameter.h" // calculates stress of descriptors from gradient of projected density matrices -void LCAO_Deepks::cal_gvepsl(const int nat) +void LCAO_Deepks::cal_gvepsl(const int nat, const std::vector& gevdm) { ModuleBase::TITLE("LCAO_Deepks", "cal_gvepsl"); - // preconditions - this->cal_gvdm(nat); if (!gdmepsl_vector.empty()) { gdmepsl_vector.erase(gdmepsl_vector.begin(), gdmepsl_vector.end()); @@ -76,13 +74,13 @@ void LCAO_Deepks::cal_gvepsl(const int nat) // einsum for each inl: // gdmepsl_vector : b:npol * a:inl(projector) * m:nm * n:nm - // gevdm_vector : a:inl * v:nm (descriptor) * m:nm (pdm, dim1) * n:nm + // gevdm : a:inl * v:nm (descriptor) * m:nm (pdm, dim1) * n:nm // (pdm, dim2) gvepsl_vector : b:npol * a:inl(projector) * // m:nm(descriptor) std::vector gvepsl_vector; for (int nl = 0; nl < nlmax; ++nl) { - gvepsl_vector.push_back(at::einsum("bamn, avmn->bav", {this->gdmepsl_vector[nl], this->gevdm_vector[nl]})); + gvepsl_vector.push_back(at::einsum("bamn, avmn->bav", {this->gdmepsl_vector[nl], gevdm[nl]})); } // cat nv-> \sum_nl(nv) = \sum_nl(nm_nl)=des_per_atom @@ -96,13 +94,14 @@ void LCAO_Deepks::cal_gvepsl(const int nat) return; } -// dDescriptor / dprojected density matrix -void LCAO_Deepks::cal_gvdm(const int nat) +// d(Descriptor) / d(projected density matrix) +// Dimension is different for each inl, so there's a vector of tensors +void LCAO_Deepks::cal_gevdm(const int nat, std::vector& gevdm) { - ModuleBase::TITLE("LCAO_Deepks", "cal_gvdm"); - if (!gevdm_vector.empty()) + ModuleBase::TITLE("LCAO_Deepks", "cal_gevdm"); + if (!gevdm.empty()) { - gevdm_vector.erase(gevdm_vector.begin(), gevdm_vector.end()); + gevdm.erase(gevdm.begin(), gevdm.end()); } // cal gevdm(d(EigenValue(D))/dD) int nlmax = inlmax / nat; @@ -114,7 +113,7 @@ void LCAO_Deepks::cal_gvdm(const int nat) int inl = iat * nlmax + nl; int nm = 2 * this->inl_l[inl] + 1; // repeat each block for nm times in an additional dimension - torch::Tensor tmp_x = this->pdm_tensor[inl].reshape({nm, nm}).unsqueeze(0).repeat({nm, 1, 1}); + torch::Tensor tmp_x = this->pdm[inl].reshape({nm, nm}).unsqueeze(0).repeat({nm, 1, 1}); // torch::Tensor tmp_y = std::get<0>(torch::symeig(tmp_x, true)); torch::Tensor tmp_y = std::get<0>(torch::linalg::eigh(tmp_x, "U")); torch::Tensor tmp_yshell = torch::eye(nm, torch::TensorOptions().dtype(torch::kFloat64)); @@ -134,9 +133,9 @@ void LCAO_Deepks::cal_gvdm(const int nat) avmmv.push_back(tmp_res[0]); } torch::Tensor avmm = torch::stack(avmmv, 0); // nat*nv**nm*nm - this->gevdm_vector.push_back(avmm); + gevdm.push_back(avmm); } - assert(this->gevdm_vector.size() == nlmax); + assert(gevdm.size() == nlmax); return; } @@ -367,7 +366,8 @@ void LCAO_Deepks::prepare_gevdm(const int nat, const LCAO_Orbitals& orb) int mmax = 2 * this->lmaxd + 1; this->gevdm_tensor = torch::zeros({nat, nlmax, mmax, mmax, mmax}, torch::TensorOptions().dtype(torch::kFloat64)); - this->cal_gvdm(nat); + std::vector gevdm; + this->cal_gevdm(nat, gevdm); int nl = 0; for (int L0 = 0; L0 <= orb.Alpha[0].getLmax(); ++L0) @@ -383,7 +383,7 @@ void LCAO_Deepks::prepare_gevdm(const int nat, const LCAO_Orbitals& orb) { for (int n = 0; n < nm; ++n) { - this->gevdm_tensor[iat][nl][v][m][n] = this->gevdm_vector[nl][iat][v][m][n]; + this->gevdm_tensor[iat][nl][v][m][n] = gevdm[nl][iat][v][m][n]; } } } diff --git a/source/module_hamilt_lcao/module_deepks/cal_descriptor.cpp b/source/module_hamilt_lcao/module_deepks/cal_descriptor.cpp index 6cc5b2a6a1..020892ddc0 100644 --- a/source/module_hamilt_lcao/module_deepks/cal_descriptor.cpp +++ b/source/module_hamilt_lcao/module_deepks/cal_descriptor.cpp @@ -27,7 +27,7 @@ void LCAO_Deepks::cal_descriptor_equiv(const int nat) for (int iat = 0; iat < nat; iat++) { auto tmp = torch::zeros(des_per_atom, torch::kFloat64); - std::memcpy(tmp.data_ptr(), pdm[iat], sizeof(double) * tmp.numel()); + std::memcpy(tmp.data_ptr(), pdm[iat].data_ptr(), sizeof(double) * tmp.numel()); this->d_tensor.push_back(tmp); } @@ -45,40 +45,17 @@ void LCAO_Deepks::cal_descriptor(const int nat) { return; } - // init pdm_tensor and d_tensor - torch::Tensor tmp; - - // if pdm_tensor and d_tensor is not empty, clear it !! + // init d_tensor + // if d_tensor is not empty, clear it !! if (!this->d_tensor.empty()) { this->d_tensor.erase(this->d_tensor.begin(), this->d_tensor.end()); } - if (!this->pdm_tensor.empty()) - { - this->pdm_tensor.erase(this->pdm_tensor.begin(), - this->pdm_tensor.end()); - } - for (int inl = 0; inl < this->inlmax; ++inl) { const int nm = 2 * inl_l[inl] + 1; - tmp = torch::ones({nm, nm}, - torch::TensorOptions().dtype(torch::kFloat64)); - - for (int m1 = 0; m1 < nm; ++m1) - { - for (int m2 = 0; m2 < nm; ++m2) - { - tmp.index_put_({m1, m2}, this->pdm[inl][m1 * nm + m2]); - } - } - - // torch::Tensor tmp = torch::from_blob(this->pdm[inl], { nm, nm }, - // torch::requires_grad()); - - tmp.requires_grad_(true); - this->pdm_tensor.push_back(tmp); + this->pdm[inl].requires_grad_(true); this->d_tensor.push_back(torch::ones({nm}, torch::requires_grad(true))); } @@ -87,9 +64,9 @@ void LCAO_Deepks::cal_descriptor(const int nat) { { torch::Tensor vd; std::tuple d_v(this->d_tensor[inl], vd); - // d_v = torch::symeig(pdm_tensor[inl], /*eigenvalues=*/true, + // d_v = torch::symeig(pdm[inl], /*eigenvalues=*/true, // /*upper=*/true); - d_v = torch::linalg::eigh(pdm_tensor[inl], /*uplo*/ "U"); + d_v = torch::linalg::eigh(pdm[inl], /*uplo*/ "U"); d_tensor[inl] = std::get<0>(d_v); } ModuleBase::timer::tick("LCAO_Deepks", "cal_descriptor"); @@ -149,7 +126,7 @@ void LCAO_Deepks::check_descriptor(const UnitCell& ucell, const std::string& out << " n_descriptor " << this->des_per_atom << std::endl; for (int i = 0; i < this->des_per_atom; i++) { - ofs << this->pdm[iat][i] << " "; + ofs << this->pdm[iat][i].item() << " "; if (i % 8 == 7) { ofs << std::endl; diff --git a/source/module_hamilt_lcao/module_deepks/cal_gedm.cpp b/source/module_hamilt_lcao/module_deepks/cal_gedm.cpp index 4b3ca70999..7fe9a59fe1 100644 --- a/source/module_hamilt_lcao/module_deepks/cal_gedm.cpp +++ b/source/module_hamilt_lcao/module_deepks/cal_gedm.cpp @@ -108,7 +108,6 @@ void LCAO_Deepks::cal_gedm(const int nat) { return; } - // using this->pdm_tensor ModuleBase::TITLE("LCAO_Deepks", "cal_gedm"); // forward @@ -125,7 +124,7 @@ void LCAO_Deepks::cal_gedm(const int nat) { std::vector gedm_shell; gedm_shell.push_back(torch::ones_like(ec[0])); this->gedm_tensor = torch::autograd::grad(ec, - this->pdm_tensor, + this->pdm, gedm_shell, /*retain_grad=*/true, /*create_graph=*/false, diff --git a/source/module_hamilt_lcao/module_deepks/cal_gvx.cpp b/source/module_hamilt_lcao/module_deepks/cal_gvx.cpp index 6744285505..ac1ba9705a 100644 --- a/source/module_hamilt_lcao/module_deepks/cal_gvx.cpp +++ b/source/module_hamilt_lcao/module_deepks/cal_gvx.cpp @@ -16,10 +16,9 @@ // calculates gradient of descriptors from gradient of projected density // matrices -void LCAO_Deepks::cal_gvx(const int nat) { +void LCAO_Deepks::cal_gvx(const int nat, const std::vector& gevdm) +{ ModuleBase::TITLE("LCAO_Deepks", "cal_gvx"); - // preconditions - this->cal_gvdm(nat); if (!gdmr_vector.empty()) { @@ -85,14 +84,14 @@ void LCAO_Deepks::cal_gvx(const int nat) { // einsum for each inl: // gdmr_vector : b:nat(derivative) * x:3 * a:inl(projector) * m:nm * - // n:nm gevdm_vector : a:inl * v:nm (descriptor) * m:nm (pdm, dim1) * + // n:nm gevdm : a:inl * v:nm (descriptor) * m:nm (pdm, dim1) * // n:nm (pdm, dim2) gvx_vector : b:nat(derivative) * x:3 * // a:inl(projector) * m:nm(descriptor) std::vector gvx_vector; for (int nl = 0; nl < nlmax; ++nl) { gvx_vector.push_back( at::einsum("bxamn, avmn->bxav", - {this->gdmr_vector[nl], this->gevdm_vector[nl]})); + {this->gdmr_vector[nl], gevdm[nl]})); } // cat nv-> \sum_nl(nv) = \sum_nl(nm_nl)=des_per_atom diff --git a/source/module_hamilt_lcao/module_deepks/orbital_precalc.cpp b/source/module_hamilt_lcao/module_deepks/deepks_orbpre.cpp similarity index 82% rename from source/module_hamilt_lcao/module_deepks/orbital_precalc.cpp rename to source/module_hamilt_lcao/module_deepks/deepks_orbpre.cpp index 8154746d1b..0b7805cb69 100644 --- a/source/module_hamilt_lcao/module_deepks/orbital_precalc.cpp +++ b/source/module_hamilt_lcao/module_deepks/deepks_orbpre.cpp @@ -1,10 +1,10 @@ #ifdef __DEEPKS -/// 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 +/// cal_orbital_precalc : orbital_precalc is used for training with orbital label, +/// which equals gvdm * orbital_pdm, +/// orbital_pdm[nks,Inl,nm,nm] = dm_hl * overlap * overlap -#include "LCAO_deepks.h" +#include "deepks_orbpre.h" #include "LCAO_deepks_io.h" // mohan add 2024-07-22 #include "module_base/blas_connector.h" #include "module_base/constants.h" @@ -13,31 +13,31 @@ #include "module_hamilt_lcao/module_hcontainer/atom_pair.h" #include "module_parameter/parameter.h" -// calculates orbital_precalc[1,NAt,NDscrpt] = gvdm * orbital_pdm_shell; -// orbital_pdm_shell[Inl,nm*nm] = dm_hl * overlap * overlap; +// calculates orbital_precalc[nks,NAt,NDscrpt] = gvdm * orbital_pdm; +// orbital_pdm[nks,Inl,nm,nm] = dm_hl * overlap * overlap; template -void LCAO_Deepks::cal_orbital_precalc(const std::vector& dm_hl, - const int lmaxd, - const int inlmax, - const int nat, - const int nks, - const int* inl_l, - const std::vector>& kvec_d, - const std::vector*> phialpha, - const ModuleBase::IntArray* inl_index, - const UnitCell& ucell, - const LCAO_Orbitals& orb, - const Parallel_Orbitals& pv, - const Grid_Driver& GridD) +void DeePKS_domain::cal_orbital_precalc(const std::vector& dm_hl, + const int lmaxd, + const int inlmax, + const int nat, + const int nks, + const int* inl_l, + const std::vector>& kvec_d, + const std::vector*> phialpha, + const std::vector 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) { - ModuleBase::TITLE("LCAO_Deepks", "cal_orbital_precalc"); - ModuleBase::timer::tick("LCAO_Deepks", "calc_orbital_precalc"); - - this->cal_gvdm(nat); + ModuleBase::TITLE("DeePKS_domain", "cal_orbital_precalc"); + ModuleBase::timer::tick("DeePKS_domain", "calc_orbital_precalc"); const double Rcut_Alpha = orb.Alpha[0].getRcut(); - this->init_orbital_pdm_shell(nks); + torch::Tensor orbital_pdm = torch::zeros({nks, inlmax, (2 * lmaxd + 1) , (2 * lmaxd + 1)}, torch::dtype(torch::kFloat64)); for (int T0 = 0; T0 < ucell.ntype; T0++) { @@ -256,7 +256,8 @@ void LCAO_Deepks::cal_orbital_precalc(const std::vector& dm_hl, for (int ik = 0; ik < nks; ik++) { - // do dot of g_1dmt and s_1t to get orbital_pdm_shell + // do dot of g_1dmt and s_1t to get orbital_pdm + const double* p_g1dmt = g_1dmt.data() + ik * row_size; int ib = 0, index = 0, inc = 1; @@ -272,7 +273,7 @@ void LCAO_Deepks::cal_orbital_precalc(const std::vector& dm_hl, { for (int m2 = 0; m2 < nm; ++m2) // m1 = 1 for s, 3 for p, 5 for d { - orbital_pdm_shell[ik][inl][m1 * nm + m2] + orbital_pdm[ik][inl][m1][m2] += ddot_(&row_size, p_g1dmt + index * row_size * nks, &inc, @@ -293,17 +294,17 @@ void LCAO_Deepks::cal_orbital_precalc(const std::vector& dm_hl, { for (int inl = 0; inl < inlmax; inl++) { - Parallel_Reduce::reduce_all(this->orbital_pdm_shell[iks][inl], + auto tensor_slice = orbital_pdm[iks][inl]; + Parallel_Reduce::reduce_all(tensor_slice.data_ptr(), (2 * lmaxd + 1) * (2 * lmaxd + 1)); } } #endif - // transfer orbital_pdm_shell to orbital_pdm_shell_vector - + // transfer orbital_pdm [nks,inl,nm,nm] to orbital_pdm_vector [nl,[nks,nat,nm,nm]] int nlmax = inlmax / nat; - std::vector orbital_pdm_shell_vector; + std::vector orbital_pdm_vector; for (int nl = 0; nl < nlmax; ++nl) { std::vector kammv; @@ -320,7 +321,7 @@ void LCAO_Deepks::cal_orbital_precalc(const std::vector& dm_hl, { for (int m2 = 0; m2 < nm; ++m2) // m1 = 1 for s, 3 for p, 5 for d { - mmv.push_back(this->orbital_pdm_shell[iks][inl][m1 * nm + m2]); + mmv.push_back(orbital_pdm[iks][inl][m1][m2].item()); } } torch::Tensor mm @@ -332,26 +333,25 @@ void LCAO_Deepks::cal_orbital_precalc(const std::vector& dm_hl, kammv.push_back(amm); } torch::Tensor kamm = torch::stack(kammv, 0); - orbital_pdm_shell_vector.push_back(kamm); + orbital_pdm_vector.push_back(kamm); } - assert(orbital_pdm_shell_vector.size() == nlmax); + assert(orbital_pdm_vector.size() == nlmax); // einsum for each nl: std::vector orbital_precalc_vector; for (int nl = 0; nl < nlmax; ++nl) { orbital_precalc_vector.push_back( - at::einsum("kamn, avmn->kav", {orbital_pdm_shell_vector[nl], this->gevdm_vector[nl]})); + at::einsum("kamn, avmn->kav", {orbital_pdm_vector[nl], gevdm[nl]})); } - this->orbital_precalc_tensor = torch::cat(orbital_precalc_vector, -1); - this->del_orbital_pdm_shell(nks); + orbital_precalc = torch::cat(orbital_precalc_vector, -1); ModuleBase::timer::tick("LCAO_Deepks", "calc_orbital_precalc"); return; } -template void LCAO_Deepks::cal_orbital_precalc( +template void DeePKS_domain::cal_orbital_precalc( const std::vector& dm_hl, const int lmaxd, const int inlmax, @@ -360,13 +360,15 @@ template void LCAO_Deepks::cal_orbital_precalc( const int* inl_l, const std::vector>& kvec_d, const std::vector*> phialpha, + const std::vector gevdm, const ModuleBase::IntArray* inl_index, const UnitCell& ucell, const LCAO_Orbitals& orb, const Parallel_Orbitals& pv, - const Grid_Driver& GridD); + const Grid_Driver& GridD, + torch::Tensor& orbital_precalc); -template void LCAO_Deepks::cal_orbital_precalc, ModuleBase::ComplexMatrix>( +template void DeePKS_domain::cal_orbital_precalc, ModuleBase::ComplexMatrix>( const std::vector& dm_hl, const int lmaxd, const int inlmax, @@ -375,9 +377,11 @@ template void LCAO_Deepks::cal_orbital_precalc, ModuleBase: const int* inl_l, const std::vector>& kvec_d, const std::vector*> phialpha, + const std::vector gevdm, const ModuleBase::IntArray* inl_index, const UnitCell& ucell, const LCAO_Orbitals& orb, const Parallel_Orbitals& pv, - const Grid_Driver& GridD); + const Grid_Driver& GridD, + torch::Tensor& orbital_precalc); #endif diff --git a/source/module_hamilt_lcao/module_deepks/deepks_orbpre.h b/source/module_hamilt_lcao/module_deepks/deepks_orbpre.h new file mode 100644 index 0000000000..a320ffba71 --- /dev/null +++ b/source/module_hamilt_lcao/module_deepks/deepks_orbpre.h @@ -0,0 +1,46 @@ +#ifndef DEEPKS_ORBPRE_H +#define DEEPKS_ORBPRE_H + +#ifdef __DEEPKS + +#include "module_base/complexmatrix.h" +#include "module_base/intarray.h" +#include "module_base/matrix.h" +#include "module_base/timer.h" +#include "module_basis/module_ao/parallel_orbitals.h" +#include "module_basis/module_nao/two_center_integrator.h" +#include "module_cell/module_neighbor/sltk_grid_driver.h" +#include "module_elecstate/module_dm/density_matrix.h" +#include "module_hamilt_lcao/module_hcontainer/hcontainer.h" + +#include +#include + +namespace DeePKS_domain +{ + //------------------------ + // deepks_orbpre.cpp + //------------------------ + + // This file contains one subroutine for calculating orbital_precalc, + // which is defind as gvdm * dm_hl * overlap * overlap + +template +void cal_orbital_precalc(const std::vector& dm_hl, + const int lmaxd, + const int inlmax, + const int nat, + const int nks, + const int* inl_l, + const std::vector>& kvec_d, + const std::vector*> phialpha, + const std::vector 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); +} +#endif +#endif diff --git a/source/module_hamilt_lcao/module_deepks/test/LCAO_deepks_test.cpp b/source/module_hamilt_lcao/module_deepks/test/LCAO_deepks_test.cpp index 1414b4bbea..33a9454747 100644 --- a/source/module_hamilt_lcao/module_deepks/test/LCAO_deepks_test.cpp +++ b/source/module_hamilt_lcao/module_deepks/test/LCAO_deepks_test.cpp @@ -1,6 +1,8 @@ #include "LCAO_deepks_test.h" #define private public #include "module_parameter/parameter.h" +#include +#include #undef private namespace Test_Deepks { @@ -170,7 +172,9 @@ void test_deepks::check_descriptor() void test_deepks::check_gvx() { - this->ld.cal_gvx(ucell.nat); + std::vector gevdm; + this->ld.cal_gevdm(ucell.nat, gevdm); + this->ld.cal_gvx(ucell.nat, gevdm); this->ld.check_gvx(ucell.nat); for (int ia = 0; ia < ucell.nat; ia++) diff --git a/source/module_hamilt_lcao/module_deepks/v_delta_precalc.cpp b/source/module_hamilt_lcao/module_deepks/v_delta_precalc.cpp index 489692ed1a..c3cea5eb76 100644 --- a/source/module_hamilt_lcao/module_deepks/v_delta_precalc.cpp +++ b/source/module_hamilt_lcao/module_deepks/v_delta_precalc.cpp @@ -30,7 +30,8 @@ void LCAO_Deepks::cal_v_delta_precalc(const int nlocal, // timeval t_start; // gettimeofday(&t_start,NULL); - this->cal_gvdm(nat); + std::vector gevdm; + this->cal_gevdm(nat, gevdm); const double Rcut_Alpha = orb.Alpha[0].getRcut(); this->init_v_delta_pdm_shell(nks, nlocal); @@ -266,13 +267,13 @@ void LCAO_Deepks::cal_v_delta_precalc(const int nlocal, if constexpr (std::is_same::value) { v_delta_precalc_vector.push_back( - at::einsum("kxyamn, avmn->kxyav", {v_delta_pdm_shell_vector[nl], this->gevdm_vector[nl]})); + at::einsum("kxyamn, avmn->kxyav", {v_delta_pdm_shell_vector[nl], gevdm[nl]})); } else { - torch::Tensor gevdm_vector_complex = this->gevdm_vector[nl].to(torch::kComplexDouble); + torch::Tensor gevdm_complex = gevdm[nl].to(torch::kComplexDouble); v_delta_precalc_vector.push_back( - at::einsum("kxyamn, avmn->kxyav", {v_delta_pdm_shell_vector[nl], gevdm_vector_complex})); + at::einsum("kxyamn, avmn->kxyav", {v_delta_pdm_shell_vector[nl], gevdm_complex})); } } From f85c16d5192fdb2f0740e01c22d3192fca90660b Mon Sep 17 00:00:00 2001 From: ErjieWu Date: Tue, 31 Dec 2024 17:55:53 +0800 Subject: [PATCH 2/2] clang-format adjustment. --- .../hamilt_lcaodft/FORCE_STRESS.cpp | 2 +- .../module_deepks/LCAO_deepks_interface.cpp | 20 +- .../module_deepks/LCAO_deepks_io.cpp | 489 +++++++++--------- .../module_deepks/LCAO_deepks_io.h | 185 ++++--- .../module_deepks/LCAO_deepks_pdm.cpp | 20 +- .../module_deepks/cal_descriptor.cpp | 118 ++--- .../module_deepks/cal_gedm.cpp | 104 ++-- .../module_deepks/cal_gvx.cpp | 84 ++- .../module_deepks/deepks_orbpre.cpp | 23 +- .../module_deepks/deepks_orbpre.h | 16 +- .../module_deepks/test/LCAO_deepks_test.cpp | 3 +- 11 files changed, 521 insertions(+), 543 deletions(-) diff --git a/source/module_hamilt_lcao/hamilt_lcaodft/FORCE_STRESS.cpp b/source/module_hamilt_lcao/hamilt_lcaodft/FORCE_STRESS.cpp index 42954019f7..bc98636d3a 100644 --- a/source/module_hamilt_lcao/hamilt_lcaodft/FORCE_STRESS.cpp +++ b/source/module_hamilt_lcao/hamilt_lcaodft/FORCE_STRESS.cpp @@ -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 diff --git a/source/module_hamilt_lcao/module_deepks/LCAO_deepks_interface.cpp b/source/module_hamilt_lcao/module_deepks/LCAO_deepks_interface.cpp index 4f8e5010d3..7207a64a27 100644 --- a/source/module_hamilt_lcao/module_deepks/LCAO_deepks_interface.cpp +++ b/source/module_hamilt_lcao/module_deepks/LCAO_deepks_interface.cpp @@ -117,7 +117,21 @@ void LCAO_Deepks_Interface::out_deepks_labels(const double& etot, torch::Tensor orbital_precalc; std::vector gevdm; ld->cal_gevdm(nat, gevdm); - DeePKS_domain::cal_orbital_precalc(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_orbital_precalc(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(dm_bandgap, *h_delta, o_delta, *ParaV, nks); // save obase and orbital_precalc @@ -139,8 +153,8 @@ void LCAO_Deepks_Interface::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! diff --git a/source/module_hamilt_lcao/module_deepks/LCAO_deepks_io.cpp b/source/module_hamilt_lcao/module_deepks/LCAO_deepks_io.cpp index 8eb065b1bf..fdf3ba7e38 100644 --- a/source/module_hamilt_lcao/module_deepks/LCAO_deepks_io.cpp +++ b/source/module_hamilt_lcao/module_deepks/LCAO_deepks_io.cpp @@ -1,50 +1,48 @@ -//wenfei 2022-1-11 -//This file contains subroutines that contains interface with libnpy +// wenfei 2022-1-11 +// This file contains subroutines that contains interface with libnpy #include "module_parameter/parameter.h" -//since many arrays must be saved in numpy format -//It also contains subroutines for printing density matrices -//which is used in unit tests - -//There are 2 subroutines for printing density matrices: -//1. print_dm : for gamma only -//2. print_dm_k : for multi-k - -//There are 4 subroutines in this file that prints to npy file: -//1. save_npy_d : descriptor ->dm_eig.npy -//2. save_npy_gvx : gvx ->grad_vx.npy -//3. save_npy_e : energy -//4. save_npy_f : force -//5. save_npy_s : stress -//6. save_npy_o : bandgap -//7. save_npy_orbital_precalc : orbital_precalc -//8. save_npy_h : Hamiltonian -//9. save_npy_v_delta_precalc : v_delta_precalc -//10. save_npy_phialpha : phialpha -//11. save_npy_gevdm : grav_evdm , can use phialpha and gevdm to calculate v_delta_precalc +// since many arrays must be saved in numpy format +// It also contains subroutines for printing density matrices +// which is used in unit tests + +// There are 2 subroutines for printing density matrices: +// 1. print_dm : for gamma only +// 2. print_dm_k : for multi-k + +// There are 4 subroutines in this file that prints to npy file: +// 1. save_npy_d : descriptor ->dm_eig.npy +// 2. save_npy_gvx : gvx ->grad_vx.npy +// 3. save_npy_e : energy +// 4. save_npy_f : force +// 5. save_npy_s : stress +// 6. save_npy_o : bandgap +// 7. save_npy_orbital_precalc : orbital_precalc +// 8. save_npy_h : Hamiltonian +// 9. save_npy_v_delta_precalc : v_delta_precalc +// 10. save_npy_phialpha : phialpha +// 11. save_npy_gevdm : grav_evdm , can use phialpha and gevdm to calculate v_delta_precalc #ifdef __DEEPKS -#include #include "LCAO_deepks_io.h" #include "npy.hpp" +#include + template -void LCAO_deepks_io::print_dm(const int nks, - const int nlocal, - const int nrow, - const std::vector>& dm) +void LCAO_deepks_io::print_dm(const int nks, const int nlocal, const int nrow, const std::vector>& dm) { std::stringstream ss; - for(int ik=0;ik npy_gedm; - std::vector dshape = {static_cast(nat), - static_cast(des_per_atom)}; + std::vector dshape = {static_cast(nat), static_cast(des_per_atom)}; std::string gedm_file = "gedm.npy"; @@ -75,57 +71,55 @@ void LCAO_deepks_io::load_npy_gedm(const int nat, for (int iat = 0; iat < nat; iat++) { - for(int ides = 0; ides < des_per_atom; ides++) + for (int ides = 0; ides < des_per_atom; ides++) { - gedm[iat][ides] = npy_gedm[iat*des_per_atom + ides] * 2.0; //Ha to Ry + gedm[iat][ides] = npy_gedm[iat * des_per_atom + ides] * 2.0; // Ha to Ry } } - //load ec.npy + // load ec.npy std::vector npy_ec; - std::vector eshape = { 1ul }; + std::vector eshape = {1ul}; std::string ec_file = "ec.npy"; npy::LoadArrayFromNumpy(ec_file, eshape, npy_ec); - e_delta = npy_ec[0] * 2.0; //Ha to Ry + e_delta = npy_ec[0] * 2.0; // Ha to Ry } #ifdef __MPI - for(int iat = 0; iat < nat; iat++) + for (int iat = 0; iat < nat; iat++) { MPI_Bcast(gedm[iat], des_per_atom, MPI_DOUBLE, 0, MPI_COMM_WORLD); } MPI_Bcast(&e_delta, 1, MPI_DOUBLE, 0, MPI_COMM_WORLD); #endif - } - -//saves descriptor into dm_eig.npy -void LCAO_deepks_io::save_npy_d(const int nat, +// saves descriptor into dm_eig.npy +void LCAO_deepks_io::save_npy_d(const int nat, const int des_per_atom, const int inlmax, - const int *inl_l, + const int* inl_l, const bool deepks_equiv, - const std::vector &d_tensor, + const std::vector& d_tensor, const std::string& out_dir, const int rank) { ModuleBase::TITLE("LCAO_deepks_io", "save_npy_d"); - if(rank!=0) - { - return; - } + if (rank != 0) + { + return; + } - //save descriptor in .npy format - // deepks_equiv was PARAM.inp.deepks_equiv - if(!deepks_equiv) + // save descriptor in .npy format + // deepks_equiv was PARAM.inp.deepks_equiv + if (!deepks_equiv) { std::vector npy_des; - for (int inl = 0;inl < inlmax;++inl) + for (int inl = 0; inl < inlmax; ++inl) { - int nm = 2*inl_l[inl] + 1; - for(int im=0;im npy_des; - for(int iat = 0; iat < nat; iat ++) + for (int iat = 0; iat < nat; iat++) { - for(int i = 0; i < des_per_atom; i++) + for (int i = 0; i < des_per_atom; i++) { npy_des.push_back(d_tensor[iat].index({i}).item().toDouble()); } @@ -153,90 +147,85 @@ void LCAO_deepks_io::save_npy_d(const int nat, if (rank == 0) { std::string file_dm_eig = out_dir + "deepks_dm_eig.npy"; - //std::string file_dm_eig = "dm_eig.npy"; + // std::string file_dm_eig = "dm_eig.npy"; npy::SaveArrayAsNumpy(file_dm_eig, false, 2, dshape, npy_des); - } + } } return; } - -//saves gvx into grad_vx.npy -void LCAO_deepks_io::save_npy_gvx(const int nat, +// saves gvx into grad_vx.npy +void LCAO_deepks_io::save_npy_gvx(const int nat, const int des_per_atom, - const torch::Tensor &gvx_tensor, - const std::string &out_dir, + const torch::Tensor& gvx_tensor, + const std::string& out_dir, const int rank) { ModuleBase::TITLE("LCAO_deepks_io", "save_npy_gvx"); - if(rank!=0) - { - return; - } + if (rank != 0) + { + return; + } - assert(nat>0); + assert(nat > 0); - //save grad_vx.npy (when force label is in use) - //unit: /Bohr - const long unsigned gshape[] - = {static_cast(nat), - 3UL, - static_cast(nat), - static_cast(des_per_atom)}; + // save grad_vx.npy (when force label is in use) + // unit: /Bohr + const long unsigned gshape[] = {static_cast(nat), + 3UL, + static_cast(nat), + static_cast(des_per_atom)}; std::vector npy_gvx; - for (int ibt = 0;ibt < nat;++ibt) + for (int ibt = 0; ibt < nat; ++ibt) { - for (int i = 0;i < 3;i++) + for (int i = 0; i < 3; i++) { - for (int iat = 0;iat < nat;++iat) + for (int iat = 0; iat < nat; ++iat) { - for(int p=0;p(nat), static_cast(des_per_atom)}; - //save grad_vepsl.npy (when stress label is in use) - //unit: none - const long unsigned gshape[] = {6UL, - static_cast(nat), - static_cast(des_per_atom)}; - std::vector npy_gvepsl; - for (int i = 0;i < 6;i++) + for (int i = 0; i < 6; i++) { - for (int ibt = 0;ibt < nat;++ibt) + for (int ibt = 0; ibt < nat; ++ibt) { - for(int p=0;p npy_e; npy_e.push_back(e); npy::SaveArrayAsNumpy(e_file, false, 1, eshape, npy_e); return; } - -//saves force in numpy format -void LCAO_deepks_io::save_npy_f(const ModuleBase::matrix &f, - const std::string &f_file, - const int nat, - const int rank) +// saves force in numpy format +void LCAO_deepks_io::save_npy_f(const ModuleBase::matrix& f, const std::string& f_file, const int nat, const int rank) { ModuleBase::TITLE("LCAO_deepks_io", "save_npy_f"); - if(rank!=0) - { - return; - } + if (rank != 0) + { + return; + } - assert(nat>0); + assert(nat > 0); - //save f_base - //caution: unit: Rydberg/Bohr + // save f_base + // caution: unit: Rydberg/Bohr const long unsigned fshape[] = {static_cast(nat), 3}; std::vector npy_f; - for (int iat = 0;iat < nat;++iat) + for (int iat = 0; iat < nat; ++iat) { - for (int i = 0;i < 3;i++) + for (int i = 0; i < 3; i++) { npy_f.push_back(f(iat, i)); } @@ -297,77 +279,73 @@ void LCAO_deepks_io::save_npy_f(const ModuleBase::matrix &f, return; } - -//saves stress in numpy format -void LCAO_deepks_io::save_npy_s(const ModuleBase::matrix &stress, - const std::string &s_file, - const double &omega, +// saves stress in numpy format +void LCAO_deepks_io::save_npy_s(const ModuleBase::matrix& stress, + const std::string& s_file, + const double& omega, const int rank) { ModuleBase::TITLE("LCAO_deepks_io", "save_npy_s"); - if(rank!=0) - { - return; - } + if (rank != 0) + { + return; + } - const long unsigned sshape[] = { 6 }; + const long unsigned sshape[] = {6}; std::vector npy_s; - for (int ipol = 0;ipol < 3;++ipol) + for (int ipol = 0; ipol < 3; ++ipol) { - for (int jpol = ipol;jpol < 3;jpol++) + for (int jpol = ipol; jpol < 3; jpol++) { - npy_s.push_back(stress(ipol, jpol)*omega); + npy_s.push_back(stress(ipol, jpol) * omega); } } npy::SaveArrayAsNumpy(s_file, false, 1, sshape, npy_s); return; } - -void LCAO_deepks_io::save_npy_o(const std::vector& bandgap, - const std::string& o_file, +void LCAO_deepks_io::save_npy_o(const std::vector& bandgap, + const std::string& o_file, const int nks, const int rank) { - ModuleBase::TITLE("LCAO_deepks_io", "save_npy_o"); - if(rank!=0) - { - return; - } - - //save o_base - const long unsigned oshape[] = {static_cast(nks), 1 }; + ModuleBase::TITLE("LCAO_deepks_io", "save_npy_o"); + if (rank != 0) + { + return; + } + + // save o_base + const long unsigned oshape[] = {static_cast(nks), 1}; npy::SaveArrayAsNumpy(o_file, false, 2, oshape, bandgap); return; } - -void LCAO_deepks_io::save_npy_orbital_precalc(const int nat, - const int nks, +void LCAO_deepks_io::save_npy_orbital_precalc(const int nat, + const int nks, const int des_per_atom, const torch::Tensor& orbital_precalc, const std::string& out_dir, const int rank) { ModuleBase::TITLE("LCAO_deepks_io", "save_npy_orbital_precalc"); - if(rank!=0) - { - return; - } + if (rank != 0) + { + return; + } - //save orbital_precalc.npy (when bandgap label is in use) - //unit: a.u. - const long unsigned gshape[] = {static_cast(nks), - static_cast(nat), - static_cast(des_per_atom)}; + // save orbital_precalc.npy (when bandgap label is in use) + // unit: a.u. + const long unsigned gshape[] + = {static_cast(nks), static_cast(nat), static_cast(des_per_atom)}; std::vector npy_orbital_precalc; for (int iks = 0; iks < nks; ++iks) { - for (int iat = 0;iat < nat;++iat) + for (int iat = 0; iat < nat; ++iat) { - for(int p=0; p -void LCAO_deepks_io::save_npy_h(const std::vector &hamilt, - const std::string &h_file, +void LCAO_deepks_io::save_npy_h(const std::vector& hamilt, + const std::string& h_file, const int nlocal, const int nks, const int rank) { ModuleBase::TITLE("LCAO_deepks_io", "save_npy_h"); - if(rank!=0) - { - return; - } + if (rank != 0) + { + return; + } - const long unsigned hshape[] = {static_cast(nks), - static_cast(nlocal), - static_cast(nlocal) }; + const long unsigned hshape[] + = {static_cast(nks), static_cast(nlocal), static_cast(nlocal)}; std::vector npy_h; - for(int k=0; k -void LCAO_deepks_io::save_npy_v_delta_precalc(const int nat, +void LCAO_deepks_io::save_npy_v_delta_precalc(const int nat, const int nks, - const int nlocal, + const int nlocal, const int des_per_atom, const torch::Tensor& v_delta_precalc_tensor, const std::string& out_dir, const int rank) { ModuleBase::TITLE("LCAO_deepks_io", "save_npy_v_delta_precalc"); - if(rank!=0) - { - return; - } + if (rank != 0) + { + return; + } - // timeval t_start; + // timeval t_start; // gettimeofday(&t_start,NULL); - //save v_delta_precalc.npy (when v_delta label is in use) - //unit: a.u. + // save v_delta_precalc.npy (when v_delta label is in use) + // unit: a.u. const long unsigned gshape[] = {static_cast(nks), static_cast(nlocal), static_cast(nlocal), static_cast(nat), static_cast(des_per_atom)}; - std::vector npy_v_delta_precalc; + std::vector npy_v_delta_precalc; for (int iks = 0; iks < nks; ++iks) { for (int mu = 0; mu < nlocal; ++mu) { for (int nu = 0; nu < nlocal; ++nu) { - for (int iat = 0;iat < nat;++iat) + for (int iat = 0; iat < nat; ++iat) { - for(int p=0; p::value) { - npy_v_delta_precalc.push_back(v_delta_precalc_tensor.index({iks, mu, nu, iat, p }).item().toDouble()); + npy_v_delta_precalc.push_back( + v_delta_precalc_tensor.index({iks, mu, nu, iat, p}).item().toDouble()); } else { - std::complex value(torch::real(v_delta_precalc_tensor.index({iks, mu, nu, iat, p})).item(), - torch::imag(v_delta_precalc_tensor.index({iks, mu, nu, iat, p})).item()); + std::complex value( + torch::real(v_delta_precalc_tensor.index({iks, mu, nu, iat, p})).item(), + torch::imag(v_delta_precalc_tensor.index({iks, mu, nu, iat, p})).item()); npy_v_delta_precalc.push_back(value); } } - } + } } } } @@ -469,53 +448,54 @@ void LCAO_deepks_io::save_npy_v_delta_precalc(const int nat, } template -void LCAO_deepks_io::save_npy_phialpha(const int nat, +void LCAO_deepks_io::save_npy_phialpha(const int nat, const int nks, const int nlocal, const int inlmax, const int lmaxd, - const torch::Tensor &phialpha_tensor, + const torch::Tensor& phialpha_tensor, const std::string& out_dir, const int rank) { ModuleBase::TITLE("LCAO_deepks_io", "save_npy_phialpha"); - if(rank!=0) - { - return; - } - - //save phialpha.npy (when v_delta label == 2) - //unit: a.u. - const int nlmax = inlmax/nat; - const int mmax = 2*lmaxd+1; + if (rank != 0) + { + return; + } + + // save phialpha.npy (when v_delta label == 2) + // unit: a.u. + const int nlmax = inlmax / nat; + const int mmax = 2 * lmaxd + 1; const long unsigned gshape[] = {static_cast(nat), static_cast(nlmax), static_cast(nks), static_cast(nlocal), static_cast(mmax)}; std::vector npy_phialpha; - for(int iat=0; iat< nat ; iat++) + for (int iat = 0; iat < nat; iat++) { - for(int nl = 0; nl < nlmax; nl++) + for (int nl = 0; nl < nlmax; nl++) { - for (int iks = 0; iks < nks ; iks++) + for (int iks = 0; iks < nks; iks++) { - for(int mu = 0; mu < nlocal ; mu++) + for (int mu = 0; mu < nlocal; mu++) { - for(int m=0; m< mmax; m++) + for (int m = 0; m < mmax; m++) { if constexpr (std::is_same::value) { - npy_phialpha.push_back(phialpha_tensor.index({ iat,nl, iks, mu, m }).item().toDouble()); + npy_phialpha.push_back(phialpha_tensor.index({iat, nl, iks, mu, m}).item().toDouble()); } else { - std::complex value(torch::real(phialpha_tensor.index({ iat, nl, iks, mu, m })).item(), - torch::imag(phialpha_tensor.index({ iat, nl, iks, mu, m })).item()); + std::complex value( + torch::real(phialpha_tensor.index({iat, nl, iks, mu, m})).item(), + torch::imag(phialpha_tensor.index({iat, nl, iks, mu, m})).item()); npy_phialpha.push_back(value); } } - } + } } } } @@ -524,7 +504,6 @@ void LCAO_deepks_io::save_npy_phialpha(const int nat, return; } - void LCAO_deepks_io::save_npy_gevdm(const int nat, const int inlmax, const int lmaxd, @@ -533,37 +512,37 @@ void LCAO_deepks_io::save_npy_gevdm(const int nat, const int rank) { ModuleBase::TITLE("LCAO_deepks_io", "save_npy_gevdm"); - if(rank!=0) - { - return; - } + if (rank != 0) + { + return; + } - assert(nat>0); + assert(nat > 0); - //save grad_evdm.npy (when v_delta label == 2) - //unit: a.u. - const int nlmax = inlmax/nat; - const int mmax = 2*lmaxd+1; + // save grad_evdm.npy (when v_delta label == 2) + // unit: a.u. + const int nlmax = inlmax / nat; + const int mmax = 2 * lmaxd + 1; const long unsigned gshape[] = {static_cast(nat), static_cast(nlmax), static_cast(mmax), static_cast(mmax), static_cast(mmax)}; std::vector npy_gevdm; - for(int iat=0; iat< nat ; iat++) + for (int iat = 0; iat < nat; iat++) { - for(int nl = 0; nl < nlmax; nl++) + for (int nl = 0; nl < nlmax; nl++) { - for(int v=0; v< mmax; v++) + for (int v = 0; v < mmax; v++) { - for(int m=0; m< mmax; m++) + for (int m = 0; m < mmax; m++) { - for(int n=0; n< mmax; n++) + for (int n = 0; n < mmax; n++) { - npy_gevdm.push_back(gevdm_tensor.index({ iat, nl, v, m, n}).item().toDouble()); - } + npy_gevdm.push_back(gevdm_tensor.index({iat, nl, v, m, n}).item().toDouble()); + } } - } + } } } const std::string file_gevdm = out_dir + "deepks_gevdm.npy"; @@ -571,51 +550,51 @@ void LCAO_deepks_io::save_npy_gevdm(const int nat, return; } - -template void LCAO_deepks_io::print_dm(const int nks, +template void LCAO_deepks_io::print_dm(const int nks, const int nlocal, const int nrow, const std::vector>& dm); -template void LCAO_deepks_io::print_dm>(const int nks, +template void LCAO_deepks_io::print_dm>(const int nks, const int nlocal, const int nrow, const std::vector>>& dm); -template void LCAO_deepks_io::save_npy_h(const std::vector &hamilt, - const std::string &h_file, +template void LCAO_deepks_io::save_npy_h(const std::vector& hamilt, + const std::string& h_file, const int nlocal, const int nks, const int rank); -template void LCAO_deepks_io::save_npy_h>(const std::vector &hamilt, - const std::string &h_file, +template void LCAO_deepks_io::save_npy_h>(const std::vector& hamilt, + const std::string& h_file, const int nlocal, const int nks, const int rank); -template void LCAO_deepks_io::save_npy_v_delta_precalc(const int nat, +template void LCAO_deepks_io::save_npy_v_delta_precalc(const int nat, const int nks, - const int nlocal, + const int nlocal, const int des_per_atom, const torch::Tensor& v_delta_precalc_tensor, const std::string& out_dir, const int rank); -template void LCAO_deepks_io::save_npy_v_delta_precalc>(const int nat, - const int nks, - const int nlocal, - const int des_per_atom, - const torch::Tensor& v_delta_precalc_tensor, - const std::string& out_dir, - const int rank); +template void LCAO_deepks_io::save_npy_v_delta_precalc>( + const int nat, + const int nks, + const int nlocal, + const int des_per_atom, + const torch::Tensor& v_delta_precalc_tensor, + const std::string& out_dir, + const int rank); template void LCAO_deepks_io::save_npy_phialpha(const int nat, const int nks, const int nlocal, const int inlmax, const int lmaxd, - const torch::Tensor &phialpha_tensor, + const torch::Tensor& phialpha_tensor, const std::string& out_dir, const int rank); @@ -624,7 +603,7 @@ template void LCAO_deepks_io::save_npy_phialpha>(const int const int nlocal, const int inlmax, const int lmaxd, - const torch::Tensor &phialpha_tensor, + const torch::Tensor& phialpha_tensor, const std::string& out_dir, const int rank); diff --git a/source/module_hamilt_lcao/module_deepks/LCAO_deepks_io.h b/source/module_hamilt_lcao/module_deepks/LCAO_deepks_io.h index 7ca700853d..8c537f1e46 100644 --- a/source/module_hamilt_lcao/module_deepks/LCAO_deepks_io.h +++ b/source/module_hamilt_lcao/module_deepks/LCAO_deepks_io.h @@ -1,144 +1,137 @@ -#ifndef LCAO_DEEPKS_IO_H +#ifndef LCAO_DEEPKS_IO_H #define LCAO_DEEPKS_IO_H #ifdef __DEEPKS -#include -#include -#include "module_base/tool_title.h" -#include "module_base/matrix.h" #include "module_base/complexmatrix.h" +#include "module_base/matrix.h" +#include "module_base/tool_title.h" +#include #include #include +#include namespace LCAO_deepks_io { - /// This file contains subroutines that contains interface with libnpy - /// since many arrays must be saved in numpy format - /// It also contains subroutines for printing density matrices - /// which is used in unit tests - - /// There are 2 subroutines for printing and loading .npy file: - /// 1. print_dm : print density matrices - /// 2. load_npy_gedm : load gedm from .npy file - - /// others print quantities in .npy format - - /// 3. save_npy_d : descriptor -> deepks_dm_eig.npy - /// 4. save_npy_e : energy - /// 5. save_npy_f : force - /// 6. save_npy_gvx : gvx -> deepks_gradvx.npy - /// 7. save_npy_s : stress - /// 8. save_npy_gvepsl : gvepsl -> deepks_gvepsl.npy - /// 9. save_npy_o: orbital - /// 10. save_npy_orbital_precalc: orbital_precalc -> deepks_orbpre.npy - /// 11. save_npy_h : Hamiltonian - /// 12. save_npy_v_delta_precalc : v_delta_precalc -> deepks_vdpre.npy - /// 13. save_npy_phialpha : phialpha -> deepks_phialpha.npy - /// 14. save_npy_gevdm : grav_evdm -> deepks_gevdm.npy, can use phialpha and gevdm to calculate v_delta_precalc +/// This file contains subroutines that contains interface with libnpy +/// since many arrays must be saved in numpy format +/// It also contains subroutines for printing density matrices +/// which is used in unit tests + +/// There are 2 subroutines for printing and loading .npy file: +/// 1. print_dm : print density matrices +/// 2. load_npy_gedm : load gedm from .npy file + +/// others print quantities in .npy format + +/// 3. save_npy_d : descriptor -> deepks_dm_eig.npy +/// 4. save_npy_e : energy +/// 5. save_npy_f : force +/// 6. save_npy_gvx : gvx -> deepks_gradvx.npy +/// 7. save_npy_s : stress +/// 8. save_npy_gvepsl : gvepsl -> deepks_gvepsl.npy +/// 9. save_npy_o: orbital +/// 10. save_npy_orbital_precalc: orbital_precalc -> deepks_orbpre.npy +/// 11. save_npy_h : Hamiltonian +/// 12. save_npy_v_delta_precalc : v_delta_precalc -> deepks_vdpre.npy +/// 13. save_npy_phialpha : phialpha -> deepks_phialpha.npy +/// 14. save_npy_gevdm : grav_evdm -> deepks_gevdm.npy, can use phialpha and gevdm to calculate v_delta_precalc /// print density matrices template -void print_dm(const int nks, - const int nlocal, - const int nrow, - const std::vector>& dm); +void print_dm(const int nks, const int nlocal, const int nrow, const std::vector>& dm); -void load_npy_gedm(const int nat, - const int des_per_atom, - double** gedm, - double& e_delta, - const int rank); +void load_npy_gedm(const int nat, const int des_per_atom, double** gedm, double& e_delta, const int rank); /// save descriptor void save_npy_d(const int nat, - const int des_per_atom, - const int inlmax, - const int* inl_l, - const bool deepks_equiv, - const std::vector &d_tensor, - const std::string& out_dir, - const int rank); + const int des_per_atom, + const int inlmax, + const int* inl_l, + const bool deepks_equiv, + const std::vector& d_tensor, + const std::string& out_dir, + const int rank); // save energy -void save_npy_e(const double &e, /**<[in] \f$E_{base}\f$ or \f$E_{tot}\f$, in Ry*/ - const std::string &e_file, - const int rank); +void save_npy_e(const double& e, /**<[in] \f$E_{base}\f$ or \f$E_{tot}\f$, in Ry*/ + const std::string& e_file, + const int rank); // save force and gvx -void save_npy_f(const ModuleBase::matrix &f, /**<[in] \f$F_{base}\f$ or \f$F_{tot}\f$, in Ry/Bohr*/ - const std::string &f_file, - const int nat, - const int rank); +void save_npy_f(const ModuleBase::matrix& f, /**<[in] \f$F_{base}\f$ or \f$F_{tot}\f$, in Ry/Bohr*/ + const std::string& f_file, + const int nat, + const int rank); void save_npy_gvx(const int nat, - const int des_per_atom, - const torch::Tensor &gvx_tensor, - const std::string& out_dir, - const int rank); + const int des_per_atom, + const torch::Tensor& gvx_tensor, + const std::string& out_dir, + const int rank); // save stress and gvepsl -void save_npy_s(const ModuleBase::matrix &stress, /**<[in] \f$S_{base}\f$ or \f$S_{tot}\f$, in Ry/Bohr^3*/ - const std::string &s_file, - const double &omega, - const int rank); +void save_npy_s(const ModuleBase::matrix& stress, /**<[in] \f$S_{base}\f$ or \f$S_{tot}\f$, in Ry/Bohr^3*/ + const std::string& s_file, + const double& omega, + const int rank); void save_npy_gvepsl(const int nat, - const int des_per_atom, - const torch::Tensor &gvepsl_tensor, - const std::string& out_dir, - const int rank); + const int des_per_atom, + const torch::Tensor& gvepsl_tensor, + const std::string& out_dir, + const int rank); /// save orbital and orbital_precalc void save_npy_o(const std::vector& bandgap, /**<[in] \f$E_{base}\f$ or \f$E_{tot}\f$, in Ry*/ - const std::string &o_file, - const int nks, - const int rank); + const std::string& o_file, + const int nks, + const int rank); -void save_npy_orbital_precalc(const int nat, - const int nks, - const int des_per_atom, - const torch::Tensor& orbital_precalc, - const std::string& out_dir, - const int rank); +void save_npy_orbital_precalc(const int nat, + const int nks, + const int des_per_atom, + const torch::Tensor& orbital_precalc, + const std::string& out_dir, + const int rank); // save Hamiltonian and v_delta_precalc(for deepks_v_delta==1)/phialpha+gevdm(for deepks_v_delta==2) template -void save_npy_h(const std::vector &hamilt, - const std::string &h_file, - const int nlocal, - const int nks, - const int rank); +void save_npy_h(const std::vector& hamilt, + const std::string& h_file, + const int nlocal, + const int nks, + const int rank); template void save_npy_v_delta_precalc(const int nat, - const int nks, - const int nlocal, - const int des_per_atom, - const torch::Tensor& v_delta_precalc_tensor, - const std::string& out_dir, - const int rank); + const int nks, + const int nlocal, + const int des_per_atom, + const torch::Tensor& v_delta_precalc_tensor, + const std::string& out_dir, + const int rank); template void save_npy_phialpha(const int nat, - const int nks, - const int nlocal, - const int inlmax, - const int lmaxd, - const torch::Tensor &phialpha_tensor, - const std::string& out_dir, - const int rank); + const int nks, + const int nlocal, + const int inlmax, + const int lmaxd, + const torch::Tensor& phialpha_tensor, + const std::string& out_dir, + const int rank); // Always real, no need for template now void save_npy_gevdm(const int nat, - const int inlmax, - const int lmaxd, - const torch::Tensor& gevdm_tensor, - const std::string& out_dir, - const int rank); -}; + const int inlmax, + const int lmaxd, + const torch::Tensor& gevdm_tensor, + const std::string& out_dir, + const int rank); +}; // namespace LCAO_deepks_io #endif #endif diff --git a/source/module_hamilt_lcao/module_deepks/LCAO_deepks_pdm.cpp b/source/module_hamilt_lcao/module_deepks/LCAO_deepks_pdm.cpp index 1d961d84e6..58e6b315c9 100644 --- a/source/module_hamilt_lcao/module_deepks/LCAO_deepks_pdm.cpp +++ b/source/module_hamilt_lcao/module_deepks/LCAO_deepks_pdm.cpp @@ -71,7 +71,7 @@ void LCAO_Deepks::read_projected_DM(bool read_pdm_file, bool is_equiv, const Num } } } - + this->init_pdm = true; } } @@ -330,10 +330,10 @@ void LCAO_Deepks::cal_projected_DM(const elecstate::DensityMatrix* d { int ind = m1 * nm + m2; pdm[inl][m1][m2] += ddot_(&row_size, - g_1dmt.data() + index * row_size, - &inc, - s_1t.data() + index * row_size, - &inc); + g_1dmt.data() + index * row_size, + &inc, + s_1t.data() + index * row_size, + &inc); index++; } } @@ -354,10 +354,10 @@ void LCAO_Deepks::cal_projected_DM(const elecstate::DensityMatrix* d for (int jproj = 0; jproj < nproj; jproj++) { pdm[iat][iproj * nproj + jproj] += ddot_(&row_size, - g_1dmt.data() + index * row_size, - &inc, - s_1t.data() + index * row_size, - &inc); + g_1dmt.data() + index * row_size, + &inc, + s_1t.data() + index * row_size, + &inc); index++; } } @@ -367,7 +367,7 @@ void LCAO_Deepks::cal_projected_DM(const elecstate::DensityMatrix* d } // T0 #ifdef __MPI - for(int inl=0; inl(), pdm_size); diff --git a/source/module_hamilt_lcao/module_deepks/cal_descriptor.cpp b/source/module_hamilt_lcao/module_deepks/cal_descriptor.cpp index 020892ddc0..e58473ec24 100644 --- a/source/module_hamilt_lcao/module_deepks/cal_descriptor.cpp +++ b/source/module_hamilt_lcao/module_deepks/cal_descriptor.cpp @@ -1,6 +1,6 @@ -///1. cal_descriptor : obtains descriptors which are eigenvalues of pdm -/// by calling torch::linalg::eigh -///2. check_descriptor : prints descriptor for checking +/// 1. cal_descriptor : obtains descriptors which are eigenvalues of pdm +/// by calling torch::linalg::eigh +/// 2. check_descriptor : prints descriptor for checking #ifdef __DEEPKS @@ -13,18 +13,18 @@ #include "module_hamilt_lcao/module_hcontainer/atom_pair.h" #include "module_parameter/parameter.h" -void LCAO_Deepks::cal_descriptor_equiv(const int nat) +void LCAO_Deepks::cal_descriptor_equiv(const int nat) { ModuleBase::TITLE("LCAO_Deepks", "cal_descriptor_equiv"); ModuleBase::timer::tick("LCAO_Deepks", "cal_descriptor_equiv"); // a rather unnecessary way of writing this, but I'll do it for now - if (!this->d_tensor.empty()) + if (!this->d_tensor.empty()) { this->d_tensor.erase(this->d_tensor.begin(), this->d_tensor.end()); } - for (int iat = 0; iat < nat; iat++) + for (int iat = 0; iat < nat; iat++) { auto tmp = torch::zeros(des_per_atom, torch::kFloat64); std::memcpy(tmp.data_ptr(), pdm[iat].data_ptr(), sizeof(double) * tmp.numel()); @@ -35,11 +35,12 @@ void LCAO_Deepks::cal_descriptor_equiv(const int nat) } // calculates descriptors from projected density matrices -void LCAO_Deepks::cal_descriptor(const int nat) { +void LCAO_Deepks::cal_descriptor(const int nat) +{ ModuleBase::TITLE("LCAO_Deepks", "cal_descriptor"); ModuleBase::timer::tick("LCAO_Deepks", "cal_descriptor"); - if (PARAM.inp.deepks_equiv) + if (PARAM.inp.deepks_equiv) { this->cal_descriptor_equiv(nat); return; @@ -47,12 +48,12 @@ void LCAO_Deepks::cal_descriptor(const int nat) { // init d_tensor // if d_tensor is not empty, clear it !! - if (!this->d_tensor.empty()) + if (!this->d_tensor.empty()) { this->d_tensor.erase(this->d_tensor.begin(), this->d_tensor.end()); } - for (int inl = 0; inl < this->inlmax; ++inl) + for (int inl = 0; inl < this->inlmax; ++inl) { const int nm = 2 * inl_l[inl] + 1; this->pdm[inl].requires_grad_(true); @@ -60,7 +61,7 @@ void LCAO_Deepks::cal_descriptor(const int nat) { } // cal d_tensor - for (int inl = 0; inl < inlmax; ++inl) + for (int inl = 0; inl < inlmax; ++inl) { torch::Tensor vd; std::tuple d_v(this->d_tensor[inl], vd); @@ -73,65 +74,64 @@ void LCAO_Deepks::cal_descriptor(const int nat) { return; } - -void LCAO_Deepks::check_descriptor(const UnitCell& ucell, const std::string& out_dir) { +void LCAO_Deepks::check_descriptor(const UnitCell& ucell, const std::string& out_dir) +{ ModuleBase::TITLE("LCAO_Deepks", "check_descriptor"); - if (GlobalV::MY_RANK != 0) + if (GlobalV::MY_RANK != 0) { return; } // mohan updated 2024-07-25 std::string file = out_dir + "deepks_desc.dat"; - + std::ofstream ofs(file.c_str()); - ofs << std::setprecision(10); - if (!PARAM.inp.deepks_equiv) - { - for (int it = 0; it < ucell.ntype; it++) - { - for (int ia = 0; ia < ucell.atoms[it].na; ia++) - { - int iat = ucell.itia2iat(it, ia); - ofs << ucell.atoms[it].label << " atom_index " << ia + 1 - << " n_descriptor " << this->des_per_atom << std::endl; - int id = 0; - for (int inl = 0; inl < inlmax / ucell.nat; inl++) - { - int nm = 2 * inl_l[inl] + 1; - for (int im = 0; im < nm; im++) - { - const int ind = iat * inlmax / ucell.nat + inl; - ofs << d_tensor[ind].index({im}).item().toDouble() - << " "; - - if (id % 8 == 7) - { - ofs << std::endl; - } - id++; - } - } - ofs << std::endl << std::endl; - } - } - } - else - { - for (int iat = 0; iat < ucell.nat; iat++) - { - const int it = ucell.iat2it[iat]; - ofs << ucell.atoms[it].label << " atom_index " << iat + 1 - << " n_descriptor " << this->des_per_atom << std::endl; - for (int i = 0; i < this->des_per_atom; i++) + ofs << std::setprecision(10); + if (!PARAM.inp.deepks_equiv) + { + for (int it = 0; it < ucell.ntype; it++) + { + for (int ia = 0; ia < ucell.atoms[it].na; ia++) + { + int iat = ucell.itia2iat(it, ia); + ofs << ucell.atoms[it].label << " atom_index " << ia + 1 << " n_descriptor " << this->des_per_atom + << std::endl; + int id = 0; + for (int inl = 0; inl < inlmax / ucell.nat; inl++) + { + int nm = 2 * inl_l[inl] + 1; + for (int im = 0; im < nm; im++) + { + const int ind = iat * inlmax / ucell.nat + inl; + ofs << d_tensor[ind].index({im}).item().toDouble() << " "; + + if (id % 8 == 7) + { + ofs << std::endl; + } + id++; + } + } + ofs << std::endl << std::endl; + } + } + } + else + { + for (int iat = 0; iat < ucell.nat; iat++) + { + const int it = ucell.iat2it[iat]; + ofs << ucell.atoms[it].label << " atom_index " << iat + 1 << " n_descriptor " << this->des_per_atom + << std::endl; + for (int i = 0; i < this->des_per_atom; i++) { ofs << this->pdm[iat][i].item() << " "; - if (i % 8 == 7) - { - ofs << std::endl; - } - } + if (i % 8 == 7) + { + ofs << std::endl; + } + } ofs << std::endl << std::endl; } } diff --git a/source/module_hamilt_lcao/module_deepks/cal_gedm.cpp b/source/module_hamilt_lcao/module_deepks/cal_gedm.cpp index 7fe9a59fe1..d35972205b 100644 --- a/source/module_hamilt_lcao/module_deepks/cal_gedm.cpp +++ b/source/module_hamilt_lcao/module_deepks/cal_gedm.cpp @@ -14,13 +14,13 @@ #include "module_hamilt_lcao/module_hcontainer/atom_pair.h" #include "module_parameter/parameter.h" +inline void generate_py_files(const int lmaxd, const int nmaxd, const std::string& out_dir) +{ -inline void generate_py_files(const int lmaxd, const int nmaxd, const std::string &out_dir) { - - if (GlobalV::MY_RANK != 0) - { - return; - } + if (GlobalV::MY_RANK != 0) + { + return; + } std::ofstream ofs("cal_gedm.py"); ofs << "import torch" << std::endl; @@ -52,12 +52,15 @@ inline void generate_py_files(const int lmaxd, const int nmaxd, const std::strin ofs.open("basis.yaml"); ofs << "proj_basis:" << std::endl; - for (int l = 0; l < lmaxd + 1; l++) { + for (int l = 0; l < lmaxd + 1; l++) + { ofs << " - - " << l << std::endl; ofs << " - ["; - for (int i = 0; i < nmaxd + 1; i++) { + for (int i = 0; i < nmaxd + 1; i++) + { ofs << "0"; - if (i != nmaxd) { + if (i != nmaxd) + { ofs << ", "; } } @@ -65,22 +68,23 @@ inline void generate_py_files(const int lmaxd, const int nmaxd, const std::strin } } -void LCAO_Deepks::cal_gedm_equiv(const int nat) { +void LCAO_Deepks::cal_gedm_equiv(const int nat) +{ ModuleBase::TITLE("LCAO_Deepks", "cal_gedm_equiv"); - LCAO_deepks_io::save_npy_d( - nat, - this->des_per_atom, - this->inlmax, - this->inl_l, - PARAM.inp.deepks_equiv, - this->d_tensor, - PARAM.globalv.global_out_dir, - GlobalV::MY_RANK); // libnpy needed + LCAO_deepks_io::save_npy_d(nat, + this->des_per_atom, + this->inlmax, + this->inl_l, + PARAM.inp.deepks_equiv, + this->d_tensor, + PARAM.globalv.global_out_dir, + GlobalV::MY_RANK); // libnpy needed generate_py_files(this->lmaxd, this->nmaxd, PARAM.globalv.global_out_dir); - if (GlobalV::MY_RANK == 0) { + if (GlobalV::MY_RANK == 0) + { std::string cmd = "python cal_gedm.py " + PARAM.inp.deepks_model; int stat = std::system(cmd.c_str()); assert(stat == 0); @@ -88,21 +92,17 @@ void LCAO_Deepks::cal_gedm_equiv(const int nat) { MPI_Barrier(MPI_COMM_WORLD); - LCAO_deepks_io::load_npy_gedm( - nat, - this->des_per_atom, - this->gedm, - this->E_delta, - GlobalV::MY_RANK); + LCAO_deepks_io::load_npy_gedm(nat, this->des_per_atom, this->gedm, this->E_delta, GlobalV::MY_RANK); std::string cmd = "rm -f cal_gedm.py basis.yaml ec.npy gedm.npy"; std::system(cmd.c_str()); } // obtain from the machine learning model dE_delta/dDescriptor -void LCAO_Deepks::cal_gedm(const int nat) { +void LCAO_Deepks::cal_gedm(const int nat) +{ - if (PARAM.inp.deepks_equiv) + if (PARAM.inp.deepks_equiv) { this->cal_gedm_equiv(nat); return; @@ -114,11 +114,10 @@ void LCAO_Deepks::cal_gedm(const int nat) { std::vector inputs; // input_dim:(natom, des_per_atom) - inputs.push_back( - torch::cat(this->d_tensor, 0).reshape({1, nat, this->des_per_atom})); + inputs.push_back(torch::cat(this->d_tensor, 0).reshape({1, nat, this->des_per_atom})); std::vector ec; ec.push_back(module.forward(inputs).toTensor()); // Hartree - this->E_delta = ec[0].item().toDouble() * 2; // Ry; *2 is for Hartree to Ry + this->E_delta = ec[0].item().toDouble() * 2; // Ry; *2 is for Hartree to Ry // cal gedm std::vector gedm_shell; @@ -131,39 +130,40 @@ void LCAO_Deepks::cal_gedm(const int nat) { /*allow_unused=*/true); // gedm_tensor(Hartree) to gedm(Ry) - for (int inl = 0; inl < inlmax; ++inl) { + for (int inl = 0; inl < inlmax; ++inl) + { int nm = 2 * inl_l[inl] + 1; - for (int m1 = 0; m1 < nm; ++m1) { - for (int m2 = 0; m2 < nm; ++m2) { + for (int m1 = 0; m1 < nm; ++m1) + { + for (int m2 = 0; m2 < nm; ++m2) + { int index = m1 * nm + m2; //*2 is for Hartree to Ry - this->gedm[inl][index] - = this->gedm_tensor[inl].index({m1, m2}).item().toDouble() - * 2; + this->gedm[inl][index] = this->gedm_tensor[inl].index({m1, m2}).item().toDouble() * 2; } } } return; } -void LCAO_Deepks::check_gedm() +void LCAO_Deepks::check_gedm() { std::ofstream ofs("gedm.dat"); - for (int inl = 0; inl < inlmax; inl++) - { - int nm = 2 * inl_l[inl] + 1; - for (int m1 = 0; m1 < nm; ++m1) - { - for (int m2 = 0; m2 < nm; ++m2) - { - int index = m1 * nm + m2; - //*2 is for Hartree to Ry - ofs << this->gedm[inl][index] << " "; - } - } - ofs << std::endl; - } + for (int inl = 0; inl < inlmax; inl++) + { + int nm = 2 * inl_l[inl] + 1; + for (int m1 = 0; m1 < nm; ++m1) + { + for (int m2 = 0; m2 < nm; ++m2) + { + int index = m1 * nm + m2; + //*2 is for Hartree to Ry + ofs << this->gedm[inl][index] << " "; + } + } + ofs << std::endl; + } } #endif diff --git a/source/module_hamilt_lcao/module_deepks/cal_gvx.cpp b/source/module_hamilt_lcao/module_deepks/cal_gvx.cpp index ac1ba9705a..cec0f1d84a 100644 --- a/source/module_hamilt_lcao/module_deepks/cal_gvx.cpp +++ b/source/module_hamilt_lcao/module_deepks/cal_gvx.cpp @@ -1,4 +1,4 @@ -/// 1. cal_gvx : gvx is used for training with force label, which is gradient of descriptors, +/// 1. cal_gvx : gvx is used for training with force label, which is gradient of descriptors, /// calculated by d(des)/dX = d(pdm)/dX * d(des)/d(pdm) = gdmx * gvdm /// using einsum /// 2. check_gvx : prints gvx into gvx.dat for checking @@ -16,59 +16,54 @@ // calculates gradient of descriptors from gradient of projected density // matrices -void LCAO_Deepks::cal_gvx(const int nat, const std::vector& gevdm) +void LCAO_Deepks::cal_gvx(const int nat, const std::vector& gevdm) { ModuleBase::TITLE("LCAO_Deepks", "cal_gvx"); - if (!gdmr_vector.empty()) + if (!gdmr_vector.empty()) { gdmr_vector.erase(gdmr_vector.begin(), gdmr_vector.end()); } // gdmr_vector : nat(derivative) * 3 * inl(projector) * nm * nm - if (GlobalV::MY_RANK == 0) + if (GlobalV::MY_RANK == 0) { // make gdmx as tensor int nlmax = this->inlmax / nat; - for (int nl = 0; nl < nlmax; ++nl) + for (int nl = 0; nl < nlmax; ++nl) { std::vector bmmv; - for (int ibt = 0; ibt < nat; ++ibt) + for (int ibt = 0; ibt < nat; ++ibt) { std::vector xmmv; - for (int i = 0; i < 3; ++i) + for (int i = 0; i < 3; ++i) { std::vector ammv; - for (int iat = 0; iat < nat; ++iat) + for (int iat = 0; iat < nat; ++iat) { int inl = iat * nlmax + nl; int nm = 2 * this->inl_l[inl] + 1; std::vector mmv; - for (int m1 = 0; m1 < nm; ++m1) + for (int m1 = 0; m1 < nm; ++m1) { - for (int m2 = 0; m2 < nm; ++m2) + for (int m2 = 0; m2 < nm; ++m2) { - if (i == 0) - { - mmv.push_back( - this->gdmx[ibt][inl][m1 * nm + m2]); - } - if (i == 1) - { - mmv.push_back( - this->gdmy[ibt][inl][m1 * nm + m2]); - } - if (i == 2) - { - mmv.push_back( - this->gdmz[ibt][inl][m1 * nm + m2]); - } + if (i == 0) + { + mmv.push_back(this->gdmx[ibt][inl][m1 * nm + m2]); + } + if (i == 1) + { + mmv.push_back(this->gdmy[ibt][inl][m1 * nm + m2]); + } + if (i == 2) + { + mmv.push_back(this->gdmz[ibt][inl][m1 * nm + m2]); + } } } // nm^2 - torch::Tensor mm - = torch::tensor( - mmv, - torch::TensorOptions().dtype(torch::kFloat64)).reshape({nm, nm}); // nm*nm + torch::Tensor mm = torch::tensor(mmv, torch::TensorOptions().dtype(torch::kFloat64)) + .reshape({nm, nm}); // nm*nm ammv.push_back(mm); } torch::Tensor amm = torch::stack(ammv, 0); // nat*nm*nm @@ -88,10 +83,9 @@ void LCAO_Deepks::cal_gvx(const int nat, const std::vector& gevdm // n:nm (pdm, dim2) gvx_vector : b:nat(derivative) * x:3 * // a:inl(projector) * m:nm(descriptor) std::vector gvx_vector; - for (int nl = 0; nl < nlmax; ++nl) { - gvx_vector.push_back( - at::einsum("bxamn, avmn->bxav", - {this->gdmr_vector[nl], gevdm[nl]})); + for (int nl = 0; nl < nlmax; ++nl) + { + gvx_vector.push_back(at::einsum("bxamn, avmn->bxav", {this->gdmr_vector[nl], gevdm[nl]})); } // cat nv-> \sum_nl(nv) = \sum_nl(nm_nl)=des_per_atom @@ -107,7 +101,8 @@ void LCAO_Deepks::cal_gvx(const int nat, const std::vector& gevdm return; } -void LCAO_Deepks::check_gvx(const int nat) { +void LCAO_Deepks::check_gvx(const int nat) +{ std::stringstream ss; std::ofstream ofs_x; std::ofstream ofs_y; @@ -117,7 +112,8 @@ void LCAO_Deepks::check_gvx(const int nat) { ofs_y << std::setprecision(12); ofs_z << std::setprecision(12); - for (int ia = 0; ia < nat; ia++) { + for (int ia = 0; ia < nat; ia++) + { ss.str(""); ss << "gvx_" << ia << ".dat"; ofs_x.open(ss.str().c_str()); @@ -132,20 +128,16 @@ void LCAO_Deepks::check_gvx(const int nat) { ofs_y << std::setprecision(10); ofs_z << std::setprecision(10); - for (int ib = 0; ib < nat; ib++) { - for (int inl = 0; inl < inlmax / nat; inl++) { + for (int ib = 0; ib < nat; ib++) + { + for (int inl = 0; inl < inlmax / nat; inl++) + { int nm = 2 * inl_l[inl] + 1; { const int ind = ib * inlmax / nat + inl; - ofs_x - << gvx_tensor.index({ia, 0, ib, inl}).item().toDouble() - << " "; - ofs_y - << gvx_tensor.index({ia, 1, ib, inl}).item().toDouble() - << " "; - ofs_z - << gvx_tensor.index({ia, 2, ib, inl}).item().toDouble() - << " "; + ofs_x << gvx_tensor.index({ia, 0, ib, inl}).item().toDouble() << " "; + ofs_y << gvx_tensor.index({ia, 1, ib, inl}).item().toDouble() << " "; + ofs_z << gvx_tensor.index({ia, 2, ib, inl}).item().toDouble() << " "; } } ofs_x << std::endl; diff --git a/source/module_hamilt_lcao/module_deepks/deepks_orbpre.cpp b/source/module_hamilt_lcao/module_deepks/deepks_orbpre.cpp index 0b7805cb69..8d48b981c0 100644 --- a/source/module_hamilt_lcao/module_deepks/deepks_orbpre.cpp +++ b/source/module_hamilt_lcao/module_deepks/deepks_orbpre.cpp @@ -5,6 +5,7 @@ /// orbital_pdm[nks,Inl,nm,nm] = dm_hl * overlap * overlap #include "deepks_orbpre.h" + #include "LCAO_deepks_io.h" // mohan add 2024-07-22 #include "module_base/blas_connector.h" #include "module_base/constants.h" @@ -37,7 +38,8 @@ void DeePKS_domain::cal_orbital_precalc(const std::vector& dm_hl, const double Rcut_Alpha = orb.Alpha[0].getRcut(); - torch::Tensor orbital_pdm = torch::zeros({nks, inlmax, (2 * lmaxd + 1) , (2 * lmaxd + 1)}, torch::dtype(torch::kFloat64)); + torch::Tensor orbital_pdm + = torch::zeros({nks, inlmax, (2 * lmaxd + 1), (2 * lmaxd + 1)}, torch::dtype(torch::kFloat64)); for (int T0 = 0; T0 < ucell.ntype; T0++) { @@ -257,7 +259,7 @@ void DeePKS_domain::cal_orbital_precalc(const std::vector& dm_hl, for (int ik = 0; ik < nks; ik++) { // do dot of g_1dmt and s_1t to get orbital_pdm - + const double* p_g1dmt = g_1dmt.data() + ik * row_size; int ib = 0, index = 0, inc = 1; @@ -273,12 +275,11 @@ void DeePKS_domain::cal_orbital_precalc(const std::vector& dm_hl, { for (int m2 = 0; m2 < nm; ++m2) // m1 = 1 for s, 3 for p, 5 for d { - orbital_pdm[ik][inl][m1][m2] - += ddot_(&row_size, - p_g1dmt + index * row_size * nks, - &inc, - s_1t.data() + index * row_size, - &inc); + orbital_pdm[ik][inl][m1][m2] += ddot_(&row_size, + p_g1dmt + index * row_size * nks, + &inc, + s_1t.data() + index * row_size, + &inc); index++; } } @@ -295,8 +296,7 @@ void DeePKS_domain::cal_orbital_precalc(const std::vector& dm_hl, for (int inl = 0; inl < inlmax; inl++) { auto tensor_slice = orbital_pdm[iks][inl]; - Parallel_Reduce::reduce_all(tensor_slice.data_ptr(), - (2 * lmaxd + 1) * (2 * lmaxd + 1)); + Parallel_Reduce::reduce_all(tensor_slice.data_ptr(), (2 * lmaxd + 1) * (2 * lmaxd + 1)); } } #endif @@ -342,8 +342,7 @@ void DeePKS_domain::cal_orbital_precalc(const std::vector& dm_hl, std::vector orbital_precalc_vector; for (int nl = 0; nl < nlmax; ++nl) { - orbital_precalc_vector.push_back( - at::einsum("kamn, avmn->kav", {orbital_pdm_vector[nl], gevdm[nl]})); + orbital_precalc_vector.push_back(at::einsum("kamn, avmn->kav", {orbital_pdm_vector[nl], gevdm[nl]})); } orbital_precalc = torch::cat(orbital_precalc_vector, -1); diff --git a/source/module_hamilt_lcao/module_deepks/deepks_orbpre.h b/source/module_hamilt_lcao/module_deepks/deepks_orbpre.h index a320ffba71..8b5ebe4998 100644 --- a/source/module_hamilt_lcao/module_deepks/deepks_orbpre.h +++ b/source/module_hamilt_lcao/module_deepks/deepks_orbpre.h @@ -1,5 +1,5 @@ -#ifndef DEEPKS_ORBPRE_H -#define DEEPKS_ORBPRE_H +#ifndef DEEPKS_ORBPRE_H +#define DEEPKS_ORBPRE_H #ifdef __DEEPKS @@ -18,12 +18,12 @@ namespace DeePKS_domain { - //------------------------ - // deepks_orbpre.cpp - //------------------------ +//------------------------ +// deepks_orbpre.cpp +//------------------------ - // This file contains one subroutine for calculating orbital_precalc, - // which is defind as gvdm * dm_hl * overlap * overlap +// This file contains one subroutine for calculating orbital_precalc, +// which is defind as gvdm * dm_hl * overlap * overlap template void cal_orbital_precalc(const std::vector& dm_hl, @@ -41,6 +41,6 @@ void cal_orbital_precalc(const std::vector& dm_hl, const Parallel_Orbitals& pv, const Grid_Driver& GridD, torch::Tensor& orbital_precalc); -} +} // namespace DeePKS_domain #endif #endif diff --git a/source/module_hamilt_lcao/module_deepks/test/LCAO_deepks_test.cpp b/source/module_hamilt_lcao/module_deepks/test/LCAO_deepks_test.cpp index 33a9454747..d1914d5ca4 100644 --- a/source/module_hamilt_lcao/module_deepks/test/LCAO_deepks_test.cpp +++ b/source/module_hamilt_lcao/module_deepks/test/LCAO_deepks_test.cpp @@ -1,8 +1,9 @@ #include "LCAO_deepks_test.h" #define private public #include "module_parameter/parameter.h" -#include + #include +#include #undef private namespace Test_Deepks {