Feature: Support ML EXX for training script.

Author: sunliang98

source/source_pw/hamilt_ofdft/ml_tools/data.cpp

@@ -208,7 +208,7 @@ void Data::init_data(const int nkernel, const int ndata, const int fftdim, const
         if (this->load_tanhxi[ik]){
             this->tanhxi[ik]    = torch::zeros({ndata, fftdim, fftdim, fftdim}).to(device);
         }
-        if (this->load_tanhxi_nl[ik{
+        if (this->load_tanhxi_nl[ik]){
             this->tanhxi_nl[ik] = torch::zeros({ndata, fftdim, fftdim, fftdim}).to(device);
         }
         if (this->load_tanh_pnl[ik]){
@@ -319,7 +319,16 @@ void Data::load_data_(
     enhancement.resize_({this->nx_tot, 1});
     pauli.resize_({nx_tot, 1});
 
-    this->tau_tf = this->cTF * torch::pow(this->rho, 5./3.);
+    if (input.energy_type == "kedf")
+    {
+        this->tau_exp = 5. / 3.;
+        this->tau_lda = this->cTF * torch::pow(this->rho, this->tau_exp);
+    }
+    else if (input.energy_type == "exx")
+    {
+        this->tau_exp = 4. / 3.;
+        this->tau_lda = this->cDirac * torch::pow(this->rho, this->tau_exp);
+    }
     // Input::print("load_data done");
 }
 

source/source_pw/hamilt_ofdft/ml_tools/data.h

@@ -17,7 +17,7 @@ class Data
     // =========== data ===========
     torch::Tensor rho;
     torch::Tensor nablaRho;
-    torch::Tensor tau_tf;
+    torch::Tensor tau_lda; // energy density of LDA, i.e. TF for KEDF, Dirac term for EXX
     // semi-local descriptors
     torch::Tensor gamma;
     torch::Tensor p;
@@ -67,7 +67,9 @@ class Data
     void init_data(const int nkernel, const int ndata, const int fftdim, const torch::Device device);
     void load_data_(Input &input, const int ndata, const int fftdim, std::string *dir);
 
-    const double cTF = 3.0/10.0 * std::pow(3*std::pow(M_PI, 2.0), 2.0/3.0) * 2; // 10/3*(3*pi^2)^{2/3}, multiply by 2 to convert unit from Hartree to Ry, finally in Ry*Bohr^(-2)
+    const double cTF = 3. /10. * std::pow(3. * std::pow(M_PI, 2.), 2. / 3.) * 2.; // 10/3*(3*pi^2)^{2/3}, multiply by 2 to convert unit from Hartree to Ry, finally in Ry*Bohr^(-2)
+    const double cDirac = - 3. /4. * std::pow(3. / M_PI, 1./3.) * 2.; // -3/4*(3/pi)^{1/3}, multiply by 2 to convert unit from Hartree to Ry, finally in Ry*Bohr^(-2)
+    double tau_exp = 5. / 3.; // 5/3 for TF KEDF, and 4/3 for Dirac term
 
   public:
     void loadTensor(std::string file,

source/source_pw/hamilt_ofdft/ml_tools/input.cpp

@@ -277,6 +277,10 @@ void Input::readInput()
         {
             this->read_value(ifs, this->device_type);
         }
+        else if (strcmp("energy_type", word) == 0)
+        {
+            this->read_value(ifs, this->energy_type);
+        }
     }
 
     std::cout << "Read nnINPUT done" << std::endl;

source/source_pw/hamilt_ofdft/ml_tools/input.h

@@ -73,6 +73,7 @@ class Input
     double lr_end = 1e-4;
     int lr_fre = 5000;
     double exponent = 5.; // exponent of weight rho^{exponent/3.}
+    std::string energy_type = "kedf"; // kedf or exx
 
     // output
     int dump_fre = 1;