From 9231adecdf8ccf4eac28ebbeab69460d4f2e9757 Mon Sep 17 00:00:00 2001 From: ErjieWu Date: Tue, 24 Mar 2026 09:47:03 +0800 Subject: [PATCH 1/2] Remove the use of std::system in DeePKS_equiv. --- .../module_deepks/LCAO_deepks_interface.cpp | 2 +- .../module_deepks/deepks_basic.cpp | 152 +++++++++++------- .../source_lcao/module_deepks/deepks_basic.h | 1 + .../module_deepks/test/LCAO_deepks_test.cpp | 1 + .../module_operator_lcao/deepks_lcao.cpp | 1 + 5 files changed, 94 insertions(+), 63 deletions(-) diff --git a/source/source_lcao/module_deepks/LCAO_deepks_interface.cpp b/source/source_lcao/module_deepks/LCAO_deepks_interface.cpp index d847714dd9..10480253e6 100644 --- a/source/source_lcao/module_deepks/LCAO_deepks_interface.cpp +++ b/source/source_lcao/module_deepks/LCAO_deepks_interface.cpp @@ -135,7 +135,7 @@ void LCAO_Deepks_Interface::out_deepks_labels(const double& etot, // new gedm is also useful in cal_f_delta, so it should be ld->gedm if (PARAM.inp.deepks_equiv) { - DeePKS_domain::cal_edelta_gedm_equiv(nat, deepks_param, descriptor, ld->gedm, E_delta, rank); + DeePKS_domain::cal_edelta_gedm_equiv(nat, deepks_param, descriptor, ld->model_deepks, ld->gedm, E_delta, rank); } else { diff --git a/source/source_lcao/module_deepks/deepks_basic.cpp b/source/source_lcao/module_deepks/deepks_basic.cpp index edcc74f0b1..561396c81e 100644 --- a/source/source_lcao/module_deepks/deepks_basic.cpp +++ b/source/source_lcao/module_deepks/deepks_basic.cpp @@ -9,7 +9,9 @@ #include "source_base/timer.h" #include "source_io/module_parameter/parameter.h" -#include // use system command +#ifdef __MPI +#include +#endif // d(Descriptor) / d(projected density matrix) // Dimension is different for each inl, so there's a vector of tensors @@ -85,57 +87,10 @@ void DeePKS_domain::load_model(const std::string& model_file, torch::jit::script return; } -inline void generate_py_files(const DeePKS_Param& deepks_param, const std::string& out_dir) -{ - std::ofstream ofs("cal_edelta_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 < deepks_param.lmaxd + 1; l++) - { - ofs << " - - " << l << std::endl; - ofs << " - ["; - for (int i = 0; i < deepks_param.nmaxd + 1; i++) - { - ofs << "0"; - if (i != deepks_param.nmaxd) - { - ofs << ", "; - } - } - ofs << "]" << std::endl; - } -} - void DeePKS_domain::cal_edelta_gedm_equiv(const int nat, const DeePKS_Param& deepks_param, const std::vector& descriptor, + torch::jit::script::Module& model_deepks, double** gedm, double& E_delta, const int rank) @@ -143,24 +98,97 @@ void DeePKS_domain::cal_edelta_gedm_equiv(const int nat, ModuleBase::TITLE("DeePKS_domain", "cal_edelta_gedm_equiv"); ModuleBase::timer::start("DeePKS_domain", "cal_edelta_gedm_equiv"); - const std::string file_d = PARAM.globalv.global_out_dir + "deepks_dm_eig.npy"; - LCAO_deepks_io::save_npy_d(nat, PARAM.inp.deepks_equiv, deepks_param, descriptor, file_d, - rank); // libnpy needed - if (rank == 0) { - generate_py_files(deepks_param, PARAM.globalv.global_out_dir); - std::string cmd = "python cal_edelta_gedm.py " + PARAM.inp.deepks_model; - int stat = std::system(cmd.c_str()); - assert(stat == 0); - } + int basis_size = 0; + for (int l = 0; l <= deepks_param.lmaxd; ++l) + { + basis_size += (2 * l + 1) * deepks_param.nmaxd; + } + assert(deepks_param.des_per_atom == basis_size * basis_size); + + torch::Tensor dm_eig = torch::cat(descriptor, 0).reshape({1, nat, deepks_param.des_per_atom}); + dm_eig = dm_eig.to(torch::kFloat64).requires_grad_(true); + torch::Tensor dm = dm_eig.reshape({1, nat, basis_size, basis_size}); + + std::vector ovlp_shells; + ovlp_shells.reserve((deepks_param.lmaxd + 1) * deepks_param.nmaxd); + int offset = 0; + for (int l = 0; l <= deepks_param.lmaxd; ++l) + { + const int nm = 2 * l + 1; + for (int n = 0; n < deepks_param.nmaxd; ++n) + { + torch::Tensor po = torch::zeros({basis_size, 1, nm}, torch::TensorOptions().dtype(torch::kFloat64)); + auto accessor = po.accessor(); + for (int m = 0; m < nm; ++m) + { + accessor[offset + m][0][m] = 1.0; + } + ovlp_shells.push_back(po); + offset += nm; + } + } + assert(offset == basis_size); + + std::vector dm_flat; + dm_flat.reserve(ovlp_shells.size()); + for (const auto& po : ovlp_shells) + { + // Equivalent to python: + // torch.einsum('rap,...rs,saq->...apq', po, dm, po) + torch::Tensor pdm_shell = torch::einsum("rap,...rs,saq->...apq", {po, dm, po}); + dm_flat.push_back(pdm_shell.squeeze(-3)); + } + + c10::List model_input; + for (const auto& pdm_shell : dm_flat) + { + model_input.push_back(pdm_shell); + } + + std::vector inputs; + inputs.emplace_back(model_input); + + torch::Tensor ec; + try + { + ec = model_deepks.forward(inputs).toTensor(); // Hartree + } + catch (const c10::Error& e) + { + ModuleBase::WARNING_QUIT("DeePKS_domain::cal_edelta_gedm_equiv", + "Failed to evaluate equivariant DeePKS model in C++."); + throw; + } - MPI_Barrier(MPI_COMM_WORLD); + E_delta = ec.item() * 2.0; // Hartree to Ry - LCAO_deepks_io::load_npy_gedm(nat, deepks_param.des_per_atom, gedm, E_delta, rank); + std::vector grad_outputs{torch::ones_like(ec)}; + std::vector grad_inputs{dm_eig}; + torch::Tensor gedm_tensor = torch::autograd::grad({ec}, grad_inputs, grad_outputs, + /*retain_graph=*/false, + /*create_graph=*/false, + /*allow_unused=*/false)[0]; - std::string cmd = "rm -f cal_edelta_gedm.py basis.yaml ec.npy gedm.npy"; - std::system(cmd.c_str()); + torch::Tensor gedm_nat = gedm_tensor.reshape({nat, deepks_param.des_per_atom}); + auto accessor = gedm_nat.accessor(); + for (int iat = 0; iat < nat; ++iat) + { + for (int ides = 0; ides < deepks_param.des_per_atom; ++ides) + { + gedm[iat][ides] = accessor[iat][ides] * 2.0; // Hartree to Ry + } + } + } + +#ifdef __MPI + for (int iat = 0; iat < nat; ++iat) + { + MPI_Bcast(gedm[iat], deepks_param.des_per_atom, MPI_DOUBLE, 0, MPI_COMM_WORLD); + } + MPI_Bcast(&E_delta, 1, MPI_DOUBLE, 0, MPI_COMM_WORLD); +#endif ModuleBase::timer::end("DeePKS_domain", "cal_edelta_gedm_equiv"); return; diff --git a/source/source_lcao/module_deepks/deepks_basic.h b/source/source_lcao/module_deepks/deepks_basic.h index 768b0aab1f..a98db86a4a 100644 --- a/source/source_lcao/module_deepks/deepks_basic.h +++ b/source/source_lcao/module_deepks/deepks_basic.h @@ -49,6 +49,7 @@ void check_gedm(const DeePKS_Param& deepks_param, double** gedm); void cal_edelta_gedm_equiv(const int nat, const DeePKS_Param& deepks_param, const std::vector& descriptor, + torch::jit::script::Module& model_deepks, double** gedm, double& E_delta, const int rank); diff --git a/source/source_lcao/module_deepks/test/LCAO_deepks_test.cpp b/source/source_lcao/module_deepks/test/LCAO_deepks_test.cpp index 68718bd9b6..517b9f9940 100644 --- a/source/source_lcao/module_deepks/test/LCAO_deepks_test.cpp +++ b/source/source_lcao/module_deepks/test/LCAO_deepks_test.cpp @@ -320,6 +320,7 @@ void test_deepks::check_edelta(std::vector& descriptor) DeePKS_domain::cal_edelta_gedm_equiv(ucell.nat, this->ld.deepks_param, descriptor, + this->ld.model_deepks, this->ld.gedm, this->ld.E_delta, 0); // 0 for rank diff --git a/source/source_lcao/module_operator_lcao/deepks_lcao.cpp b/source/source_lcao/module_operator_lcao/deepks_lcao.cpp index ef198d3834..040de01e0b 100644 --- a/source/source_lcao/module_operator_lcao/deepks_lcao.cpp +++ b/source/source_lcao/module_operator_lcao/deepks_lcao.cpp @@ -173,6 +173,7 @@ void hamilt::DeePKS>::contributeHR() DeePKS_domain::cal_edelta_gedm_equiv(this->ucell->nat, this->ld->deepks_param, descriptor, + this->ld->model_deepks, this->ld->gedm, this->ld->E_delta, GlobalV::MY_RANK); From 7a8e397ba534baee0b2aeecb77fb5265686ec339 Mon Sep 17 00:00:00 2001 From: ErjieWu Date: Tue, 24 Mar 2026 10:19:48 +0800 Subject: [PATCH 2/2] Fix a dimension match problem. --- .../source_lcao/module_deepks/LCAO_deepks.cpp | 5 +++ .../module_deepks/deepks_basic.cpp | 33 +++++++++++++++---- .../source_lcao/module_deepks/deepks_param.h | 1 + 3 files changed, 32 insertions(+), 7 deletions(-) diff --git a/source/source_lcao/module_deepks/LCAO_deepks.cpp b/source/source_lcao/module_deepks/LCAO_deepks.cpp index deb888e9d0..cd64cc9850 100644 --- a/source/source_lcao/module_deepks/LCAO_deepks.cpp +++ b/source/source_lcao/module_deepks/LCAO_deepks.cpp @@ -76,6 +76,11 @@ void LCAO_Deepks::init(const LCAO_Orbitals& orb, this->deepks_param.lmaxd = lm; this->deepks_param.nmaxd = nm; + this->deepks_param.nchi_d_l.assign(lm + 1, 0); + for (int l = 0; l <= lm; ++l) + { + this->deepks_param.nchi_d_l[l] = orb.Alpha[0].getNchi(l); + } ofs << " lmax of descriptor = " << deepks_param.lmaxd << std::endl; ofs << " nmax of descriptor = " << deepks_param.nmaxd << std::endl; diff --git a/source/source_lcao/module_deepks/deepks_basic.cpp b/source/source_lcao/module_deepks/deepks_basic.cpp index 561396c81e..f0531012cd 100644 --- a/source/source_lcao/module_deepks/deepks_basic.cpp +++ b/source/source_lcao/module_deepks/deepks_basic.cpp @@ -9,6 +9,8 @@ #include "source_base/timer.h" #include "source_io/module_parameter/parameter.h" +#include + #ifdef __MPI #include #endif @@ -100,24 +102,37 @@ void DeePKS_domain::cal_edelta_gedm_equiv(const int nat, if (rank == 0) { - int basis_size = 0; - for (int l = 0; l <= deepks_param.lmaxd; ++l) + const int basis_size + = static_cast(std::llround(std::sqrt(static_cast(deepks_param.des_per_atom)))); + if (basis_size * basis_size != deepks_param.des_per_atom) { - basis_size += (2 * l + 1) * deepks_param.nmaxd; + ModuleBase::WARNING_QUIT("DeePKS_domain::cal_edelta_gedm_equiv", + "Invalid des_per_atom for equivariant DeePKS: it must be a perfect square."); } - assert(deepks_param.des_per_atom == basis_size * basis_size); torch::Tensor dm_eig = torch::cat(descriptor, 0).reshape({1, nat, deepks_param.des_per_atom}); dm_eig = dm_eig.to(torch::kFloat64).requires_grad_(true); torch::Tensor dm = dm_eig.reshape({1, nat, basis_size, basis_size}); + if (static_cast(deepks_param.nchi_d_l.size()) != deepks_param.lmaxd + 1) + { + ModuleBase::WARNING_QUIT( + "DeePKS_domain::cal_edelta_gedm_equiv", + "Invalid nchi_d_l in DeePKS parameters: expected size lmaxd + 1 for equivariant shell construction."); + } + std::vector ovlp_shells; - ovlp_shells.reserve((deepks_param.lmaxd + 1) * deepks_param.nmaxd); + int total_shells = 0; + for (int l = 0; l <= deepks_param.lmaxd; ++l) + { + total_shells += deepks_param.nchi_d_l[l]; + } + ovlp_shells.reserve(total_shells); int offset = 0; for (int l = 0; l <= deepks_param.lmaxd; ++l) { const int nm = 2 * l + 1; - for (int n = 0; n < deepks_param.nmaxd; ++n) + for (int n = 0; n < deepks_param.nchi_d_l[l]; ++n) { torch::Tensor po = torch::zeros({basis_size, 1, nm}, torch::TensorOptions().dtype(torch::kFloat64)); auto accessor = po.accessor(); @@ -129,7 +144,11 @@ void DeePKS_domain::cal_edelta_gedm_equiv(const int nat, offset += nm; } } - assert(offset == basis_size); + if (offset != basis_size) + { + ModuleBase::WARNING_QUIT("DeePKS_domain::cal_edelta_gedm_equiv", + "Invalid shell layout: accumulated shell offset does not match basis size."); + } std::vector dm_flat; dm_flat.reserve(ovlp_shells.size()); diff --git a/source/source_lcao/module_deepks/deepks_param.h b/source/source_lcao/module_deepks/deepks_param.h index 75e161b9da..13eee68566 100644 --- a/source/source_lcao/module_deepks/deepks_param.h +++ b/source/source_lcao/module_deepks/deepks_param.h @@ -14,6 +14,7 @@ struct DeePKS_Param int inlmax = 0; int n_descriptor = 0; int des_per_atom = 0; + std::vector nchi_d_l; std::vector inl2l; ModuleBase::IntArray* inl_index = nullptr; };