pw_basis.h

Go to the documentation of this file.

#ifndef PWBASIS_H
#define PWBASIS_H
 
#include "source_base/macros.h"
#include "source_base/module_device/memory_op.h"
#include "source_base/matrix.h"
#include "source_base/matrix3.h"
#include "source_base/vector3.h"
#include <complex>
#include "module_fft/fft_bundle.h"
#include <cstring>
#ifdef __MPI
#include "mpi.h"
#endif
 
namespace ModulePW
{
 
class PW_Basis
{
 
public:
    std::string classname;
    PW_Basis();
    PW_Basis(std::string device_, std::string precision_);
    virtual ~PW_Basis();
    //Init mpi parameters
#ifdef __MPI
    void initmpi(
        const int poolnproc_in, // Number of processors in this pool
        const int poolrank_in, // Rank in this pool
        MPI_Comm pool_world_in //Comm world for pw_basis
    );
#endif
 
    //Init the grids for FFT
    virtual void initgrids(
        const double lat0_in, //unit length (unit in bohr)
        const ModuleBase::Matrix3 latvec_in, // Unitcell lattice vectors (unit in lat0)
        const double gridecut //unit in Ry, ecut to set up grids
    );
    //Init the grids for FFT
    virtual void initgrids(
        const double lat0_in,
        const ModuleBase::Matrix3 latvec_in, // Unitcell lattice vectors
        const int nx_in, int ny_in, int nz_in
    );
 
    //Init some parameters
    void initparameters(
        const bool gamma_only_in,
        const double pwecut_in, //unit in Ry, ecut to decides plane waves
        const int distribution_type_in = 1,
        const bool xprime_in = true
    );
 
    //Set parameters about full planewave, only used in OFDFT for now.
    void setfullpw(
        const bool inpt_full_pw = false,
        const int inpt_full_pw_dim = 0
    );
//===============================================
//                 distribution maps
//===============================================
public:
#ifdef __MPI
    MPI_Comm pool_world=MPI_COMM_NULL;
#endif
 
    int *ig2isz=nullptr; // map ig to (is, iz).
    int *ig2ixyz_gpu = nullptr;
    int *istot2ixy=nullptr; // istot2ixy[is]: iy + ix * ny of is^th stick among all sticks.
    int *is2fftixy=nullptr, * d_is2fftixy = nullptr; // is2fftixy[is]: iy + ix * ny of is^th stick among sticks on current proc.
    int *fftixy2ip=nullptr; // fftixy2ip[iy + ix * fftny]: ip of proc which contains stick on (ix, iy). if no stick: -1
    int nst=0; //num. of sticks in current proc.
    int *nst_per=nullptr;// nst on each core
    int nstnz=0; // nst * nz
    int nstot=0; //num. of sticks in total.
    int npw=0; //num. of plane waves in current proc.
    int *npw_per=nullptr; //npw on each core
    int npwtot=0; // total num. of plane waves in all proc. in this pool
 
    //real space
    int nrxx=0; //num. of real space grids
    int *startz=nullptr; //startz[ip]: starting z plane in the ip-th proc. in current POOL_WORLD
    int *numz=nullptr; //numz[ip]: num. of z planes in the ip-th proc. in current POOL_WORLD
    int *numg=nullptr; //numg[ip] :  nst_per[poolrank] * numz[ip]
    int *numr=nullptr; //numr[ip] :  numz[poolrank] * nst_per[ip]
    int *startg=nullptr;  // startg[ip] = numg[ip-1] + startg[ip-1]
    int *startr=nullptr;  // startr[ip] = numr[ip-1] + startr[ip-1]
    int startz_current=0;
    int nplane=0; //num. of planes in current proc.
 
    ModuleBase::Vector3<double> *gdirect=nullptr;       //(= *G1d) ; // ig = new Vector igc[npw]
    ModuleBase::Vector3<double> *gcar=nullptr;              //G vectors in cartesian corrdinate
    double *gg=nullptr;         // modulus (G^2) of G vectors [npw]
    //gg[ng]=ig[ng]*GGT*ig[ng]/(lat0*lat0)=g[ng]*g[ng] (/lat0*lat0)
    // gg_global dimension: [cutgg_num_now] (save memory skill is used)
    int ig_gge0=-1;    //ig when gg == 0
 
    //distribute plane waves and grids and set up fft
    void setuptransform();
 
protected:
    int *startnsz_per=nullptr;//useless intermediate variable// startnsz_per[ip]: starting is * nz stick in the ip^th proc.
 
    //distribute plane waves to different processors
    void distribute_g();
 
    //distribute real-space grids to different processors
    virtual void distribute_r();
 
    //prepare for MPI_Alltoall
    void getstartgr();
 
 
public:
    //collect gdirect, gcar, gg
    void collect_local_pw();
 
public:
    int ngg=0; //number of different modulus (G^2) of G vectors
    int *ig2igg=nullptr;//[npw] map ig to igg(<ngg: the index of G^2)
    double *gg_uniq=nullptr; //[ngg] modulus (G^2) of G vectors of igg, each gg of igg is unique.
    //collect gg_uniq
    void collect_uniqgg();
 
 
public:
  bool gamma_only = false; 
  bool full_pw = false; 
  double ggecut = 0.0;  
  double gridecut_lat = 0.0;  
  double lat0 = 1;            
  double tpiba = 1;           
  double tpiba2 = 1;          
  ModuleBase::Matrix3 latvec; 
  ModuleBase::Matrix3 G;      
  ModuleBase::Matrix3 GT;     
  ModuleBase::Matrix3 GGT;    
  double omega = 1.0;         
  int distribution_type = 1;  
  int full_pw_dim = 0; 
  int poolnproc = 1;
  int poolrank = 0;
 
protected:
    //distribute plane waves to different processors
    //method 1: first consider number of plane waves
    void distribution_method1();
    // Distribute sticks to cores in method 1.
    void divide_sticks_1(
        int* st_i,          // x or x + fftnx (if x < 0) of stick.
        int* st_j,          // y or y + fftny (if y < 0) of stick.
        int* st_length     // the stick on (ix, iy) consists of st_length[ix*fftny+iy] planewaves.
    );
 
    //method 2: first consider number of sticks
    void distribution_method2();
    // Distribute sticks to cores in method 2.
    void divide_sticks_2();
 
    //Count the total number of planewaves (tot_npw) and sticks (this->nstot) (in distributeg method1 and method2)
    void count_pw_st(
        int* st_length2D, // the number of planewaves that belong to the stick located on (x, y).
        int* st_bottom2D  // the z-coordinate of the bottom of stick on (x, y).
    );
 
    //get ig2isz and is2fftixy
    void get_ig2isz_is2fftixy(
        int* st_bottom,     // minimum z of stick, stored in 1d array with tot_nst elements.
        int* st_length     // the stick on (x, y) consists of st_length[x*fftny+y] planewaves.
    );
 
    //Collect the x, y indexs, length of the sticks (in distributeg method1)
    void collect_st(
        int* st_length2D,                               // the number of planewaves that belong to the stick located on (x, y), stored in 2d x-y plane.
        int* st_bottom2D,                               // the z-coordinate of the bottom of stick on (x, y), stored in 2d x-y plane.
        int* st_i,                                      // x or x + fftnx (if x < 0) of stick.
        int* st_j,                                      // y or y + fftny (if y < 0) of stick.
        int* st_length                                 // number of planewaves in stick, stored in 1d array with tot_nst elements.
    );
 
    //get istot2ixy
    void get_istot2ixy(
        int* st_i,          // x or x + fftnx (if x < 0) of stick.
        int* st_j          // y or y + fftny (if y < 0) of stick.
    );
 
    //Create the maps from ixy to (in method 2)
    void create_maps(
        int* st_length2D  // the number of planewaves that belong to the stick located on (x, y), stored in 2d x-y plane.
    );
 
//===============================================
//                  FFT
//===============================================
public:
    // FFT dimensions for wave functions.
    int fftnx=0, fftny=0, fftnz=0, fftnxyz=0, fftnxy=0;
    int nx=0, ny=0, nz=0, nxyz=0, nxy=0; // Gamma_only: fftny = int(ny/2)-1 , others: fftny = ny
    int liy=0, riy=0;// liy: the left edge of the pw ball; riy: the right edge of the pw ball in the y direction
    int lix=0, rix=0;// lix: the left edge of the pw ball; rix: the right edge of the pw ball in the x direction
    bool xprime = true; // true: when do recip2real, x-fft will be done last and when doing real2recip, x-fft will be
                        // done first; false: y-fft For gamma_only, true: we use half x; false: we use half y
    int ng_xeq0 = 0; //only used when xprime = true, number of g whose gx = 0
    int nmaxgr = 0;  // Gamma_only: max between npw and (nrxx+1)/2, others: max between npw and nrxx
                     // Thus std::complex<double>[nmaxgr] is able to contain either reciprocal or real data
    // FFT ft;
    FFT_Bundle fft_bundle;
    //The position of pointer in and out can be equal(in-place transform) or different(out-of-place transform).
    
    template <typename FPTYPE>
    void real2recip(const FPTYPE* in,
                    std::complex<FPTYPE>* out,
                    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 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 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 bool add = false,
                    const FPTYPE factor = 1.0) const; // in:(nz, ns)  ; out(nplane,nx*ny)
    
    template <typename FPTYPE>
    void real2recip_gpu(const FPTYPE* in,
                    std::complex<FPTYPE>* out,
                    const bool add = false,
                    const FPTYPE factor = 1.0) const; // in:(nplane,nx*ny)  ; out(nz, ns)
    template <typename FPTYPE>
    void real2recip_gpu(const std::complex<FPTYPE>* in,
                    std::complex<FPTYPE>* out,
                    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,
                    FPTYPE* out,
                    const bool add = false,
                    const FPTYPE factor = 1.0) const; // in:(nz, ns)  ; out(nplane,nx*ny)
    template <typename FPTYPE>
    void recip2real_gpu(const std::complex<FPTYPE>* in,
                    std::complex<FPTYPE>* out,
                    const bool add = false,
                    const FPTYPE factor = 1.0) const; // in:(nz, ns)  ; out(nplane,nx*ny)
 
    template <typename TK,
              typename TR,
              typename Device,
              typename std::enable_if<!std::is_same<TK, typename GetTypeReal<TK>::type>::value
                                          && (std::is_same<TR, typename GetTypeReal<TK>::type>::value
                                              || std::is_same<TR, TK>::value)
                                          && std::is_same<Device, base_device::DEVICE_CPU>::value,
                                      int>::type
              = 0>
    void recip_to_real(TK* in, TR* out, const bool add = false, const typename GetTypeReal<TK>::type factor = 1.0) const
    {
        this->recip2real(in, out, add, factor);
    };
 
    template <typename TK,
              typename TR,
              typename Device,
              typename std::enable_if<!std::is_same<TK, typename GetTypeReal<TK>::type>::value
                                          && (std::is_same<TR, typename GetTypeReal<TK>::type>::value
                                              || std::is_same<TR, TK>::value)
                                          && std::is_same<Device, base_device::DEVICE_GPU>::value,
                                      int>::type
              = 0>
    void recip_to_real(TK* in,
                       TR* out,
                       const bool add = false,
                       const typename GetTypeReal<TK>::type factor = 1.0) const
                       {
                        this->recip2real_gpu(in,out,add,factor);
                       };
 
    // template <typename FPTYPE,
    //         typename Device,
    //         typename std::enable_if<!std::is_same<FPTYPE, typename GetTypeReal<FPTYPE>::type>::value, int>::type = 0>
    // void recip_to_real(FPTYPE* in,
    //                    FPTYPE* out,
    //                    const bool add = false,
    //                    const typename GetTypeReal<FPTYPE>::type factor = 1.0) const;
    template <typename TR,typename TK,typename Device,
            typename std::enable_if<!std::is_same<TK, typename GetTypeReal<TK>::type>::value
                    && (std::is_same<TR, typename GetTypeReal<TK>::type>::value || std::is_same<TR, TK>::value)
                    && std::is_same<Device, base_device::DEVICE_CPU>::value ,int>::type = 0>
    void real_to_recip(TR* in,
                       TK* out,
                       const bool add = false,
                       const typename GetTypeReal<TK>::type factor = 1.0) const
    {
        this->real2recip(in, out, add, factor);
    }
 
    template <typename TR, typename TK, typename Device,
            typename std::enable_if<!std::is_same<TK, typename GetTypeReal<TK>::type>::value
                    && (std::is_same<TR, typename GetTypeReal<TK>::type>::value || std::is_same<TR, TK>::value)
                    && std::is_same<Device, base_device::DEVICE_GPU>::value ,int>::type = 0>
    void real_to_recip(TR* in,
                       TK* out,
                       const bool add = false,
                       const typename GetTypeReal<TK>::type factor = 1.0) const
                       {
                        this->real2recip_gpu(in,out,add,factor);
                       };
 
  protected:
    //gather planes and scatter sticks of all processors
    template <typename T>
    void gatherp_scatters(std::complex<T>* in, std::complex<T>* out) const;
 
    // gather sticks of and scatter planes of all processors
    template <typename T>
    void gathers_scatterp(std::complex<T>* in, std::complex<T>* out) const;
 
  public:
    //get fftixy2is;
    void getfftixy2is(int * fftixy2is) const;
 
    using resmem_int_op = base_device::memory::resize_memory_op<int, base_device::DEVICE_GPU>;
    using delmem_int_op = base_device::memory::delete_memory_op<int, base_device::DEVICE_GPU>;
    using syncmem_int_h2d_op = base_device::memory::synchronize_memory_op<int, base_device::DEVICE_GPU, base_device::DEVICE_CPU>;
    // using default_device_cpu = base_device::DEVICE_CPU;
    
    void set_device(std::string device_);
    void set_precision(std::string precision_);
 
    std::string get_device() const { return device; }
    std::string get_precision() const { return precision; }
 
protected:
 
  std::string device = "cpu";       
  std::string precision = "double"; 
  bool double_data_ = true;         
  bool float_data_ = false;         
};
}
#endif // PWBASIS_H
#include "pw_basis_sup.h"
#include "pw_basis_big.h" //temporary it will be removed