diff --git a/source/Makefile.Objects b/source/Makefile.Objects index 4301c0943a..5d01dd1839 100644 --- a/source/Makefile.Objects +++ b/source/Makefile.Objects @@ -193,21 +193,20 @@ OBJS_CELL=atom_pseudo.o\ read_atom_species.o\ OBJS_DEEPKS=LCAO_deepks.o\ + deepks_basic.o\ + deepks_descriptor.o\ deepks_force.o\ deepks_fpre.o\ deepks_spre.o\ - deepks_descriptor.o\ deepks_orbital.o\ deepks_orbpre.o\ + deepks_vdelta.o\ deepks_vdpre.o\ deepks_hmat.o\ + deepks_pdm.o\ + deepks_phialpha.o\ LCAO_deepks_io.o\ - LCAO_deepks_pdm.o\ - LCAO_deepks_phialpha.o\ - LCAO_deepks_torch.o\ - LCAO_deepks_vdelta.o\ LCAO_deepks_interface.o\ - cal_gedm.o\ OBJS_ELECSTAT=elecstate.o\ diff --git a/source/module_esolver/esolver_ks_lcao.cpp b/source/module_esolver/esolver_ks_lcao.cpp index 8c87cc352b..1fc3b1cdc3 100644 --- a/source/module_esolver/esolver_ks_lcao.cpp +++ b/source/module_esolver/esolver_ks_lcao.cpp @@ -225,9 +225,16 @@ void ESolver_KS_LCAO::before_all_runners(UnitCell& ucell, const Input_pa if (PARAM.inp.deepks_scf) { // load the DeePKS model from deep neural network - GlobalC::ld.load_model(PARAM.inp.deepks_model); + DeePKS_domain::load_model(PARAM.inp.deepks_model, GlobalC::ld.model_deepks); // read pdm from file for NSCF or SCF-restart, do it only once in whole calculation - GlobalC::ld.read_projected_DM((PARAM.inp.init_chg == "file"), PARAM.inp.deepks_equiv, *orb_.Alpha); + DeePKS_domain::read_pdm((PARAM.inp.init_chg == "file"), + PARAM.inp.deepks_equiv, + GlobalC::ld.init_pdm, + GlobalC::ld.inlmax, + GlobalC::ld.lmaxd, + GlobalC::ld.inl_l, + *orb_.Alpha, + GlobalC::ld.pdm); } #endif @@ -928,9 +935,7 @@ void ESolver_KS_LCAO::after_scf(UnitCell& ucell, const int istep) // 1) calculate the kinetic energy density tau, sunliang 2024-09-18 if (PARAM.inp.out_elf[0] > 0) { - elecstate::lcao_cal_tau(&(this->GG), - &(this->GK), - this->pelec->charge); + elecstate::lcao_cal_tau(&(this->GG), &(this->GK), this->pelec->charge); } //! 2) call after_scf() of ESolver_KS @@ -1047,7 +1052,6 @@ void ESolver_KS_LCAO::after_scf(UnitCell& ucell, const int istep) std::shared_ptr ld_shared_ptr(&GlobalC::ld, [](LCAO_Deepks*) {}); LCAO_Deepks_Interface LDI(ld_shared_ptr); - ModuleBase::timer::tick("ESolver_KS_LCAO", "out_deepks_labels"); LDI.out_deepks_labels(this->pelec->f_en.etot, this->pelec->klist->get_nks(), ucell.nat, @@ -1061,8 +1065,6 @@ void ESolver_KS_LCAO::after_scf(UnitCell& ucell, const int istep) *(this->psi), dynamic_cast*>(this->pelec)->get_DM(), p_ham_deepks); - - ModuleBase::timer::tick("ESolver_KS_LCAO", "out_deepks_labels"); } #endif diff --git a/source/module_esolver/lcao_before_scf.cpp b/source/module_esolver/lcao_before_scf.cpp index 8aad327da7..724b4e06e3 100644 --- a/source/module_esolver/lcao_before_scf.cpp +++ b/source/module_esolver/lcao_before_scf.cpp @@ -211,13 +211,19 @@ void ESolver_KS_LCAO::before_scf(UnitCell& ucell, const int istep) { const Parallel_Orbitals* pv = &this->pv; // allocate , phialpha is different every ion step, so it is allocated here - GlobalC::ld.allocate_phialpha(PARAM.inp.cal_force, ucell, orb_, this->gd); + DeePKS_domain::allocate_phialpha(PARAM.inp.cal_force, ucell, orb_, this->gd, pv, GlobalC::ld.phialpha); // build and save at beginning - GlobalC::ld.build_phialpha(PARAM.inp.cal_force, ucell, orb_, this->gd, *(two_center_bundle_.overlap_orb_alpha)); + DeePKS_domain::build_phialpha(PARAM.inp.cal_force, + ucell, + orb_, + this->gd, + pv, + *(two_center_bundle_.overlap_orb_alpha), + GlobalC::ld.phialpha); if (PARAM.inp.deepks_out_unittest) { - GlobalC::ld.check_phialpha(PARAM.inp.cal_force, ucell, orb_, this->gd); + DeePKS_domain::check_phialpha(PARAM.inp.cal_force, ucell, orb_, this->gd, pv, GlobalC::ld.phialpha); } } #endif diff --git a/source/module_esolver/lcao_others.cpp b/source/module_esolver/lcao_others.cpp index 3c8c0ff879..2ba5d9c37c 100644 --- a/source/module_esolver/lcao_others.cpp +++ b/source/module_esolver/lcao_others.cpp @@ -217,13 +217,19 @@ void ESolver_KS_LCAO::others(UnitCell& ucell, const int istep) { const Parallel_Orbitals* pv = &this->pv; // allocate , phialpha is different every ion step, so it is allocated here - GlobalC::ld.allocate_phialpha(PARAM.inp.cal_force, ucell, orb_, this->gd); + DeePKS_domain::allocate_phialpha(PARAM.inp.cal_force, ucell, orb_, this->gd, pv, GlobalC::ld.phialpha); // build and save at beginning - GlobalC::ld.build_phialpha(PARAM.inp.cal_force, ucell, orb_, this->gd, *(two_center_bundle_.overlap_orb_alpha)); + DeePKS_domain::build_phialpha(PARAM.inp.cal_force, + ucell, + orb_, + this->gd, + pv, + *(two_center_bundle_.overlap_orb_alpha), + GlobalC::ld.phialpha); if (PARAM.inp.deepks_out_unittest) { - GlobalC::ld.check_phialpha(PARAM.inp.cal_force, ucell, orb_, this->gd); + DeePKS_domain::check_phialpha(PARAM.inp.cal_force, ucell, orb_, this->gd, pv, GlobalC::ld.phialpha); } } #endif diff --git a/source/module_hamilt_lcao/hamilt_lcaodft/FORCE_STRESS.cpp b/source/module_hamilt_lcao/hamilt_lcaodft/FORCE_STRESS.cpp index bfeb7b9f7f..1d64bbe049 100644 --- a/source/module_hamilt_lcao/hamilt_lcaodft/FORCE_STRESS.cpp +++ b/source/module_hamilt_lcao/hamilt_lcaodft/FORCE_STRESS.cpp @@ -500,87 +500,16 @@ void Force_Stress_LCAO::getForceStress(UnitCell& ucell, if (PARAM.inp.deepks_out_labels) // not parallelized yet { const std::string file_ftot = PARAM.globalv.global_out_dir + "deepks_ftot.npy"; - LCAO_deepks_io::save_npy_f(fcs, file_ftot, ucell.nat, - GlobalV::MY_RANK); // Ty/Bohr, F_tot + LCAO_deepks_io::save_npy_f(fcs, file_ftot, GlobalV::MY_RANK); // Ry/Bohr, F_tot + const std::string file_fbase = PARAM.globalv.global_out_dir + "deepks_fbase.npy"; if (PARAM.inp.deepks_scf) { - const std::string file_fbase = PARAM.globalv.global_out_dir + "deepks_fbase.npy"; - LCAO_deepks_io::save_npy_f(fcs - fvnl_dalpha, - file_fbase, - ucell.nat, - GlobalV::MY_RANK); // Ry/Bohr, F_base - - if (!PARAM.inp.deepks_equiv) // training with force label not supported by equivariant version now - { - torch::Tensor gdmx; - if (PARAM.globalv.gamma_only_local) - { - const std::vector>& dm_gamma - = dynamic_cast*>(pelec)->get_DM()->get_DMK_vector(); - - DeePKS_domain::cal_gdmx(GlobalC::ld.lmaxd, - GlobalC::ld.inlmax, - kv.get_nks(), - kv.kvec_d, - GlobalC::ld.phialpha, - GlobalC::ld.inl_index, - dm_gamma, - ucell, - orb, - pv, - gd, - gdmx); - } - else - { - const std::vector>>& dm_k - = dynamic_cast>*>(pelec) - ->get_DM() - ->get_DMK_vector(); - - DeePKS_domain::cal_gdmx(GlobalC::ld.lmaxd, - GlobalC::ld.inlmax, - kv.get_nks(), - kv.kvec_d, - GlobalC::ld.phialpha, - GlobalC::ld.inl_index, - dm_k, - ucell, - orb, - pv, - gd, - gdmx); - } - std::vector gevdm; - GlobalC::ld.cal_gevdm(ucell.nat, gevdm); - torch::Tensor gvx; - DeePKS_domain::cal_gvx(ucell.nat, - GlobalC::ld.inlmax, - GlobalC::ld.des_per_atom, - GlobalC::ld.inl_l, - gevdm, - gdmx, - gvx); - - if (PARAM.inp.deepks_out_unittest) - { - DeePKS_domain::check_gdmx(gdmx); - DeePKS_domain::check_gvx(gvx); - } - - LCAO_deepks_io::save_npy_gvx(ucell.nat, - GlobalC::ld.des_per_atom, - gvx, - PARAM.globalv.global_out_dir, - GlobalV::MY_RANK); - } + LCAO_deepks_io::save_npy_f(fcs - fvnl_dalpha, file_fbase, GlobalV::MY_RANK); // Ry/Bohr, F_base } else { - const std::string file_fbase = PARAM.globalv.global_out_dir + "deepks_fbase.npy"; - LCAO_deepks_io::save_npy_f(fcs, file_fbase, ucell.nat, - GlobalV::MY_RANK); // no scf, F_base=F_tot + LCAO_deepks_io::save_npy_f(fcs, file_fbase, GlobalV::MY_RANK); // no scf, F_base=F_tot } } #endif @@ -758,80 +687,18 @@ void Force_Stress_LCAO::getForceStress(UnitCell& ucell, ucell.omega, GlobalV::MY_RANK); // change to energy unit Ry when printing, S_tot, w/ model - // wenfei add 2021/11/2 + const std::string file_sbase = PARAM.globalv.global_out_dir + "deepks_sbase.npy"; if (PARAM.inp.deepks_scf) { - const std::string file_sbase = PARAM.globalv.global_out_dir + "deepks_sbase.npy"; LCAO_deepks_io::save_npy_s(scs - svnl_dalpha, file_sbase, ucell.omega, GlobalV::MY_RANK); // change to energy unit Ry when printing, S_base; - - if (!PARAM.inp.deepks_equiv) // training with stress label not supported by equivariant version now - { - torch::Tensor gdmepsl; - if (PARAM.globalv.gamma_only_local) - { - const std::vector>& dm_gamma - = dynamic_cast*>(pelec)->get_DM()->get_DMK_vector(); - - DeePKS_domain::cal_gdmepsl(GlobalC::ld.lmaxd, - GlobalC::ld.inlmax, - kv.get_nks(), - kv.kvec_d, - GlobalC::ld.phialpha, - GlobalC::ld.inl_index, - dm_gamma, - ucell, - orb, - pv, - gd, - gdmepsl); - } - else - { - const std::vector>>& dm_k - = dynamic_cast>*>(pelec) - ->get_DM() - ->get_DMK_vector(); - - DeePKS_domain::cal_gdmepsl(GlobalC::ld.lmaxd, - GlobalC::ld.inlmax, - kv.get_nks(), - kv.kvec_d, - GlobalC::ld.phialpha, - GlobalC::ld.inl_index, - dm_k, - ucell, - orb, - pv, - gd, - gdmepsl); - } - - std::vector gevdm; - GlobalC::ld.cal_gevdm(ucell.nat, gevdm); - torch::Tensor gvepsl; - DeePKS_domain::cal_gvepsl(ucell.nat, - GlobalC::ld.inlmax, - GlobalC::ld.des_per_atom, - GlobalC::ld.inl_l, - gevdm, - gdmepsl, - gvepsl); - - if (PARAM.inp.deepks_out_unittest) - { - DeePKS_domain::check_gdmepsl(gdmepsl); - DeePKS_domain::check_gvepsl(gvepsl); - } - - LCAO_deepks_io::save_npy_gvepsl(ucell.nat, - GlobalC::ld.des_per_atom, - gvepsl, - PARAM.globalv.global_out_dir, - GlobalV::MY_RANK); // unitless, grad_vepsl - } + } + else + { + LCAO_deepks_io::save_npy_s(scs, file_sbase, ucell.omega, + GlobalV::MY_RANK); // sbase = stot } } #endif diff --git a/source/module_hamilt_lcao/hamilt_lcaodft/FORCE_gamma.cpp b/source/module_hamilt_lcao/hamilt_lcaodft/FORCE_gamma.cpp index cdcf7bb4cc..506a1ada42 100644 --- a/source/module_hamilt_lcao/hamilt_lcaodft/FORCE_gamma.cpp +++ b/source/module_hamilt_lcao/hamilt_lcaodft/FORCE_gamma.cpp @@ -252,15 +252,23 @@ void Force_LCAO::ftable(const bool isforce, if (PARAM.inp.deepks_scf) { // when deepks_scf is on, the init pdm should be same as the out pdm, so we should not recalculate the pdm - // GlobalC::ld.cal_projected_DM(dm, ucell, orb, gd); - DeePKS_domain::cal_descriptor(ucell.nat, GlobalC::ld.inlmax, GlobalC::ld.inl_l, GlobalC::ld.pdm, descriptor, GlobalC::ld.des_per_atom); - GlobalC::ld.cal_gedm(ucell.nat, descriptor); + DeePKS_domain::cal_gedm(ucell.nat, + GlobalC::ld.lmaxd, + GlobalC::ld.nmaxd, + GlobalC::ld.inlmax, + GlobalC::ld.des_per_atom, + GlobalC::ld.inl_l, + descriptor, + GlobalC::ld.pdm, + GlobalC::ld.model_deepks, + GlobalC::ld.gedm, + GlobalC::ld.E_delta); const int nks = 1; DeePKS_domain::cal_f_delta(dm_gamma, @@ -302,32 +310,8 @@ void Force_LCAO::ftable(const bool isforce, } #ifdef __DEEPKS - // It seems these test should not all be here, should be moved in the future - // Also, these test are not in multi-k case now if (PARAM.inp.deepks_scf && PARAM.inp.deepks_out_unittest) { - const int nks = 1; // 1 for gamma-only - LCAO_deepks_io::print_dm(nks, PARAM.globalv.nlocal, this->ParaV->nrow, dm_gamma); - - GlobalC::ld.check_projected_dm(); - - DeePKS_domain::check_descriptor(GlobalC::ld.inlmax, - GlobalC::ld.des_per_atom, - GlobalC::ld.inl_l, - ucell, - PARAM.globalv.global_out_dir, - descriptor); - - GlobalC::ld.check_gedm(); - - GlobalC::ld.cal_e_delta_band(dm_gamma, nks); - - std::ofstream ofs("E_delta_bands.dat"); - ofs << std::setprecision(10) << GlobalC::ld.e_delta_band; - - std::ofstream ofs1("E_delta.dat"); - ofs1 << std::setprecision(10) << GlobalC::ld.E_delta; - DeePKS_domain::check_f_delta(ucell.nat, fvnl_dalpha, svnl_dalpha); } #endif diff --git a/source/module_hamilt_lcao/hamilt_lcaodft/FORCE_k.cpp b/source/module_hamilt_lcao/hamilt_lcaodft/FORCE_k.cpp index e8a7bd5c49..ac44ee4194 100644 --- a/source/module_hamilt_lcao/hamilt_lcaodft/FORCE_k.cpp +++ b/source/module_hamilt_lcao/hamilt_lcaodft/FORCE_k.cpp @@ -347,8 +347,6 @@ void Force_LCAO>::ftable(const bool isforce, const std::vector>>& dm_k = dm->get_DMK_vector(); // when deepks_scf is on, the init pdm should be same as the out pdm, so we should not recalculate the pdm - // GlobalC::ld.cal_projected_DM(dm, ucell, orb, gd); - std::vector descriptor; DeePKS_domain::cal_descriptor(ucell.nat, GlobalC::ld.inlmax, @@ -356,7 +354,17 @@ void Force_LCAO>::ftable(const bool isforce, GlobalC::ld.pdm, descriptor, GlobalC::ld.des_per_atom); - GlobalC::ld.cal_gedm(ucell.nat, descriptor); + DeePKS_domain::cal_gedm(ucell.nat, + GlobalC::ld.lmaxd, + GlobalC::ld.nmaxd, + GlobalC::ld.inlmax, + GlobalC::ld.des_per_atom, + GlobalC::ld.inl_l, + descriptor, + GlobalC::ld.pdm, + GlobalC::ld.model_deepks, + GlobalC::ld.gedm, + GlobalC::ld.E_delta); DeePKS_domain::cal_f_delta>(dm_k, ucell, @@ -399,6 +407,13 @@ void Force_LCAO>::ftable(const bool isforce, #endif } +#ifdef __DEEPKS + if (PARAM.inp.deepks_scf && PARAM.inp.deepks_out_unittest) + { + DeePKS_domain::check_f_delta(ucell.nat, fvnl_dalpha, svnl_dalpha); + } +#endif + ModuleBase::timer::tick("Force_LCAO", "ftable"); return; } diff --git a/source/module_hamilt_lcao/hamilt_lcaodft/operator_lcao/deepks_lcao.cpp b/source/module_hamilt_lcao/hamilt_lcaodft/operator_lcao/deepks_lcao.cpp index 00de8f1149..09410cf101 100644 --- a/source/module_hamilt_lcao/hamilt_lcaodft/operator_lcao/deepks_lcao.cpp +++ b/source/module_hamilt_lcao/hamilt_lcaodft/operator_lcao/deepks_lcao.cpp @@ -160,7 +160,18 @@ void hamilt::DeePKS>::contributeHR() { ModuleBase::timer::tick("DeePKS", "contributeHR"); - GlobalC::ld.cal_projected_DM(this->DM, *this->ucell, *ptr_orb_, *(this->gd)); + DeePKS_domain::cal_pdm(GlobalC::ld.init_pdm, + GlobalC::ld.inlmax, + GlobalC::ld.lmaxd, + GlobalC::ld.inl_l, + GlobalC::ld.inl_index, + this->DM, + GlobalC::ld.phialpha, + *this->ucell, + *ptr_orb_, + *(this->gd), + *(this->hR->get_paraV()), + GlobalC::ld.pdm); std::vector descriptor; DeePKS_domain::cal_descriptor(this->ucell->nat, @@ -169,7 +180,17 @@ void hamilt::DeePKS>::contributeHR() GlobalC::ld.pdm, descriptor, GlobalC::ld.des_per_atom); - GlobalC::ld.cal_gedm(this->ucell->nat, descriptor); + DeePKS_domain::cal_gedm(this->ucell->nat, + GlobalC::ld.lmaxd, + GlobalC::ld.nmaxd, + GlobalC::ld.inlmax, + GlobalC::ld.des_per_atom, + GlobalC::ld.inl_l, + descriptor, + GlobalC::ld.pdm, + GlobalC::ld.model_deepks, + GlobalC::ld.gedm, + GlobalC::ld.E_delta); // // recalculate the H_V_delta // if (this->H_V_delta == nullptr) @@ -444,14 +465,17 @@ void hamilt::DeePKS>::cal_HR_IJR(const double* hr_i } } -inline void get_h_delta_k(int ik, double*& h_delta_k) +template +inline void get_h_delta_k(int ik, TK*& h_delta_k) { - h_delta_k = GlobalC::ld.H_V_delta[ik].data(); - return; -} -inline void get_h_delta_k(int ik, std::complex*& h_delta_k) -{ - h_delta_k = GlobalC::ld.H_V_delta_k[ik].data(); + if constexpr (std::is_same::value) + { + h_delta_k = GlobalC::ld.H_V_delta[ik].data(); + } + else + { + h_delta_k = GlobalC::ld.H_V_delta_k[ik].data(); + } return; } @@ -463,7 +487,7 @@ void hamilt::DeePKS>::contributeHk(int ik) ModuleBase::timer::tick("DeePKS", "contributeHk"); TK* h_delta_k = nullptr; - get_h_delta_k(ik, h_delta_k); + get_h_delta_k(ik, h_delta_k); // set SK to zero and then calculate SK for each k vector ModuleBase::GlobalFunc::ZEROS(h_delta_k, this->hsk->get_size()); diff --git a/source/module_hamilt_lcao/hamilt_lcaodft/operator_lcao/deepks_lcao.h b/source/module_hamilt_lcao/hamilt_lcaodft/operator_lcao/deepks_lcao.h index a2ad622ec5..08e19a6c9a 100644 --- a/source/module_hamilt_lcao/hamilt_lcaodft/operator_lcao/deepks_lcao.h +++ b/source/module_hamilt_lcao/hamilt_lcaodft/operator_lcao/deepks_lcao.h @@ -58,8 +58,6 @@ class DeePKS> : public OperatorLCAO private: elecstate::DensityMatrix* DM; - // LCAO_Deepks* ld = nullptr; - const UnitCell* ucell = nullptr; Grid_Driver* gridD = nullptr; @@ -73,6 +71,8 @@ class DeePKS> : public OperatorLCAO #ifdef __DEEPKS + LCAO_Deepks* ld = nullptr; + /** * @brief initialize HR, search the nearest neighbor atoms * HContainer is used to store the DeePKS real space Hamiltonian correction with specific atom-pairs diff --git a/source/module_hamilt_lcao/module_deepks/CMakeLists.txt b/source/module_hamilt_lcao/module_deepks/CMakeLists.txt index 2db7d5956f..e5e7b05ce9 100644 --- a/source/module_hamilt_lcao/module_deepks/CMakeLists.txt +++ b/source/module_hamilt_lcao/module_deepks/CMakeLists.txt @@ -1,21 +1,20 @@ if(ENABLE_DEEPKS) list(APPEND objects LCAO_deepks.cpp + deepks_basic.cpp deepks_descriptor.cpp deepks_force.cpp deepks_fpre.cpp deepks_spre.cpp deepks_orbital.cpp deepks_orbpre.cpp + deepks_vdelta.cpp deepks_vdpre.cpp deepks_hmat.cpp + deepks_pdm.cpp + deepks_phialpha.cpp LCAO_deepks_io.cpp - LCAO_deepks_pdm.cpp - LCAO_deepks_phialpha.cpp - LCAO_deepks_torch.cpp - LCAO_deepks_vdelta.cpp LCAO_deepks_interface.cpp - cal_gedm.cpp ) add_library( diff --git a/source/module_hamilt_lcao/module_deepks/LCAO_deepks.cpp b/source/module_hamilt_lcao/module_deepks/LCAO_deepks.cpp index cbfb5eb4cb..8d9e363d47 100644 --- a/source/module_hamilt_lcao/module_deepks/LCAO_deepks.cpp +++ b/source/module_hamilt_lcao/module_deepks/LCAO_deepks.cpp @@ -189,7 +189,6 @@ void LCAO_Deepks::init_index(const int ntype, void LCAO_Deepks::allocate_V_delta(const int nat, const int nks) { ModuleBase::TITLE("LCAO_Deepks", "allocate_V_delta"); - nks_V_delta = nks; // initialize the H matrix H_V_delta if (PARAM.globalv.gamma_only_local) @@ -232,7 +231,16 @@ void LCAO_Deepks::allocate_V_delta(const int nat, const int nks) template void LCAO_Deepks::dpks_cal_e_delta_band(const std::vector>& dm, const int nks) { - this->cal_e_delta_band(dm, nks); + std::vector> h_delta; + if constexpr (std::is_same::value) + { + h_delta = this->H_V_delta; + } + else + { + h_delta = this->H_V_delta_k; + } + DeePKS_domain::cal_e_delta_band(dm, h_delta, nks, this->pv, this->e_delta_band); } template void LCAO_Deepks::dpks_cal_e_delta_band(const std::vector>& dm, const int nks); diff --git a/source/module_hamilt_lcao/module_deepks/LCAO_deepks.h b/source/module_hamilt_lcao/module_deepks/LCAO_deepks.h index d2cc460ddb..5b3041e72d 100644 --- a/source/module_hamilt_lcao/module_deepks/LCAO_deepks.h +++ b/source/module_hamilt_lcao/module_deepks/LCAO_deepks.h @@ -3,13 +3,17 @@ #ifdef __DEEPKS +#include "deepks_basic.h" #include "deepks_descriptor.h" #include "deepks_force.h" #include "deepks_fpre.h" #include "deepks_hmat.h" #include "deepks_orbital.h" #include "deepks_orbpre.h" +#include "deepks_pdm.h" +#include "deepks_phialpha.h" #include "deepks_spre.h" +#include "deepks_vdelta.h" #include "deepks_vdpre.h" #include "module_base/complexmatrix.h" #include "module_base/intarray.h" @@ -18,13 +22,11 @@ #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 "module_io/winput.h" #include #include -#include /// /// The LCAO_Deepks contains subroutines for implementation of the DeePKS method in atomic basis. @@ -90,18 +92,21 @@ class LCAO_Deepks // private variables //------------------- // private: - public: // change to public to reconstuct the code, 2024-07-22 by mohan - int lmaxd = 0; // max l of descirptors - int nmaxd = 0; //#. descriptors per l - int inlmax = 0; // tot. number {i,n,l} - atom, n, l - int nat_gdm = 0; - int nks_V_delta = 0; + public: // change to public to reconstuct the code, 2024-07-22 by mohan + int lmaxd = 0; // max l of descirptors + int nmaxd = 0; //#. descriptors per l + int inlmax = 0; // tot. number {i,n,l} - atom, n, l + int n_descriptor; // natoms * des_per_atom, size of descriptor(projector) basis set + int des_per_atom; // \sum_L{Nchi(L)*(2L+1)} + int* inl_l; // inl_l[inl_index] = l of descriptor with inl_index + ModuleBase::IntArray* alpha_index; // seems not used in the code + ModuleBase::IntArray* inl_index; // caoyu add 2021-05-07 - bool init_pdm = false; // for DeePKS NSCF calculation + bool init_pdm = false; // for DeePKS NSCF calculation, set init_pdm to skip the calculation of pdm in SCF iteration // deep neural network module that provides corrected Hamiltonian term and - // related derivatives. - torch::jit::script::Module module; + // related derivatives. Used in cal_gedm. + torch::jit::script::Module model_deepks; // saves and its derivatives // index 0 for itself and index 1-3 for derivatives over x,y,z @@ -112,21 +117,8 @@ class LCAO_Deepks // [nat][nlm*nlm] for equivariant version std::vector pdm; - // gedm:dE/dD, [tot_Inl][2l+1][2l+1] (E: Hartree) - std::vector gedm_tensor; - /// dE/dD, autograd from loaded model(E: Ry) - double** gedm; //[tot_Inl][2l+1][2l+1] - - /// size of descriptor(projector) basis set - int n_descriptor; - - // \sum_L{Nchi(L)*(2L+1)} - int des_per_atom; - - ModuleBase::IntArray* alpha_index; // seems not used in the code - ModuleBase::IntArray* inl_index; // caoyu add 2021-05-07 - int* inl_l; // inl_l[inl_index] = l of descriptor with inl_index + double** gedm; //[tot_Inl][(2l+1)*(2l+1)] // HR status, // true : HR should be calculated @@ -167,141 +159,14 @@ class LCAO_Deepks /// Allocate memory for correction to Hamiltonian void allocate_V_delta(const int nat, const int nks = 1); - 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); - - //------------------- - // LCAO_deepks_phialpha.cpp - //------------------- - - // E.Wu 2024-12-24 - // This file contains 3 subroutines: - // 1. allocate_phialpha, which allocates memory for phialpha - // 2. build_phialpha, which calculates the overlap - // between atomic basis and projector alpha : - // which will be used in calculating pdm, gdmx, H_V_delta, F_delta; - // 3. check_phialpha, which prints the results into .dat files - // for checking - - public: - // calculates - void allocate_phialpha(const bool& cal_deri, - const UnitCell& ucell, - const LCAO_Orbitals& orb, - const Grid_Driver& GridD); - - void build_phialpha(const bool& cal_deri /**< [in] 0 for 2-center intergration, 1 for its derivation*/, - const UnitCell& ucell, - const LCAO_Orbitals& orb, - const Grid_Driver& GridD, - const TwoCenterIntegrator& overlap_orb_alpha); - - void check_phialpha(const bool& cal_deri /**< [in] 0 for 2-center intergration, 1 for its derivation*/, - const UnitCell& ucell, - const LCAO_Orbitals& orb, - const Grid_Driver& GridD); - - //------------------- - // LCAO_deepks_pdm.cpp - //------------------- - - // This file contains subroutines for calculating pdm, - // which is defind as sum_mu,nu rho_mu,nu (); - // as well as gdmx, which is the gradient of pdm, defined as - // sum_mu,nu rho_mu,nu d/dX() - - // It also contains subroutines for printing pdm and gdmx - // for checking purpose - - // There are 2 subroutines in this file: - // 1. cal_projected_DM, which is used for calculating pdm - // 2. check_projected_dm, which prints pdm to descriptor.dat - - public: - /** - * @brief calculate projected density matrix: - * pdm = sum_i,occ - * 3 cases to skip calculation of pdm: - * 1. NSCF calculation of DeePKS, init_chg = file and pdm has been read - * 2. SCF calculation of DeePKS with init_chg = file and pdm has been read for restarting SCF - * 3. Relax/Cell-Relax/MD calculation, non-first step will use the convergence pdm from the last step as initial - * pdm - */ - template - void cal_projected_DM(const elecstate::DensityMatrix* dm, - const UnitCell& ucell, - const LCAO_Orbitals& orb, - const Grid_Driver& GridD); - - void check_projected_dm(); - - /** - * @brief set init_pdm to skip the calculation of pdm in SCF iteration - */ - void set_init_pdm(bool ipdm) - { - this->init_pdm = ipdm; - } - /** - * @brief read pdm from file, do it only once in whole calculation - */ - void read_projected_DM(bool read_pdm_file, bool is_equiv, const Numerical_Orbital& alpha); - - //------------------- - // LCAO_deepks_vdelta.cpp - //------------------- - - // This file contains subroutines related to V_delta, which is the deepks contribution to Hamiltonian - // defined as |alpha>V(D)& dm/**<[in] density matrix*/); - template - void cal_e_delta_band(const std::vector>& dm /**<[in] density matrix*/, const int nks); - - //! a temporary interface for cal_e_delta_band and cal_e_delta_band_k + //! a temporary interface for cal_e_delta_band template void dpks_cal_e_delta_band(const std::vector>& dm, const int nks); - public: - //------------------- - // LCAO_deepks_torch.cpp - //------------------- - - // This file contains interfaces with libtorch, - // including loading of model and calculating gradients - // as well as subroutines that prints the results for checking - - // The file contains 8 subroutines: - // 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)& descriptor); - void check_gedm(); - void cal_gedm_equiv(const int nat, const std::vector& descriptor); - - // calculate gevdm - void cal_gevdm(const int nat, std::vector& gevdm); - 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); + const Parallel_Orbitals* pv; }; 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 a6547f2187..8215e634ca 100644 --- a/source/module_hamilt_lcao/module_deepks/LCAO_deepks_interface.cpp +++ b/source/module_hamilt_lcao/module_deepks/LCAO_deepks_interface.cpp @@ -34,13 +34,40 @@ void LCAO_Deepks_Interface::out_deepks_labels(const double& etot, // define TH for different types using TH = std::conditional_t::value, ModuleBase::matrix, ModuleBase::ComplexMatrix>; + // These variables are frequently used in the following code + const int inlmax = ld->inlmax; + const int lmaxd = ld->lmaxd; + const int des_per_atom = ld->des_per_atom; + const int* inl_l = ld->inl_l; + const ModuleBase::IntArray* inl_index = ld->inl_index; + const std::vector*> phialpha = ld->phialpha; + std::vector pdm = ld->pdm; + const int my_rank = GlobalV::MY_RANK; const int nspin = PARAM.inp.nspin; + // Used for deepks_bandgap == 1 and deepks_v_delta > 0 + std::vector>* h_delta = nullptr; + if constexpr (std::is_same::value) + { + h_delta = &ld->H_V_delta; + } + else + { + h_delta = &ld->H_V_delta_k; + } + // calculating deepks correction to bandgap and save the results if (PARAM.inp.deepks_out_labels) { - // mohan updated 2024-07-25 + // Used for deepks_scf == 1 + std::vector gevdm; + if (PARAM.inp.deepks_scf) + { + DeePKS_domain::cal_gevdm(nat, inlmax, inl_l, pdm, gevdm); + } + + // Energy Part const std::string file_etot = PARAM.globalv.global_out_dir + "deepks_etot.npy"; const std::string file_ebase = PARAM.globalv.global_out_dir + "deepks_ebase.npy"; @@ -57,16 +84,55 @@ void LCAO_Deepks_Interface::out_deepks_labels(const double& etot, LCAO_deepks_io::save_npy_e(etot, file_ebase, my_rank); } - std::vector>* h_delta = nullptr; - if constexpr (std::is_same::value) + // Force Part + if (PARAM.inp.cal_force) { - h_delta = &ld->H_V_delta; + if (PARAM.inp.deepks_scf + && !PARAM.inp.deepks_equiv) // training with force label not supported by equivariant version now + { + std::vector> dm_vec = dm->get_DMK_vector(); + torch::Tensor gdmx; + DeePKS_domain::cal_gdmx< + TK>(lmaxd, inlmax, nks, kvec_d, phialpha, inl_index, dm_vec, ucell, orb, *ParaV, GridD, gdmx); + + torch::Tensor gvx; + DeePKS_domain::cal_gvx(ucell.nat, inlmax, des_per_atom, inl_l, gevdm, gdmx, gvx); + const std::string file_gradvx = PARAM.globalv.global_out_dir + "deepks_gradvx.npy"; + LCAO_deepks_io::save_tensor2npy(file_gradvx, gvx, my_rank); + + if (PARAM.inp.deepks_out_unittest) + { + DeePKS_domain::check_gdmx(gdmx); + DeePKS_domain::check_gvx(gvx); + } + } } - else + + // Stress Part + if (PARAM.inp.cal_stress) { - h_delta = &ld->H_V_delta_k; + if (PARAM.inp.deepks_scf + && !PARAM.inp.deepks_equiv) // training with stress label not supported by equivariant version now + { + std::vector> dm_vec = dm->get_DMK_vector(); + torch::Tensor gdmepsl; + DeePKS_domain::cal_gdmepsl< + TK>(lmaxd, inlmax, nks, kvec_d, phialpha, inl_index, dm_vec, ucell, orb, *ParaV, GridD, gdmepsl); + + torch::Tensor gvepsl; + DeePKS_domain::cal_gvepsl(ucell.nat, inlmax, des_per_atom, inl_l, gevdm, gdmepsl, gvepsl); + const std::string file_gvepsl = PARAM.globalv.global_out_dir + "deepks_gvepsl.npy"; + LCAO_deepks_io::save_tensor2npy(file_gvepsl, gvepsl, my_rank); + + if (PARAM.inp.deepks_out_unittest) + { + DeePKS_domain::check_gdmepsl(gdmepsl); + DeePKS_domain::check_gvepsl(gvepsl); + } + } } + // Bandgap Part if (PARAM.inp.deepks_bandgap) { const int nocc = (PARAM.inp.nelec + 1) / 2; @@ -115,18 +181,16 @@ void LCAO_Deepks_Interface::out_deepks_labels(const double& etot, // 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, + lmaxd, + inlmax, nat, nks, - ld->inl_l, + inl_l, kvec_d, - ld->phialpha, + phialpha, gevdm, - ld->inl_index, + inl_index, ucell, orb, *ParaV, @@ -135,12 +199,9 @@ void LCAO_Deepks_Interface::out_deepks_labels(const double& etot, 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, - orbital_precalc, - PARAM.globalv.global_out_dir, - my_rank); + const std::string file_orbpre = PARAM.globalv.global_out_dir + "deepks_orbpre.npy"; + LCAO_deepks_io::save_tensor2npy(file_orbpre, orbital_precalc, my_rank); + const std::string file_obase = PARAM.globalv.global_out_dir + "deepks_obase.npy"; std::vector o_base(nks); for (int iks = 0; iks < nks; ++iks) @@ -156,7 +217,7 @@ void LCAO_Deepks_Interface::out_deepks_labels(const double& etot, } // end deepks_scf == 0 } // end bandgap label - // save H(R) matrix + // H(R) matrix part, not realized now if (true) // should be modified later! { const std::string file_hr = PARAM.globalv.global_out_dir + "deepks_hr.npy"; @@ -165,6 +226,7 @@ void LCAO_Deepks_Interface::out_deepks_labels(const double& etot, // How to save H(R)? } + // H(k) matrix part if (PARAM.inp.deepks_v_delta) { std::vector h_tot(nks); @@ -209,70 +271,38 @@ void LCAO_Deepks_Interface::out_deepks_labels(const double& etot, if (PARAM.inp.deepks_v_delta == 1) // v_delta_precalc storage method 1 { - std::vector gevdm; - ld->cal_gevdm(nat, gevdm); - torch::Tensor v_delta_precalc; DeePKS_domain::cal_v_delta_precalc(nlocal, - ld->lmaxd, - ld->inlmax, + lmaxd, + inlmax, nat, nks, - ld->inl_l, + inl_l, kvec_d, - ld->phialpha, + phialpha, gevdm, - ld->inl_index, + inl_index, ucell, orb, *ParaV, GridD, v_delta_precalc); - LCAO_deepks_io::save_npy_v_delta_precalc(nat, - nks, - nlocal, - ld->des_per_atom, - v_delta_precalc, - PARAM.globalv.global_out_dir, - my_rank); + const std::string file_vdpre = PARAM.globalv.global_out_dir + "deepks_vdpre.npy"; + LCAO_deepks_io::save_tensor2npy(file_vdpre, v_delta_precalc, my_rank); } else if (PARAM.inp.deepks_v_delta == 2) // v_delta_precalc storage method 2 { torch::Tensor phialpha_out; - DeePKS_domain::prepare_phialpha(nlocal, - ld->lmaxd, - ld->inlmax, - nat, - nks, - kvec_d, - ld->phialpha, - ucell, - orb, - *ParaV, - GridD, - phialpha_out); - - LCAO_deepks_io::save_npy_phialpha(nat, - nks, - nlocal, - ld->inlmax, - ld->lmaxd, - phialpha_out, - PARAM.globalv.global_out_dir, - my_rank); - std::vector gevdm; - ld->cal_gevdm(nat, gevdm); + DeePKS_domain::prepare_phialpha< + TK>(nlocal, lmaxd, inlmax, nat, nks, kvec_d, phialpha, ucell, orb, *ParaV, GridD, phialpha_out); + const std::string file_phialpha = PARAM.globalv.global_out_dir + "deepks_phialpha.npy"; + LCAO_deepks_io::save_tensor2npy(file_phialpha, phialpha_out, my_rank); torch::Tensor gevdm_out; - DeePKS_domain::prepare_gevdm(nat, ld->lmaxd, ld->inlmax, orb, gevdm, gevdm_out); - - LCAO_deepks_io::save_npy_gevdm(nat, - ld->inlmax, - ld->lmaxd, - gevdm_out, - PARAM.globalv.global_out_dir, - my_rank); + DeePKS_domain::prepare_gevdm(nat, lmaxd, inlmax, orb, gevdm, gevdm_out); + const std::string file_gevdm = PARAM.globalv.global_out_dir + "deepks_gevdm.npy"; + LCAO_deepks_io::save_tensor2npy(file_gevdm, gevdm_out, my_rank); } } else // deepks_scf == 0 @@ -292,31 +322,23 @@ void LCAO_Deepks_Interface::out_deepks_labels(const double& etot, // when deepks_scf is on, the init pdm should be same as the out pdm, so we should not recalculate the pdm if (!PARAM.inp.deepks_scf) { - ld->cal_projected_DM(dm, ucell, orb, GridD); + DeePKS_domain::cal_pdm< + TK>(ld->init_pdm, inlmax, lmaxd, inl_l, inl_index, dm, phialpha, ucell, orb, GridD, *ParaV, pdm); } - ld->check_projected_dm(); // print out the projected dm for NSCF calculaiton + DeePKS_domain::check_pdm(inlmax, inl_l, pdm); // print out the projected dm for NSCF calculaiton std::vector descriptor; - DeePKS_domain::cal_descriptor(nat, - ld->inlmax, - ld->inl_l, - ld->pdm, - descriptor, - ld->des_per_atom); // final descriptor - DeePKS_domain::check_descriptor(ld->inlmax, - ld->des_per_atom, - ld->inl_l, - ucell, - PARAM.globalv.global_out_dir, - descriptor); + DeePKS_domain::cal_descriptor(nat, inlmax, inl_l, pdm, descriptor, + des_per_atom); // final descriptor + DeePKS_domain::check_descriptor(inlmax, des_per_atom, inl_l, ucell, PARAM.globalv.global_out_dir, descriptor); if (PARAM.inp.deepks_out_labels) { LCAO_deepks_io::save_npy_d(nat, - ld->des_per_atom, - ld->inlmax, - ld->inl_l, + des_per_atom, + inlmax, + inl_l, PARAM.inp.deepks_equiv, descriptor, PARAM.globalv.global_out_dir, @@ -327,12 +349,25 @@ void LCAO_Deepks_Interface::out_deepks_labels(const double& etot, /// print out deepks information to the screen if (PARAM.inp.deepks_scf) { - ld->cal_e_delta_band(dm->get_DMK_vector(), nks); + DeePKS_domain::cal_e_delta_band(dm->get_DMK_vector(), *h_delta, nks, ParaV, ld->e_delta_band); std::cout << "E_delta_band = " << std::setprecision(8) << ld->e_delta_band << " Ry" << " = " << std::setprecision(8) << ld->e_delta_band * ModuleBase::Ry_to_eV << " eV" << std::endl; std::cout << "E_delta_NN = " << std::setprecision(8) << ld->E_delta << " Ry" << " = " << std::setprecision(8) << ld->E_delta * ModuleBase::Ry_to_eV << " eV" << std::endl; + if (PARAM.inp.deepks_out_unittest) + { + LCAO_deepks_io::print_dm(nks, PARAM.globalv.nlocal, ParaV->nrow, dm->get_DMK_vector()); + + DeePKS_domain::check_gedm(inlmax, inl_l, ld->gedm); + + std::ofstream ofs("E_delta_bands.dat"); + ofs << std::setprecision(10) << ld->e_delta_band; + + std::ofstream ofs1("E_delta.dat"); + ofs1 << std::setprecision(10) << ld->E_delta; + } } + ModuleBase::timer::tick("LCAO_Deepks_Interface", "out_deepks_labels"); } template class LCAO_Deepks_Interface; diff --git a/source/module_hamilt_lcao/module_deepks/LCAO_deepks_interface.h b/source/module_hamilt_lcao/module_deepks/LCAO_deepks_interface.h index 695f942990..13075bcf4e 100644 --- a/source/module_hamilt_lcao/module_deepks/LCAO_deepks_interface.h +++ b/source/module_hamilt_lcao/module_deepks/LCAO_deepks_interface.h @@ -6,6 +6,7 @@ #include "module_base/complexmatrix.h" #include "module_base/matrix.h" #include "module_hamilt_lcao/hamilt_lcaodft/hamilt_lcao.h" + #include template @@ -26,11 +27,9 @@ class LCAO_Deepks_Interface /// @param[in] orb /// @param[in] GridD /// @param[in] ParaV - /// @param[in] psi /// @param[in] psid - /// @param[in] dm_gamma - /// @param[in] dm_k - // for Gamma-only + /// @param[in] dm + /// @param[in] p_ham void out_deepks_labels(const double& etot, const int& nks, const int& nat, 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 c0b1b3cd93..28e5268891 100644 --- a/source/module_hamilt_lcao/module_deepks/LCAO_deepks_io.cpp +++ b/source/module_hamilt_lcao/module_deepks/LCAO_deepks_io.cpp @@ -1,26 +1,4 @@ -// 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 #ifdef __DEEPKS @@ -156,89 +134,6 @@ void LCAO_deepks_io::save_npy_d(const int nat, return; } -// 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, - const std::string& out_dir, - const int rank) -{ - ModuleBase::TITLE("LCAO_deepks_io", "save_npy_gvx"); - - if (rank != 0) - { - return; - } - - 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)}; - - std::vector npy_gvx; - auto accessor = gvx.accessor(); - for (int ibt = 0; ibt < nat; ++ibt) - { - for (int i = 0; i < 3; i++) - { - for (int iat = 0; iat < nat; ++iat) - { - for (int p = 0; p < des_per_atom; ++p) - { - npy_gvx.push_back(accessor[ibt][i][iat][p]); - } - } - } - } - - std::string file_gradvx = out_dir + "deepks_gradvx.npy"; - npy::SaveArrayAsNumpy(file_gradvx, false, 4, gshape, npy_gvx); - return; -} - -// saves gvx into grad_vepsl.npy -void LCAO_deepks_io::save_npy_gvepsl(const int nat, - const int des_per_atom, - const torch::Tensor& gvepsl, - const std::string& out_dir, - const int rank) -{ - ModuleBase::TITLE("LCAO_deepks_io", "save_npy_gvepsl"); - - if (rank != 0) - { - return; - } - - // 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; - auto accessor = gvepsl.accessor(); - - for (int i = 0; i < 6; i++) - { - for (int ibt = 0; ibt < nat; ++ibt) - { - - for (int p = 0; p < des_per_atom; ++p) - { - npy_gvepsl.push_back(accessor[i][ibt][p]); - } - } - } - - // change the name from grad_vepsl.npy to deepks_gvepsl.npy - const std::string file = out_dir + "deepks_gvepsl.npy"; - npy::SaveArrayAsNumpy(file, false, 3, gshape, npy_gvepsl); - return; -} - // saves energy in numpy format void LCAO_deepks_io::save_npy_e(const double& e, const std::string& e_file, const int rank) { @@ -248,16 +143,14 @@ void LCAO_deepks_io::save_npy_e(const double& e, const std::string& e_file, cons return; } - // save e_base + // save energy in .npy format const long unsigned eshape[] = {1}; - std::vector npy_e; - npy_e.push_back(e); - npy::SaveArrayAsNumpy(e_file, false, 1, eshape, npy_e); + npy::SaveArrayAsNumpy(e_file, false, 1, eshape, &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) +void LCAO_deepks_io::save_npy_f(const ModuleBase::matrix& f, const std::string& f_file, const int rank) { ModuleBase::TITLE("LCAO_deepks_io", "save_npy_f"); @@ -266,20 +159,9 @@ void LCAO_deepks_io::save_npy_f(const ModuleBase::matrix& f, const std::string& return; } - assert(nat > 0); - - // 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 i = 0; i < 3; i++) - { - npy_f.push_back(f(iat, i)); - } - } - npy::SaveArrayAsNumpy(f_file, false, 2, fshape, npy_f); + // save force in unit: Rydberg/Bohr + const long unsigned fshape[] = {static_cast(f.nr), static_cast(f.nc)}; + npy::SaveArrayAsNumpy(f_file, false, 2, fshape, f.c); return; } @@ -326,42 +208,6 @@ void LCAO_deepks_io::save_npy_o(const std::vector& bandgap, return; } -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; - } - - // 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; - auto accessor = orbital_precalc.accessor(); - for (int iks = 0; iks < nks; ++iks) - { - for (int iat = 0; iat < nat; ++iat) - { - for (int p = 0; p < des_per_atom; ++p) - { - npy_orbital_precalc.push_back(accessor[iks][iat][p]); - } - } - } - - const std::string file_orbpre = out_dir + "deepks_orbpre.npy"; - npy::SaveArrayAsNumpy(file_orbpre, false, 3, gshape, npy_orbital_precalc); - return; -} - template void LCAO_deepks_io::save_npy_h(const std::vector& hamilt, const std::string& h_file, @@ -394,169 +240,61 @@ void LCAO_deepks_io::save_npy_h(const std::vector& hamilt, return; } -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, - const std::string& out_dir, - const int rank) +void LCAO_deepks_io::save_matrix2npy(const std::string& file_name, + const ModuleBase::matrix& matrix, + const int rank, + const double& scale) { - ModuleBase::TITLE("LCAO_deepks_io", "save_npy_v_delta_precalc"); + ModuleBase::TITLE("LCAO_deepks_io", "save_matrix2npy"); + if (rank != 0) { return; } - // timeval t_start; - // gettimeofday(&t_start,NULL); - // 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; - auto accessor - = v_delta_precalc - .accessor::value, double, c10::complex>, 5>(); - for (int iks = 0; iks < nks; ++iks) + const unsigned long shape[] = {static_cast(matrix.nr), static_cast(matrix.nc)}; + + std::vector scaled_data(matrix.nr * matrix.nc); + for (int i = 0; i < matrix.nr * matrix.nc; ++i) { - for (int mu = 0; mu < nlocal; ++mu) - { - for (int nu = 0; nu < nlocal; ++nu) - { - for (int iat = 0; iat < nat; ++iat) - { - for (int p = 0; p < des_per_atom; ++p) - { - if constexpr (std::is_same::value) - { - npy_v_delta_precalc.push_back(accessor[iks][mu][nu][iat][p]); - } - else - { - c10::complex tmp_c10 = accessor[iks][mu][nu][iat][p]; - std::complex tmp = std::complex(tmp_c10.real(), tmp_c10.imag()); - npy_v_delta_precalc.push_back(tmp); - } - } - } - } - } + scaled_data[i] = matrix.c[i] * scale; } - const std::string file_vdpre = out_dir + "deepks_vdpre.npy"; - npy::SaveArrayAsNumpy(file_vdpre, false, 5, gshape, npy_v_delta_precalc); + + npy::SaveArrayAsNumpy(file_name, false, 2, shape, scaled_data.data()); return; } -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 std::string& out_dir, - const int rank) +template +void LCAO_deepks_io::save_tensor2npy(const std::string& file_name, const torch::Tensor& tensor, 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; - const long unsigned gshape[] = {static_cast(nat), - static_cast(nlmax), - static_cast(nks), - static_cast(nlocal), - static_cast(mmax)}; - std::vector npy_phialpha; - auto accessor - = phialpha_tensor - .accessor::value, double, c10::complex>, 5>(); - for (int iat = 0; iat < nat; iat++) + ModuleBase::TITLE("LCAO_deepks_io", "save_tensor2npy"); + const int dim = tensor.dim(); + std::vector shape(dim); + for (int i = 0; i < dim; i++) { - for (int nl = 0; nl < nlmax; nl++) - { - for (int iks = 0; iks < nks; iks++) - { - for (int mu = 0; mu < nlocal; mu++) - { - for (int m = 0; m < mmax; m++) - { - if constexpr (std::is_same::value) - { - npy_phialpha.push_back(accessor[iat][nl][iks][mu][m]); - } - else - { - c10::complex tmp_c10 = accessor[iat][nl][iks][mu][m]; - std::complex tmp = std::complex(tmp_c10.real(), tmp_c10.imag()); - npy_phialpha.push_back(tmp); - } - } - } - } - } + shape[i] = tensor.size(i); } - const std::string file_phialpha = out_dir + "deepks_phialpha.npy"; - npy::SaveArrayAsNumpy(file_phialpha, false, 5, gshape, npy_phialpha); - return; -} -void LCAO_deepks_io::save_npy_gevdm(const int nat, - const int inlmax, - const int lmaxd, - const torch::Tensor& gevdm, - const std::string& out_dir, - const int rank) -{ - ModuleBase::TITLE("LCAO_deepks_io", "save_npy_gevdm"); - if (rank != 0) + std::vector data(tensor.numel()); + + if constexpr (std::is_same::value) { - return; + std::memcpy(data.data(), tensor.data_ptr(), tensor.numel() * sizeof(double)); } - - 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; - const long unsigned gshape[] = {static_cast(nat), - static_cast(nlmax), - static_cast(mmax), - static_cast(mmax), - static_cast(mmax)}; - std::vector npy_gevdm; - auto accessor = gevdm.accessor(); - for (int iat = 0; iat < nat; iat++) + else { - for (int nl = 0; nl < nlmax; nl++) + auto tensor_data = tensor.data_ptr>(); + for (size_t i = 0; i < tensor.numel(); ++i) { - for (int v = 0; v < mmax; v++) - { - for (int m = 0; m < mmax; m++) - { - for (int n = 0; n < mmax; n++) - { - npy_gevdm.push_back(accessor[iat][nl][v][m][n]); - } - } - } + data[i] = std::complex(tensor_data[i].real(), tensor_data[i].imag()); } } - const std::string file_gevdm = out_dir + "deepks_gevdm.npy"; - npy::SaveArrayAsNumpy(file_gevdm, false, 5, gshape, npy_gevdm); - return; + + npy::SaveArrayAsNumpy(file_name, false, shape.size(), shape.data(), data); } template void LCAO_deepks_io::print_dm(const int nks, @@ -581,39 +319,11 @@ template void LCAO_deepks_io::save_npy_h>(const std::vector const int nks, 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 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 std::string& out_dir, - const int rank); +template void LCAO_deepks_io::save_tensor2npy(const std::string& file_name, + const torch::Tensor& tensor, + const int rank); +template void LCAO_deepks_io::save_tensor2npy>(const std::string& file_name, + const torch::Tensor& tensor, + const int rank); #endif 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 8bd3134007..3d27c1df75 100644 --- a/source/module_hamilt_lcao/module_deepks/LCAO_deepks_io.h +++ b/source/module_hamilt_lcao/module_deepks/LCAO_deepks_io.h @@ -26,18 +26,14 @@ namespace LCAO_deepks_io /// 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 +/// 1. save_npy_d : descriptor -> deepks_dm_eig.npy +/// 2. save_npy_e : energy +/// 3. save_npy_f : force +/// 4. save_npy_s : stress +/// 5. save_npy_o: orbital +/// 6. save_npy_h : Hamiltonian +/// 7. save_matrix2npy : ModuleBase::matrix -> .npy +/// 8. save_tensor2npy : torch::Tensor -> .npy /// print density matrices template @@ -60,43 +56,23 @@ 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 +// save force 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, - const std::string& out_dir, - const int rank); - -// save stress and gvepsl +// save stress 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, - const std::string& out_dir, - const int rank); - -/// save orbital and orbital_precalc +/// save orbital 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); -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, @@ -105,32 +81,13 @@ void save_npy_h(const std::vector& hamilt, 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, - const std::string& out_dir, - const int rank); +void save_matrix2npy(const std::string& file_name, + const ModuleBase::matrix& matrix, + const int rank, + const double& scale = 1.0); -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); - -// Always real, no need for template now -void save_npy_gevdm(const int nat, - const int inlmax, - const int lmaxd, - const torch::Tensor& gevdm, - const std::string& out_dir, - const int rank); +template +void save_tensor2npy(const std::string& file_name, const torch::Tensor& tensor, const int rank); }; // namespace LCAO_deepks_io #endif diff --git a/source/module_hamilt_lcao/module_deepks/LCAO_deepks_torch.cpp b/source/module_hamilt_lcao/module_deepks/LCAO_deepks_torch.cpp deleted file mode 100644 index 07b5a3fe7f..0000000000 --- a/source/module_hamilt_lcao/module_deepks/LCAO_deepks_torch.cpp +++ /dev/null @@ -1,84 +0,0 @@ -// This file contains interfaces with libtorch, -// including loading of model and calculating gradients -// as well as subroutines that prints the results for checking - -// The file contains 3 subroutines: - -// cal_gevdm : d(des)/d(pdm) -// calculated using torch::autograd::grad -// load_model : loads model for applying V_delta - -#ifdef __DEEPKS - -#include "LCAO_deepks.h" -#include "LCAO_deepks_io.h" // mohan add 2024-07-22 -#include "module_base/blas_connector.h" -#include "module_base/constants.h" -#include "module_base/libm/libm.h" -#include "module_base/parallel_reduce.h" -#include "module_hamilt_lcao/module_hcontainer/atom_pair.h" -#include "module_parameter/parameter.h" - -// 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_gevdm"); - if (!gevdm.empty()) - { - gevdm.erase(gevdm.begin(), gevdm.end()); - } - // cal gevdm(d(EigenValue(D))/dD) - int nlmax = inlmax / nat; - for (int nl = 0; nl < nlmax; ++nl) - { - std::vector avmmv; - for (int iat = 0; iat < nat; ++iat) - { - 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[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)); - std::vector tmp_rpt; // repeated-pdm-tensor (x) - std::vector tmp_rdt; // repeated-d-tensor (y) - std::vector tmp_gst; // gvx-shell - tmp_rpt.push_back(tmp_x); - tmp_rdt.push_back(tmp_y); - tmp_gst.push_back(tmp_yshell); - std::vector tmp_res; - tmp_res = torch::autograd::grad(tmp_rdt, - tmp_rpt, - tmp_gst, - false, - false, - /*allow_unused*/ true); // nm(v)**nm*nm - avmmv.push_back(tmp_res[0]); - } - torch::Tensor avmm = torch::stack(avmmv, 0); // nat*nv**nm*nm - gevdm.push_back(avmm); - } - assert(gevdm.size() == nlmax); - return; -} - -void LCAO_Deepks::load_model(const std::string& deepks_model) -{ - ModuleBase::TITLE("LCAO_Deepks", "load_model"); - - try - { - this->module = torch::jit::load(deepks_model); - } - catch (const c10::Error& e) - - { - std::cerr << "error loading the model" << std::endl; - return; - } - return; -} - -#endif diff --git a/source/module_hamilt_lcao/module_deepks/LCAO_deepks_vdelta.cpp b/source/module_hamilt_lcao/module_deepks/LCAO_deepks_vdelta.cpp deleted file mode 100644 index 9bf6b08c21..0000000000 --- a/source/module_hamilt_lcao/module_deepks/LCAO_deepks_vdelta.cpp +++ /dev/null @@ -1,79 +0,0 @@ -//This file contains subroutines related to V_delta, which is the deepks contribution to Hamiltonian -//defined as |alpha>V(D) -void LCAO_Deepks::cal_e_delta_band(const std::vector>& dm, const int nks) -{ - ModuleBase::TITLE("LCAO_Deepks", "cal_e_delta_band"); - ModuleBase::timer::tick("LCAO_Deepks","cal_e_delta_band"); - TK e_delta_band_tmp = TK(0); - for (int i = 0; i < PARAM.globalv.nlocal; ++i) - { - for (int j = 0; j < PARAM.globalv.nlocal; ++j) - { - const int mu = pv->global2local_row(j); - const int nu = pv->global2local_col(i); - - if (mu >= 0 && nu >= 0) - { - int iic; - if (ModuleBase::GlobalFunc::IS_COLUMN_MAJOR_KS_SOLVER(PARAM.inp.ks_solver)) - { - iic = mu + nu * pv->nrow; - } - else - { - iic = mu * pv->ncol + nu; - } - if constexpr (std::is_same::value) - { - for (int is = 0; is < dm.size(); ++is) //dm.size() == PARAM.inp.nspin - { - e_delta_band_tmp += dm[is][nu * this->pv->nrow + mu] * this->H_V_delta[0][iic]; - } - } - else - { - for(int ik = 0; ik < nks; ik++) - { - e_delta_band_tmp += dm[ik][nu * this->pv->nrow + mu] * this->H_V_delta_k[ik][iic]; - } - } - - } - } - } - if constexpr (std::is_same::value) - { - this->e_delta_band = e_delta_band_tmp; - } - else - { - this->e_delta_band = e_delta_band_tmp.real(); - } -#ifdef __MPI - Parallel_Reduce::reduce_all(this->e_delta_band); -#endif - ModuleBase::timer::tick("LCAO_Deepks","cal_e_delta_band"); - return; -} - -template void LCAO_Deepks::cal_e_delta_band(const std::vector>& dm, const int nks); -template void LCAO_Deepks::cal_e_delta_band>(const std::vector>>& dm, const int nks); - -#endif diff --git a/source/module_hamilt_lcao/module_deepks/cal_gedm.cpp b/source/module_hamilt_lcao/module_deepks/cal_gedm.cpp deleted file mode 100644 index 3f0b2cf025..0000000000 --- a/source/module_hamilt_lcao/module_deepks/cal_gedm.cpp +++ /dev/null @@ -1,170 +0,0 @@ -/// cal_gedm : calculates d(E_delta)/d(pdm) -/// this is the term V(D) that enters the expression H_V_delta = |alpha>V(D)& descriptor) -{ - 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, - descriptor, - 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) - { - std::string cmd = "python cal_gedm.py " + PARAM.inp.deepks_model; - int stat = std::system(cmd.c_str()); - assert(stat == 0); - } - - MPI_Barrier(MPI_COMM_WORLD); - - 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, const std::vector& descriptor) -{ - - if (PARAM.inp.deepks_equiv) - { - this->cal_gedm_equiv(nat, descriptor); - return; - } - - ModuleBase::TITLE("LCAO_Deepks", "cal_gedm"); - - // forward - std::vector inputs; - - // input_dim:(natom, des_per_atom) - inputs.push_back(torch::cat(descriptor, 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() * 2; // Ry; *2 is for Hartree to Ry - - // cal gedm - std::vector gedm_shell; - gedm_shell.push_back(torch::ones_like(ec[0])); - this->gedm_tensor = torch::autograd::grad(ec, - this->pdm, - gedm_shell, - /*retain_grad=*/true, - /*create_graph=*/false, - /*allow_unused=*/true); - - // gedm_tensor(Hartree) to gedm(Ry) - for (int inl = 0; inl < inlmax; ++inl) - { - int nm = 2 * inl_l[inl] + 1; - auto accessor = this->gedm_tensor[inl].accessor(); - 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] = accessor[m1][m2] * 2; - } - } - } - return; -} - -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; - } -} - -#endif diff --git a/source/module_hamilt_lcao/module_deepks/deepks_basic.cpp b/source/module_hamilt_lcao/module_deepks/deepks_basic.cpp new file mode 100644 index 0000000000..0ff45425a9 --- /dev/null +++ b/source/module_hamilt_lcao/module_deepks/deepks_basic.cpp @@ -0,0 +1,239 @@ +// This file contains interfaces with libtorch, +// including loading of model and calculating gradients +// as well as subroutines that prints the results for checking + +#ifdef __DEEPKS +#include "deepks_basic.h" +#include "module_parameter/parameter.h" + +// d(Descriptor) / d(projected density matrix) +// Dimension is different for each inl, so there's a vector of tensors +void DeePKS_domain::cal_gevdm(const int nat, + const int inlmax, + const int* inl_l, + const std::vector& pdm, + std::vector& gevdm) +{ + ModuleBase::TITLE("DeePKS_domain", "cal_gevdm"); + // cal gevdm(d(EigenValue(D))/dD) + int nlmax = inlmax / nat; + for (int nl = 0; nl < nlmax; ++nl) + { + std::vector avmmv; + for (int iat = 0; iat < nat; ++iat) + { + int inl = iat * nlmax + nl; + int nm = 2 * inl_l[inl] + 1; + // repeat each block for nm times in an additional dimension + torch::Tensor tmp_x = 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)); + std::vector tmp_rpt; // repeated-pdm-tensor (x) + std::vector tmp_rdt; // repeated-d-tensor (y) + std::vector tmp_gst; // gvx-shell + tmp_rpt.push_back(tmp_x); + tmp_rdt.push_back(tmp_y); + tmp_gst.push_back(tmp_yshell); + std::vector tmp_res; + tmp_res = torch::autograd::grad(tmp_rdt, + tmp_rpt, + tmp_gst, + false, + false, + /*allow_unused*/ true); // nm(v)**nm*nm + avmmv.push_back(tmp_res[0]); + } + torch::Tensor avmm = torch::stack(avmmv, 0); // nat*nv**nm*nm + gevdm.push_back(avmm); + } + assert(gevdm.size() == nlmax); + return; +} + +void DeePKS_domain::load_model(const std::string& model_file, torch::jit::script::Module& model) +{ + ModuleBase::TITLE("DeePKS_domain", "load_model"); + + try + { + model = torch::jit::load(model_file); + } + catch (const c10::Error& e) + { + std::cerr << "error loading the model" << std::endl; + return; + } + return; +} + +inline void generate_py_files(const int lmaxd, const int nmaxd, const std::string& out_dir) +{ + + if (GlobalV::MY_RANK != 0) + { + return; + } + + std::ofstream ofs("cal_gedm.py"); + ofs << "import torch" << std::endl; + ofs << "import numpy as np" << std::endl << std::endl; + ofs << "import sys" << std::endl; + + ofs << "from deepks.scf.enn.scf import BasisInfo" << std::endl; + ofs << "from deepks.iterate.template_abacus import t_make_pdm" << std::endl; + ofs << "from deepks.utils import load_yaml" << std::endl << std::endl; + + ofs << "basis = load_yaml('basis.yaml')['proj_basis']" << std::endl; + ofs << "model = torch.jit.load(sys.argv[1])" << std::endl; + ofs << "dm_eig = np.expand_dims(np.load('" << out_dir << "dm_eig.npy'),0)" << std::endl; + ofs << "dm_eig = torch.tensor(dm_eig, " + "dtype=torch.float64,requires_grad=True)" + << std::endl + << std::endl; + + ofs << "dm_flat,basis_info = t_make_pdm(dm_eig,basis)" << std::endl; + ofs << "ec = model(dm_flat.double())" << std::endl; + ofs << "gedm = " + "torch.autograd.grad(ec,dm_eig,grad_outputs=torch.ones_like(ec))[0]" + << std::endl + << std::endl; + + ofs << "np.save('ec.npy',ec.double().detach().numpy())" << std::endl; + ofs << "np.save('gedm.npy',gedm.double().numpy())" << std::endl; + ofs.close(); + + ofs.open("basis.yaml"); + ofs << "proj_basis:" << std::endl; + for (int l = 0; l < lmaxd + 1; l++) + { + ofs << " - - " << l << std::endl; + ofs << " - ["; + for (int i = 0; i < nmaxd + 1; i++) + { + ofs << "0"; + if (i != nmaxd) + { + ofs << ", "; + } + } + ofs << "]" << std::endl; + } +} + +void DeePKS_domain::cal_gedm_equiv(const int nat, + const int lmaxd, + const int nmaxd, + const int inlmax, + const int des_per_atom, + const int* inl_l, + const std::vector& descriptor, + double** gedm, + double& E_delta) +{ + ModuleBase::TITLE("DeePKS_domain", "cal_gedm_equiv"); + + LCAO_deepks_io::save_npy_d(nat, + des_per_atom, + inlmax, + inl_l, + PARAM.inp.deepks_equiv, + descriptor, + PARAM.globalv.global_out_dir, + GlobalV::MY_RANK); // libnpy needed + + generate_py_files(lmaxd, nmaxd, PARAM.globalv.global_out_dir); + + 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); + } + + MPI_Barrier(MPI_COMM_WORLD); + + LCAO_deepks_io::load_npy_gedm(nat, des_per_atom, gedm, 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 +// E_delta is also calculated here +void DeePKS_domain::cal_gedm(const int nat, + const int lmaxd, + const int nmaxd, + const int inlmax, + const int des_per_atom, + const int* inl_l, + const std::vector& descriptor, + const std::vector& pdm, + torch::jit::script::Module& model_deepks, + double** gedm, + double& E_delta) +{ + if (PARAM.inp.deepks_equiv) + { + DeePKS_domain::cal_gedm_equiv(nat, lmaxd, nmaxd, inlmax, des_per_atom, inl_l, descriptor, gedm, E_delta); + return; + } + ModuleBase::TITLE("DeePKS_domain", "cal_gedm"); + + // forward + std::vector inputs; + + // input_dim:(natom, des_per_atom) + inputs.push_back(torch::cat(descriptor, 0).reshape({1, nat, des_per_atom})); + std::vector ec; + ec.push_back(model_deepks.forward(inputs).toTensor()); // Hartree + E_delta = ec[0].item() * 2; // Ry; *2 is for Hartree to Ry + + // cal gedm + std::vector gedm_shell; + gedm_shell.push_back(torch::ones_like(ec[0])); + std::vector gedm_tensor = torch::autograd::grad(ec, + pdm, + gedm_shell, + /*retain_grad=*/true, + /*create_graph=*/false, + /*allow_unused=*/true); + + // gedm_tensor(Hartree) to gedm(Ry) + for (int inl = 0; inl < inlmax; ++inl) + { + int nm = 2 * inl_l[inl] + 1; + auto accessor = gedm_tensor[inl].accessor(); + for (int m1 = 0; m1 < nm; ++m1) + { + for (int m2 = 0; m2 < nm; ++m2) + { + int index = m1 * nm + m2; + gedm[inl][index] = accessor[m1][m2] * 2; //*2 is for Hartree to Ry + } + } + } + return; +} + +void DeePKS_domain::check_gedm(const int inlmax, const int* inl_l, double** 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 << gedm[inl][index] << " "; + } + } + ofs << std::endl; + } +} + +#endif diff --git a/source/module_hamilt_lcao/module_deepks/deepks_basic.h b/source/module_hamilt_lcao/module_deepks/deepks_basic.h new file mode 100644 index 0000000000..9bbaa5b553 --- /dev/null +++ b/source/module_hamilt_lcao/module_deepks/deepks_basic.h @@ -0,0 +1,62 @@ +#ifndef DEEPKS_BASIC_H +#define DEEPKS_BASIC_H + +#ifdef __DEEPKS +#include "LCAO_deepks_io.h" +#include "module_base/parallel_reduce.h" +#include "module_base/tool_title.h" + +#include +#include + +namespace DeePKS_domain +{ +//------------------------ +// deepks_basic.cpp +//------------------------ + +// The file contains 2 subroutines: +// 1. load_model : loads model for applying V_delta +// 2. cal_gevdm : d(des)/d(pdm), calculated using torch::autograd::grad +// 3. cal_gedm : calculates d(E_delta)/d(pdm) +// this is the term V(D) that enters the expression H_V_delta = |alpha>V(D)& pdm, + std::vector& gevdm); + +/// calculate partial of energy correction to descriptors +void cal_gedm(const int nat, + const int lmaxd, + const int nmaxd, + const int inlmax, + const int des_per_atom, + const int* inl_l, + const std::vector& descriptor, + const std::vector& pdm, + torch::jit::script::Module& model_deepks, + double** gedm, + double& E_delta); +void check_gedm(const int inlmax, const int* inl_l, double** gedm); +void cal_gedm_equiv(const int nat, + const int lmaxd, + const int nmaxd, + const int inlmax, + const int des_per_atom, + const int* inl_l, + const std::vector& descriptor, + double** gedm, + double& E_delta); + +} // namespace DeePKS_domain +#endif +#endif \ No newline at end of file diff --git a/source/module_hamilt_lcao/module_deepks/deepks_force.cpp b/source/module_hamilt_lcao/module_deepks/deepks_force.cpp index c7eea52036..be2f716661 100644 --- a/source/module_hamilt_lcao/module_deepks/deepks_force.cpp +++ b/source/module_hamilt_lcao/module_deepks/deepks_force.cpp @@ -112,7 +112,7 @@ void DeePKS_domain::cal_f_delta(const std::vector>& dm, hamilt::AtomPair dm_pair(ibt1, ibt2, dRx, dRy, dRz, &pv); - dm_pair.allocate(nullptr, 1); + dm_pair.allocate(nullptr, true); if constexpr (std::is_same::value) // for gamma-only { @@ -318,7 +318,7 @@ void DeePKS_domain::cal_f_delta(const std::vector>& dm, // prints forces and stress from DeePKS (LCAO) void DeePKS_domain::check_f_delta(const int nat, ModuleBase::matrix& f_delta, ModuleBase::matrix& svnl_dalpha) { - ModuleBase::TITLE("LCAO_Deepks", "check_F_delta"); + ModuleBase::TITLE("DeePKS_domain", "check_F_delta"); std::ofstream ofs("F_delta.dat"); ofs << std::setprecision(10); diff --git a/source/module_hamilt_lcao/module_deepks/deepks_orbpre.cpp b/source/module_hamilt_lcao/module_deepks/deepks_orbpre.cpp index 774599b6a2..f49f9f586b 100644 --- a/source/module_hamilt_lcao/module_deepks/deepks_orbpre.cpp +++ b/source/module_hamilt_lcao/module_deepks/deepks_orbpre.cpp @@ -341,7 +341,7 @@ void DeePKS_domain::cal_orbital_precalc(const std::vector& dm_hl, } orbital_precalc = torch::cat(orbital_precalc_vector, -1); - ModuleBase::timer::tick("LCAO_Deepks", "calc_orbital_precalc"); + ModuleBase::timer::tick("DeePKS_domain", "calc_orbital_precalc"); return; } diff --git a/source/module_hamilt_lcao/module_deepks/LCAO_deepks_pdm.cpp b/source/module_hamilt_lcao/module_deepks/deepks_pdm.cpp similarity index 74% rename from source/module_hamilt_lcao/module_deepks/LCAO_deepks_pdm.cpp rename to source/module_hamilt_lcao/module_deepks/deepks_pdm.cpp index 7ac566c85f..dbe3684766 100644 --- a/source/module_hamilt_lcao/module_deepks/LCAO_deepks_pdm.cpp +++ b/source/module_hamilt_lcao/module_deepks/deepks_pdm.cpp @@ -9,39 +9,45 @@ // for checking purpose // There are 3 subroutines in this file: -// 1. read_projected_DM, which reads pdm from file -// 2. cal_projected_DM, which is used for calculating pdm -// 3. check_projected_dm, which prints pdm to descriptor.dat +// 1. read_pdm, which reads pdm from file +// 2. cal_pdm, which is used for calculating pdm +// 3. check_pdm, which prints pdm to descriptor.dat #ifdef __DEEPKS -#include "LCAO_deepks.h" +#include "deepks_pdm.h" #include "module_base/constants.h" #include "module_base/libm/libm.h" #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) +void DeePKS_domain::read_pdm(bool read_pdm_file, + bool is_equiv, + bool& init_pdm, + const int inlmax, + const int lmaxd, + const int* inl_l, + const Numerical_Orbital& alpha, + std::vector& pdm) { - if (read_pdm_file && !this->init_pdm) // for DeePKS NSCF calculation + if (read_pdm_file && !init_pdm) // for DeePKS NSCF calculation { 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"); + ModuleBase::WARNING_QUIT("DeePKS_domain::read_pdm", "Cannot find the file deepks_projdm.dat"); } if (!is_equiv) { - for (int inl = 0; inl < this->inlmax; inl++) + for (int inl = 0; inl < inlmax; inl++) { - int nm = this->inl_l[inl] * 2 + 1; - auto accessor = this->pdm[inl].accessor(); + int nm = inl_l[inl] * 2 + 1; + auto accessor = pdm[inl].accessor(); for (int m1 = 0; m1 < nm; m1++) { for (int m2 = 0; m2 < nm; m2++) @@ -57,14 +63,14 @@ void LCAO_Deepks::read_projected_DM(bool read_pdm_file, bool is_equiv, const Num { int pdm_size = 0; int nproj = 0; - for (int il = 0; il < this->lmaxd + 1; il++) + for (int il = 0; il < lmaxd + 1; il++) { nproj += (2 * il + 1) * alpha.getNchi(il); } pdm_size = nproj * nproj; - for (int inl = 0; inl < this->inlmax; inl++) + for (int inl = 0; inl < inlmax; inl++) { - auto accessor = this->pdm[inl].accessor(); + auto accessor = pdm[inl].accessor(); for (int ind = 0; ind < pdm_size; ind++) { double c; @@ -74,52 +80,60 @@ void LCAO_Deepks::read_projected_DM(bool read_pdm_file, bool is_equiv, const Num } } - this->init_pdm = true; + init_pdm = true; } } // this subroutine performs the calculation of projected density matrices // pdm_m,m'=\sum_{mu,nu} rho_{mu,nu} template -void LCAO_Deepks::cal_projected_DM(const elecstate::DensityMatrix* dm, - const UnitCell& ucell, - const LCAO_Orbitals& orb, - const Grid_Driver& GridD) +void DeePKS_domain::cal_pdm(bool& init_pdm, + const int inlmax, + const int lmaxd, + const int* inl_l, + const ModuleBase::IntArray* inl_index, + const elecstate::DensityMatrix* dm, + const std::vector*> phialpha, + const UnitCell& ucell, + const LCAO_Orbitals& orb, + const Grid_Driver& GridD, + const Parallel_Orbitals& pv, + std::vector& pdm) { - ModuleBase::TITLE("LCAO_Deepks", "cal_projected_DM"); + ModuleBase::TITLE("DeePKS_domain", "cal_pdm"); // if pdm has been initialized, skip the calculation - if (this->init_pdm) + if (init_pdm) { - this->init_pdm = false; + init_pdm = false; return; } if (!PARAM.inp.deepks_equiv) { - for (int inl = 0; inl < this->inlmax; inl++) + for (int inl = 0; inl < inlmax; inl++) { - int nm = this->inl_l[inl] * 2 + 1; - this->pdm[inl] = torch::zeros({nm, nm}, torch::kFloat64); + int nm = inl_l[inl] * 2 + 1; + 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++) + for (int il = 0; il < 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); + pdm[inl] = torch::zeros({pdm_size}, torch::kFloat64); } } - ModuleBase::timer::tick("LCAO_Deepks", "cal_projected_DM"); + ModuleBase::timer::tick("DeePKS_domain", "cal_pdm"); const double Rcut_Alpha = orb.Alpha[0].getRcut(); for (int T0 = 0; T0 < ucell.ntype; T0++) @@ -141,7 +155,7 @@ void LCAO_Deepks::cal_projected_DM(const elecstate::DensityMatrix* d { for (int N0 = 0; N0 < orb.Alpha[0].getNchi(L0); ++N0) { - const int inl = this->inl_index[T0](I0, L0, N0); + const int inl = inl_index[T0](I0, L0, N0); const int nm = 2 * L0 + 1; for (int m1 = 0; m1 < nm; ++m1) // m1 = 1 for s, 3 for p, 5 for d @@ -159,7 +173,7 @@ void LCAO_Deepks::cal_projected_DM(const elecstate::DensityMatrix* d else { int nproj = 0; - for (int il = 0; il < this->lmaxd + 1; il++) + for (int il = 0; il < lmaxd + 1; il++) { nproj += (2 * il + 1) * orb.Alpha[0].getNchi(il); } @@ -192,13 +206,13 @@ void LCAO_Deepks::cal_projected_DM(const elecstate::DensityMatrix* d ModuleBase::Vector3 dR1(GridD.getBox(ad1).x, GridD.getBox(ad1).y, GridD.getBox(ad1).z); if constexpr (std::is_same>::value) { - if (this->phialpha[0]->find_matrix(iat, ibt1, dR1.x, dR1.y, dR1.z) == nullptr) + if (phialpha[0]->find_matrix(iat, ibt1, dR1.x, dR1.y, dR1.z) == nullptr) { continue; } } - auto row_indexes = pv->get_indexes_row(ibt1); + auto row_indexes = pv.get_indexes_row(ibt1); const int row_size = row_indexes.size(); if (row_size == 0) { @@ -210,7 +224,7 @@ void LCAO_Deepks::cal_projected_DM(const elecstate::DensityMatrix* d std::vector g_1dmt(trace_alpha_size * row_size, 0.0); for (int irow = 0; irow < row_size; irow++) { - hamilt::BaseMatrix* row_ptr = this->phialpha[0]->find_matrix(iat, ibt1, dR1); + hamilt::BaseMatrix* row_ptr = phialpha[0]->find_matrix(iat, ibt1, dR1); for (int i = 0; i < trace_alpha_size; i++) { @@ -230,7 +244,7 @@ void LCAO_Deepks::cal_projected_DM(const elecstate::DensityMatrix* d ModuleBase::Vector3 dR2(GridD.getBox(ad2).x, GridD.getBox(ad2).y, GridD.getBox(ad2).z); if constexpr (std::is_same>::value) { - if (this->phialpha[0]->find_matrix(iat, ibt2, dR2.x, dR2.y, dR2.z) == nullptr) + if (phialpha[0]->find_matrix(iat, ibt2, dR2.x, dR2.y, dR2.z) == nullptr) { continue; } @@ -244,7 +258,7 @@ void LCAO_Deepks::cal_projected_DM(const elecstate::DensityMatrix* d continue; } - auto col_indexes = pv->get_indexes_col(ibt2); + auto col_indexes = pv.get_indexes_col(ibt2); const int col_size = col_indexes.size(); if (col_size == 0) { @@ -255,7 +269,7 @@ void LCAO_Deepks::cal_projected_DM(const elecstate::DensityMatrix* d // no possible to unexist key for (int icol = 0; icol < col_size; icol++) { - hamilt::BaseMatrix* col_ptr = this->phialpha[0]->find_matrix(iat, ibt2, dR2); + hamilt::BaseMatrix* col_ptr = phialpha[0]->find_matrix(iat, ibt2, dR2); for (int i = 0; i < trace_alpha_size; i++) { s_2t[i * col_size + icol] = col_ptr->get_value(col_indexes[icol], trace_alpha_col[i]); @@ -323,10 +337,10 @@ void LCAO_Deepks::cal_projected_DM(const elecstate::DensityMatrix* d { for (int N0 = 0; N0 < orb.Alpha[0].getNchi(L0); ++N0) { - const int inl = this->inl_index[T0](I0, L0, N0); + const int inl = inl_index[T0](I0, L0, N0); const int nm = 2 * L0 + 1; - auto accessor = this->pdm[inl].accessor(); + auto accessor = pdm[inl].accessor(); for (int m1 = 0; m1 < nm; ++m1) // m1 = 1 for s, 3 for p, 5 for d { for (int m2 = 0; m2 < nm; ++m2) // m1 = 1 for s, 3 for p, 5 for d @@ -345,10 +359,10 @@ void LCAO_Deepks::cal_projected_DM(const elecstate::DensityMatrix* d } else { - auto accessor = this->pdm[iat].accessor(); + auto accessor = pdm[iat].accessor(); int index = 0, inc = 1; int nproj = 0; - for (int il = 0; il < this->lmaxd + 1; il++) + for (int il = 0; il < lmaxd + 1; il++) { nproj += (2 * il + 1) * orb.Alpha[0].getNchi(il); } @@ -376,11 +390,11 @@ void LCAO_Deepks::cal_projected_DM(const elecstate::DensityMatrix* d Parallel_Reduce::reduce_all(pdm[inl].data_ptr(), pdm_size); } #endif - ModuleBase::timer::tick("LCAO_Deepks", "cal_projected_DM"); + ModuleBase::timer::tick("DeePKS_domain", "cal_pdm"); return; } -void LCAO_Deepks::check_projected_dm() +void DeePKS_domain::check_pdm(const int inlmax, const int* inl_l, const std::vector& pdm) { const std::string file_projdm = PARAM.globalv.global_out_dir + "pdm.dat"; std::ofstream ofs(file_projdm.c_str()); @@ -388,7 +402,7 @@ void LCAO_Deepks::check_projected_dm() ofs << std::setprecision(10); for (int inl = 0; inl < inlmax; inl++) { - const int nm = 2 * this->inl_l[inl] + 1; + const int nm = 2 * inl_l[inl] + 1; auto accessor = pdm[inl].accessor(); for (int m1 = 0; m1 < nm; m1++) { @@ -401,15 +415,31 @@ void LCAO_Deepks::check_projected_dm() } } -template void LCAO_Deepks::cal_projected_DM(const elecstate::DensityMatrix* dm, - const UnitCell& ucell, - const LCAO_Orbitals& orb, - const Grid_Driver& GridD); - -template void LCAO_Deepks::cal_projected_DM>( +template void DeePKS_domain::cal_pdm(bool& init_pdm, + const int inlmax, + const int lmaxd, + const int* inl_l, + const ModuleBase::IntArray* inl_index, + const elecstate::DensityMatrix* dm, + const std::vector*> phialpha, + const UnitCell& ucell, + const LCAO_Orbitals& orb, + const Grid_Driver& GridD, + const Parallel_Orbitals& pv, + std::vector& pdm); + +template void DeePKS_domain::cal_pdm>( + bool& init_pdm, + const int inlmax, + const int lmaxd, + const int* inl_l, + const ModuleBase::IntArray* inl_index, const elecstate::DensityMatrix, double>* dm, + const std::vector*> phialpha, const UnitCell& ucell, const LCAO_Orbitals& orb, - const Grid_Driver& GridD); + const Grid_Driver& GridD, + const Parallel_Orbitals& pv, + std::vector& pdm); #endif diff --git a/source/module_hamilt_lcao/module_deepks/deepks_pdm.h b/source/module_hamilt_lcao/module_deepks/deepks_pdm.h new file mode 100644 index 0000000000..fe4a7e0e32 --- /dev/null +++ b/source/module_hamilt_lcao/module_deepks/deepks_pdm.h @@ -0,0 +1,64 @@ +#ifndef DEEPKS_PDM_H +#define DEEPKS_PDM_H + +#ifdef __DEEPKS + +#include "module_base/complexmatrix.h" +#include "module_base/matrix.h" +#include "module_base/timer.h" +#include "module_basis/module_ao/parallel_orbitals.h" +#include "module_elecstate/module_dm/density_matrix.h" +#include "module_hamilt_lcao/module_hcontainer/hcontainer.h" + +#include +#include + +namespace DeePKS_domain +{ +//------------------- +// deepks_pdm.cpp +//------------------- + +// This file contains subroutines for calculating pdm, +// which is defind as sum_mu,nu rho_mu,nu (); +// It also contains subroutines for printing pdm for checking purpose + +// There are 3 subroutines in this file: +// 1. read_pdm, which reads pdm from file +// 2. cal_pdm, which is used for calculating pdm +// 3. check_pdm, which prints pdm to descriptor.dat + +// read pdm from file, do it only once in whole calculation +void read_pdm(bool read_pdm_file, + bool is_equiv, + bool& init_pdm, + const int inlmax, + const int lmaxd, + const int* inl_l, + const Numerical_Orbital& alpha, + std::vector& pdm); + +// calculate projected density matrix: pdm = sum_i,occ +// 3 cases to skip calculation of pdm: +// - NSCF calculation of DeePKS, init_chg = file and pdm has been read +// - SCF calculation of DeePKS with init_chg = file and pdm has been read for restarting SCF +// - Relax/Cell-Relax/MD calculation, non-first step will use the convergence pdm from the last step as initial pdm +template +void cal_pdm(bool& init_pdm, + const int inlmax, + const int lmaxd, + const int* inl_l, + const ModuleBase::IntArray* inl_index, + const elecstate::DensityMatrix* dm, + const std::vector*> phialpha, + const UnitCell& ucell, + const LCAO_Orbitals& orb, + const Grid_Driver& GridD, + const Parallel_Orbitals& pv, + std::vector& pdm); + +void check_pdm(const int inlmax, const int* inl_l, const std::vector& pdm); +} // namespace DeePKS_domain + +#endif +#endif diff --git a/source/module_hamilt_lcao/module_deepks/LCAO_deepks_phialpha.cpp b/source/module_hamilt_lcao/module_deepks/deepks_phialpha.cpp similarity index 83% rename from source/module_hamilt_lcao/module_deepks/LCAO_deepks_phialpha.cpp rename to source/module_hamilt_lcao/module_deepks/deepks_phialpha.cpp index c73dbaffa3..236b7fb9c2 100644 --- a/source/module_hamilt_lcao/module_deepks/LCAO_deepks_phialpha.cpp +++ b/source/module_hamilt_lcao/module_deepks/deepks_phialpha.cpp @@ -8,21 +8,24 @@ #ifdef __DEEPKS -#include "LCAO_deepks.h" +#include "deepks_phialpha.h" + #include "module_base/timer.h" #include "module_base/vector3.h" #include "module_parameter/parameter.h" -void LCAO_Deepks::allocate_phialpha(const bool& cal_deri, - const UnitCell& ucell, - const LCAO_Orbitals& orb, - const Grid_Driver& GridD) +void DeePKS_domain::allocate_phialpha(const bool& cal_deri, + const UnitCell& ucell, + const LCAO_Orbitals& orb, + const Grid_Driver& GridD, + const Parallel_Orbitals* pv, + std::vector*>& phialpha) { - ModuleBase::TITLE("LCAO_Deepks", "allocate_phialpha"); + ModuleBase::TITLE("DeePKS_domain", "allocate_phialpha"); - this->phialpha.resize(cal_deri ? 4 : 1); + phialpha.resize(cal_deri ? 4 : 1); - this->phialpha[0] = new hamilt::HContainer(pv); // phialpha is always real + phialpha[0] = new hamilt::HContainer(pv); // phialpha is always real // Do not use fix_gamma, since it may find wrong matrix for gamma-only case in DeePKS // cutoff for alpha is same for all types of atoms @@ -63,31 +66,33 @@ void LCAO_Deepks::allocate_phialpha(const bool& cal_deri, hamilt::AtomPair pair(iat, ibt, R_index, pv); // Notice: in AtomPair, the usage is set_size(ncol, nrow) pair.set_size(nw_alpha, atom1->nw * PARAM.globalv.npol); - this->phialpha[0]->insert_pair(pair); + phialpha[0]->insert_pair(pair); } } } - this->phialpha[0]->allocate(nullptr, true); + phialpha[0]->allocate(nullptr, true); // whether to calculate the derivative of phialpha if (cal_deri) { for (int i = 1; i < 4; ++i) { - this->phialpha[i] = new hamilt::HContainer(*this->phialpha[0], nullptr); // copy constructor + phialpha[i] = new hamilt::HContainer(*phialpha[0], nullptr); // copy constructor } } return; } -void LCAO_Deepks::build_phialpha(const bool& cal_deri, - const UnitCell& ucell, - const LCAO_Orbitals& orb, - const Grid_Driver& GridD, - const TwoCenterIntegrator& overlap_orb_alpha) +void DeePKS_domain::build_phialpha(const bool& cal_deri, + const UnitCell& ucell, + const LCAO_Orbitals& orb, + const Grid_Driver& GridD, + const Parallel_Orbitals* pv, + const TwoCenterIntegrator& overlap_orb_alpha, + std::vector*>& phialpha) { - ModuleBase::TITLE("LCAO_Deepks", "build_phialpha"); - ModuleBase::timer::tick("LCAO_Deepks", "build_phialpha"); + ModuleBase::TITLE("DeePKS_domain", "build_phialpha"); + ModuleBase::timer::tick("DeePKS_domain", "build_phialpha"); // cutoff for alpha is same for all types of atoms const double Rcut_Alpha = orb.Alpha[0].getRcut(); @@ -126,13 +131,13 @@ void LCAO_Deepks::build_phialpha(const bool& cal_deri, continue; } - double* data_pointer = this->phialpha[0]->data(iat, ibt, R); + double* data_pointer = phialpha[0]->data(iat, ibt, R); std::vector grad_pointer(3); if (cal_deri) { for (int i = 0; i < 3; ++i) { - grad_pointer[i] = this->phialpha[i + 1]->data(iat, ibt, R); + grad_pointer[i] = phialpha[i + 1]->data(iat, ibt, R); } } @@ -192,17 +197,19 @@ void LCAO_Deepks::build_phialpha(const bool& cal_deri, } } - ModuleBase::timer::tick("LCAO_Deepks", "build_phialpha"); + ModuleBase::timer::tick("DeePKS_domain", "build_phialpha"); return; } -void LCAO_Deepks::check_phialpha(const bool& cal_deri, - const UnitCell& ucell, - const LCAO_Orbitals& orb, - const Grid_Driver& GridD) +void DeePKS_domain::check_phialpha(const bool& cal_deri, + const UnitCell& ucell, + const LCAO_Orbitals& orb, + const Grid_Driver& GridD, + const Parallel_Orbitals* pv, + std::vector*>& phialpha) { - ModuleBase::TITLE("LCAO_Deepks", "check_phialpha"); - ModuleBase::timer::tick("LCAO_Deepks", "check_phialpha"); + ModuleBase::TITLE("DeePKS_domain", "check_phialpha"); + ModuleBase::timer::tick("DeePKS_domain", "check_phialpha"); const double Rcut_Alpha = orb.Alpha[0].getRcut(); // same for all types of atoms @@ -280,13 +287,13 @@ void LCAO_Deepks::check_phialpha(const bool& cal_deri, ofs_z << "R : " << R[0] << " " << R[1] << " " << R[2] << std::endl; } - const double* data_pointer = this->phialpha[0]->data(iat, ibt, R); + const double* data_pointer = phialpha[0]->data(iat, ibt, R); std::vector grad_pointer(3, nullptr); if (cal_deri) { - grad_pointer[0] = this->phialpha[1]->data(iat, ibt, R); - grad_pointer[1] = this->phialpha[2]->data(iat, ibt, R); - grad_pointer[2] = this->phialpha[3]->data(iat, ibt, R); + grad_pointer[0] = phialpha[1]->data(iat, ibt, R); + grad_pointer[1] = phialpha[2]->data(iat, ibt, R); + grad_pointer[2] = phialpha[3]->data(iat, ibt, R); } for (int iw1 = 0; iw1 < nw1_tot; ++iw1) @@ -334,7 +341,7 @@ void LCAO_Deepks::check_phialpha(const bool& cal_deri, } // end I0 } // end T0 - ModuleBase::timer::tick("LCAO_Deepks", "check_phialpha"); + ModuleBase::timer::tick("DeePKS_domain", "check_phialpha"); return; } diff --git a/source/module_hamilt_lcao/module_deepks/deepks_phialpha.h b/source/module_hamilt_lcao/module_deepks/deepks_phialpha.h new file mode 100644 index 0000000000..6c5ea9172a --- /dev/null +++ b/source/module_hamilt_lcao/module_deepks/deepks_phialpha.h @@ -0,0 +1,52 @@ +#ifndef DEEPKS_PHIALPHA_H +#define DEEPKS_PHIALPHA_H + +#ifdef __DEEPKS + +#include "module_base/complexmatrix.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_hamilt_lcao/module_hcontainer/hcontainer.h" + +#include +#include + +namespace DeePKS_domain +{ +// This file contains 3 subroutines: +// 1. allocate_phialpha, which allocates memory for phialpha +// 2. build_phialpha, which calculates the overlap +// between atomic basis and projector alpha : +// which will be used in calculating pdm, gdmx, H_V_delta, F_delta; +// 3. check_phialpha, which prints the results into .dat files +// for checking + +// calculates +void allocate_phialpha(const bool& cal_deri, + const UnitCell& ucell, + const LCAO_Orbitals& orb, + const Grid_Driver& GridD, + const Parallel_Orbitals* pv, + std::vector*>& phialpha); + +void build_phialpha(const bool& cal_deri /**< [in] 0 for 2-center intergration, 1 for its derivation*/, + const UnitCell& ucell, + const LCAO_Orbitals& orb, + const Grid_Driver& GridD, + const Parallel_Orbitals* pv, + const TwoCenterIntegrator& overlap_orb_alpha, + std::vector*>& phialpha); + +void check_phialpha(const bool& cal_deri /**< [in] 0 for 2-center intergration, 1 for its derivation*/, + const UnitCell& ucell, + const LCAO_Orbitals& orb, + const Grid_Driver& GridD, + const Parallel_Orbitals* pv, + std::vector*>& phialpha); +} // namespace DeePKS_domain + +#endif +#endif diff --git a/source/module_hamilt_lcao/module_deepks/deepks_vdelta.cpp b/source/module_hamilt_lcao/module_deepks/deepks_vdelta.cpp new file mode 100644 index 0000000000..1d6ee3ecae --- /dev/null +++ b/source/module_hamilt_lcao/module_deepks/deepks_vdelta.cpp @@ -0,0 +1,92 @@ +// This file contains subroutines related to V_delta, which is the deepks contribution to Hamiltonian +// defined as |alpha>V(D) +void DeePKS_domain::cal_e_delta_band(const std::vector>& dm, + const std::vector>& H_V_delta, + const int nks, + const Parallel_Orbitals* pv, + double& e_delta_band) +{ + ModuleBase::TITLE("DeePKS_domain", "cal_e_delta_band"); + ModuleBase::timer::tick("DeePKS_domain", "cal_e_delta_band"); + TK e_delta_band_tmp = TK(0); + for (int i = 0; i < PARAM.globalv.nlocal; ++i) + { + for (int j = 0; j < PARAM.globalv.nlocal; ++j) + { + const int mu = pv->global2local_row(j); + const int nu = pv->global2local_col(i); + + if (mu >= 0 && nu >= 0) + { + int iic; + if (ModuleBase::GlobalFunc::IS_COLUMN_MAJOR_KS_SOLVER(PARAM.inp.ks_solver)) + { + iic = mu + nu * pv->nrow; + } + else + { + iic = mu * pv->ncol + nu; + } + if constexpr (std::is_same::value) + { + for (int is = 0; is < dm.size(); ++is) // dm.size() == PARAM.inp.nspin + { + e_delta_band_tmp += dm[is][nu * pv->nrow + mu] * H_V_delta[0][iic]; + } + } + else + { + for (int ik = 0; ik < nks; ik++) + { + e_delta_band_tmp += dm[ik][nu * pv->nrow + mu] * H_V_delta[ik][iic]; + } + } + } + } + } + if constexpr (std::is_same::value) + { + e_delta_band = e_delta_band_tmp; + } + else + { + e_delta_band = e_delta_band_tmp.real(); + } +#ifdef __MPI + Parallel_Reduce::reduce_all(e_delta_band); +#endif + ModuleBase::timer::tick("DeePKS_domain", "cal_e_delta_band"); + return; +} + +template void DeePKS_domain::cal_e_delta_band(const std::vector>& dm, + const std::vector>& H_V_delta, + const int nks, + const Parallel_Orbitals* pv, + double& e_delta_band); +template void DeePKS_domain::cal_e_delta_band>( + const std::vector>>& dm, + const std::vector>>& H_V_delta, + const int nks, + const Parallel_Orbitals* pv, + double& e_delta_band); + +#endif diff --git a/source/module_hamilt_lcao/module_deepks/deepks_vdelta.h b/source/module_hamilt_lcao/module_deepks/deepks_vdelta.h new file mode 100644 index 0000000000..0022d49a4e --- /dev/null +++ b/source/module_hamilt_lcao/module_deepks/deepks_vdelta.h @@ -0,0 +1,25 @@ +#ifndef DEEPKS_VDELTA_H +#define DEEPKS_VDELTA_H + +#ifdef __DEEPKS +#include "module_basis/module_ao/parallel_orbitals.h" + +namespace DeePKS_domain +{ +//------------------------ +// deepks_vdelta.cpp +//------------------------ + +// This file contains 1 subroutine for calculating e_delta_bands +// 1. cal_e_delta_band : calculates e_delta_bands + +/// calculate tr(\rho V_delta) +template +void cal_e_delta_band(const std::vector>& dm, + const std::vector>& H_V_delta, + const int nks, + const Parallel_Orbitals* pv, + double& e_delta_band); +} // namespace DeePKS_domain +#endif +#endif \ No newline at end of file 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 107675e278..493469072b 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 @@ -36,11 +36,22 @@ void test_deepks::check_phialpha() } GlobalC::ld.init(ORB, ucell.nat, ucell.ntype, kv.nkstot, ParaO, na); - GlobalC::ld.allocate_phialpha(PARAM.input.cal_force, ucell, ORB, Test_Deepks::GridD); - - GlobalC::ld.build_phialpha(PARAM.input.cal_force, ucell, ORB, Test_Deepks::GridD, overlap_orb_alpha_); - - GlobalC::ld.check_phialpha(PARAM.input.cal_force, ucell, ORB, Test_Deepks::GridD); + DeePKS_domain::allocate_phialpha(PARAM.input.cal_force, + ucell, + ORB, + Test_Deepks::GridD, + &ParaO, + GlobalC::ld.phialpha); + + DeePKS_domain::build_phialpha(PARAM.input.cal_force, + ucell, + ORB, + Test_Deepks::GridD, + &ParaO, + overlap_orb_alpha_, + GlobalC::ld.phialpha); + + DeePKS_domain::check_phialpha(PARAM.input.cal_force, ucell, ORB, Test_Deepks::GridD, &ParaO, GlobalC::ld.phialpha); this->compare_with_ref("phialpha.dat", "phialpha_ref.dat"); this->compare_with_ref("dphialpha_x.dat", "dphialpha_x_ref.dat"); @@ -145,16 +156,38 @@ void test_deepks::check_pdm() this->read_dm(); this->set_dm_new(); this->set_p_elec_DM(); - GlobalC::ld.cal_projected_DM(p_elec_DM, ucell, ORB, Test_Deepks::GridD); + DeePKS_domain::cal_pdm(GlobalC::ld.init_pdm, + GlobalC::ld.inlmax, + GlobalC::ld.lmaxd, + GlobalC::ld.inl_l, + GlobalC::ld.inl_index, + p_elec_DM, + GlobalC::ld.phialpha, + ucell, + ORB, + Test_Deepks::GridD, + ParaO, + GlobalC::ld.pdm); } else { this->read_dm_k(kv.nkstot); this->set_dm_k_new(); this->set_p_elec_DM_k(); - GlobalC::ld.cal_projected_DM(p_elec_DM_k, ucell, ORB, Test_Deepks::GridD); + DeePKS_domain::cal_pdm(GlobalC::ld.init_pdm, + GlobalC::ld.inlmax, + GlobalC::ld.lmaxd, + GlobalC::ld.inl_l, + GlobalC::ld.inl_index, + p_elec_DM_k, + GlobalC::ld.phialpha, + ucell, + ORB, + Test_Deepks::GridD, + ParaO, + GlobalC::ld.pdm); } - GlobalC::ld.check_projected_dm(); + DeePKS_domain::check_pdm(GlobalC::ld.inlmax, GlobalC::ld.inl_l, GlobalC::ld.pdm); this->compare_with_ref("pdm.dat", "pdm_ref.dat"); } @@ -283,7 +316,7 @@ void test_deepks::check_descriptor(std::vector& descriptor) void test_deepks::check_gvx(torch::Tensor& gdmx) { std::vector gevdm; - GlobalC::ld.cal_gevdm(ucell.nat, gevdm); + DeePKS_domain::cal_gevdm(ucell.nat, GlobalC::ld.inlmax, GlobalC::ld.inl_l, GlobalC::ld.pdm, gevdm); torch::Tensor gvx; DeePKS_domain::cal_gvx(ucell.nat, GlobalC::ld.inlmax, @@ -321,7 +354,7 @@ void test_deepks::check_gvx(torch::Tensor& gdmx) void test_deepks::check_gvepsl(torch::Tensor& gdmepsl) { std::vector gevdm; - GlobalC::ld.cal_gevdm(ucell.nat, gevdm); + DeePKS_domain::cal_gevdm(ucell.nat, GlobalC::ld.inlmax, GlobalC::ld.inl_l, GlobalC::ld.pdm, gevdm); torch::Tensor gvepsl; DeePKS_domain::cal_gvepsl(ucell.nat, GlobalC::ld.inlmax, @@ -346,7 +379,7 @@ void test_deepks::check_gvepsl(torch::Tensor& gdmepsl) void test_deepks::check_edelta(std::vector& descriptor) { - GlobalC::ld.load_model("model.ptg"); + DeePKS_domain::load_model("model.ptg", GlobalC::ld.model_deepks); if (PARAM.sys.gamma_only_local) { GlobalC::ld.allocate_V_delta(ucell.nat, 1); // 1 for gamma-only @@ -355,14 +388,24 @@ void test_deepks::check_edelta(std::vector& descriptor) { GlobalC::ld.allocate_V_delta(ucell.nat, kv.nkstot); } - GlobalC::ld.cal_gedm(ucell.nat, descriptor); + DeePKS_domain::cal_gedm(ucell.nat, + GlobalC::ld.lmaxd, + GlobalC::ld.nmaxd, + GlobalC::ld.inlmax, + GlobalC::ld.des_per_atom, + GlobalC::ld.inl_l, + descriptor, + GlobalC::ld.pdm, + GlobalC::ld.model_deepks, + GlobalC::ld.gedm, + GlobalC::ld.E_delta); std::ofstream ofs("E_delta.dat"); ofs << std::setprecision(10) << GlobalC::ld.E_delta << std::endl; ofs.close(); this->compare_with_ref("E_delta.dat", "E_delta_ref.dat"); - GlobalC::ld.check_gedm(); + DeePKS_domain::check_gedm(GlobalC::ld.inlmax, GlobalC::ld.inl_l, GlobalC::ld.gedm); this->compare_with_ref("gedm.dat", "gedm_ref.dat"); } @@ -412,12 +455,12 @@ void test_deepks::check_e_deltabands() if (PARAM.sys.gamma_only_local) { this->cal_H_V_delta(); - GlobalC::ld.cal_e_delta_band(dm_new, 1); + GlobalC::ld.dpks_cal_e_delta_band(dm_new, 1); } else { this->cal_H_V_delta_k(); - GlobalC::ld.cal_e_delta_band(dm_k_new, kv.nkstot); + GlobalC::ld.dpks_cal_e_delta_band(dm_k_new, kv.nkstot); } std::ofstream ofs("E_delta_bands.dat"); diff --git a/source/module_hamilt_lcao/module_deepks/test/Makefile.Objects b/source/module_hamilt_lcao/module_deepks/test/Makefile.Objects index f1a13cfd2d..5e6debbcdf 100644 --- a/source/module_hamilt_lcao/module_deepks/test/Makefile.Objects +++ b/source/module_hamilt_lcao/module_deepks/test/Makefile.Objects @@ -7,13 +7,13 @@ HEADERS= *.h OBJS_MAIN=klist_1.o\ parallel_orbitals.o\ +deepks_basic.o\ +deepks_force.o\ +deepks_vdelta.o\ +deepks_pdm.o\ +deepks_phialpha.o\ LCAO_deepks.o\ -LCAO_deepks_fdelta.o\ LCAO_deepks_io.o\ -LCAO_deepks_pdm.o\ -LCAO_deepks_phialpha.o\ -LCAO_deepks_vdelta.o\ -LCAO_deepks_torch.o\ OBJS_IO=output.o\