onsite_proj_tools.h

Go to the documentation of this file.

#ifndef MODULEHAMILTPW_ONSITEPROJTOOLS_H
#define MODULEHAMILTPW_ONSITEPROJTOOLS_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_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 Onsite_Proj_tools
{
  public:
    Onsite_Proj_tools(
        const pseudopot_cell_vnl* nlpp_in,
        const UnitCell* ucell_in,
        const psi::Psi<std::complex<FPTYPE>, Device>* psi_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& ekb
    );
 
    // a more general constructor is in the following
    Onsite_Proj_tools(
        const std::vector<int>& nproj,     // number of projectors for each atom type
        const std::vector<int>& lproj,
        const ModuleBase::realArray& tab,  // radials' spherical bessel transform
        const ModuleBase::matrix& nhtol,
        std::complex<FPTYPE>* vkb_buf,
        const UnitCell* ucell_in,
        const psi::Psi<std::complex<FPTYPE>, Device>* psi_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& ekb
    );
 
    ~Onsite_Proj_tools();
 
    void cal_becp(int ik, int npm, std::complex<FPTYPE>* becp_in = nullptr, const std::complex<FPTYPE>* ppsi_in = nullptr);
    void cal_dbecp_s(int ik, int npm, int ipol, int jpol);
    void cal_dbecp_f(int ik, int npm, int ipol);
 
    void cal_force_dftu(
        int ik,
        int npm,
        FPTYPE* force,
        const int* orbital_corr,
        const std::complex<FPTYPE>* vu,
        const int size_vu,
        const FPTYPE* h_wg
    );
 
    void cal_force_dspin(
        int ik,
        int npm,
        FPTYPE* force,
        const ModuleBase::Vector3<double>* lambda,
        const FPTYPE* h_wg
    );
 
    // return stress(i,j) value
    double cal_stress_dftu(
        int ik,
        int npm,
        const int* orbital_corr,
        const std::complex<FPTYPE>* vu,
        const int size_vu,
        const FPTYPE* h_wg
    );
 
    // return stress(i,j) value
    double cal_stress_dspin(
        int ik,
        int npm,
        const ModuleBase::Vector3<double>* lambda,
        const FPTYPE* h_wg
    );
 
 
    std::complex<FPTYPE>* get_becp() const { return becp; }
    std::complex<FPTYPE>* get_dbecp() const { return dbecp; }
 
  private:
    void allocate_memory(
        const ModuleBase::matrix& wg,
        const ModuleBase::matrix& ekb,
        const std::vector<int>& nproj,
        const std::vector<int>& nch
    );
    void delete_memory();
 
  private:
    const Structure_Factor* sf_ = nullptr;
    const pseudopot_cell_vnl* nlpp_ = nullptr;
    const UnitCell* ucell_ = nullptr;
    const psi::Psi<std::complex<FPTYPE>, Device>* psi_;
    const K_Vectors* kv_ = nullptr;
    const ModulePW::PW_Basis_K* wfc_basis_ = nullptr;
 
    Device* ctx = {};
    base_device::DEVICE_CPU* cpu_ctx = {};
    base_device::AbacusDevice_t device = {};
    int nkb = 0;
    int nbands = 0;
    int deeq_dims[4] = {0, 0, 0, 0};    // deeq can be something other than that in pseudopotentials
    int deeq_nc_dims[4] = {0, 0, 0, 0};
 
    int current_ik = -1;
 
    int max_nh = 0;
    int max_npw = 0;
    int ntype = 0;
    bool nondiagonal = false;
    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;
    std::vector<int> nproj;
 
    int* gcar_zero_indexes = nullptr;
    int gcar_zero_counts[3] = {0, 0, 0};
    std::complex<FPTYPE>* vkb_save = nullptr;
    void transfer_gcar(int npw, int npw_max, const FPTYPE* gcar_in);
    void save_vkb(int npw, int ipol);
    void revert_vkb(int npw, int ipol);
 
    const ModuleBase::realArray* tabtpr = nullptr;
    const ModuleBase::matrix* nhtol = nullptr;
    int lprojmax = -1;
 
    FPTYPE* d_wg = nullptr;
    FPTYPE* d_ekb = nullptr;
    FPTYPE* gcar = nullptr;
 
    FPTYPE* deeq = nullptr;
    std::complex<FPTYPE>* deeq_nc = 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