LCAO_deepks_io.h

Go to the documentation of this file.

#ifndef LCAO_DEEPKS_IO_H
#define LCAO_DEEPKS_IO_H
 
#ifdef __MLALGO
 
#include "deepks_param.h"
#include "source_base/complexmatrix.h"
#include "source_base/matrix.h"
#include "source_base/tool_title.h"
 
#include <string>
#include <torch/script.h>
#include <torch/torch.h>
#include <vector>
 
namespace LCAO_deepks_io
{
 
 
 
 
 
constexpr double Ry2Hartree = 0.5;
 
template <typename TK>
void print_dm(const int nks, const int nlocal, const int nrow, const std::vector<std::vector<TK>>& dm);
 
void load_npy_gedm(const int nat, const int des_per_atom, double** gedm, double& e_delta, const int rank);
 
void save_npy_d(const int nat,
                const bool deepks_equiv,
                const DeePKS_Param& deepks_param,
                const std::vector<torch::Tensor>& descriptor,
                const std::string& dm_eig_file,
                const int rank);
 
// save energy
void save_npy_e(const double& e, 
                const std::string& e_file,
                const int rank,
                const double unit_scale = Ry2Hartree);
 
// save Hamiltonian
template <typename TK, typename TH>
void save_npy_h(const std::vector<TH>& hamilt,
                const std::string& h_file,
                const int nlocal,
                const int nks,
                const int rank,
                const double unit_scale = Ry2Hartree);
 
void save_matrix2npy(const std::string& file_name,
                     const ModuleBase::matrix& matrix,
                     const int rank,
                     const double& scale = 1.0,
                     const char mode = 'N',
                     const double unit_scale = Ry2Hartree);
 
template <typename T>
void save_tensor2npy(const std::string& file_name, const torch::Tensor& tensor, const int rank);
}; // namespace LCAO_deepks_io
 
#endif
#endif