psi.h

Go to the documentation of this file.

#ifndef PSI_H
#define PSI_H
 
#include "source_base/module_device/memory_op.h"
#include "source_base/module_device/types.h"
 
#include <tuple>
#include <vector>
 
namespace psi
{
 
// structure for getting range of Psi
// two display method: k index first or bands index first
struct Range
{
    bool k_first;
    size_t index_1;
    size_t range_1;
    size_t range_2;
    // this is simple constructor for hPsi return
    Range(const size_t range_in);
    // constructor 2
    Range(const bool k_first_in, const size_t index_1_in, const size_t range_1_in, const size_t range_2_in);
};
 
// there is the structure of electric wavefunction coefficient
// the basic operations defined in the Operator Class
template <typename T, typename Device = base_device::DEVICE_CPU>
class Psi
{
  public:
    // Constructor 0: basic
    Psi();
 
    // Constructor 1:
    Psi(const int nk_in, const int nbd_in, const int nbs_in, const std::vector<int>& ngk_in, const bool k_first_in);
 
    // Constructor 2-1: initialize a new psi from the given psi_in
    Psi(const Psi& psi_in);
 
    // Constructor 2-2: initialize a new psi from the given psi_in with a different class template
    // in this case, psi_in may have a different device type.
    template <typename T_in, typename Device_in = Device>
    Psi(const Psi<T_in, Device_in>& psi_in);
 
    // Constructor 3-1: 2D Psi version
    // used in hsolver-pw function pointer and somewhere.
    Psi(T* psi_pointer,
        const int nk_in,
        const int nbd_in,
        const int nbs_in,
        const int current_nbasis_in,
        const bool k_first_in = true);
 
    // Constructor 3-2: 2D Psi version
    Psi(const int nk_in, const int nbd_in, const int nbs_in, const int current_nbasis_in, const bool k_first_in);
 
    // Destructor for deleting the psi array manually
    ~Psi();
 
    // set psi value func 1
    void set_all_psi(const T* another_pointer, const std::size_t size_in);
 
    // set psi value func 2
    void zero_out();
 
    // size_t size() const {return this->psi.size();}
    size_t size() const;
 
    // copy assignment operator
    Psi& operator=(const Psi& psi_in);
 
    // allocate psi for three dimensions
    void resize(const int nks_in, const int nbands_in, const int nbasis_in);
 
    // get the pointer for the 1st index
    T* get_pointer() const;
 
    // get the pointer for the 2nd index (iband for k_first = true, ik for k_first = false)
    T* get_pointer(const int& ikb) const;
 
    // interface to get three dimension size
    const int& get_nk() const;
    const int& get_nbands() const;
    const int& get_nbasis() const;
 
    void fix_k(const int ik) const;
    void fix_b(const int ib) const;
    void fix_kb(const int ik, const int ib) const;
 
    T& operator()(const int ikb1, const int ikb2, const int ibasis) const;
    T& operator()(const int ikb2, const int ibasis) const;
    // use operator "(ibasis)" to reach target element for current k and current band
    T& operator()(const int ibasis) const;
 
    // return current k-point index
    int get_current_k() const;
    // return current band index
    int get_current_b() const;
    // return current ngk for PW base
    int get_current_nbas() const;
 
    const int& get_ngk(const int ik_in) const;
 
    const int* get_ngk_pointer() const;
 
    // return k_first
    const bool& get_k_first() const;
 
    // return device type of psi
    const Device* get_device() const;
 
    // return psi_bias
    const int& get_psi_bias() const;
 
    const int& get_current_ngk() const;
 
    // solve Range: return(pointer of begin, number of bands or k-points)
    std::tuple<const T*, int> to_range(const Range& range) const;
 
    int get_npol() const;
 
  private:
    T* psi = nullptr; // avoid using C++ STL
 
    Device* ctx = {}; // an context identifier for obtaining the device variable
 
    // dimensions
    int nk = 1;     // number of k points
    int nbands = 1; // number of bands
    int nbasis = 1; // number of basis
 
    mutable int current_k = 0;      // current k point
    mutable int current_b = 0;      // current band index
    mutable int current_nbasis = 1; // current number of basis of current_k
 
    // current pointer for getting the psi
    mutable T* psi_current = nullptr;
    // psi_current = psi + psi_bias;
    mutable int psi_bias = 0;
 
    const int* ngk = nullptr;
 
    bool k_first = true;
 
    bool allocate_inside = true; 
 
#ifdef __DSP
    using delete_memory_op = base_device::memory::delete_memory_op_mt<T, Device>;
    using resize_memory_op = base_device::memory::resize_memory_op_mt<T, Device>;
#else
    using delete_memory_op = base_device::memory::delete_memory_op<T, Device>;
    using resize_memory_op = base_device::memory::resize_memory_op<T, Device>;
#endif
    using set_memory_op = base_device::memory::set_memory_op<T, Device>;
    using synchronize_memory_op = base_device::memory::synchronize_memory_op<T, Device, Device>;
};
 
} // end of namespace psi
 
#endif