hcontainer.h

Go to the documentation of this file.

#ifndef W_ABACUS_DEVELOP_ABACUS_DEVELOP_SOURCE_MODULE_HAMILT_LCAO_MODULE_HCONTAINER_HCONTAINER_H
#define W_ABACUS_DEVELOP_ABACUS_DEVELOP_SOURCE_MODULE_HAMILT_LCAO_MODULE_HCONTAINER_HCONTAINER_H
 
#include "atom_pair.h"
#include "source_base/vector3.h"
#include "source_cell/unitcell.h"
 
#include <set>
#include <vector>
 
namespace hamilt
{
 
template <typename T>
class HContainer
{
  public:
    // Destructor of class HContainer
    ~HContainer();
 
    HContainer();
 
    HContainer(const HContainer<T>& HR_in, T* data_array = nullptr);
    
    // copy assignment, not allowed
    HContainer& operator=(const HContainer<T>& HR_in) = delete;
    
    // move constructor
    HContainer(HContainer<T>&& HR_in) noexcept;
    // move assignment
    HContainer& operator=(HContainer<T>&& HR_in) noexcept;
 
    // simple constructor
    HContainer(int natom);
 
    // use unitcell to initialize atom_pairs
    HContainer(const UnitCell& ucell_, const Parallel_Orbitals* paraV = nullptr);
 
    HContainer(const Parallel_Orbitals* paraV, T* data_pointer = nullptr, const std::vector<int>* ijr_info = nullptr);
 
    void set_paraV(const Parallel_Orbitals* paraV_in)
    {
        this->paraV = paraV_in;
        for (auto& ap : atom_pairs) ap.set_paraV(paraV_in);
    };
    const Parallel_Orbitals* get_paraV() const { return this->paraV; };
 
    void allocate(T* data_array = nullptr, bool if_zero = false);
 
    void set_zero();
 
    void insert_pair(const AtomPair<T>& atom_ij);
 
    AtomPair<T>* find_pair(int i, int j) const;
 
    BaseMatrix<T>* find_matrix(int i, int j, int rx, int ry, int rz);
    const BaseMatrix<T>* find_matrix(int i, int j, int rx, int ry, int rz) const;
    BaseMatrix<T>* find_matrix(int i, int j, const ModuleBase::Vector3<int>& R_index);
    const BaseMatrix<T>* find_matrix(int i, int j, const ModuleBase::Vector3<int>& R_index) const;
 
    int find_matrix_offset(int i, int j, int rx, int ry, int rz) const;
    int find_matrix_offset(int i, int j, const ModuleBase::Vector3<int>& R_index) const;
 
    AtomPair<T>& get_atom_pair(int i, int j) const;
    AtomPair<T>& get_atom_pair(int index) const;
 
    // T& operator()(int atom_i, int atom_j, int rx_in, int ry_in, int rz_in, int mu, int nu) const;
 
    void add(const HContainer<T>& other);
 
    // save atom-pair pointers into this->tmp_atom_pairs for selected R index
    bool fix_R(int rx_in, int ry_in, int rz_in) const;
 
    void unfix_R() const;
 
    void fix_gamma();
 
    // interface for call a R loop for HContainer
    // it can return a new R-index with (rx,ry,rz) for each loop
    // if index==0, a new loop of R will be initialized
    void loop_R(const size_t& index, int& rx, int& ry, int& rz) const;
 
    size_t size_R_loop() const;
 
    
    int find_R(const int& rx_in, const int& ry_in, const int& rz_in) const;
    int find_R(const ModuleBase::Vector3<int>& R_in) const;
 
    size_t size_atom_pairs() const;
 
    T* data(int i, int j) const;
 
    T* data(int i, int j, int* R) const;
 
    int get_current_R() const;
 
    bool is_gamma_only() const;
 
    size_t get_memory_size() const;
 
    size_t get_nnr() const
    {
        size_t sum = 0;
        for (int iap = 0; iap < this->atom_pairs.size(); ++iap)
        {
            sum += this->atom_pairs[iap].get_R_size() * this->atom_pairs[iap].get_size();
        }
        return sum;
    }
 
    std::vector<int> get_ijr_info() const;
 
    void insert_ijrs(const std::vector<int>* ijrs);
 
    void insert_ijrs(const std::vector<int>* ijrs, const UnitCell& ucell, const int npol=1);
    
    T* get_wrapper() const
    {
        return this->wrapper_pointer;
    }
 
    void shape_synchron(const HContainer<T>& other);
 
    const std::vector<std::vector<int>>& get_sparse_ap() const
    {
        return sparse_ap;
    }
 
    const std::vector<std::vector<int>>& get_sparse_ap_index() const
    {
        return sparse_ap_index;
    }
 
    int get_nbasis() const
    {
        return paraV->get_global_row_size();
    }
 
  private:
    // i-j atom pairs, sorted by matrix of (atom_i, atom_j)
    std::vector<AtomPair<T>> atom_pairs;
 
    // sparse table for (atom_i, atom_j)->index of atom_pairs
    std::vector<std::vector<int>> sparse_ap;
    std::vector<std::vector<int>> sparse_ap_index;
 
    mutable std::vector<const AtomPair<T>*> tmp_atom_pairs;
    // it contains 3 index of cell, size of R_index is three times of values.
    mutable std::vector<ModuleBase::Vector3<int>> tmp_R_index;
    // current index of R in tmp_atom_pairs, -1 means not initialized
    mutable int current_R = -1;
 
    bool gamma_only = false;
 
    T* wrapper_pointer = nullptr;
    bool allocated = false;
 
    const Parallel_Orbitals* paraV = nullptr;
 
    // int multiple = 1;
    // int current_multiple = 0;
};
 
} // namespace hamilt
 
#endif