diff --git a/source/source_estate/module_pot/pot_ml_exx.h b/source/source_estate/module_pot/pot_ml_exx.h index 41a4b63c0ba..fa95bb90d07 100644 --- a/source/source_estate/module_pot/pot_ml_exx.h +++ b/source/source_estate/module_pot/pot_ml_exx.h @@ -5,7 +5,7 @@ #ifdef __MLALGO #include "pot_base.h" -#include "source_pw/module_ofdft/ml_tools/nn_of.h" +#include "source_pw/module_ofdft/nn_of.h" #include "source_io/module_ml/cal_mlkedf_descriptors.h" #include "source_pw/module_ofdft/ml_base.h" diff --git a/source/source_pw/module_ofdft/CMakeLists.txt b/source/source_pw/module_ofdft/CMakeLists.txt index 8ec09bd72d2..8326faa0b4a 100644 --- a/source/source_pw/module_ofdft/CMakeLists.txt +++ b/source/source_pw/module_ofdft/CMakeLists.txt @@ -26,7 +26,7 @@ if(ENABLE_MLALGO) kedf_ml_label.cpp ml_base.cpp ml_base_pot.cpp - ml_tools/nn_of.cpp + nn_of.cpp ) add_library( diff --git a/source/source_pw/module_ofdft/ml_base.h b/source/source_pw/module_ofdft/ml_base.h index bace3c1745b..41c08d22327 100644 --- a/source/source_pw/module_ofdft/ml_base.h +++ b/source/source_pw/module_ofdft/ml_base.h @@ -8,7 +8,7 @@ #include #ifdef __MLALGO -#include "source_pw/module_ofdft/ml_tools/nn_of.h" +#include "source_pw/module_ofdft/nn_of.h" #include "source_io/module_ml/cal_mlkedf_descriptors.h" // The ML_Base class encapsulates common functionality for Machine Learning based diff --git a/source/source_pw/module_ofdft/ml_tools/CMakeLists.txt b/source/source_pw/module_ofdft/ml_tools/CMakeLists.txt deleted file mode 100644 index 96f06172167..00000000000 --- a/source/source_pw/module_ofdft/ml_tools/CMakeLists.txt +++ /dev/null @@ -1,28 +0,0 @@ -cmake_minimum_required(VERSION 3.0 FATAL_ERROR) -project(nnof) - -#list(APPEND CMAKE_PREFIX_PATH /path/to/Torch) -#list(APPEND libnpy_INCLUDE_DIR /path/to/libnpy) - -find_package(Torch REQUIRED) -include_directories(${TORCH_INCLUDE_DIRS}) - -find_path(libnpy_SOURCE_DIR npy.hpp HINTS ${libnpy_INCLUDE_DIR}) -if(NOT libnpy_SOURCE_DIR) - include(FetchContent) - FetchContent_Declare( - libnpy - GIT_REPOSITORY https://github.com/llohse/libnpy.git - GIT_SHALLOW TRUE - GIT_PROGRESS TRUE) - FetchContent_MakeAvailable(libnpy) -else() - include_directories(${libnpy_INCLUDE_DIR}) -endif() -include_directories(${libnpy_SOURCE_DIR}/include) - -add_executable(nnof main.cpp data.cpp nn_of.cpp grid.cpp input.cpp kernel.cpp pauli_potential.cpp train_kedf.cpp) -target_link_libraries(nnof "${TORCH_LIBRARIES}") -set_property(TARGET nnof PROPERTY CXX_STANDARD 14) - -#cmake -B build -Dlibnpy_INCLUDE_DIR=/path/to/libnpy/ -DTorch_DIR=/path/to/Torch/ diff --git a/source/source_pw/module_ofdft/ml_tools/data.cpp b/source/source_pw/module_ofdft/ml_tools/data.cpp deleted file mode 100644 index d5a658415f1..00000000000 --- a/source/source_pw/module_ofdft/ml_tools/data.cpp +++ /dev/null @@ -1,373 +0,0 @@ -#include "./data.h" -#include "npy.hpp" - -Data::~Data() -{ - delete[] load_gammanl; - delete[] load_pnl; - delete[] load_qnl; - delete[] load_xi; - delete[] load_tanhxi; - delete[] load_tanhxi_nl; - delete[] load_tanh_pnl; - delete[] load_tanh_qnl; - delete[] load_tanhp_nl; - delete[] load_tanhq_nl; -} - -void Data::load_data(Input &input, const int ndata, std::string *dir, const torch::Device device) -{ - if (ndata <= 0) { return; -} - this->init_label(input); - this->init_data(input.nkernel, ndata, input.fftdim, device); - this->load_data_(input, ndata, input.fftdim, dir); - std::cout << "enhancement mean: " << this->enhancement_mean << std::endl; - std::cout << "exponent: " << input.exponent << std::endl; - std::cout << "tau mean: " << this->tau_mean << std::endl; - std::cout << "pauli potential mean: " << this->pauli_mean << std::endl; - std::cout << "Load data done" << std::endl; -} - -torch::Tensor Data::get_data(std::string parameter, const int ikernel) const { - if (parameter == "gamma"){ - return this->gamma.reshape({this->nx_tot}); - } - if (parameter == "p"){ - return this->p.reshape({this->nx_tot}); - } - if (parameter == "q"){ - return this->q.reshape({this->nx_tot}); - } - if (parameter == "tanhp"){ - return this->tanhp.reshape({this->nx_tot}); - } - if (parameter == "tanhq"){ - return this->tanhq.reshape({this->nx_tot}); - } - if (parameter == "gammanl"){ - return this->gammanl[ikernel].reshape({this->nx_tot}); - } - if (parameter == "pnl"){ - return this->pnl[ikernel].reshape({this->nx_tot}); - } - if (parameter == "qnl"){ - return this->qnl[ikernel].reshape({this->nx_tot}); - } - if (parameter == "xi"){ - return this->xi[ikernel].reshape({this->nx_tot}); - } - if (parameter == "tanhxi"){ - return this->tanhxi[ikernel].reshape({this->nx_tot}); - } - if (parameter == "tanhxi_nl"){ - return this->tanhxi_nl[ikernel].reshape({this->nx_tot}); - } - if (parameter == "tanh_pnl"){ - return this->tanh_pnl[ikernel].reshape({this->nx_tot}); - } - if (parameter == "tanh_qnl"){ - return this->tanh_qnl[ikernel].reshape({this->nx_tot}); - } - if (parameter == "tanhp_nl"){ - return this->tanhp_nl[ikernel].reshape({this->nx_tot}); - } - if (parameter == "tanhq_nl"){ - return this->tanhq_nl[ikernel].reshape({this->nx_tot}); - } - return torch::zeros({}); -} - -void Data::init_label(Input &input) -{ - // Input::print("init_label"); - this->load_gammanl = new bool[input.nkernel]; - this->load_pnl = new bool[input.nkernel]; - this->load_qnl = new bool[input.nkernel]; - this->load_xi = new bool[input.nkernel]; - this->load_tanhxi = new bool[input.nkernel]; - this->load_tanhxi_nl = new bool[input.nkernel]; - this->load_tanh_pnl = new bool[input.nkernel]; - this->load_tanh_qnl = new bool[input.nkernel]; - this->load_tanhp_nl = new bool[input.nkernel]; - this->load_tanhq_nl = new bool[input.nkernel]; - - bool load_gammanl_tot = false; - bool load_pnl_tot = false; - bool load_qnl_tot = false; - // bool load_xi_tot = false; - // bool load_tanhxi_tot = false; - // bool load_tanhxi_nl_tot = false; - bool load_tanh_pnl_tot = false; - bool load_tanh_qnl_tot = false; - bool load_tanhp_nl_tot = false; - bool load_tanhq_nl_tot = false; - - for (int ik = 0; ik < input.nkernel; ++ik) - { - this->load_gammanl[ik] = input.ml_gammanl[ik]; - this->load_pnl[ik] = input.ml_pnl[ik]; - this->load_qnl[ik] = input.ml_qnl[ik]; - this->load_tanhxi_nl[ik] = input.ml_tanhxi_nl[ik]; - this->load_tanhxi[ik] = input.ml_tanhxi[ik] || input.ml_tanhxi_nl[ik]; - this->load_xi[ik] = input.ml_xi[ik] || this->load_tanhxi[ik]; - this->load_tanh_pnl[ik] = input.ml_tanh_pnl[ik]; - this->load_tanh_qnl[ik] = input.ml_tanh_qnl[ik]; - this->load_tanhp_nl[ik] = input.ml_tanhp_nl[ik]; - this->load_tanhq_nl[ik] = input.ml_tanhq_nl[ik]; - // this->load_pnl[ik] = input.ml_pnl[ik] || input.ml_tanh_pnl[ik]; - - load_gammanl_tot = load_gammanl_tot || this->load_gammanl[ik]; - load_pnl_tot = load_pnl_tot || this->load_pnl[ik]; - load_qnl_tot = load_qnl_tot || this->load_qnl[ik]; - // load_xi_tot = load_xi_tot || this->load_xi[ik]; - // load_tanhxi_tot = load_tanhxi_tot || this->load_tanhxi[ik]; - // load_tanhxi_nl_tot = load_tanhxi_nl_tot || this->load_tanhxi_nl[ik]; - load_tanh_pnl_tot = load_tanh_pnl_tot || this->load_tanh_pnl[ik]; - load_tanh_qnl_tot = load_tanh_qnl_tot || this->load_tanh_qnl[ik]; - load_tanhp_nl_tot = load_tanhp_nl_tot || this->load_tanhp_nl[ik]; - load_tanhq_nl_tot = load_tanhq_nl_tot || this->load_tanhq_nl[ik]; - - // std::cout << "load_gammanl " << this->load_gammanl[ik] << std::endl; - // std::cout << "load_pnl " << this->load_pnl[ik] << std::endl; - // std::cout << "load_qnl " << this->load_qnl[ik] << std::endl; - // std::cout << "load_tanhxi_nl " << this->load_tanhxi_nl[ik] << std::endl; - // std::cout << "load_tanhxi " << this->load_tanhxi[ik] << std::endl; - // std::cout << "load_xi " << this->load_xi[ik] << std::endl; - // std::cout << "load_tanh_pnl " << this->load_tanh_pnl[ik] << std::endl; - // std::cout << "load_tanh_qnl " << this->load_tanh_qnl[ik] << std::endl; - // std::cout << "load_tanhp_nl " << this->load_tanhp_nl[ik] << std::endl; - // std::cout << "load_tanhq_nl " << this->load_tanhq_nl[ik] << std::endl; - } - this->load_gamma = input.ml_gamma || load_gammanl_tot; - this->load_tanhp = input.ml_tanhp || load_tanhp_nl_tot || load_tanh_pnl_tot; - this->load_tanhq = input.ml_tanhq || load_tanhq_nl_tot || load_tanh_qnl_tot; - this->load_p = input.ml_p || this->load_tanhp || load_pnl_tot; - this->load_q = input.ml_q || this->load_tanhq || load_qnl_tot; - // Input::print("init_label done"); -} - -void Data::init_data(const int nkernel, const int ndata, const int fftdim, const torch::Device device) -{ - // Input::print("init_data"); - this->nx = std::pow(fftdim, 3); - this->nx_tot = this->nx * ndata; - - this->rho = torch::zeros({ndata, fftdim, fftdim, fftdim}).to(device); - if (this->load_p){ - this->nablaRho = torch::zeros({ndata, 3, fftdim, fftdim, fftdim}).to(device); - } - - this->enhancement = torch::zeros({ndata, fftdim, fftdim, fftdim}).to(device); - this->enhancement_mean = torch::zeros(ndata).to(device); - this->tau_mean = torch::zeros(ndata).to(device); - this->pauli = torch::zeros({ndata, fftdim, fftdim, fftdim}).to(device); - this->pauli_mean = torch::zeros(ndata).to(device); - - if (this->load_gamma){ - this->gamma = torch::zeros({ndata, fftdim, fftdim, fftdim}).to(device); - } - if (this->load_p){ - this->p = torch::zeros({ndata, fftdim, fftdim, fftdim}).to(device); - } - if (this->load_q){ - this->q = torch::zeros({ndata, fftdim, fftdim, fftdim}).to(device); - } - if (this->load_tanhp){ - this->tanhp = torch::zeros({ndata, fftdim, fftdim, fftdim}).to(device); - } - if (this->load_tanhq){ - this->tanhq = torch::zeros({ndata, fftdim, fftdim, fftdim}).to(device); - } - - for (int ik = 0; ik < nkernel; ++ik) - { - this->gammanl.push_back(torch::zeros({}).to(device)); - this->pnl.push_back(torch::zeros({}).to(device)); - this->qnl.push_back(torch::zeros({}).to(device)); - this->xi.push_back(torch::zeros({}).to(device)); - this->tanhxi.push_back(torch::zeros({}).to(device)); - this->tanhxi_nl.push_back(torch::zeros({}).to(device)); - this->tanh_pnl.push_back(torch::zeros({}).to(device)); - this->tanh_qnl.push_back(torch::zeros({}).to(device)); - this->tanhp_nl.push_back(torch::zeros({}).to(device)); - this->tanhq_nl.push_back(torch::zeros({}).to(device)); - - if (this->load_gammanl[ik]){ - this->gammanl[ik] = torch::zeros({ndata, fftdim, fftdim, fftdim}).to(device); - } - if (this->load_pnl[ik]){ - this->pnl[ik] = torch::zeros({ndata, fftdim, fftdim, fftdim}).to(device); - } - if (this->load_qnl[ik]){ - this->qnl[ik] = torch::zeros({ndata, fftdim, fftdim, fftdim}).to(device); - } - if (this->load_xi[ik]){ - this->xi[ik] = torch::zeros({ndata, fftdim, fftdim, fftdim}).to(device); - } - if (this->load_tanhxi[ik]){ - this->tanhxi[ik] = torch::zeros({ndata, fftdim, fftdim, fftdim}).to(device); - } - if (this->load_tanhxi_nl[ik]){ - this->tanhxi_nl[ik] = torch::zeros({ndata, fftdim, fftdim, fftdim}).to(device); - } - if (this->load_tanh_pnl[ik]){ - this->tanh_pnl[ik] = torch::zeros({ndata, fftdim, fftdim, fftdim}).to(device); - } - if (this->load_tanh_qnl[ik]){ - this->tanh_qnl[ik] = torch::zeros({ndata, fftdim, fftdim, fftdim}).to(device); - } - if (this->load_tanhp_nl[ik]){ - this->tanhp_nl[ik] = torch::zeros({ndata, fftdim, fftdim, fftdim}).to(device); - } - if (this->load_tanhq_nl[ik]){ - this->tanhq_nl[ik] = torch::zeros({ndata, fftdim, fftdim, fftdim}).to(device); - } - } - - // Input::print("init_data done"); -} - -void Data::load_data_( - Input &input, - const int ndata, - const int fftdim, - std::string *dir -) -{ - // Input::print("load_data"); - if (ndata <= 0){ - return; - } - - std::vector cshape = {(long unsigned) nx}; - std::vector container(nx); - bool fortran_order = false; - - for (int idata = 0; idata < ndata; ++idata) - { - this->loadTensor(dir[idata] + "/rho.npy", idata, fftdim, rho); - if (this->load_gamma){ - this->loadTensor(dir[idata] + "/gamma.npy", idata, fftdim, gamma); - } - if (this->load_p){ - this->loadTensor(dir[idata] + "/p.npy", idata, fftdim, p); - this->loadVector(dir[idata] + "/nablaRhox.npy", container); - // npy::LoadArrayFromNumpy(dir[idata] + "/nablaRhox.npy", cshape, fortran_order, container); - nablaRho[idata][0] = torch::tensor(container).reshape({fftdim, fftdim, fftdim}); - this->loadVector(dir[idata] + "/nablaRhoy.npy", container); - nablaRho[idata][1] = torch::tensor(container).reshape({fftdim, fftdim, fftdim}); - this->loadVector(dir[idata] + "/nablaRhoz.npy", container); - nablaRho[idata][2] = torch::tensor(container).reshape({fftdim, fftdim, fftdim}); - } - if (this->load_q){ - this->loadTensor(dir[idata] + "/q.npy", idata, fftdim, q); - } - if (this->load_tanhp){ - this->loadTensor(dir[idata] + "/tanhp.npy", idata, fftdim, tanhp); - } - if (this->load_tanhq){ - this->loadTensor(dir[idata] + "/tanhq.npy", idata, fftdim, tanhq); - } - - for (int ik = 0; ik < input.nkernel; ++ik) - { - int ktype = input.kernel_type[ik]; - double kscaling = input.kernel_scaling[ik]; - - if (this->load_gammanl[ik]){ - this->loadTensor(dir[idata] + this->file_name("gammanl", ktype, kscaling), idata, fftdim, gammanl[ik]); - } - if (this->load_pnl[ik]){ - this->loadTensor(dir[idata] + this->file_name("pnl", ktype, kscaling), idata, fftdim, pnl[ik]); - } - if (this->load_qnl[ik]){ - this->loadTensor(dir[idata] + this->file_name("qnl", ktype, kscaling), idata, fftdim, qnl[ik]); - } - if (this->load_xi[ik]){ - this->loadTensor(dir[idata] + this->file_name("xi", ktype, kscaling), idata, fftdim, xi[ik]); - } - if (this->load_tanhxi[ik]){ - this->loadTensor(dir[idata] + this->file_name("tanhxi", ktype, kscaling), idata, fftdim, tanhxi[ik]); - } - if (this->load_tanhxi_nl[ik]){ - this->loadTensor(dir[idata] + this->file_name("tanhxi_nl", ktype, kscaling), idata, fftdim, tanhxi_nl[ik]); - } - if (this->load_tanh_pnl[ik]){ - this->loadTensor(dir[idata] + this->file_name("tanh_pnl", ktype, kscaling), idata, fftdim, tanh_pnl[ik]); - } - if (this->load_tanh_qnl[ik]){ - this->loadTensor(dir[idata] + this->file_name("tanh_qnl", ktype, kscaling), idata, fftdim, tanh_qnl[ik]); - } - if (this->load_tanhp_nl[ik]){ - this->loadTensor(dir[idata] + this->file_name("tanhp_nl", ktype, kscaling), idata, fftdim, tanhp_nl[ik]); - } - if (this->load_tanhq_nl[ik]){ - this->loadTensor(dir[idata] + this->file_name("tanhq_nl", ktype, kscaling), idata, fftdim, tanhq_nl[ik]); - } - } - - this->loadTensor(dir[idata] + "/enhancement.npy", idata, fftdim, enhancement); - enhancement_mean[idata] = torch::mean(enhancement[idata]); - tau_mean[idata] = torch::mean(torch::pow(rho[idata], input.exponent/3.) * enhancement[idata]); - - if (input.loss == "potential" || input.loss == "both" || input.loss == "both_new") - { - this->loadTensor(dir[idata] + "/pauli.npy", idata, fftdim, pauli); - pauli_mean[idata] = torch::mean(pauli[idata]); - } - } - enhancement.resize_({this->nx_tot, 1}); - pauli.resize_({nx_tot, 1}); - - 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"); -} - -void Data::loadTensor( - std::string file, - const int index, - const int fftdim, - torch::Tensor &data -) -{ - npy::npy_data d = npy::read_npy(file); - data[index] = torch::tensor(d.data).reshape({fftdim, fftdim, fftdim}); -} - -void Data::loadVector(std::string file, - std::vector &data) -{ - npy::npy_data d = npy::read_npy(file); - data = d.data; -} - -void Data::dumpTensor(const torch::Tensor &data, std::string filename, int nx) -{ - std::vector v(nx); - for (int ir = 0; ir < nx; ++ir){ - v[ir] = data[ir].item(); - } - npy::npy_data_ptr d; - d.data_ptr = v.data(); - d.shape = {(long unsigned) nx}; - d.fortran_order = false; - npy::write_npy(filename, d); - std::cout << "Dumping " << filename << " done" << std::endl; -} - -std::string Data::file_name(std::string parameter, const int kernel_type, const double kernel_scaling) -{ - std::stringstream ss; - ss << "/" << parameter << "_" << kernel_type << "_" << kernel_scaling << ".npy"; - return ss.str(); -} diff --git a/source/source_pw/module_ofdft/ml_tools/data.h b/source/source_pw/module_ofdft/ml_tools/data.h deleted file mode 100644 index dbd04545dec..00000000000 --- a/source/source_pw/module_ofdft/ml_tools/data.h +++ /dev/null @@ -1,85 +0,0 @@ -#ifndef DATA_H -#define DATA_H - -#include "./input.h" - -#include - -class Data -{ - public: - // --------- load the data from .npy files ------ - ~Data(); - - int nx = 0; - int nx_tot = 0; - - // =========== data =========== - torch::Tensor rho; - torch::Tensor nablaRho; - 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; - torch::Tensor q; - torch::Tensor tanhp; - torch::Tensor tanhq; - // non-local descriptors - std::vector gammanl = {}; - std::vector pnl = {}; - std::vector qnl = {}; - std::vector xi = {}; - std::vector tanhxi = {}; - std::vector tanhxi_nl = {}; - std::vector tanh_pnl = {}; - std::vector tanh_qnl = {}; - std::vector tanhp_nl = {}; - std::vector tanhq_nl = {}; - // target - torch::Tensor enhancement; - torch::Tensor pauli; - torch::Tensor enhancement_mean; - torch::Tensor tau_mean; // mean Pauli energy - torch::Tensor pauli_mean; - - // =========== label =========== - bool load_gamma = false; - bool load_p = false; - bool load_q = false; - bool load_tanhp = false; - bool load_tanhq = false; - bool* load_gammanl = nullptr; - bool* load_pnl = nullptr; - bool* load_qnl = nullptr; - bool* load_xi = nullptr; - bool* load_tanhxi = nullptr; - bool* load_tanhxi_nl = nullptr; - bool* load_tanh_pnl = nullptr; - bool* load_tanh_qnl = nullptr; - bool* load_tanhp_nl = nullptr; - bool* load_tanhq_nl = nullptr; - - void load_data(Input &input, const int ndata, std::string *dir, const torch::Device device); - torch::Tensor get_data(std::string parameter, const int ikernel) const; - - private: - void init_label(Input &input); - 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. / 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, - const int index, - const int fftdim, - torch::Tensor &data); - void loadVector(std::string file, - std::vector &data); - // -------- dump Tensor into .npy files --------- - void dumpTensor(const torch::Tensor &data, std::string filename, int nx); - std::string file_name(std::string parameter, const int kernel_type, const double kernel_scaling); -}; -#endif \ No newline at end of file diff --git a/source/source_pw/module_ofdft/ml_tools/grid.cpp b/source/source_pw/module_ofdft/ml_tools/grid.cpp deleted file mode 100644 index d36094de7e6..00000000000 --- a/source/source_pw/module_ofdft/ml_tools/grid.cpp +++ /dev/null @@ -1,95 +0,0 @@ -#include "./grid.h" - -void Grid::initGrid(const int fftdim, - const int ndata, - const std::string *cell, - const double *a, - const torch::Device device, - double *volume) -{ - this->initGrid_(fftdim, ndata, cell, a, device, volume, this->fft_grid, this->fft_gg); - std::cout << "Init grid done" << std::endl; -} - -void Grid::initGrid_(const int fftdim, - const int ndata, - const std::string *cell, - const double *a, - const torch::Device device, - double *volume, - std::vector> &grid, - std::vector &gg) -{ - this->fft_grid = std::vector>(ndata); - this->fft_gg = std::vector(ndata); - for (int i = 0; i < ndata; ++i) - { - this->fft_grid[i] = std::vector(3); - } - - for (int it = 0; it < ndata; ++it) - { - if (cell[it] == "sc"){ - this->initScRecipGrid(fftdim, a[it], it, device, volume, grid, gg); - } - else if (cell[it] == "fcc"){ - this->initFccRecipGrid(fftdim, a[it], it, device, volume, grid, gg); - } - else if (cell[it] == "bcc"){ - this->initBccRecipGrid(fftdim, a[it], it, device, volume, grid, gg); - } - } -} - -void Grid::initScRecipGrid(const int fftdim, - const double a, - const int index, - const torch::Device device, - double *volume, - std::vector> &grid, - std::vector &gg) -{ - volume[index] = std::pow(a, 3); - torch::Tensor fre = torch::fft::fftfreq(fftdim, a / fftdim).to(device) * 2. * M_PI; - grid[index] = torch::meshgrid({fre, fre, fre}); - gg[index] = grid[index][0] * grid[index][0] + grid[index][1] * grid[index][1] + grid[index][2] * grid[index][2]; -} - -void Grid::initFccRecipGrid(const int fftdim, - const double a, - const int index, - const torch::Device device, - double *volume, - std::vector> &grid, - std::vector &gg) -{ - // std::cout << "init grid" << std::endl; - volume[index] = std::pow(a, 3) / 4.; - double coef = 1. / sqrt(2.); - // std::cout << "fftfreq" << std::endl; - torch::Tensor fre = torch::fft::fftfreq(fftdim, a * coef / fftdim).to(device) * 2. * M_PI; - auto originalGrid = torch::meshgrid({fre, fre, fre}); - grid[index][0] = coef * (-originalGrid[0] + originalGrid[1] + originalGrid[2]); - grid[index][1] = coef * (originalGrid[0] - originalGrid[1] + originalGrid[2]); - grid[index][2] = coef * (originalGrid[0] + originalGrid[1] - originalGrid[2]); - // std::cout << "gg" << std::endl; - gg[index] = grid[index][0] * grid[index][0] + grid[index][1] * grid[index][1] + grid[index][2] * grid[index][2]; -} - -void Grid::initBccRecipGrid(const int fftdim, - const double a, - const int index, - const torch::Device device, - double *volume, - std::vector> &grid, - std::vector &gg) -{ - volume[index] = std::pow(a, 3) / 2.; - double coef = sqrt(3.) / 2.; - torch::Tensor fre = torch::fft::fftfreq(fftdim, a * coef / fftdim).to(device) * 2. * M_PI; - auto originalGrid = torch::meshgrid({fre, fre, fre}); - grid[index][0] = coef * (originalGrid[1] + originalGrid[2]); - grid[index][1] = coef * (originalGrid[0] + originalGrid[2]); - grid[index][2] = coef * (originalGrid[0] + originalGrid[1]); - gg[index] = grid[index][0] * grid[index][0] + grid[index][1] * grid[index][1] + grid[index][2] * grid[index][2]; -} \ No newline at end of file diff --git a/source/source_pw/module_ofdft/ml_tools/grid.h b/source/source_pw/module_ofdft/ml_tools/grid.h deleted file mode 100644 index 9d615a0a5f2..00000000000 --- a/source/source_pw/module_ofdft/ml_tools/grid.h +++ /dev/null @@ -1,51 +0,0 @@ -#ifndef GRID_H -#define GRID_H - -#include - -class Grid -{ - public: - void initGrid(const int fftdim, - const int ndata, - const std::string *cell, - const double *a, - const torch::Device device, - double *volume); - - // fft grid - std::vector> fft_grid; // ntrain*3*fftdim*fftdim*fftdim - std::vector fft_gg; - - private: - void initGrid_(const int fftdim, - const int ndata, - const std::string *cell, - const double *a, - const torch::Device device, - double *volume, - std::vector> &grid, - std::vector &gg); - void initScRecipGrid(const int fftdim, - const double a, - const int index, - const torch::Device device, - double *volume, - std::vector> &grid, - std::vector &gg); - void initFccRecipGrid(const int fftdim, - const double a, - const int index, - const torch::Device device, - double *volume, - std::vector> &grid, - std::vector &gg); - void initBccRecipGrid(const int fftdim, - const double a, - const int index, - const torch::Device device, - double *volume, - std::vector> &grid, - std::vector &gg); -}; -#endif \ No newline at end of file diff --git a/source/source_pw/module_ofdft/ml_tools/input.cpp b/source/source_pw/module_ofdft/ml_tools/input.cpp deleted file mode 100644 index e3a7c8641c0..00000000000 --- a/source/source_pw/module_ofdft/ml_tools/input.cpp +++ /dev/null @@ -1,287 +0,0 @@ -#include "./input.h" - -void Input::readInput() -{ - std::ifstream ifs("nnINPUT", std::ios::in); - if (!ifs) - { - std::cout << " Can't find the nnINPUT file." << std::endl; - exit(0); - } - - std::string word; - int ierr = 0; - - ifs.rdstate(); - while (ifs.good()) - { - ifs >> word; - if (ifs.eof()) - break; - - if (strcmp("fftdim", word) == 0) - { - this->read_value(ifs, this->fftdim); - } - else if (strcmp("nbatch", word) == 0) - { - this->read_value(ifs, this->nbatch); - } - else if (strcmp("ntrain", word) == 0) - { - this->read_value(ifs, this->ntrain); - this->train_dir = new std::string[this->ntrain]; - this->train_cell = new std::string[this->ntrain]; - this->train_a = new double[this->ntrain]; - } - else if (strcmp("nvalidation", word) == 0) - { - this->read_value(ifs, this->nvalidation); - if (this->nvalidation > 0) - { - this->validation_dir = new std::string[this->nvalidation]; - this->validation_cell = new std::string[this->nvalidation]; - this->validation_a = new double[this->nvalidation]; - } - } - else if (strcmp("train_dir", word) == 0) - { - this->read_values(ifs, this->ntrain, this->train_dir); - } - else if (strcmp("train_cell", word) == 0) - { - this->read_values(ifs, this->ntrain, this->train_cell); - } - else if (strcmp("train_a", word) == 0) - { - this->read_values(ifs, this->ntrain, this->train_a); - } - else if (strcmp("validation_dir", word) == 0) - { - this->read_values(ifs, this->nvalidation, this->validation_dir); - } - else if (strcmp("validation_cell", word) == 0 && this->nvalidation > 0) - { - this->read_values(ifs, this->nvalidation, this->validation_cell); - } - else if (strcmp("validation_a", word) == 0 && this->nvalidation > 0) - { - this->read_values(ifs, this->nvalidation, this->validation_a); - } - else if (strcmp("loss", word) == 0) - { - this->read_value(ifs, this->loss); - } - else if (strcmp("exponent", word) == 0) - { - this->read_value(ifs, this->exponent); - } - else if (strcmp("nepoch", word) == 0) - { - this->read_value(ifs, this->nepoch); - } - else if (strcmp("lr_start", word) == 0) - { - this->read_value(ifs, this->lr_start); - } - else if (strcmp("lr_end", word) == 0) - { - this->read_value(ifs, this->lr_end); - } - else if (strcmp("lr_fre", word) == 0) - { - this->read_value(ifs, this->lr_fre); - } - else if (strcmp("dump_fre", word) == 0) - { - this->read_value(ifs, this->dump_fre); - } - else if (strcmp("print_fre", word) == 0) - { - this->read_value(ifs, this->print_fre); - } - else if (strcmp("gamma", word) == 0) - { - this->read_value(ifs, this->ml_gamma); - } - else if (strcmp("p", word) == 0) - { - this->read_value(ifs, this->ml_p); - } - else if (strcmp("q", word) == 0) - { - this->read_value(ifs, this->ml_q); - } - else if (strcmp("gammanl", word) == 0) - { - this->read_values(ifs, this->nkernel, this->ml_gammanl); - } - else if (strcmp("pnl", word) == 0) - { - this->read_values(ifs, this->nkernel, this->ml_pnl); - } - else if (strcmp("qnl", word) == 0) - { - this->read_values(ifs, this->nkernel, this->ml_qnl); - } - else if (strcmp("xi", word) == 0) - { - this->read_values(ifs, this->nkernel, this->ml_xi); - } - else if (strcmp("tanhxi", word) == 0) - { - this->read_values(ifs, this->nkernel, this->ml_tanhxi); - } - else if (strcmp("tanhxi_nl", word) == 0) - { - this->read_values(ifs, this->nkernel, this->ml_tanhxi_nl); - } - else if (strcmp("tanhp", word) == 0) - { - this->read_value(ifs, this->ml_tanhp); - } - else if (strcmp("tanhq", word) == 0) - { - this->read_value(ifs, this->ml_tanhq); - } - else if (strcmp("tanh_pnl", word) == 0) - { - this->read_values(ifs, this->nkernel, this->ml_tanh_pnl); - } - else if (strcmp("tanh_qnl", word) == 0) - { - this->read_values(ifs, this->nkernel, this->ml_tanh_qnl); - } - else if (strcmp("tanhp_nl", word) == 0) - { - this->read_values(ifs, this->nkernel, this->ml_tanhp_nl); - } - else if (strcmp("tanhq_nl", word) == 0) - { - this->read_values(ifs, this->nkernel, this->ml_tanhq_nl); - } - else if (strcmp("chi_xi", word) == 0) - { - this->read_values(ifs, this->nkernel, this->chi_xi); - } - else if (strcmp("chi_p", word) == 0) - { - this->read_value(ifs, this->chi_p); - } - else if (strcmp("chi_q", word) == 0) - { - this->read_value(ifs, this->chi_q); - } - else if (strcmp("chi_pnl", word) == 0) - { - this->read_values(ifs, this->nkernel, this->chi_pnl); - } - else if (strcmp("chi_qnl", word) == 0) - { - this->read_values(ifs, this->nkernel, this->chi_qnl); - } - else if (strcmp("feg_limit", word) == 0) - { - this->read_value(ifs, this->feg_limit); - } - else if (strcmp("change_step", word) == 0) - { - this->read_value(ifs, this->change_step); - } - else if (strcmp("coef_e", word) == 0) - { - this->read_value(ifs, this->coef_e); - } - else if (strcmp("coef_p", word) == 0) - { - this->read_value(ifs, this->coef_p); - } - else if (strcmp("coef_feg_e", word) == 0) - { - this->read_value(ifs, this->coef_feg_e); - } - else if (strcmp("coef_feg_p", word) == 0) - { - this->read_value(ifs, this->coef_feg_p); - } - else if (strcmp("check_pot", word) == 0) - { - this->read_value(ifs, this->check_pot); - } - else if (strcmp("nnode", word) == 0) - { - this->read_value(ifs, this->nnode); - } - else if (strcmp("nlayer", word) == 0) - { - this->read_value(ifs, this->nlayer); - } - else if (strcmp("nkernel", word) == 0) - { - this->read_value(ifs, this->nkernel); - this->ml_gammanl = new bool[this->nkernel]; - this->ml_pnl = new bool[this->nkernel]; - this->ml_qnl = new bool[this->nkernel]; - this->ml_xi = new bool[this->nkernel]; - this->ml_tanhxi = new bool[this->nkernel]; - this->ml_tanhxi_nl = new bool[this->nkernel]; - this->ml_tanh_pnl = new bool[this->nkernel]; - this->ml_tanh_qnl = new bool[this->nkernel]; - this->ml_tanhp_nl = new bool[this->nkernel]; - this->ml_tanhq_nl = new bool[this->nkernel]; - this->chi_xi = new double[this->nkernel]; - this->chi_pnl = new double[this->nkernel]; - this->chi_qnl = new double[this->nkernel]; - this->kernel_type = new int[this->nkernel]; - this->kernel_scaling = new double[this->nkernel]; - this->yukawa_alpha = new double[this->nkernel]; - this->kernel_file = new std::string[this->nkernel]; - for (int ik = 0; ik < this->nkernel; ++ik) - { - this->ml_gammanl[ik] = 0; - this->ml_pnl[ik] = 0; - this->ml_qnl[ik] = 0; - this->ml_xi[ik] = 0; - this->ml_tanhxi[ik] = 0; - this->ml_tanhxi_nl[ik] = 0; - this->ml_tanh_pnl[ik] = 0; - this->ml_tanh_qnl[ik] = 0; - this->ml_tanhp_nl[ik] = 0; - this->ml_tanhq_nl[ik] = 0; - this->chi_xi[ik] = 1.; - this->chi_pnl[ik] = 1.; - this->chi_qnl[ik] = 1.; - this->kernel_type[ik] = 1; - this->kernel_scaling[ik] = 1.; - this->yukawa_alpha[ik] = 1.; - this->kernel_file[ik] = "none"; - } - } - else if (strcmp("kernel_type", word) == 0) - { - this->read_values(ifs, this->nkernel, this->kernel_type); - } - else if (strcmp("yukawa_alpha", word) == 0) - { - this->read_values(ifs, this->nkernel, this->yukawa_alpha); - } - else if (strcmp("kernel_scaling", word) == 0) - { - this->read_values(ifs, this->nkernel, this->kernel_scaling); - } - else if (strcmp("kernel_file", word) == 0) - { - this->read_values(ifs, this->nkernel, this->kernel_file); - } - else if (strcmp("device_type", word) == 0) - { - 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; -} diff --git a/source/source_pw/module_ofdft/ml_tools/input.h b/source/source_pw/module_ofdft/ml_tools/input.h deleted file mode 100644 index e456c2b38ac..00000000000 --- a/source/source_pw/module_ofdft/ml_tools/input.h +++ /dev/null @@ -1,135 +0,0 @@ -#ifndef INPUT_H -#define INPUT_H - -#include - -class Input -{ - // ---------- read in the settings from nnINPUT -------- - public: - Input(){}; - ~Input() - { - delete[] this->train_dir; - delete[] this->train_cell; - delete[] this->train_a; - delete[] this->validation_dir; - delete[] this->validation_cell; - delete[] this->validation_a; - - delete[] this->ml_gammanl; - delete[] this->ml_pnl; - delete[] this->ml_qnl; - delete[] this->ml_xi; - delete[] this->ml_tanhxi; - delete[] this->ml_tanhxi_nl; - delete[] this->ml_tanh_pnl; - delete[] this->ml_tanh_qnl; - delete[] this->ml_tanhp_nl; - delete[] this->ml_tanhq_nl; - delete[] this->chi_xi; - delete[] this->chi_pnl; - delete[] this->chi_qnl; - delete[] this->kernel_type; - delete[] this->kernel_scaling; - delete[] this->yukawa_alpha; - delete[] this->kernel_file; - - }; - - void readInput(); - - template static void read_value(std::ifstream &ifs, T &var) - { - ifs >> var; - ifs.ignore(150, '\n'); - return; - } - - template static void read_values(std::ifstream &ifs, const int length, T *var) - { - for (int i = 0; i < length; ++i) - { - ifs >> var[i]; - } - ifs.ignore(150, '\n'); - return; - } - - // training - int fftdim = 0; - int nbatch = 0; - int ntrain = 1; - int nvalidation = 0; - std::string *train_dir = nullptr; - std::string *train_cell = nullptr; - double *train_a = nullptr; - std::string *validation_dir = nullptr; - std::string *validation_cell = nullptr; - double *validation_a = nullptr; - std::string loss = "both"; - int nepoch = 1000; - double lr_start = 0.01; // learning rate 2023-02-24 - 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; - int print_fre = 1; - - // descriptors - // semi-local descriptors - bool ml_gamma = false; - bool ml_p = false; - bool ml_q = false; - bool ml_tanhp = false; - bool ml_tanhq = false; - double chi_p = 1.; - double chi_q = 1.; - // non-local descriptors - bool* ml_gammanl = nullptr; - bool* ml_pnl = nullptr; - bool* ml_qnl = nullptr; - bool* ml_xi = nullptr; - bool* ml_tanhxi = nullptr; - bool* ml_tanhxi_nl = nullptr; - bool* ml_tanh_pnl = nullptr; - bool* ml_tanh_qnl = nullptr; - bool* ml_tanhp_nl = nullptr; - bool* ml_tanhq_nl = nullptr; - double* chi_xi = nullptr; - double* chi_pnl = nullptr; - double* chi_qnl = nullptr; - - int feg_limit = 0; // Free Electron Gas - int change_step = 0; // when feg_limit=3, change the output of net after change_step - - // coefficients in loss function - double coef_e = 1.; - double coef_p = 1.; - double coef_feg_e = 1.; - double coef_feg_p = 1.; - - // size of nn - int nnode = 10; - int nlayer = 3; - - // kernel - int nkernel = 1; - int* kernel_type = nullptr; - double* kernel_scaling = nullptr; - double* yukawa_alpha = nullptr; - std::string* kernel_file = nullptr; - - // GPU - std::string device_type = "gpu"; - bool check_pot = false; - - static void print(std::string message) - { - std::cout << message << std::endl; - } -}; -#endif \ No newline at end of file diff --git a/source/source_pw/module_ofdft/ml_tools/kernel.cpp b/source/source_pw/module_ofdft/ml_tools/kernel.cpp deleted file mode 100644 index 53038ef58a6..00000000000 --- a/source/source_pw/module_ofdft/ml_tools/kernel.cpp +++ /dev/null @@ -1,215 +0,0 @@ -#include "./kernel.h" - -void Kernel::fill_kernel(const int fftdim, - const int ndata, - const torch::Tensor &rho, - const double *volume, - const std::string *cell, - const torch::Device device, - const std::vector &fft_gg) -{ - double rho0 = 0.; - double tkF = 0.; - double eta = 0.; - this->kernel = std::vector(ndata); - - // read in the kernel - if (this->kernel_type == 3 || this->kernel_type == 4) // 3 for TKK Al, and 4 for TKK Si - { - this->read_kernel(fftdim, ndata, rho, volume, cell, device, fft_gg); - } - else - { - for (int id = 0; id < ndata; ++id) - { - rho0 = torch::sum(rho[id]).item() / std::pow(fftdim, 3); - std::cout << "There are " << rho0 * volume[id] << " electrons in " << cell[id] << " strcture." << std::endl; - tkF = 2. * std::pow(3. * std::pow(M_PI, 2) * rho0, 1. / 3.); - this->kernel[id] = torch::zeros({fftdim, fftdim, fftdim}).to(device); - for (int ix = 0; ix < fftdim; ++ix) - { - for (int iy = 0; iy < fftdim; ++iy) - { - for (int iz = 0; iz < fftdim; ++iz) - { - eta = sqrt(fft_gg[id][ix][iy][iz].item()) / tkF; - eta = eta * this->scaling; - if (this->kernel_type == 1) - { - this->kernel[id][ix][iy][iz] = this->wt_kernel(eta); - // this->kernel[id][ix][iy][iz] = std::pow(1. / this->scaling, 3) * this->wt_kernel(eta); - // if (ix == 26 && iy == 0 && iz == 26) - // { - // std::cout << "kernel1 " << this->kernel[id][ix][iy][iz] << std::endl; - // std::cout << "eta " << eta << std::endl; - // std::cout << "tkF " << tkF << std::endl; - // } - } - else if (this->kernel_type == 2) - { - this->kernel[id][ix][iy][iz] = this->yukawa_kernel(eta, this->yukawa_alpha); - } - } - } - } - } - } - std::cout << "Fill kernel done" << std::endl; -} - -double Kernel::wt_kernel(double eta, double tf_weight, double vw_weight) const -{ - if (eta < 0.) - { - return 0.; - } - // limit for small eta - else if (eta < 1e-10) - { - return 1. - tf_weight + eta * eta * (1./3. - 3. * vw_weight); - } - // around the singularity - else if (std::abs(eta - 1.) < 1e-10) - { - return 2. - tf_weight - 3. * vw_weight + 20. * (eta - 1); - } - // Taylor expansion for high eta - else if (eta > 3.65) - { - double eta2 = eta * eta; - double invEta2 = 1. / eta2; - double LindG = 3. * (1. - vw_weight) * eta2 - -tf_weight-0.6 - + invEta2 * (-0.13714285714285712 - + invEta2 * (-6.39999999999999875E-2 - + invEta2 * (-3.77825602968460128E-2 - + invEta2 * (-2.51824061652633074E-2 - + invEta2 * (-1.80879839616166146E-2 - + invEta2 * (-1.36715733124818332E-2 - + invEta2 * (-1.07236045520990083E-2 - + invEta2 * (-8.65192783339199453E-3 - + invEta2 * (-7.1372762502456763E-3 - + invEta2 * (-5.9945117538835746E-3 - + invEta2 * (-5.10997527675418131E-3 - + invEta2 * (-4.41060829979912465E-3 - + invEta2 * (-3.84763737842981233E-3 - + invEta2 * (-3.38745061493813488E-3 - + invEta2 * (-3.00624946457977689E-3))))))))))))))); - return LindG; - } - else - { - return 1. / (0.5 + 0.25 * (1. - eta * eta) / eta * std::log((1 + eta)/std::abs(1 - eta))) - - 3. * vw_weight * eta * eta - tf_weight; - } -} - -double Kernel::yukawa_kernel(double eta, double alpha) const -{ - return (eta == 0 && alpha == 0) ? 0. : M_PI / (eta * eta + alpha * alpha / 4.); -} - -// Read kernel from file -void Kernel::read_kernel(const int fftdim, - const int ndata, - const torch::Tensor &rho, - const double *volume, - const std::string *cell, - const torch::Device device, - const std::vector &fft_gg) -{ - std::ifstream ifs(kernel_file.c_str(), std::ios::in); - - if (!ifs) - { - std::cout << " Can't find the kernel file " << kernel_file << std::endl; - exit(0); - } - - std::cout << "Read WT kernel from " << kernel_file << std::endl; - - int kineType = 0; - double kF_in = 0.; - double rho0_in = 0.; - int nq_in = 0; - double maxEta_in = 0.; - - ifs >> kineType; - ifs >> kF_in; - ifs >> rho0_in; - ifs >> nq_in; - - double *eta_in = new double[nq_in]; - double *w0_in = new double[nq_in]; - - for (int iq = 0; iq < nq_in; ++iq) - { - ifs >> eta_in[iq] >> w0_in[iq]; - } - - maxEta_in = eta_in[nq_in-1]; - - double rho0 = 0.; - double tkF = 0.; - double eta = 0.; - for (int id = 0; id < ndata; ++id) - { - rho0 = torch::sum(rho[id]).item() / std::pow(fftdim, 3); - std::cout << "There are " << rho0 * volume[id] << " electrons in " << cell[id] << " strcture." << std::endl; - tkF = 2. * std::pow(3. * std::pow(M_PI, 2) * rho0, 1. / 3.); - this->kernel[id] = torch::zeros({fftdim, fftdim, fftdim}).to(device); - - double eta = 0.; - double maxEta = 0.; - int ind1 = 0; - int ind2 = 0; - int ind_mid = 0; - double fac1 = 0.; - double fac2 = 0.; - - for (int ix = 0; ix < fftdim; ++ix) - { - for (int iy = 0; iy < fftdim; ++iy) - { - for (int iz = 0; iz < fftdim; ++iz) - { - eta = sqrt(fft_gg[id][ix][iy][iz].item()) / tkF; - eta = eta * this->scaling; - maxEta = std::max(eta, maxEta); - - if (eta <= eta_in[0]) { - this->kernel[id][ix][iy][iz] = w0_in[0]; - } else if (eta > maxEta_in) { - this->kernel[id][ix][iy][iz] = w0_in[nq_in-1]; - } else - { - ind1 = 1; - ind2 = nq_in; - while (ind1 < ind2 - 1) - { - ind_mid = (ind1 + ind2)/2; - if (eta > eta_in[ind_mid]) - { - ind1 = ind_mid; - } - else - { - ind2 = ind_mid; - } - } - fac1 = (eta_in[ind2] - eta)/(eta_in[ind2] - eta_in[ind1]); - fac2 = (eta - eta_in[ind1])/(eta_in[ind2] - eta_in[ind1]); - this->kernel[id][ix][iy][iz] = fac1 * w0_in[ind1] + fac2 * w0_in[ind2]; - // this->kernel[id][ix][iy][iz] *= std::pow(1. / this->scaling, 3); - } - } - } - } - if (maxEta > maxEta_in) { std::cout << "Please increase the maximal eta value in KEDF kernel file" << std::endl; -} - } - - - delete[] eta_in; - delete[] w0_in; -} \ No newline at end of file diff --git a/source/source_pw/module_ofdft/ml_tools/kernel.h b/source/source_pw/module_ofdft/ml_tools/kernel.h deleted file mode 100644 index e9e614c5225..00000000000 --- a/source/source_pw/module_ofdft/ml_tools/kernel.h +++ /dev/null @@ -1,43 +0,0 @@ -#ifndef KERNEL_H -#define KERNEL_H - -#include - -class Kernel -{ - // ------------ fill the kernel in reciprocal space ---------- - public: - Kernel(){}; - - void set_para(const int kernel_type_in, const double scaling_in, const double yukawa_alpha_in, const std::string &kernel_file_in) - { - this->kernel_type = kernel_type_in; - this->scaling = scaling_in; - this->yukawa_alpha = yukawa_alpha_in; - this->kernel_file = kernel_file_in; - } - - int kernel_type = 0; // 1: WT, 2: Yukawa - double scaling = 0.; - double yukawa_alpha = 0.; - std::string kernel_file = "none"; - std::vector kernel; - - void fill_kernel(const int fftdim, - const int ndata, - const torch::Tensor &rho, - const double *volume, - const std::string *cell, - const torch::Device device, - const std::vector &fft_gg); - double wt_kernel(double eta, double tf_weight = 1., double vw_weight = 1.) const; - double yukawa_kernel(double eta, double alpha) const; - void read_kernel(const int fftdim, - const int ndata, - const torch::Tensor &rho, - const double *volume, - const std::string *cell, - const torch::Device device, - const std::vector &fft_gg); -}; -#endif \ No newline at end of file diff --git a/source/source_pw/module_ofdft/ml_tools/main.cpp b/source/source_pw/module_ofdft/ml_tools/main.cpp deleted file mode 100644 index bfd6c4ff24d..00000000000 --- a/source/source_pw/module_ofdft/ml_tools/main.cpp +++ /dev/null @@ -1,20 +0,0 @@ -#include "./train_kedf.h" - -int main() -{ - torch::set_default_dtype(caffe2::TypeMeta::fromScalarType(torch::kDouble)); - auto output = torch::get_default_dtype(); - std::cout << "Default type: " << output << std::endl; - - Train_KEDF train; - train.input.readInput(); - if (train.input.check_pot) - { - train.potTest(); - } - else - { - train.init(); - train.train(); - } -} \ No newline at end of file diff --git a/source/source_pw/module_ofdft/ml_tools/pauli_potential.cpp b/source/source_pw/module_ofdft/ml_tools/pauli_potential.cpp deleted file mode 100644 index aee2e352072..00000000000 --- a/source/source_pw/module_ofdft/ml_tools/pauli_potential.cpp +++ /dev/null @@ -1,634 +0,0 @@ -#include "./pauli_potential.h" - -void PauliPotential::init(const Input &input, - const int ninput, - const std::vector &descriptor_type, - const std::vector &kernel_index) -{ - this->fftdim = input.fftdim; - std::cout << "descriptor_type " << descriptor_type << std::endl; - std::cout << "kernel_index " << kernel_index << std::endl; - - // this->descriptor_type = descriptor_type; - // this->kernel_index = kernel_index; - - this->chi_xi = input.chi_xi; - this->chi_p = input.chi_p; - this->chi_q = input.chi_q; - this->chi_pnl = input.chi_pnl; - this->chi_qnl = input.chi_qnl; - - this->descriptor2kernel = {{"gamma", {}}, - {"p", {}}, - {"q", {}}, - {"tanhp", {}}, - {"tanhq", {}}, - {"gammanl", {}}, - {"pnl", {}}, - {"qnl", {}}, - {"xi", {}}, - {"tanhxi", {}}, - {"tanhxi_nl", {}}, - {"tanh_pnl", {}}, - {"tanh_qnl", {}}, - {"tanhp_nl", {}}, - {"tanhq_nl", {}}}; - this->descriptor2index = this->descriptor2kernel; - - for (int i = 0; i < ninput; ++i) - { - this->descriptor2kernel[descriptor_type[i]].push_back(kernel_index[i]); - std::cout << "this->descriptor2kernel[descriptor_type[i]] " << this->descriptor2kernel[descriptor_type[i]] - << std::endl; - this->descriptor2index[descriptor_type[i]].push_back(i); - std::cout << "this->descriptor2index[descriptor_type[i]] " << this->descriptor2index[descriptor_type[i]] - << std::endl; - } - std::cout << "descriptor2index " << descriptor2index << std::endl; - std::cout << "descriptor2kernel " << descriptor2kernel << std::endl; - - this->ml_gamma = this->descriptor2index["gamma"].size() > 0; - this->ml_p = this->descriptor2index["p"].size() > 0; - this->ml_q = this->descriptor2index["q"].size() > 0; - this->ml_tanhp = this->descriptor2index["tanhp"].size() > 0; - this->ml_tanhq = this->descriptor2index["tanhq"].size() > 0; - this->ml_gammanl = this->descriptor2index["gammanl"].size() > 0; - this->ml_pnl = this->descriptor2index["pnl"].size() > 0; - this->ml_qnl = this->descriptor2index["qnl"].size() > 0; - this->ml_xi = this->descriptor2index["xi"].size() > 0; - this->ml_tanhxi = this->descriptor2index["tanhxi"].size() > 0; - this->ml_tanhxi_nl = this->descriptor2index["tanhxi_nl"].size() > 0; - this->ml_tanh_pnl = this->descriptor2index["tanh_pnl"].size() > 0; - this->ml_tanh_qnl = this->descriptor2index["tanh_qnl"].size() > 0; - this->ml_tanhp_nl = this->descriptor2index["tanhp_nl"].size() > 0; - this->ml_tanhq_nl = this->descriptor2index["tanhq_nl"].size() > 0; - - if (input.energy_type == "kedf") - { - this->cLDA = this->cTF; - this->tau_exp = 5. / 3.; - } - else if (input.energy_type == "exx") - { - this->cLDA = this->cDirac; - this->tau_exp = 4. / 3.; - } -} - -torch::Tensor PauliPotential::get_potential(const int istru, - const Data &data, - const torch::Tensor &F, - const torch::Tensor &gradient, - const Kernel *kernels, - const Grid &grid) -{ - // Input::print("get potential begin"); - this->istru = istru; - torch::Tensor potential = F * this->tau_exp; - - // semi-local potential terms - if (this->ml_gamma) { - potential += this->potGammaTerm(data.gamma[istru], gradient); - } - if (this->ml_p) { - potential += this->potPTerm1(data.p[istru], gradient); - } - if (this->ml_q) { - potential += this->potQTerm1(data.q[istru], gradient); - } - if (this->ml_xi) { - potential += this->potXiTerm1(data.rho[istru], data.xi, gradient); - } - if (this->ml_tanhxi) { - potential += this->potTanhxiTerm1(data.rho[istru], data.xi, data.tanhxi, gradient); - } - if (this->ml_tanhp) { - potential += this->potTanhpTerm1(data.p[istru], data.tanhp[istru], gradient); - } - if (this->ml_tanhq) { - potential += this->potTanhqTerm1(data.q[istru], data.tanhq[istru], gradient); - } - potential *= data.tau_lda[istru] / data.rho[istru]; - - // non-local potential terms - if (this->ml_gammanl) { - potential += this->potGammanlTerm(data.rho[istru], data.gamma[istru], kernels, data.tau_lda[istru], gradient); - } - if (this->ml_p || this->ml_pnl) { - potential += this->potPPnlTerm(data.rho[istru], - data.nablaRho[istru], - data.p[istru], - kernels, - data.tau_lda[istru], - gradient, - grid.fft_grid[istru]); - } - if (this->ml_q || this->ml_qnl) { - potential += this->potQQnlTerm(data.rho[istru], - data.q[istru], - kernels, - data.tau_lda[istru], - gradient, - grid.fft_gg[istru]); - } - if (this->ml_xi) { - potential += this->potXinlTerm(data.rho[istru], kernels, data.tau_lda[istru], gradient); - } - if (this->ml_tanhxi) { - potential += this->potTanhxinlTerm(data.rho[istru], data.tanhxi, kernels, data.tau_lda[istru], gradient); - } - if (this->ml_tanhxi_nl) { - potential - += this->potTanhxi_nlTerm(data.rho[istru], data.xi, data.tanhxi, kernels, data.tau_lda[istru], gradient); - } - if ((this->ml_tanhp || this->ml_tanhp_nl) && !this->ml_tanh_pnl) { - potential += this->potTanhpTanhp_nlTerm(data.rho[istru], - data.nablaRho[istru], - data.p[istru], - data.tanhp[istru], - kernels, - data.tau_lda[istru], - gradient, - grid.fft_grid[istru]); - } - if ((this->ml_tanhq || this->ml_tanhq_nl) && !this->ml_tanh_qnl) { - potential += this->potTanhqTanhq_nlTerm(data.rho[istru], - data.q[istru], - data.tanhq[istru], - kernels, - data.tau_lda[istru], - gradient, - grid.fft_gg[istru]); - } - if (this->ml_tanh_pnl) { - potential += this->potTanhpTanh_pnlTerm(data.rho[istru], - data.nablaRho[istru], - data.p[istru], - data.tanhp[istru], - data.tanh_pnl, - kernels, - data.tau_lda[istru], - gradient, - grid.fft_grid[istru]); - } - if (this->ml_tanh_qnl) { - potential += this->potTanhqTanh_qnlTerm(data.rho[istru], - data.q[istru], - data.tanhq[istru], - data.tanh_qnl, - kernels, - data.tau_lda[istru], - gradient, - grid.fft_gg[istru]); - } - - // Input::print("get potential done"); - return potential; -} - -torch::Tensor PauliPotential::potGammaTerm(const torch::Tensor &gamma, const torch::Tensor &gradient) -{ - // std::cout << "potGammaTerm" << std::endl; - return 1. / 3. * gamma - * gradient.index({"...", this->descriptor2index["gamma"][0]}) - .reshape({this->fftdim, this->fftdim, this->fftdim}); -} - -torch::Tensor PauliPotential::potPTerm1(const torch::Tensor &p, const torch::Tensor &gradient) -{ - // std::cout << "potPTerm1" << std::endl; - return -8. / 3. * p - * gradient.index({"...", this->descriptor2index["p"][0]}) - .reshape({this->fftdim, this->fftdim, this->fftdim}); -} - -torch::Tensor PauliPotential::potQTerm1(const torch::Tensor &q, const torch::Tensor &gradient) -{ - // std::cout << "potQTerm1" << std::endl; - return -5. / 3. * q - * gradient.index({"...", this->descriptor2index["q"][0]}) - .reshape({this->fftdim, this->fftdim, this->fftdim}); -} - -torch::Tensor PauliPotential::potGammanlTerm(const torch::Tensor &rho, - const torch::Tensor &gamma, - const Kernel *kernels, - // const torch::Tensor &kernel, - const torch::Tensor &tauLDA, - const torch::Tensor &gradient) -{ - // std::cout << "potGmmamnlTerm" << std::endl; - torch::Tensor result = torch::zeros_like(gamma); - for (int ik = 0; ik < this->descriptor2kernel["gammanl"].size(); ++ik) - { - int d2k = this->descriptor2kernel["gammanl"][ik]; - int d2i = this->descriptor2index["gammanl"][ik]; - result += 1. / 3. * gamma / rho - * torch::real( - torch::fft::ifftn(torch::fft::fftn(gradient.index({"...", d2i}) - .reshape({this->fftdim, this->fftdim, this->fftdim}) - * tauLDA) - * kernels[d2k].kernel[istru])); - } - return result; -} - -torch::Tensor PauliPotential::potPPnlTerm(const torch::Tensor &rho, - const torch::Tensor &nablaRho, - const torch::Tensor &p, - const Kernel *kernels, - // const torch::Tensor &kernel, - const torch::Tensor &tauLDA, - const torch::Tensor &gradient, - const std::vector &grid) -{ - // std::cout << "potPPnlTerm" << std::endl; - torch::Tensor dFdpnl_nl = torch::zeros_like(p); - for (int ik = 0; ik < this->descriptor2index["pnl"].size(); ++ik) - { - int d2k = this->descriptor2kernel["pnl"][ik]; - int d2i = this->descriptor2index["pnl"][ik]; - dFdpnl_nl - += torch::real(torch::fft::ifftn(torch::fft::fftn(gradient.index({"...", d2i}) - .reshape({this->fftdim, this->fftdim, this->fftdim}) - * tauLDA) - * kernels[d2k].kernel[istru])); - } - - torch::Tensor temp = torch::zeros_like(nablaRho); - for (int i = 0; i < 3; ++i) - { - temp[i] = (this->ml_p) ? - this->pqcoef * 2. - * gradient.index({"...", this->descriptor2index["p"][0]}) - .reshape({this->fftdim, this->fftdim, this->fftdim}) - * nablaRho[i] * tauLDA / torch::pow(rho, 8. / 3.) - : torch::zeros_like(nablaRho[i]); - if (this->ml_pnl) { - temp[i] += -this->pqcoef * 2. * nablaRho[i] / torch::pow(rho, 8. / 3.) * dFdpnl_nl; - } - } - // std::cout << torch::slice(temp[0][0][0], 0, 0, 10); - torch::Tensor result = this->divergence(temp, grid); - - if (this->ml_pnl) { - result += -8. / 3. * p / rho * dFdpnl_nl; - } - // std::cout << torch::slice(result[0][0], 0, 20) << std::endl; - - // std::cout << "potPPnlTerm done" << std::endl; - return result; -} - -torch::Tensor PauliPotential::potQQnlTerm(const torch::Tensor &rho, - const torch::Tensor &q, - const Kernel *kernels, - // const torch::Tensor &kernel, - const torch::Tensor &tauLDA, - const torch::Tensor &gradient, - const torch::Tensor &gg) -{ - // std::cout << "potQQnlTerm" << std::endl; - torch::Tensor dFdqnl_nl = torch::zeros_like(q); - for (int ik = 0; ik < this->descriptor2index["qnl"].size(); ++ik) - { - int d2k = this->descriptor2kernel["qnl"][ik]; - int d2i = this->descriptor2index["qnl"][ik]; - dFdqnl_nl - = torch::real(torch::fft::ifftn(torch::fft::fftn(gradient.index({"...", d2i}) - .reshape({this->fftdim, this->fftdim, this->fftdim}) - * tauLDA) - * kernels[d2k].kernel[istru])); - } - - torch::Tensor temp = (this->ml_q) ? this->pqcoef - * gradient.index({"...", this->descriptor2index["q"][0]}) - .reshape({this->fftdim, this->fftdim, this->fftdim}) - * tauLDA / torch::pow(rho, 5. / 3.) - : torch::zeros_like(q); - if (this->ml_qnl) { - temp += this->pqcoef / torch::pow(rho, 5. / 3.) * dFdqnl_nl; - } - torch::Tensor result = this->Laplacian(temp, gg); - - if (this->ml_qnl) { - result += -5. / 3. * q / rho * dFdqnl_nl; - } - - // std::cout << "potQQnlTerm done" << std::endl; - return result; -} - -torch::Tensor PauliPotential::potXiTerm1(const torch::Tensor &rho, const std::vector &xi, const torch::Tensor &gradient) -{ - torch::Tensor result = torch::zeros_like(rho); - for (int ik = 0; ik < this->descriptor2kernel["xi"].size(); ++ik) - { - int d2k = this->descriptor2kernel["xi"][ik]; - int d2i = this->descriptor2index["xi"][ik]; - result += -1. / 3. * xi[d2k][istru] - * gradient.index({"...", d2i}) - .reshape({this->fftdim, this->fftdim, this->fftdim}); - } - return result; -} - -torch::Tensor PauliPotential::potTanhxiTerm1(const torch::Tensor &rho, - const std::vector &xi, - const std::vector &tanhxi, - const torch::Tensor &gradient) -{ - torch::Tensor result = torch::zeros_like(rho); - for (int ik = 0; ik < this->descriptor2kernel["tanhxi"].size(); ++ik) - { - int d2k = this->descriptor2kernel["tanhxi"][ik]; - int d2i = this->descriptor2index["tanhxi"][ik]; - result += -1. / 3. * xi[d2k][istru] * this->dtanh(tanhxi[d2k][istru], this->chi_xi[d2k]) - * gradient.index({"...", d2i}).reshape({this->fftdim, this->fftdim, this->fftdim}); - } - return result; -} - -torch::Tensor PauliPotential::potTanhpTerm1(const torch::Tensor &p, - const torch::Tensor &tanhp, - const torch::Tensor &gradient) -{ - return -8. / 3. * p * this->dtanh(tanhp, this->chi_p) - * gradient.index({"...", this->descriptor2index["tanhp"][0]}) - .reshape({this->fftdim, this->fftdim, this->fftdim}); -} - -torch::Tensor PauliPotential::potTanhqTerm1(const torch::Tensor &q, - const torch::Tensor &tanhq, - const torch::Tensor &gradient) -{ - return -5. / 3. * q * this->dtanh(tanhq, this->chi_q) - * gradient.index({"...", this->descriptor2index["tanhq"][0]}) - .reshape({this->fftdim, this->fftdim, this->fftdim}); -} - -torch::Tensor PauliPotential::potXinlTerm(const torch::Tensor &rho, - const Kernel *kernels, - // const torch::Tensor &kernel, - const torch::Tensor &tauLDA, - const torch::Tensor &gradient) -{ - torch::Tensor result = torch::zeros_like(rho); - for (int ik = 0; ik < this->descriptor2kernel["xi"].size(); ++ik) - { - int d2k = this->descriptor2kernel["xi"][ik]; - int d2i = this->descriptor2index["xi"][ik]; - result += 1. / 3. * torch::pow(rho, -2. / 3.) - * torch::real( - torch::fft::ifftn(torch::fft::fftn(gradient.index({"...", d2i}) - .reshape({this->fftdim, this->fftdim, this->fftdim}) - * tauLDA * torch::pow(rho, -1. / 3.)) - * kernels[d2k].kernel[istru])); - } - return result; -} - -torch::Tensor PauliPotential::potTanhxinlTerm(const torch::Tensor &rho, - const std::vector &tanhxi, - const Kernel *kernels, - // const torch::Tensor &kernel, - const torch::Tensor &tauLDA, - const torch::Tensor &gradient) -{ - torch::Tensor result = torch::zeros_like(rho); - - for (int ik = 0; ik < this->descriptor2kernel["tanhxi"].size(); ++ik) - { - int d2k = this->descriptor2kernel["tanhxi"][ik]; - int d2i = this->descriptor2index["tanhxi"][ik]; - result += 1. / 3. * torch::pow(rho, -2. / 3.) - * torch::real(torch::fft::ifftn( - torch::fft::fftn(gradient.index({"...", d2i}).reshape({this->fftdim, this->fftdim, this->fftdim}) - * this->dtanh(tanhxi[d2k][istru], this->chi_xi[d2k]) * tauLDA - * torch::pow(rho, -1. / 3.)) - * kernels[d2k].kernel[istru])); - } - return result; -} - -torch::Tensor PauliPotential::potTanhxi_nlTerm(const torch::Tensor &rho, - const std::vector &xi, - const std::vector &tanhxi, - const Kernel *kernels, - // const torch::Tensor &kernel, - const torch::Tensor &tauLDA, - const torch::Tensor &gradient) -{ - torch::Tensor result = torch::zeros_like(rho); - for (int ik = 0; ik < this->descriptor2kernel["tanhxi_nl"].size(); ++ik) - { - int d2k = this->descriptor2kernel["tanhxi_nl"][ik]; - int d2i = this->descriptor2index["tanhxi_nl"][ik]; - torch::Tensor dFdxi - = torch::real(torch::fft::ifftn( - torch::fft::fftn(tauLDA - * gradient.index({"...", d2i}).reshape({this->fftdim, this->fftdim, this->fftdim})) - * kernels[d2k].kernel[istru])) - * this->dtanh(tanhxi[d2k][istru], this->chi_xi[d2k]) * torch::pow(rho, -1. / 3.); - // std::cout << "tanhxi\n" << torch::slice(tanhxi[d2k][istru][0][0], 0, 0, 10) << std::endl; - // std::cout << "gradient\n" << torch::slice(gradient, 0, 0, 10) << std::endl; - // std::cout << "kernel\n" << torch::slice(kernels[d2k].kernel[istru][0][0], 0, 0, 10) << std::endl; - result += 1. / 3. * torch::pow(rho, -2. / 3.) - * (-xi[d2k][istru] * dFdxi - + torch::real(torch::fft::ifftn(torch::fft::fftn(dFdxi) * kernels[d2k].kernel[istru]))); - } - return result; -} - -torch::Tensor PauliPotential::potTanhpTanh_pnlTerm(const torch::Tensor &rho, - const torch::Tensor &nablaRho, - const torch::Tensor &p, - const torch::Tensor &tanhp, - const std::vector &tanh_pnl, - const Kernel *kernels, - // const torch::Tensor &kernel, - const torch::Tensor &tauLDA, - const torch::Tensor &gradient, - const std::vector &grid) -{ - torch::Tensor dFdpnl_nl = torch::zeros_like(tanhp); - for (int ik = 0; ik < this->descriptor2kernel["tanh_pnl"].size(); ++ik) - { - int d2k = this->descriptor2kernel["tanh_pnl"][ik]; - int d2i = this->descriptor2index["tanh_pnl"][ik]; - dFdpnl_nl += torch::real(torch::fft::ifftn( - torch::fft::fftn(gradient.index({"...", d2i}) - .reshape({this->fftdim, this->fftdim, this->fftdim}) - * this->dtanh(tanh_pnl[d2k][istru], this->chi_pnl[d2k]) - * tauLDA) - * kernels[d2k].kernel[istru])); - } - - torch::Tensor temp = torch::zeros_like(nablaRho); - for (int i = 0; i < 3; ++i) - { - temp[i] = (this->ml_tanhp) ? - this->pqcoef * 2. - * gradient.index({"...", this->descriptor2index["tanhp"][0]}) - .reshape({this->fftdim, this->fftdim, this->fftdim}) - * this->dtanh(tanhp, this->chi_p) * nablaRho[i] * tauLDA / torch::pow(rho, 8./3.) - : torch::zeros_like(nablaRho[i]); - if (this->ml_tanh_pnl) { - temp[i] += -this->pqcoef * 2. * nablaRho[i] / torch::pow(rho, 8. / 3.) * dFdpnl_nl; - } - } - torch::Tensor result = this->divergence(temp, grid); - - if (this->ml_tanh_pnl) { - result += -8. / 3. * p / rho * dFdpnl_nl; - } - - return result; -} - -torch::Tensor PauliPotential::potTanhqTanh_qnlTerm(const torch::Tensor &rho, - const torch::Tensor &q, - const torch::Tensor &tanhq, - const std::vector &tanh_qnl, - const Kernel *kernels, - // const torch::Tensor &kernel, - const torch::Tensor &tauLDA, - const torch::Tensor &gradient, - const torch::Tensor &gg) -{ - torch::Tensor dFdqnl_nl = torch::zeros_like(tanhq); - for (int ik = 0; ik < this->descriptor2kernel["tanh_qnl"].size(); ++ik) - { - int d2k = this->descriptor2kernel["tanh_qnl"][ik]; - int d2i = this->descriptor2index["tanh_qnl"][ik]; - dFdqnl_nl += torch::real(torch::fft::ifftn( - torch::fft::fftn(gradient.index({"...", d2i}) - .reshape({this->fftdim, this->fftdim, this->fftdim}) - * this->dtanh(tanh_qnl[d2k][istru], this->chi_qnl[d2k]) - * tauLDA) - * kernels[d2k].kernel[istru])); - } - - torch::Tensor temp = (this->ml_tanhq) ? this->pqcoef - * gradient.index({"...", this->descriptor2index["tanhq"][0]}) - .reshape({this->fftdim, this->fftdim, this->fftdim}) - * this->dtanh(tanhq, this->chi_q) * tauLDA / torch::pow(rho, 5. / 3.) - : torch::zeros_like(q); - if (this->ml_tanh_qnl) { - temp += this->pqcoef / torch::pow(rho, 5. / 3.) * dFdqnl_nl; - } - torch::Tensor result = this->Laplacian(temp, gg); - - if (this->ml_tanh_qnl) { - result += -5. / 3. * q / rho * dFdqnl_nl; - } - - return result; -} - -torch::Tensor PauliPotential::potTanhpTanhp_nlTerm(const torch::Tensor &rho, - const torch::Tensor &nablaRho, - const torch::Tensor &p, - const torch::Tensor &tanhp, - const Kernel *kernels, - // const torch::Tensor &kernel, - const torch::Tensor &tauLDA, - const torch::Tensor &gradient, - const std::vector &grid) -{ - torch::Tensor dFdpnl_nl = torch::zeros_like(tanhp); - for (int ik = 0; ik < this->descriptor2kernel["tanhp_nl"].size(); ++ik) - { - int d2k = this->descriptor2kernel["tanhp_nl"][ik]; - int d2i = this->descriptor2index["tanhp_nl"][ik]; - dFdpnl_nl += torch::real(torch::fft::ifftn( - torch::fft::fftn(gradient.index({"...", d2i}) - .reshape({this->fftdim, this->fftdim, this->fftdim}) - * tauLDA) - * kernels[d2k].kernel[istru])) - * this->dtanh(tanhp, this->chi_p); - } - - torch::Tensor temp = torch::zeros_like(nablaRho); - for (int i = 0; i < 3; ++i) - { - temp[i] = (this->ml_tanhp) ? - this->pqcoef * 2. - * gradient.index({"...", this->descriptor2index["tanhp"][0]}) - .reshape({this->fftdim, this->fftdim, this->fftdim}) - * this->dtanh(tanhp, this->chi_p) * nablaRho[i] * tauLDA / torch::pow(rho, 8. / 3.) - : torch::zeros_like(nablaRho[i]); - if (this->ml_tanhp_nl) { - temp[i] += -this->pqcoef * 2. * nablaRho[i] / torch::pow(rho, 8. / 3.) * dFdpnl_nl; - } - } - torch::Tensor result = this->divergence(temp, grid); - - if (this->ml_tanhp_nl) { - result += -8. / 3. * p / rho * dFdpnl_nl; - } - - return result; -} - -torch::Tensor PauliPotential::potTanhqTanhq_nlTerm(const torch::Tensor &rho, - const torch::Tensor &q, - const torch::Tensor &tanhq, - const Kernel *kernels, - // const torch::Tensor &kernel, - const torch::Tensor &tauLDA, - const torch::Tensor &gradient, - const torch::Tensor &gg) -{ - torch::Tensor dFdqnl_nl = torch::zeros_like(tanhq); - for (int ik = 0; ik < this->descriptor2kernel["tanhq_nl"].size(); ++ik) - { - int d2k = this->descriptor2kernel["tanhq_nl"][ik]; - int d2i = this->descriptor2index["tanhq_nl"][ik]; - dFdqnl_nl += torch::real(torch::fft::ifftn( - torch::fft::fftn(gradient.index({"...", d2i}) - .reshape({this->fftdim, this->fftdim, this->fftdim}) - * tauLDA) - * kernels[d2k].kernel[istru])) - * this->dtanh(tanhq, this->chi_q); - } - - torch::Tensor temp = (this->ml_tanhq) ? this->pqcoef - * gradient.index({"...", this->descriptor2index["tanhq"][0]}) - .reshape({this->fftdim, this->fftdim, this->fftdim}) - * this->dtanh(tanhq, this->chi_q) * tauLDA / torch::pow(rho, 5. / 3.) - : torch::zeros_like(q); - if (this->ml_tanhq_nl) { - temp += this->pqcoef / torch::pow(rho, 5. / 3.) * dFdqnl_nl; - } - torch::Tensor result = this->Laplacian(temp, gg); - - if (this->ml_tanhq_nl) { - result += -5. / 3. * q / rho * dFdqnl_nl; - } - - return result; -} - -torch::Tensor PauliPotential::divergence(const torch::Tensor &input, const std::vector &grid) -{ - torch::Tensor result = torch::zeros_like(input[0]); - // torch::Tensor img = torch::tensor({1.0j}); - // for (int i = 0; i < 3; ++i) - // { - // result += torch::real(torch::fft::ifftn(torch::fft::fftn(input[i]) * grid[i] * img)); - // } - for (int i = 0; i < 3; ++i) - { - result -= torch::imag(torch::fft::ifftn(torch::fft::fftn(input[i]) * grid[i])); - } - return result; -} - -torch::Tensor PauliPotential::Laplacian(const torch::Tensor &input, const torch::Tensor &gg) -{ - return torch::real(torch::fft::ifftn(torch::fft::fftn(input) * -gg)); -} - -torch::Tensor PauliPotential::dtanh(const torch::Tensor &tanhx, const double chi) -{ - return (torch::ones_like(tanhx) - tanhx * tanhx) * chi; - // return (1. - tanhx * tanhx) * chi; -} \ No newline at end of file diff --git a/source/source_pw/module_ofdft/ml_tools/pauli_potential.h b/source/source_pw/module_ofdft/ml_tools/pauli_potential.h deleted file mode 100644 index 780f3d8d1ed..00000000000 --- a/source/source_pw/module_ofdft/ml_tools/pauli_potential.h +++ /dev/null @@ -1,206 +0,0 @@ -#ifndef PAULI_POTENTIAL_H -#define PAULI_POTENTIAL_H - -#include -#include "./input.h" -#include "./data.h" -#include "./kernel.h" -#include "./grid.h" - -class PauliPotential{ - -public: - void init(const Input &input, const int ninput, const std::vector &descriptor_type, const std::vector &kernel_index); - - int fftdim = 0; - int istru = 0; - - std::map> descriptor2kernel = {}; - std::map> descriptor2index = {}; - - // semi-local descriptors - bool ml_gamma = false; - bool ml_p = false; - bool ml_q = false; - bool ml_tanhp = false; - bool ml_tanhq = false; - // non-local descriptors - bool ml_gammanl = false; - bool ml_pnl = false; - bool ml_qnl = false; - bool ml_xi = false; - bool ml_tanhxi = false; - bool ml_tanhxi_nl = false; - bool ml_tanh_pnl = false; - bool ml_tanh_qnl = false; - bool ml_tanhp_nl = false; - bool ml_tanhq_nl = false; - - double chi_p = 1.; - double chi_q = 1.; - double* chi_xi = nullptr; - double* chi_pnl = nullptr; - double* chi_qnl = nullptr; - - torch::Tensor get_potential( - const int istru, - const Data &data, - const torch::Tensor &F, - const torch::Tensor &gradient, - const Kernel *kernels, - const Grid &grid - ); -private: - - torch::Tensor potGammaTerm( - const torch::Tensor &gamma, - const torch::Tensor &gradient - ); - torch::Tensor potPTerm1( - const torch::Tensor &p, - const torch::Tensor &gradient - ); - torch::Tensor potQTerm1( - const torch::Tensor &q, - const torch::Tensor &gradient - ); - torch::Tensor potGammanlTerm( - const torch::Tensor &rho, - const torch::Tensor &gamma, - const Kernel *kernels, - // const torch::Tensor &kernel, - const torch::Tensor &tauLDA, - const torch::Tensor &gradient - ); - torch::Tensor potPPnlTerm( - const torch::Tensor &rho, - const torch::Tensor &nablaRho, - const torch::Tensor &p, - const Kernel *kernels, - // const torch::Tensor &kernel, - const torch::Tensor &tauLDA, - const torch::Tensor &gradient, - const std::vector &grid - ); - torch::Tensor potQQnlTerm( - const torch::Tensor &rho, - const torch::Tensor &q, - const Kernel *kernels, - // const torch::Tensor &kernel, - const torch::Tensor &tauLDA, - const torch::Tensor &gradient, - const torch::Tensor &gg - ); - - - torch::Tensor potXiTerm1( - const torch::Tensor &rho, - const std::vector &xi, - const torch::Tensor &gradient - ); - torch::Tensor potTanhxiTerm1( - const torch::Tensor &rho, - const std::vector &xi, - const std::vector &tanhxi, - const torch::Tensor &gradient - ); - torch::Tensor potTanhpTerm1( - const torch::Tensor &p, - const torch::Tensor &tanhp, - const torch::Tensor &gradient - ); - torch::Tensor potTanhqTerm1( - const torch::Tensor &q, - const torch::Tensor &tanhq, - const torch::Tensor &gradient - ); - torch::Tensor potXinlTerm( - const torch::Tensor &rho, - const Kernel *kernels, - // const torch::Tensor &kernel, - const torch::Tensor &tauLDA, - const torch::Tensor &gradient - ); - torch::Tensor potTanhxinlTerm( - const torch::Tensor &rho, - const std::vector &tanhxi, - const Kernel *kernels, - // const torch::Tensor &kernel, - const torch::Tensor &tauLDA, - const torch::Tensor &gradient - ); - torch::Tensor potTanhxi_nlTerm( - const torch::Tensor &rho, - const std::vector &xi, - const std::vector &tanhxi, - const Kernel *kernels, - // const torch::Tensor &kernel, - const torch::Tensor &tauLDA, - const torch::Tensor &gradient - ); - torch::Tensor potTanhpTanh_pnlTerm( - const torch::Tensor &rho, - const torch::Tensor &nablaRho, - const torch::Tensor &p, - const torch::Tensor &tanhp, - const std::vector &tanh_pnl, - const Kernel *kernels, - // const torch::Tensor &kernel, - const torch::Tensor &tauLDA, - const torch::Tensor &gradient, - const std::vector &grid - ); - torch::Tensor potTanhqTanh_qnlTerm( - const torch::Tensor &rho, - const torch::Tensor &q, - const torch::Tensor &tanhq, - const std::vector &tanh_qnl, - const Kernel *kernels, - // const torch::Tensor &kernel, - const torch::Tensor &tauLDA, - const torch::Tensor &gradient, - const torch::Tensor &gg - ); - torch::Tensor potTanhpTanhp_nlTerm( - const torch::Tensor &rho, - const torch::Tensor &nablaRho, - const torch::Tensor &p, - const torch::Tensor &tanhp, - const Kernel *kernels, - // const torch::Tensor &kernel, - const torch::Tensor &tauLDA, - const torch::Tensor &gradient, - const std::vector &grid - ); - torch::Tensor potTanhqTanhq_nlTerm( - const torch::Tensor &rho, - const torch::Tensor &q, - const torch::Tensor &tanhq, - const Kernel *kernels, - // const torch::Tensor &kernel, - const torch::Tensor &tauLDA, - const torch::Tensor &gradient, - const torch::Tensor &gg - ); - - // Tools for getting potential - torch::Tensor divergence( - const torch::Tensor &input, - const std::vector &grid - ); - torch::Tensor Laplacian( - const torch::Tensor &input, - const torch::Tensor &gg - ); - torch::Tensor dtanh( - const torch::Tensor &tanhx, - const double chi - ); - - 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 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 cLDA = cTF; - double tau_exp = 5. / 3.; - const double pqcoef = 1.0 / (4.0 * std::pow(3*std::pow(M_PI, 2.0), 2.0/3.0)); // coefficient of p and q -}; -#endif \ No newline at end of file diff --git a/source/source_pw/module_ofdft/ml_tools/train_kedf.cpp b/source/source_pw/module_ofdft/ml_tools/train_kedf.cpp deleted file mode 100644 index 6613829be75..00000000000 --- a/source/source_pw/module_ofdft/ml_tools/train_kedf.cpp +++ /dev/null @@ -1,512 +0,0 @@ -#include "./train_kedf.h" -#include -#include -#include - -Train_KEDF::~Train_KEDF() -{ - delete[] this->train_volume; - delete[] this->vali_volume; - delete[] this->kernel_train; - delete[] this->kernel_vali; -} - -void Train_KEDF::setUpFFT() -{ - this->train_volume = new double[this->input.ntrain]; - this->grid_train.initGrid( - this->input.fftdim, - this->input.ntrain, - this->input.train_cell, - this->input.train_a, - this->device, - this->train_volume - ); - this->kernel_train = new Kernel[this->input.nkernel]; - for (int ik = 0; ik < this->input.nkernel; ++ik) - { - this->kernel_train[ik].set_para(this->input.kernel_type[ik], this->input.kernel_scaling[ik], this->input.yukawa_alpha[ik], this->input.kernel_file[ik]); - this->kernel_train[ik].fill_kernel( - this->input.fftdim, - this->input.ntrain, - this->data_train.rho, - this->train_volume, - this->input.train_cell, - this->device, - this->grid_train.fft_gg - ); - } - if (this->input.nvalidation > 0){ - this->vali_volume = new double[this->input.nvalidation]; - this->grid_vali.initGrid( - this->input.fftdim, - this->input.nvalidation, - this->input.validation_cell, - this->input.validation_a, - this->device, - this->vali_volume - ); - this->kernel_vali = new Kernel[this->input.nkernel]; - for (int ik = 0; ik < this->input.nkernel; ++ik) - { - this->kernel_vali[ik].set_para(this->input.kernel_type[ik], this->input.kernel_scaling[ik], this->input.yukawa_alpha[ik], this->input.kernel_file[ik]); - this->kernel_vali[ik].fill_kernel( - this->input.fftdim, - this->input.nvalidation, - this->data_vali.rho, - this->vali_volume, - this->input.validation_cell, - this->device, - this->grid_vali.fft_gg - ); - } - } - - // this->dumpTensor(this->fft_kernel_train[0].reshape({this->data_train.nx}), "kernel_fcc.npy", this->data_train.nx); - // this->dumpTensor(this->fft_kernel_vali[0].reshape({this->data_train.nx}), "kernel_bcc.npy", this->data_train.nx); -} - -void Train_KEDF::set_device() -{ - if (this->input.device_type == "cpu") - { - std::cout << "---------- Running on CPU ----------" << std::endl; - this->device = torch::Device(torch::kCPU); - } - else if (this->input.device_type == "gpu") - { - if (torch::cuda::cudnn_is_available()) - { - std::cout << "---------- Running on GPU ----------" << std::endl; - this->device = torch::Device(torch::kCUDA); - } - else - { - std::cout << "---- Warning: GPU is unaviable -----" << std::endl; - std::cout << "---------- Running on CPU ----------" << std::endl; - this->device = torch::Device(torch::kCPU); - } - } -} - -void Train_KEDF::init_input_index() -{ - this->ninput = 0; - - // --------- semi-local descriptors --------- - if (this->input.ml_gamma){ - this->descriptor_type.push_back("gamma"); - this->kernel_index.push_back(-1); - ninput++; - } - if (this->input.ml_p){ - this->descriptor_type.push_back("p"); - this->kernel_index.push_back(-1); - ninput++; - } - if (this->input.ml_q){ - this->descriptor_type.push_back("q"); - this->kernel_index.push_back(-1); - ninput++; - } - // --------- non-local descriptors --------- - for (int ik = 0; ik < this->input.nkernel; ++ik) - { - if (this->input.ml_gammanl[ik]){ - this->descriptor_type.push_back("gammanl"); - this->kernel_index.push_back(ik); - this->ninput++; - } - if (this->input.ml_pnl[ik]){ - this->descriptor_type.push_back("pnl"); - this->kernel_index.push_back(ik); - this->ninput++; - } - if (this->input.ml_qnl[ik]){ - this->descriptor_type.push_back("qnl"); - this->kernel_index.push_back(ik); - this->ninput++; - } - if (this->input.ml_xi[ik]){ - this->descriptor_type.push_back("xi"); - this->kernel_index.push_back(ik); - this->ninput++; - } - if (this->input.ml_tanhxi[ik]){ - this->descriptor_type.push_back("tanhxi"); - this->kernel_index.push_back(ik); - this->ninput++; - } - if (this->input.ml_tanhxi_nl[ik]){ - this->descriptor_type.push_back("tanhxi_nl"); - this->kernel_index.push_back(ik); - this->ninput++; - } - } - // --------- semi-local descriptors --------- - if (this->input.ml_tanhp){ - this->descriptor_type.push_back("tanhp"); - this->kernel_index.push_back(-1); - ninput++; - } - if (this->input.ml_tanhq){ - this->descriptor_type.push_back("tanhq"); - this->kernel_index.push_back(-1); - ninput++; - } - // --------- non-local descriptors --------- - for (int ik = 0; ik < this->input.nkernel; ++ik) - { - if (this->input.ml_tanh_pnl[ik]){ - this->descriptor_type.push_back("tanh_pnl"); - this->kernel_index.push_back(ik); - this->ninput++; - } - if (this->input.ml_tanh_qnl[ik]){ - this->descriptor_type.push_back("tanh_qnl"); - this->kernel_index.push_back(ik); - this->ninput++; - } - if (this->input.ml_tanhp_nl[ik]){ - this->descriptor_type.push_back("tanhp_nl"); - this->kernel_index.push_back(ik); - this->ninput++; - } - if (this->input.ml_tanhq_nl[ik]){ - this->descriptor_type.push_back("tanhq_nl"); - this->kernel_index.push_back(ik); - this->ninput++; - } - } - - std::cout << "ninput = " << ninput << std::endl; - - if (this->input.feg_limit != 0) - { - this->feg_inpt = torch::zeros(ninput).to(device); - for (int i = 0; i < this->ninput; ++i) - { - if (this->descriptor_type[i] == "gamma") this->feg_inpt[i] = 1.; - } - - this->feg_predict = torch::zeros(1).to(device); - this->feg_dFdgamma = torch::zeros(1).to(device); - } - std::cout << "feg_limit = " << this->input.feg_limit << std::endl; -} - -void Train_KEDF::init() -{ - this->set_device(); - this->init_input_index(); - this->data_train.load_data(this->input, this->input.ntrain, this->input.train_dir, this->device); - this->data_vali.load_data(this->input, this->input.nvalidation, this->input.validation_dir, this->device); - // Input::print("LOAD DATA done"); - - this->potential.init(this->input, this->ninput, this->descriptor_type, this->kernel_index); - // Input::print("init potential done"); - if (this->input.loss == "potential" || this->input.loss == "both" || this->input.loss == "both_new") this->setUpFFT(); - // Input::print("init fft done"); - - this->nn = std::make_shared(this->data_train.nx_tot, this->data_vali.nx_tot, this->ninput, this->input.nnode, this->input.nlayer, this->device); - // Input::print("init_nn done"); - this->nn->set_data(&(this->data_train), this->descriptor_type, this->kernel_index, this->nn->inputs); - this->nn->set_data(&(this->data_vali), this->descriptor_type, this->kernel_index, this->nn->input_vali); -} - -torch::Tensor Train_KEDF::lossFunction(torch::Tensor enhancement, torch::Tensor target, torch::Tensor coef) -{ - return torch::sum(torch::pow(enhancement - target, 2))/this->data_train.nx/coef/coef; -} - -torch::Tensor Train_KEDF::lossFunction_new(torch::Tensor enhancement, torch::Tensor target, torch::Tensor weight, torch::Tensor coef) -{ - return torch::sum(torch::pow(weight * (enhancement - target), 2.))/this->data_train.nx/coef/coef; -} - - -void Train_KEDF::train() -{ - // time - double tot = 0.; - double totF = 0.; - double totFback = 0.; - double totLoss = 0.; - double totLback = 0.; - double totP = 0.; - double totStep = 0.; - std::chrono::_V2::system_clock::time_point start, startF, startFB, startL, startLB, startP, startStep, end, endF, endFB, endL, endLB, endP, endStep; - - start = std::chrono::high_resolution_clock::now(); - - std::cout << "========== Train_KEDF begin ==========" << std::endl; - // torch::Tensor target = (this->input.loss=="energy") ? this->data_train.enhancement : this->data_train.pauli; - if (this->input.loss == "potential" || this->input.loss == "both" || this->input.loss == "both_new") - { - this->data_train.pauli.resize_({this->input.ntrain, this->input.fftdim, this->input.fftdim, this->input.fftdim}); - } - - torch::optim::SGD optimizer(this->nn->parameters(), this->input.lr_start); - double update_coef = this->input.nepoch/std::log(this->input.lr_end/this->input.lr_start); // used to reduce the learning rate - - std::cout << "Epoch\tLoss\tValidation\tLoss_pot\tLoss_E\tLoss_FEG_pot\tLoss_FEG_E\n"; - double lossTrain = 0.; - double lossPot = 0.; - double lossE = 0.; - double lossFEG_pot = 0.; - double lossFEG_E = 0.; - double lossVali = 0.; - double maxLoss = 100.; - - // bool increase_coef_feg_e = false; - torch::Tensor weight = torch::pow(this->data_train.rho, this->input.exponent/3.); - for (size_t epoch = 1; epoch <= this->input.nepoch; ++epoch) - { - for (int batch_index = 0; batch_index < this->input.ntrain; ++batch_index) - { - startF = std::chrono::high_resolution_clock::now(); - - optimizer.zero_grad(); - if (this->input.loss == "energy") - { - torch::Tensor inpt = torch::slice(this->nn->inputs, 0, batch_index*this->data_train.nx, (batch_index + 1)*this->data_train.nx); - inpt.requires_grad_(true); - torch::Tensor prediction = this->nn->forward(inpt); - startL = std::chrono::high_resolution_clock::now(); - torch::Tensor loss = this->lossFunction(prediction, torch::slice(this->data_train.enhancement, 0, batch_index*this->data_train.nx, (batch_index + 1)*this->data_train.nx), this->data_train.enhancement_mean[batch_index]) * this->input.coef_e; - lossTrain = loss.item(); - endL = std::chrono::high_resolution_clock::now(); - totLoss += double(std::chrono::duration_cast(endL - startL).count()) * std::chrono::microseconds::period::num / std::chrono::microseconds::period::den; - - startLB = std::chrono::high_resolution_clock::now(); - loss.backward(); - endLB = std::chrono::high_resolution_clock::now(); - totLback += double(std::chrono::duration_cast(endLB - startLB).count()) * std::chrono::microseconds::period::num / std::chrono::microseconds::period::den; - } - else if (this->input.loss == "potential" || this->input.loss == "both" || this->input.loss == "both_new") - { - torch::Tensor inpt = torch::slice(this->nn->inputs, 0, batch_index*this->data_train.nx, (batch_index + 1)*this->data_train.nx); - inpt.requires_grad_(true); - torch::Tensor prediction = this->nn->forward(inpt); - - if (this->input.feg_limit != 3) - { - prediction = torch::softplus(prediction); - } - if (this->input.feg_limit != 0) - { - // if (this->feg_inpt.grad().numel()) this->feg_inpt.grad().zero_(); - this->feg_predict = this->nn->forward(this->feg_inpt); - // if (this->input.ml_gamma) this->feg_dFdgamma = torch::autograd::grad({this->feg_predict}, {this->feg_inpt}, - // {torch::ones_like(this->feg_predict)}, true, true)[0][this->nn_input_index["gamma"]]; - if (this->input.feg_limit == 1) prediction = prediction - torch::softplus(this->feg_predict) + 1.; - if (this->input.feg_limit == 3 && epoch < this->input.change_step) prediction = torch::softplus(prediction); - if (this->input.feg_limit == 3 && epoch >= this->input.change_step) prediction = torch::softplus(prediction - this->feg_predict + this->feg3_correct); - } - endF = std::chrono::high_resolution_clock::now(); - totF += double(std::chrono::duration_cast(endF - startF).count()) * std::chrono::microseconds::period::num / std::chrono::microseconds::period::den; - - startFB = std::chrono::high_resolution_clock::now(); - torch::Tensor gradient = torch::autograd::grad({prediction}, {inpt}, - {torch::ones_like(prediction)}, true, true)[0]; - endFB = std::chrono::high_resolution_clock::now(); - totFback += double(std::chrono::duration_cast(endFB - startFB).count()) * std::chrono::microseconds::period::num / std::chrono::microseconds::period::den; - - startL = std::chrono::high_resolution_clock::now(); - torch::Tensor pot = this->potential.get_potential(batch_index, this->data_train, prediction.reshape({this->input.fftdim, this->input.fftdim, this->input.fftdim}), gradient, this->kernel_train, this->grid_train); - torch::Tensor loss = this->lossFunction(pot, this->data_train.pauli[batch_index], this->data_train.pauli_mean[batch_index]) - * this->input.coef_p; - lossPot = loss.item(); - if (this->input.loss == "both") - { - loss = loss + this->input.coef_e * this->lossFunction(prediction, torch::slice(this->data_train.enhancement, 0, batch_index*this->data_train.nx, (batch_index + 1)*this->data_train.nx), this->data_train.enhancement_mean[batch_index]); - lossE = loss.item() - lossPot; - } - if (this->input.loss == "both_new") - { - loss = loss + this->input.coef_e * this->lossFunction_new(prediction, torch::slice(this->data_train.enhancement, 0, batch_index*this->data_train.nx, (batch_index + 1)*this->data_train.nx), weight[batch_index].reshape({this->data_train.nx, 1}), this->data_train.tau_mean[batch_index]); - lossE = loss.item() - lossPot; - } - if (this->input.feg_limit != 0) - { - if (this->input.feg_limit == 1 || this->input.feg_limit == 2) - { - loss = loss + torch::pow(this->feg_predict - 1., 2) * this->input.coef_feg_e; - lossFEG_E = loss.item() - (lossPot + lossE + lossFEG_pot); - // if (lossFEG_E/lossE < 1e-3 && increase_coef_feg_e == false) - // { - // this->input.coef_feg_e *= 2.; - // increase_coef_feg_e = true; - // std::cout << "---------ICREASE COEF FEG E--------" << std::endl; - // } - } - if (this->input.feg_limit == 3) - { - loss = loss + torch::pow(this->feg_predict - this->feg3_correct, 2) * this->input.coef_feg_e; - lossFEG_E = loss.item() - (lossPot + lossE + lossFEG_pot); - } - } - - lossTrain = loss.item(); - endL = std::chrono::high_resolution_clock::now(); - totLoss += double(std::chrono::duration_cast(endL - startL).count()) * std::chrono::microseconds::period::num / std::chrono::microseconds::period::den; - - startLB = std::chrono::high_resolution_clock::now(); - loss.backward(); - endLB = std::chrono::high_resolution_clock::now(); - totLback += double(std::chrono::duration_cast(endLB - startLB).count()) * std::chrono::microseconds::period::num / std::chrono::microseconds::period::den; - // this->dumpTensor(pot.reshape({this->data_train.nx}), "pot_fcc.npy", this->data_train.nx); - // this->dumpTensor(torch::slice(prediction, 0, batch_index*this->data_train.nx, (batch_index + 1)*this->data_train.nx).reshape({this->data_train.nx}), "F_fcc.npy", this->data_train.nx); - } - - startP = std::chrono::high_resolution_clock::now(); - if (epoch % this->input.print_fre == 0) { - if (this->input.nvalidation > 0) - { - torch::Tensor valid_pre = this->nn->forward(this->nn->input_vali); - if (this->input.feg_limit == 3) - { - valid_pre = torch::softplus(valid_pre - this->feg_predict + this->feg3_correct); - } - lossVali = this->lossFunction(valid_pre, this->data_vali.enhancement, this->data_vali.enhancement_mean).item(); - } - std::cout << std::setiosflags(std::ios::scientific) << std::setprecision(3) << epoch - << std::setw(12) << lossTrain - << std::setw(12) << lossVali - << std::setw(12) << lossPot - << std::setw(12) << lossE - << std::setw(12) << lossFEG_pot - << std::setw(12) << lossFEG_E - << std::endl; - } - - if (lossTrain > maxLoss){ - std::cout << "ERROR: too large loss: " << lossTrain << std::endl; - exit(0); - } - endP = std::chrono::high_resolution_clock::now(); - totP += double(std::chrono::duration_cast(endP - startP).count()) * std::chrono::microseconds::period::num / std::chrono::microseconds::period::den; - - startStep = std::chrono::high_resolution_clock::now(); - optimizer.step(); - endStep = std::chrono::high_resolution_clock::now(); - totStep += double(std::chrono::duration_cast(endStep - startStep).count()) * std::chrono::microseconds::period::num / std::chrono::microseconds::period::den; - } - if (epoch % this->input.dump_fre == 0) - { - std::stringstream file; - file << "model/net" << epoch << ".pt"; - torch::save(this->nn, file.str()); - } - // Reduce the learning_rate - if (epoch % this->input.lr_fre == 0) - { - for (auto &group : optimizer.param_groups()) - { - if(group.has_options()) - { - auto &options = static_cast(group.options()); - options.lr(this->input.lr_start * std::exp(epoch/update_coef)); - } - } - } - } - end = std::chrono::high_resolution_clock::now(); - - tot = double(std::chrono::duration_cast(end - start).count()) * std::chrono::microseconds::period::num / std::chrono::microseconds::period::den; - - std::cout << "=========== Done ============" << std::endl; - std::cout << std::setprecision(2) << std::setiosflags(std::ios::fixed) << "Item\t\t\tTime (s)\tPercentage" << std::endl; - std::cout.unsetf(std::ios::scientific); - std::cout << "Total\t\t\t" << tot << "\t\t" << tot/tot * 100. << " %" << std::endl; - std::cout << "Forward\t\t\t" << totF << "\t\t" << totF/tot * 100. << " %" << std::endl; - std::cout << "Enhancement back\t" << totFback << "\t\t" << totFback/tot * 100. << " %" << std::endl; - std::cout << "Loss function\t\t" << totLoss << "\t\t" << totLoss/tot * 100. << " %" << std::endl; - std::cout << "Loss backward\t\t" << totLback << "\t\t" << totLback/tot * 100. << " %" << std::endl; - std::cout << "Print\t\t\t" << totP << "\t\t" << totP/tot * 100. << " %" << std::endl; - std::cout << "Step\t\t\t" << totStep << "\t\t" << totStep/tot * 100. << " %" << std::endl; -} - -void Train_KEDF::potTest() -{ - this->set_device(); - this->init_input_index(); - this->data_train.load_data(this->input, this->input.ntrain, this->input.train_dir, this->device); - this->data_vali.load_data(this->input, this->input.nvalidation, this->input.validation_dir, this->device); - Input::print("LOAD DATA done"); - - this->potential.init(this->input, this->ninput, this->descriptor_type, this->kernel_index); - Input::print("init potential done"); - if (this->input.loss == "potential" || this->input.loss == "both" || this->input.loss == "both_new") this->setUpFFT(); - Input::print("init fft done"); - - std::chrono::_V2::system_clock::time_point start, end; - this->nn = std::make_shared(this->data_train.nx_tot, this->data_vali.nx_tot, this->ninput, this->input.nnode, this->input.nlayer, this->device); - torch::DeviceType device_type = torch::kCPU; - torch::load(this->nn, "net.pt", device_type); - // this->nn->to(this->device); - Input::print("init_nn done"); - this->nn->set_data(&(this->data_train), this->descriptor_type, this->kernel_index, this->nn->inputs); - this->nn->set_data(&(this->data_vali), this->descriptor_type, this->kernel_index, this->nn->input_vali); - - this->nn->inputs.requires_grad_(true); - - if (this->input.loss == "potential" || this->input.loss == "both" || this->input.loss == "both_new") this->data_train.pauli.resize_({this->input.ntrain, this->input.fftdim, this->input.fftdim, this->input.fftdim}); - - for (int batch_index = 0; batch_index < this->input.ntrain; ++batch_index) - { - for (int ii = 0; ii < 1; ++ii) - { - torch::Tensor inpts = torch::slice(this->nn->inputs, 0, ii*this->data_train.nx, (ii + 1)*this->data_train.nx); - inpts.requires_grad_(true); - torch::Tensor prediction = this->nn->forward(inpts); - if (this->input.feg_limit != 3) - { - prediction = torch::softplus(prediction); - } - if (this->input.feg_limit != 0) - { - // if (this->input.ml_gamma) if (this->feg_inpt[this->nn_input_index["gamma"]].grad().numel()) this->feg_inpt[this->nn_input_index["gamma"]].grad().zero_(); - if (this->feg_inpt.grad().numel()) this->feg_inpt.grad().zero_(); - this->feg_predict = this->nn->forward(this->feg_inpt); - // if (this->input.ml_gamma) this->feg_dFdgamma = torch::autograd::grad({this->feg_predict}, {this->feg_inpt[this->nn_input_index["gamma"]]}, - // {torch::ones(1)}, true, true)[0]; - // if (this->input.ml_gamma) this->feg_dFdgamma = torch::autograd::grad({this->feg_predict}, {this->feg_inpt}, - // {torch::ones_like(this->feg_predict)}, true, true)[0][this->nn_input_index["gamma"]]; - if (this->input.feg_limit == 1) prediction = prediction - torch::softplus(this->feg_predict) + 1.; - if (this->input.feg_limit == 3) prediction = torch::softplus(prediction - this->feg_predict + this->feg3_correct); - } - - start = std::chrono::high_resolution_clock::now(); - torch::Tensor gradient = torch::autograd::grad({prediction}, {inpts}, - {torch::ones_like(prediction)}, true, true)[0]; - end = std::chrono::high_resolution_clock::now(); - std::cout << "spend " << double(std::chrono::duration_cast(end - start).count()) * std::chrono::microseconds::period::num / std::chrono::microseconds::period::den << " s" << std::endl; - - std::cout << "begin potential" << std::endl; - torch::Tensor weight = torch::pow(this->data_train.rho, this->input.exponent/3.); - - torch::Tensor pot = this->potential.get_potential(ii, this->data_train, torch::slice(prediction, 0, ii*this->data_train.nx, (ii + 1)*this->data_train.nx).reshape({this->input.fftdim, this->input.fftdim, this->input.fftdim}), gradient, this->kernel_train, this->grid_train); - std::cout << "after potential" << std::endl; - - torch::Tensor loss = this->lossFunction(pot, this->data_train.pauli[ii], this->data_train.pauli_mean[ii]) * this->input.coef_p; - if (this->input.loss == "both") - { - loss = loss + this->input.coef_e * this->lossFunction(prediction, torch::slice(this->data_train.enhancement, 0, ii*this->data_train.nx, (ii + 1)*this->data_train.nx), this->data_train.enhancement_mean[ii]); - } - if (this->input.loss == "both_new") - { - loss = loss + this->input.coef_e * this->lossFunction_new(prediction, torch::slice(this->data_train.enhancement, 0, ii*this->data_train.nx, (ii + 1)*this->data_train.nx), weight[ii].reshape({this->data_train.nx, 1}), this->data_train.tau_mean[ii]); - } - std::cout << "after loss" << std::endl; - // loss = loss + this->input.coef_e * this->lossFunction(prediction, torch::slice(this->data_train.enhancement, 0, ii*this->data_train.nx, (ii + 1)*this->data_train.nx)); - double lossTrain = loss.item(); - std::cout << "loss = " << lossTrain << std::endl; - this->data_train.dumpTensor(pot.reshape({this->data_train.nx}), "potential-nnof.npy", this->data_train.nx); - this->data_train.dumpTensor(torch::slice(prediction, 0, ii*this->data_train.nx, (ii + 1)*this->data_train.nx).reshape({this->data_train.nx}), "enhancement-nnof.npy", this->data_train.nx); - // this->dumpTensor(torch::slice(this->nn->F, 0, ii*this->data_train.nx, (ii + 1)*this->data_train.nx).reshape({this->data_train.nx}), "F_fcc.npy", this->data_train.nx); - } - // std::cout << std::setiosflags(std::ios::scientific) << std::setprecision(12) << this->nn->parameters() << std::endl; - } -} - diff --git a/source/source_pw/module_ofdft/ml_tools/train_kedf.h b/source/source_pw/module_ofdft/ml_tools/train_kedf.h deleted file mode 100644 index f744de00df6..00000000000 --- a/source/source_pw/module_ofdft/ml_tools/train_kedf.h +++ /dev/null @@ -1,89 +0,0 @@ -#ifndef TRAIN_KEDF_H -#define TRAIN_KEDF_H - -#include "./data.h" -#include "./grid.h" -#include "./input.h" -#include "./kernel.h" -#include "./nn_of.h" -#include "./pauli_potential.h" - -#include - -class Train_KEDF -{ - public: - Train_KEDF(){}; - ~Train_KEDF(); - - std::shared_ptr nn; - Input input; - Grid grid_train; - Grid grid_vali; - Kernel *kernel_train = nullptr; - Kernel *kernel_vali = nullptr; - PauliPotential potential; - //----------- training set ----------- - Data data_train; - double *train_volume = nullptr; - //---------validation set ------------ - Data data_vali; - double *vali_volume = nullptr; - // ------------------------------------ - - torch::Device device = torch::Device(torch::kCUDA); - int ninput = 0; - std::vector descriptor_type = {}; - std::vector kernel_index = {}; - - // -------- free electron gas --------- - torch::Tensor feg_inpt; - torch::Tensor feg_predict; - torch::Tensor feg_dFdgamma; - - // ----------- constants --------------- - double feg3_correct = 0.541324854612918; // ln(e - 1) - 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 pqcoef = 1.0 / (4.0 * std::pow(3 * std::pow(M_PI, 2.0), 2.0 / 3.0)); // coefficient of p and q - - void train(); - void potTest(); - void setUpFFT(); - void set_device(); - void init_input_index(); - void init(); - - torch::Tensor lossFunction(torch::Tensor enhancement, torch::Tensor target, torch::Tensor coef = torch::ones(1)); - torch::Tensor lossFunction_new(torch::Tensor enhancement, - torch::Tensor target, - torch::Tensor weight, - torch::Tensor coef = torch::ones(1)); -}; - -// class OF_data : public torch::data::Dataset -// { -// private: -// torch::Tensor input; -// torch::Tensor target; - -// public: -// explicit OF_data(torch::Tensor &input, torch::Tensor &target) -// { -// this->input = input.clone(); -// this->target = target.clone(); -// } - -// torch::data::Example<> get(size_t index) override -// { -// return {this->input[index], this->target[index]}; -// } - -// torch::optional size() const override -// { -// return this->input.size(0); -// } -// }; - -#endif \ No newline at end of file diff --git a/source/source_pw/module_ofdft/ml_tools/nn_of.cpp b/source/source_pw/module_ofdft/nn_of.cpp similarity index 100% rename from source/source_pw/module_ofdft/ml_tools/nn_of.cpp rename to source/source_pw/module_ofdft/nn_of.cpp diff --git a/source/source_pw/module_ofdft/ml_tools/nn_of.h b/source/source_pw/module_ofdft/nn_of.h similarity index 100% rename from source/source_pw/module_ofdft/ml_tools/nn_of.h rename to source/source_pw/module_ofdft/nn_of.h