fs_nonlocal_tools.h

Go to the documentation of this file.

#ifndef W_ABACUS_DEVELOP_ABACUS_DEVELOP_SOURCE_source_pw_HAMILT_PWDFT_FS_NONLOCAL_TOOLS_H
#define W_ABACUS_DEVELOP_ABACUS_DEVELOP_SOURCE_source_pw_HAMILT_PWDFT_FS_NONLOCAL_TOOLS_H
 
#include "source_base/module_device/device.h"
#include "source_basis/module_pw/pw_basis_k.h"
#include "source_cell/klist.h"
#include "source_cell/unitcell.h"
#include "source_pw/module_pwdft/VNL_in_pw.h"
#include "source_pw/module_pwdft/kernels/stress_op.h"
#include "source_base/kernels/math_kernel_op.h"
#include "source_psi/psi.h"
 
#include <complex>
 
namespace hamilt
{
 
template <typename FPTYPE, typename Device>
class FS_Nonlocal_tools
{
  public:
    FS_Nonlocal_tools(const pseudopot_cell_vnl* nlpp_in,
                      const UnitCell* ucell_in,
                      const K_Vectors* kv_in,
                      const ModulePW::PW_Basis_K* wfc_basis_in,
                      const Structure_Factor* sf_in,
                      const ModuleBase::matrix& wg,
                      const ModuleBase::matrix* p_ekb);
    ~FS_Nonlocal_tools();
 
    void cal_vkb(const int& ik, const int& nbdall);
 
    void cal_becp(const int& ik, const int& npm, const std::complex<FPTYPE>* ppsi, const int& nbd0 = 0);
    
    void reduce_pool_becp(const int& npm);
 
    void cal_vkb_deri_s(const int& ik, const int& nbdall, const int& ipol, const int& jpol);
 
    void cal_dbecp_s(const int& ik, const int& npm, const std::complex<FPTYPE>* ppsi, const int& nbd0 = 0);
 
    void cal_stress(const int& ik,
                    const int& npm,
                    const bool& occ,
                    const int& ipol,
                    const int& jpol,
                    FPTYPE* stress,
                    const int& nbd0 = 0);
 
        void cal_vkb_deri_f(const int& ik, const int& nbdall, const int& ipol);
    void cal_dbecp_f(const int& ik, const int& nbdall, const int& npm, const int& ipol, const std::complex<FPTYPE>* ppsi, const int& nbd0 = 0);
    void cal_force(const int& ik, const int& nbdall, const int& npm, const bool& occ, FPTYPE* force, const int& nbd0 = 0);
 
    void revert_vkb(const int& ik, const int& ipol);
 
  private:
    void allocate_memory(const ModuleBase::matrix& wg, const ModuleBase::matrix* p_ekb);
    void delete_memory();
 
  private:
    const Structure_Factor* sf_;
    const pseudopot_cell_vnl* nlpp_;
    const UnitCell* ucell_;
    const K_Vectors* kv_;
    const ModulePW::PW_Basis_K* wfc_basis_;
 
    Device* ctx = {};
    base_device::DEVICE_CPU* cpu_ctx = {};
    base_device::AbacusDevice_t device = {};
    int nkb;
    int nbands;
 
    int max_nh = 0;
    int max_npw = 0;
    int ntype;
    bool nondiagonal;
    int pre_ik_s = -1;
    int pre_ik_f = -1;
 
    int* atom_nh = nullptr;
    int* atom_na = nullptr;
    std::vector<int> h_atom_nh;
    std::vector<int> h_atom_na;
 
    int* gcar_zero_indexes = nullptr;
    int gcar_zero_counts[3] = {0, 0, 0};
    std::complex<FPTYPE>* vkb_save = nullptr;
    void transfer_gcar(const int& npw, const int& npw_max, const FPTYPE* gcar_in);
    void save_vkb(const int& npw, const int& ipol);
 
    FPTYPE* d_wg = nullptr;
    FPTYPE* d_wk = nullptr;
    FPTYPE* d_ekb = nullptr;
    FPTYPE* gcar = nullptr;
    FPTYPE* deeq = nullptr;
    FPTYPE* kvec_c = nullptr;
    FPTYPE* qq_nt = nullptr;
    std::complex<FPTYPE>* ppcell_vkb = nullptr;
    std::vector<FPTYPE> g_plus_k;
    FPTYPE* hd_ylm = nullptr;              // (lmax + 1) * (lmax + 1) * npw
    FPTYPE* hd_ylm_deri = nullptr;         // 3 * (lmax + 1) * (lmax + 1) * npw
    FPTYPE* hd_vq = nullptr;               // this->ucell->atoms[it].ncpp.nbeta * npw
    FPTYPE* hd_vq_deri = nullptr;          // this->ucell->atoms[it].ncpp.nbeta * npw
    std::complex<FPTYPE>* hd_sk = nullptr; // this->ucell->nat * npw
    FPTYPE* d_g_plus_k = nullptr;              // npw * 5
    FPTYPE* d_pref = nullptr;                  // this->ucell->atoms[it].ncpp.nh
    FPTYPE* d_gk = nullptr;                    // this->ucell->atoms[it].ncpp.nh * npw
    FPTYPE* d_vq_tab = nullptr;                // this->ucell->atoms[it].ncpp.nbeta * npw
    std::vector<int> dvkb_indexes;             // this->ucell->atoms[it].ncpp.nh * 4
    int* d_dvkb_indexes = nullptr;             // this->ucell->atoms[it].ncpp.nh * 4
    std::complex<FPTYPE>* d_pref_in = nullptr; // this->ucell->atoms[it].ncpp.nh
 
    std::complex<FPTYPE>* dbecp = nullptr; // nbands * nkb (for stress) or nbands * nkb * 3 (for force)
    std::complex<FPTYPE>* becp = nullptr;  // nbands * nkb
 
    using gemm_op = ModuleBase::gemm_op<std::complex<FPTYPE>, Device>;
    using cal_stress_nl_op = hamilt::cal_stress_nl_op<FPTYPE, Device>;
    using cal_dbecp_noevc_nl_op = hamilt::cal_dbecp_noevc_nl_op<FPTYPE, Device>;
 
    using resmem_complex_op = base_device::memory::resize_memory_op<std::complex<FPTYPE>, Device>;
    using resmem_complex_h_op = base_device::memory::resize_memory_op<std::complex<FPTYPE>, base_device::DEVICE_CPU>;
    using setmem_complex_op = base_device::memory::set_memory_op<std::complex<FPTYPE>, Device>;
    using delmem_complex_op = base_device::memory::delete_memory_op<std::complex<FPTYPE>, Device>;
    using delmem_complex_h_op = base_device::memory::delete_memory_op<std::complex<FPTYPE>, base_device::DEVICE_CPU>;
    using syncmem_complex_h2d_op
        = base_device::memory::synchronize_memory_op<std::complex<FPTYPE>, Device, base_device::DEVICE_CPU>;
    using syncmem_complex_d2h_op
        = base_device::memory::synchronize_memory_op<std::complex<FPTYPE>, base_device::DEVICE_CPU, Device>;
 
    using resmem_var_op = base_device::memory::resize_memory_op<FPTYPE, Device>;
    using resmem_var_h_op = base_device::memory::resize_memory_op<FPTYPE, base_device::DEVICE_CPU>;
    using setmem_var_op = base_device::memory::set_memory_op<FPTYPE, Device>;
    using delmem_var_op = base_device::memory::delete_memory_op<FPTYPE, Device>;
    using delmem_var_h_op = base_device::memory::delete_memory_op<FPTYPE, base_device::DEVICE_CPU>;
    using syncmem_var_h2d_op = base_device::memory::synchronize_memory_op<FPTYPE, Device, base_device::DEVICE_CPU>;
    using syncmem_var_d2h_op = base_device::memory::synchronize_memory_op<FPTYPE, base_device::DEVICE_CPU, Device>;
 
    using resmem_int_op = base_device::memory::resize_memory_op<int, Device>;
    using delmem_int_op = base_device::memory::delete_memory_op<int, Device>;
    using syncmem_int_h2d_op = base_device::memory::synchronize_memory_op<int, Device, base_device::DEVICE_CPU>;
 
    using cal_vq_op = hamilt::cal_vq_op<FPTYPE, Device>;
    using cal_vq_deri_op = hamilt::cal_vq_deri_op<FPTYPE, Device>;
 
    using cal_vkb_op = hamilt::cal_vkb_op<FPTYPE, Device>;
    using cal_vkb_deri_op = hamilt::cal_vkb_deri_op<FPTYPE, Device>;
};
 
} // namespace hamilt
 
#endif