From 071059f93ce1aa669c869d40dfcac7738e27eec1 Mon Sep 17 00:00:00 2001 From: abacus_fixer Date: Thu, 2 Jul 2026 09:54:59 +0800 Subject: [PATCH] fix KEDF --- .../source_io/module_ctrl/ctrl_output_pw.cpp | 3 +- .../module_ml/cal_mlkedf_descriptors.cpp | 44 ++++---- .../module_ml/cal_mlkedf_descriptors.h | 2 +- .../module_ml/write_mlkedf_descriptors.cpp | 103 +++++++++--------- .../module_ml/write_mlkedf_descriptors.h | 15 ++- source/source_pw/module_ofdft/ml_base.cpp | 2 +- 6 files changed, 86 insertions(+), 83 deletions(-) diff --git a/source/source_io/module_ctrl/ctrl_output_pw.cpp b/source/source_io/module_ctrl/ctrl_output_pw.cpp index 580bc314efb..b820c721c36 100644 --- a/source/source_io/module_ctrl/ctrl_output_pw.cpp +++ b/source/source_io/module_ctrl/ctrl_output_pw.cpp @@ -370,7 +370,8 @@ void ModuleIO::ctrl_runner_pw(UnitCell& ucell, pw_wfc, pw_rho, ucell, - pelec->pot->get_eff_v(0)); + pelec->pot->get_eff_v(0), + chr.nrxx); } #endif diff --git a/source/source_io/module_ml/cal_mlkedf_descriptors.cpp b/source/source_io/module_ml/cal_mlkedf_descriptors.cpp index 22a7b794d43..bd47c6e4c5a 100644 --- a/source/source_io/module_ml/cal_mlkedf_descriptors.cpp +++ b/source/source_io/module_ml/cal_mlkedf_descriptors.cpp @@ -4,7 +4,7 @@ namespace ModuleIO { void Cal_MLKEDF_Descriptors::set_para( - const int &nx, + const int &nrxx, const double &nelec, const double &tf_weight, const double &vw_weight, @@ -23,7 +23,7 @@ void Cal_MLKEDF_Descriptors::set_para( std::ostream& ofs_running ) { - this->nx = nx; + this->nrxx = nrxx; this->nkernel = nkernel; this->chi_p = chi_p; this->chi_q = chi_q; @@ -235,7 +235,7 @@ void Cal_MLKEDF_Descriptors::divergence(double ** pinput, const ModulePW::PW_Bas { std::complex *recipContainer = new std::complex[pw_rho->npw]; std::complex img(0.0, 1.0); - ModuleBase::GlobalFunc::ZEROS(routput, this->nx); + ModuleBase::GlobalFunc::ZEROS(routput, this->nrxx); for (int i = 0; i < 3; ++i) { pw_rho->real2recip(pinput[i], recipContainer); @@ -251,7 +251,7 @@ void Cal_MLKEDF_Descriptors::divergence(double ** pinput, const ModulePW::PW_Bas void Cal_MLKEDF_Descriptors::tanh(std::vector &pinput, std::vector &routput, double chi) { - for (int i = 0; i < this->nx; ++i) + for (int i = 0; i < this->nrxx; ++i) { routput[i] = std::tanh(pinput[i] * chi); } @@ -264,7 +264,7 @@ double Cal_MLKEDF_Descriptors::dtanh(double tanhx, double chi) void Cal_MLKEDF_Descriptors::getGamma(const double * const *prho, std::vector &rgamma) { - for(int ir = 0; ir < this->nx; ++ir) + for(int ir = 0; ir < this->nrxx; ++ir) { rgamma[ir] = std::pow(prho[0][ir]/this->rho0, 1./3.); } @@ -272,7 +272,7 @@ void Cal_MLKEDF_Descriptors::getGamma(const double * const *prho, std::vector> &pnablaRho, std::vector &rp) { - for(int ir = 0; ir < this->nx; ++ir) + for(int ir = 0; ir < this->nrxx; ++ir) { rp[ir] = 0.; for (int j = 0; j < 3; ++j) @@ -294,7 +294,7 @@ void Cal_MLKEDF_Descriptors::getQ(const double * const *prho, const ModulePW::PW } pw_rho->recip2real(recipRho, rq.data()); - for (int ir = 0; ir < this->nx; ++ir) + for (int ir = 0; ir < this->nrxx; ++ir) { rq[ir] *= this->pqcoef / std::pow(prho[0][ir], 5.0/3.0); } @@ -320,7 +320,7 @@ void Cal_MLKEDF_Descriptors::getQnl(const int ikernel, std::vector &pq, // xi = gammanl/gamma void Cal_MLKEDF_Descriptors::getXi(std::vector &pgamma, std::vector &pgammanl, std::vector &rxi) { - for (int ir = 0; ir < this->nx; ++ir) + for (int ir = 0; ir < this->nrxx; ++ir) { if (pgamma[ir] == 0) { @@ -337,7 +337,7 @@ void Cal_MLKEDF_Descriptors::getXi(std::vector &pgamma, std::vector &pgamma, std::vector &pgammanl, std::vector &rtanhxi) { - for (int ir = 0; ir < this->nx; ++ir) + for (int ir = 0; ir < this->nrxx; ++ir) { if (pgamma[ir] == 0) { @@ -404,14 +404,14 @@ void Cal_MLKEDF_Descriptors::getF_KS( std::vector &rpauli ) { - double *pauliED = new double[this->nx]; // Pauli Energy Density - ModuleBase::GlobalFunc::ZEROS(pauliED, this->nx); + double *pauliED = new double[this->nrxx]; // Pauli Energy Density + ModuleBase::GlobalFunc::ZEROS(pauliED, this->nrxx); - double *pauliPot = new double[this->nx]; - ModuleBase::GlobalFunc::ZEROS(pauliPot, this->nx); + double *pauliPot = new double[this->nrxx]; + ModuleBase::GlobalFunc::ZEROS(pauliPot, this->nrxx); - std::complex *wfcr = new std::complex[this->nx]; - ModuleBase::GlobalFunc::ZEROS(wfcr, this->nx); + std::complex *wfcr = new std::complex[this->nrxx]; + ModuleBase::GlobalFunc::ZEROS(wfcr, this->nrxx); double epsilonM = pelec->ekb(0,0); assert(PARAM.inp.nspin == 1); @@ -438,14 +438,14 @@ void Cal_MLKEDF_Descriptors::getF_KS( // output one wf, to check KS equation if (ik == 0 && ibnd == 0) { - std::vector wf_real = std::vector(this->nx); - std::vector wf_imag = std::vector(this->nx); - for (int ir = 0; ir < this->nx; ++ir) + std::vector wf_real = std::vector(this->nrxx); + std::vector wf_imag = std::vector(this->nrxx); + for (int ir = 0; ir < this->nrxx; ++ir) { wf_real[ir] = wfcr[ir].real(); wf_imag[ir] = wfcr[ir].imag(); } - const long unsigned cshape[] = {(long unsigned) this->nx}; // shape of container and containernl + const long unsigned cshape[] = {(long unsigned) this->nrxx}; // shape of container and containernl } if (w1 != 0.0) @@ -474,7 +474,7 @@ void Cal_MLKEDF_Descriptors::getF_KS( pw_psi->recip2real(wfcr, wfcr, ik); - for (int ir = 0; ir < this->nx; ++ir) + for (int ir = 0; ir < this->nrxx; ++ir) { pauliED[ir] += w1 * norm(wfcr[ir]); // actually, here should be w1/2 * norm(wfcr[ir]), but we multiply 2 to convert Ha to Ry. } @@ -484,13 +484,13 @@ void Cal_MLKEDF_Descriptors::getF_KS( for (int j = 0; j < 3; ++j) { - for (int ir = 0; ir < this->nx; ++ir) + for (int ir = 0; ir < this->nrxx; ++ir) { pauliED[ir] -= nablaRho[j][ir] * nablaRho[j][ir] / (8. * pelec->charge->rho[0][ir]) * 2.; // convert Ha to Ry. } } - for (int ir = 0; ir < this->nx; ++ir) + for (int ir = 0; ir < this->nrxx; ++ir) { rF[ir] = pauliED[ir] / (this->cTF * std::pow(pelec->charge->rho[0][ir], 5./3.)); rpauli[ir] = (pauliED[ir] + pauliPot[ir])/pelec->charge->rho[0][ir] + epsilonM; diff --git a/source/source_io/module_ml/cal_mlkedf_descriptors.h b/source/source_io/module_ml/cal_mlkedf_descriptors.h index 7a0b70b69f3..357c1706130 100644 --- a/source/source_io/module_ml/cal_mlkedf_descriptors.h +++ b/source/source_io/module_ml/cal_mlkedf_descriptors.h @@ -90,7 +90,7 @@ class Cal_MLKEDF_Descriptors std::vector chi_pnl = {1.0}; std::vector chi_qnl = {1.0}; - int nx = 0; + int nrxx = 0; double dV = 0.; double rho0 = 0.; // average rho double kF = 0.; // Fermi vector kF = (3 pi^2 rho0)^(1/3) diff --git a/source/source_io/module_ml/write_mlkedf_descriptors.cpp b/source/source_io/module_ml/write_mlkedf_descriptors.cpp index 59ee758bd70..9350f28ba26 100644 --- a/source/source_io/module_ml/write_mlkedf_descriptors.cpp +++ b/source/source_io/module_ml/write_mlkedf_descriptors.cpp @@ -14,57 +14,50 @@ void Write_MLKEDF_Descriptors::generateTrainData_KS( ModulePW::PW_Basis_K *pw_psi, ModulePW::PW_Basis *pw_rho, UnitCell& ucell, - const double* veff + const double* veff, + const int nrxx ) { - std::vector> nablaRho(3, std::vector(this->cal_tool->nx, 0.)); + if (nrxx <= 0) + { + ModuleBase::WARNING_QUIT("Write_MLKEDF_Descriptors::generateTrainData_KS", "nrxx must be greater than 0"); + } + + std::vector> drho(3, std::vector(nrxx, 0.)); - this->generate_descriptor(out_dir, pelec->charge->rho, pw_rho, nablaRho); + this->generate_descriptor(out_dir, pelec->charge->rho, pw_rho, drho, nrxx); - std::vector container(this->cal_tool->nx); - std::vector containernl(this->cal_tool->nx); + std::vector enhancement(nrxx); + std::vector pauli(nrxx); - const long unsigned cshape[] = {(long unsigned) this->cal_tool->nx}; // shape of container and containernl - // enhancement factor of Pauli energy, and Pauli potential - this->cal_tool->getF_KS(psi, pelec, pw_psi, pw_rho, ucell, nablaRho, container, containernl); + this->cal_tool->getF_KS(psi, pelec, pw_psi, pw_rho, ucell, drho, enhancement, pauli); Symmetry_rho srho; - Charge* ptempRho = new Charge(); - ptempRho->nspin = PARAM.inp.nspin; - ptempRho->nrxx = this->cal_tool->nx; - ptempRho->rho_core = pelec->charge->rho_core; - ptempRho->rho = new double*[1]; - ptempRho->rho[0] = new double[this->cal_tool->nx]; - ptempRho->rhog = new std::complex*[1]; - ptempRho->rhog[0] = new std::complex[pw_rho->npw]; - - for (int ir = 0; ir < this->cal_tool->nx; ++ir){ - ptempRho->rho[0][ir] = container[ir]; - } - srho.begin(0, *ptempRho, pw_rho, ucell.symm); - for (int ir = 0; ir < this->cal_tool->nx; ++ir){ - container[ir] = ptempRho->rho[0][ir]; - } + std::vector rho_vec(nrxx); + std::vector> rhog_vec(pw_rho->npw); + double* rho_ptr = rho_vec.data(); + std::complex* rhog_ptr = rhog_vec.data(); + + std::copy(enhancement.begin(), enhancement.end(), rho_vec.begin()); + srho.begin(0, &rho_ptr, &rhog_ptr, pw_rho->npw, nullptr, pw_rho, ucell.symm); + std::copy(rho_vec.begin(), rho_vec.end(), enhancement.begin()); + + std::copy(pauli.begin(), pauli.end(), rho_vec.begin()); + srho.begin(0, &rho_ptr, &rhog_ptr, pw_rho->npw, nullptr, pw_rho, ucell.symm); + std::copy(rho_vec.begin(), rho_vec.end(), pauli.begin()); - for (int ir = 0; ir < this->cal_tool->nx; ++ir){ - ptempRho->rho[0][ir] = containernl[ir]; - } - srho.begin(0, *ptempRho, pw_rho, ucell.symm); - for (int ir = 0; ir < this->cal_tool->nx; ++ir){ - containernl[ir] = ptempRho->rho[0][ir]; - } - npy::SaveArrayAsNumpy(out_dir + "/enhancement.npy", false, 1, cshape, container); - npy::SaveArrayAsNumpy(out_dir + "/pauli.npy", false, 1, cshape, containernl); + // output data in .npy format + const long unsigned cshape[] = {(long unsigned) nrxx}; + npy::SaveArrayAsNumpy(out_dir + "/enhancement.npy", false, 1, cshape, enhancement); + npy::SaveArrayAsNumpy(out_dir + "/pauli.npy", false, 1, cshape, pauli); - for (int ir = 0; ir < this->cal_tool->nx; ++ir) + for (int ir = 0; ir < nrxx; ++ir) { - container[ir] = veff[ir]; + enhancement[ir] = veff[ir]; } - npy::SaveArrayAsNumpy(out_dir + "/veff.npy", false, 1, cshape, container); - - delete ptempRho; + npy::SaveArrayAsNumpy(out_dir + "/veff.npy", false, 1, cshape, enhancement); } void Write_MLKEDF_Descriptors::generateTrainData_KS( @@ -74,12 +67,13 @@ void Write_MLKEDF_Descriptors::generateTrainData_KS( ModulePW::PW_Basis_K *pw_psi, ModulePW::PW_Basis *pw_rho, UnitCell& ucell, - const double* veff + const double* veff, + const int nrxx ) { psi::Psi, base_device::DEVICE_CPU> psi_double(*psi); - this->generateTrainData_KS(out_dir, &psi_double, pelec, pw_psi, pw_rho, ucell, veff); + this->generateTrainData_KS(out_dir, &psi_double, pelec, pw_psi, pw_rho, ucell, veff, nrxx); } #if ((defined __CUDA) || (defined __ROCM)) @@ -90,12 +84,13 @@ void Write_MLKEDF_Descriptors::generateTrainData_KS( ModulePW::PW_Basis_K *pw_psi, ModulePW::PW_Basis *pw_rho, UnitCell& ucell, - const double* veff + const double* veff, + const int nrxx ) { psi::Psi, base_device::DEVICE_CPU> psi_cpu(*psi); - this->generateTrainData_KS(out_dir, &psi_cpu, pelec, pw_psi, pw_rho, ucell, veff); + this->generateTrainData_KS(out_dir, &psi_cpu, pelec, pw_psi, pw_rho, ucell, veff, nrxx); } void Write_MLKEDF_Descriptors::generateTrainData_KS( @@ -105,12 +100,13 @@ void Write_MLKEDF_Descriptors::generateTrainData_KS( ModulePW::PW_Basis_K *pw_psi, ModulePW::PW_Basis *pw_rho, UnitCell& ucell, - const double *veff + const double *veff, + const int nrxx ) { psi::Psi, base_device::DEVICE_CPU> psi_cpu_double(*psi); - this->generateTrainData_KS(dir, &psi_cpu_double, pelec, pw_psi, pw_rho, ucell, veff); + this->generateTrainData_KS(dir, &psi_cpu_double, pelec, pw_psi, pw_rho, ucell, veff, nrxx); } #endif @@ -118,21 +114,22 @@ void Write_MLKEDF_Descriptors::generate_descriptor( const std::string& out_dir, const double * const *prho, ModulePW::PW_Basis *pw_rho, - std::vector> &nablaRho + std::vector> &nablaRho, + const int nrxx ) { // container which will contain gamma, p, q in turn - std::vector container(this->cal_tool->nx); - std::vector new_container(this->cal_tool->nx); + std::vector container(nrxx); + std::vector new_container(nrxx); // container contains gammanl, pnl, qnl in turn - std::vector containernl(this->cal_tool->nx); - std::vector new_containernl(this->cal_tool->nx); + std::vector containernl(nrxx); + std::vector new_containernl(nrxx); - const long unsigned cshape[] = {(long unsigned) this->cal_tool->nx}; // shape of container and containernl + const long unsigned cshape[] = {(long unsigned) nrxx}; // rho - std::vector rho(this->cal_tool->nx); - for (int ir = 0; ir < this->cal_tool->nx; ++ir){ + std::vector rho(nrxx); + for (int ir = 0; ir < nrxx; ++ir){ rho[ir] = prho[0][ir]; } npy::SaveArrayAsNumpy(out_dir + "/rho.npy", false, 1, cshape, rho); @@ -236,4 +233,4 @@ std::string Write_MLKEDF_Descriptors::file_name( } -#endif \ No newline at end of file +#endif diff --git a/source/source_io/module_ml/write_mlkedf_descriptors.h b/source/source_io/module_ml/write_mlkedf_descriptors.h index 801fc401437..cd50b1ab6e8 100644 --- a/source/source_io/module_ml/write_mlkedf_descriptors.h +++ b/source/source_io/module_ml/write_mlkedf_descriptors.h @@ -30,7 +30,8 @@ class Write_MLKEDF_Descriptors ModulePW::PW_Basis_K *pw_psi, ModulePW::PW_Basis *pw_rho, UnitCell& ucell, - const double *veff + const double *veff, + const int nrxx ); void generateTrainData_KS( const std::string& out_dir, @@ -39,7 +40,8 @@ class Write_MLKEDF_Descriptors ModulePW::PW_Basis_K *pw_psi, ModulePW::PW_Basis *pw_rho, UnitCell& ucell, - const double *veff + const double *veff, + const int nrxx ); #if ((defined __CUDA) || (defined __ROCM)) @@ -50,7 +52,8 @@ class Write_MLKEDF_Descriptors ModulePW::PW_Basis_K *pw_psi, ModulePW::PW_Basis *pw_rho, UnitCell& ucell, - const double *veff + const double *veff, + const int nrxx ); void generateTrainData_KS( const std::string& dir, @@ -59,7 +62,8 @@ class Write_MLKEDF_Descriptors ModulePW::PW_Basis_K *pw_psi, ModulePW::PW_Basis *pw_rho, UnitCell& ucell, - const double *veff + const double *veff, + const int nrxx ); #endif @@ -67,7 +71,8 @@ class Write_MLKEDF_Descriptors const std::string& out_dir, const double * const *prho, ModulePW::PW_Basis *pw_rho, - std::vector> &nablaRho + std::vector> &nablaRho, + const int nrxx ); std::string file_name( diff --git a/source/source_pw/module_ofdft/ml_base.cpp b/source/source_pw/module_ofdft/ml_base.cpp index eb16ff5d940..2e04cb6231d 100644 --- a/source/source_pw/module_ofdft/ml_base.cpp +++ b/source/source_pw/module_ofdft/ml_base.cpp @@ -285,7 +285,7 @@ void ML_Base::dump_vector(std::string filename, const std::vector &data) { npy::npy_data_ptr d; d.data_ptr = data.data(); - d.shape = {(long unsigned) this->cal_tool->nx}; + d.shape = {(long unsigned) this->cal_tool->nrxx}; d.fortran_order = false; npy::write_npy(filename, d); }