fix KEDF output descriptors

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
 

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<double> *recipContainer = new std::complex<double>[pw_rho->npw];
     std::complex<double> 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<double> &pinput, std::vector<double> &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,15 +264,15 @@ double Cal_MLKEDF_Descriptors::dtanh(double tanhx, double chi)
 
 void Cal_MLKEDF_Descriptors::getGamma(const double * const *prho, std::vector<double> &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.);
     }
 }
 
 void Cal_MLKEDF_Descriptors::getP(const double * const *prho, const ModulePW::PW_Basis *pw_rho, std::vector<std::vector<double>> &pnablaRho, std::vector<double> &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<double> &pq,
 // xi = gammanl/gamma
 void Cal_MLKEDF_Descriptors::getXi(std::vector<double> &pgamma, std::vector<double> &pgammanl, std::vector<double> &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<double> &pgamma, std::vector<doub
 // tanhxi = tanh(gammanl/gamma)
 void Cal_MLKEDF_Descriptors::getTanhXi(const int ikernel, std::vector<double> &pgamma, std::vector<double> &pgammanl, std::vector<double> &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<double> &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<double> *wfcr = new std::complex<double>[this->nx];
-    ModuleBase::GlobalFunc::ZEROS(wfcr, this->nx);
+    std::complex<double> *wfcr = new std::complex<double>[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<double> wf_real = std::vector<double>(this->nx);
-                std::vector<double> wf_imag = std::vector<double>(this->nx);
-                for (int ir = 0; ir < this->nx; ++ir)
+                std::vector<double> wf_real = std::vector<double>(this->nrxx);
+                std::vector<double> wf_imag = std::vector<double>(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;

source/source_io/module_ml/cal_mlkedf_descriptors.h

@@ -90,7 +90,7 @@ class Cal_MLKEDF_Descriptors
     std::vector<double> chi_pnl = {1.0};
     std::vector<double> 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)

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<std::vector<double>> nablaRho(3, std::vector<double>(this->cal_tool->nx, 0.));
+    if (nrxx <= 0)
+    {
+        ModuleBase::WARNING_QUIT("Write_MLKEDF_Descriptors::generateTrainData_KS", "nrxx must be greater than 0");
+    }
+
+    std::vector<std::vector<double>> drho(3, std::vector<double>(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<double> container(this->cal_tool->nx);
-    std::vector<double> containernl(this->cal_tool->nx);
+    std::vector<double> enhancement(nrxx);
+    std::vector<double> 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<double>*[1];
-    ptempRho->rhog[0] = new std::complex<double>[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<double> rho_vec(nrxx);
+    std::vector<std::complex<double>> rhog_vec(pw_rho->npw);
+    double* rho_ptr = rho_vec.data();
+    std::complex<double>* 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<std::complex<double>, 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<std::complex<double>, 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,34 +100,36 @@ 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<std::complex<double>, 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
 
 void Write_MLKEDF_Descriptors::generate_descriptor(
     const std::string& out_dir,
     const double * const *prho, 
     ModulePW::PW_Basis *pw_rho,
-    std::vector<std::vector<double>> &nablaRho
+    std::vector<std::vector<double>> &nablaRho,
+    const int nrxx
 )
 {
     // container which will contain gamma, p, q in turn
-    std::vector<double> container(this->cal_tool->nx);
-    std::vector<double> new_container(this->cal_tool->nx);
+    std::vector<double> container(nrxx);
+    std::vector<double> new_container(nrxx);
     // container contains gammanl, pnl, qnl in turn
-    std::vector<double> containernl(this->cal_tool->nx);
-    std::vector<double> new_containernl(this->cal_tool->nx);
+    std::vector<double> containernl(nrxx);
+    std::vector<double> 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<double> rho(this->cal_tool->nx);
-    for (int ir = 0; ir < this->cal_tool->nx; ++ir){
+    std::vector<double> 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
+#endif

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,15 +62,17 @@ 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
 
     void generate_descriptor(
         const std::string& out_dir,
         const double * const *prho, 
         ModulePW::PW_Basis *pw_rho,
-        std::vector<std::vector<double>> &nablaRho
+        std::vector<std::vector<double>> &nablaRho,
+        const int nrxx
     );
 
     std::string file_name(

source/source_pw/module_ofdft/ml_base.cpp

@@ -285,7 +285,7 @@ void ML_Base::dump_vector(std::string filename, const std::vector<double> &data)
 {
     npy::npy_data_ptr<double> 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);
 }