pw_basis_k.h

Go to the documentation of this file.

#ifndef PWBASISK_H
#define PWBASISK_H
 
#include "pw_basis.h"
#include "source_psi/psi.h"
#include "source_base/module_device/device.h"
namespace ModulePW
{
 
class PW_Basis_K : public PW_Basis
{
 
public:
    PW_Basis_K();
    PW_Basis_K(std::string device_, std::string precision_) : PW_Basis(device_, precision_) {classname="PW_Basis_K";}
    ~PW_Basis_K();
 
    //init parameters of pw_basis_k class
    void initparameters(
        const bool gamma_only_in,
        const double ecut_in,
        const int nk_in, //number of k points in this pool
        const ModuleBase::Vector3<double> *kvec_d, // Direct coordinates of k points
        const int distribution_type_in = 1,
        const bool xprime_in = true
    );
 
  public:
    int nks=0;//number of k points in this pool
    ModuleBase::Vector3<double> *kvec_d=nullptr; // Direct coordinates of k points
    ModuleBase::Vector3<double> *kvec_c=nullptr; // Cartesian coordinates of k points
    int *npwk=nullptr; //[nks] number of plane waves of different k-points
    int npwk_max=0; //max npwk among all nks k-points, it may be smaller than npw
                  //npw cutoff: (|g|+|k|)^2, npwk in the the npw ball, thus is smaller
    double gk_ecut=0; //Energy cut off for (g+k)^2/2
 
public:
    //prepare for transforms between real and reciprocal spaces
    void setuptransform();
 
    int *igl2isz_k=nullptr, * d_igl2isz_k = nullptr; //[npwk_max*nks] map (igl,ik) to (is,iz)
    int *igl2ig_k=nullptr;//[npwk_max*nks] map (igl,ik) to ig
    int *ig2ixyz_k=nullptr; 
    std::vector<int> ig2ixyz_k_cpu; 
    double *gk2=nullptr; // modulus (G+K)^2 of G vectors [npwk_max*nks]
 
    // liuyu add 2023-09-06
    double erf_ecut=0.0;   // the value of the constant energy cutoff
    double erf_height=0.0; // the height of the energy step for reciprocal vectors
    double erf_sigma=0.0;  // the width of the energy step for reciprocal vectors
 
    //collect gdirect, gcar, gg
    void collect_local_pw(const double& erf_ecut_in = 0.0,
                          const double& erf_height_in = 0.0,
                          const double& erf_sigma_in = 0.1);
 
  private:
    float  * s_gk2 = nullptr;
    double * d_gk2 = nullptr; // modulus (G+K)^2 of G vectors [npwk_max*nks]
    //create igl2isz_k map array for fft
    void setupIndGk();
    // get ig2ixyz_k
    void get_ig2ixyz_k();
    //calculate G+K, it is a private function
    ModuleBase::Vector3<double> cal_GplusK_cartesian(const int ik, const int ig) const;
 
  public:
    template <typename FPTYPE>
    void real2recip(const FPTYPE* in,
                    std::complex<FPTYPE>* out,
                    const int ik,
                    const bool add = false,
                    const FPTYPE factor = 1.0) const; // in:(nplane,nx*ny)  ; out(nz, ns)
    template <typename FPTYPE>
    void real2recip(const std::complex<FPTYPE>* in,
                    std::complex<FPTYPE>* out,
                    const int ik,
                    const bool add = false,
                    const FPTYPE factor = 1.0) const; // in:(nplane,nx*ny)  ; out(nz, ns)
    template <typename FPTYPE>
    void recip2real(const std::complex<FPTYPE>* in,
                    FPTYPE* out,
                    const int ik,
                    const bool add = false,
                    const FPTYPE factor = 1.0) const; // in:(nz, ns)  ; out(nplane,nx*ny)
    template <typename FPTYPE>
    void recip2real(const std::complex<FPTYPE>* in,
                    std::complex<FPTYPE>* out,
                    const int ik,
                    const bool add = false,
                    const FPTYPE factor = 1.0) const; // in:(nz, ns)  ; out(nplane,nx*ny)
    #if defined(__DSP)
    template <typename FPTYPE, typename Device>
    void convolution(const Device* ctx,
                      const int ik,
                      const int size,
                      const std::complex<FPTYPE>* input,
                      const FPTYPE*               input1,
                      std::complex<FPTYPE>*       output,
                      const bool add = false,
                      const FPTYPE factor =1.0) const ;
 
    template <typename FPTYPE>
    void real2recip_dsp(const std::complex<FPTYPE>* in,
                       std::complex<FPTYPE>* out,
                       const int ik,
                       const bool add = false,
                       const FPTYPE factor = 1.0) const; // in:(nplane,nx*ny)  ; out(nz, ns)
    template <typename FPTYPE>
    void recip2real_dsp(const std::complex<FPTYPE>* in,
                       std::complex<FPTYPE>* out,
                       const int ik,
                       const bool add = false,
                       const FPTYPE factor = 1.0) const; // in:(nz, ns)  ; out(nplane,nx*ny)
    
    #endif
 
     template <typename FPTYPE, typename Device>
    void real_to_recip(const Device* ctx,
                       const std::complex<FPTYPE>* in,
                       std::complex<FPTYPE>* out,
                       const int ik,
                       const bool add = false,
                       const FPTYPE factor = 1.0) const; // in:(nplane,nx*ny)  ; out(nz, ns)
    template <typename FPTYPE, typename Device>
    void recip_to_real(const Device* ctx,
                       const std::complex<FPTYPE>* in,
                       std::complex<FPTYPE>* out,
                       const int ik,
                       const bool add = false,
                       const FPTYPE factor = 1.0) const; // in:(nz, ns)  ; out(nplane,nx*ny)
 
 
    template <typename TK,
              typename Device,
              typename std::enable_if<std::is_same<Device, base_device::DEVICE_CPU>::value, int>::type = 0>
    void real_to_recip(const TK* in,
                       TK* out,
                       const int ik,
                       const bool add = false,
                       const typename GetTypeReal<TK>::type factor = 1.0) const
    {
      #if defined(__DSP)
        this->real2recip_dsp(in, out, ik, add, factor);
      #else
        this->real2recip(in,out,ik,add,factor);
      #endif
    }
    template <typename TK,
              typename Device,
              typename std::enable_if<std::is_same<Device, base_device::DEVICE_CPU>::value, int>::type = 0>
    void recip_to_real(const TK* in,
                       TK* out,
                       const int ik,
                       const bool add = false,
                       const typename GetTypeReal<TK>::type factor = 1.0) const
    {
      
      #if defined(__DSP)
        this->recip2real_dsp(in,out,ik,add,factor);
      #else
        this->recip2real(in,out,ik,add,factor);
      #endif
    }
    template <typename FPTYPE>
    void real2recip_gpu(const std::complex<FPTYPE>* in,
                    std::complex<FPTYPE>* out,
                    const int ik,
                    const bool add = false,
                    const FPTYPE factor = 1.0) const; // in:(nplane,nx*ny)  ; out(nz, ns)
                    
    template <typename FPTYPE>
    void recip2real_gpu(const std::complex<FPTYPE>* in,
                    std::complex<FPTYPE>* out,
                    const int ik,
                    const bool add = false,
                    const FPTYPE factor = 1.0) const; // in:(nz, ns)  ; out(nplane,nx*ny)
 
    template <typename FPTYPE,
              typename Device,
              typename std::enable_if<!std::is_same<Device, base_device::DEVICE_CPU>::value, int>::type = 0>
    void real_to_recip(const FPTYPE* in,
                       FPTYPE* out,
                       const int ik,
                       const bool add = false,
                       const typename GetTypeReal<FPTYPE>::type factor = 1.0) const
    {
        this->real2recip_gpu(in, out, ik, add, factor);
    }
 
    template <typename TK,
              typename Device,
              typename std::enable_if<std::is_same<Device, base_device::DEVICE_GPU>::value, int>::type = 0>
    void recip_to_real(const TK* in,
                       TK* out,
                       const int ik,
                       const bool add = false,
                       const typename GetTypeReal<TK>::type factor = 1.0) const
    {
        this->recip2real_gpu(in, out, ik, add, factor);
    }
 
  public:
    //operator:
    //get (G+K)^2:
    double& getgk2(const int ik, const int igl) const;
    //get G
    ModuleBase::Vector3<double>& getgcar(const int ik, const int igl) const;
    //get G-direct
    ModuleBase::Vector3<double> getgdirect(const int ik, const int igl) const;
    //get (G+K)
    ModuleBase::Vector3<double> getgpluskcar(const int ik, const int igl) const;
    //get igl2isz_k
    int& getigl2isz(const int ik, const int igl) const;
    //get igl2ig_k or igk(ik,ig) in older ABACUS
    int& getigl2ig(const int ik, const int igl) const;
 
    //get ig_to_ix
    std::vector<int> get_ig2ix(const int ik) const;
    //get ig_to_iy
    std::vector<int> get_ig2iy(const int ik) const;
    //get ig_to_iz
    std::vector<int> get_ig2iz(const int ik) const;
 
    template <typename FPTYPE> FPTYPE * get_gk2_data() const;
    template <typename FPTYPE> FPTYPE * get_gcar_data() const;
    template <typename FPTYPE> FPTYPE * get_kvec_c_data() const;
 
private:
    float * s_gcar = nullptr, * s_kvec_c = nullptr;
    double * d_gcar = nullptr, * d_kvec_c = nullptr;
};
 
}
#endif //PlaneWave_K class
 
#include "./pw_basis_k_big.h" //temporary it will be removed