add ngk_vector

Author: haozhihan

source/module_psi/psi.cpp

@@ -59,6 +59,8 @@ Psi<T, Device>::Psi(const int nk_in,
     this->allocate_inside = true;
 
     this->ngk = ngk_in.data(); // modify later
+    ngk_vector = ngk_in;
+
     // This function will delete the psi array first(if psi exist), then malloc a new memory for it.
     resize_memory_op()(this->psi, nk_in * static_cast<std::size_t>(nbd_in) * nbs_in, "no_record");
 
@@ -80,7 +82,7 @@ Psi<T, Device>::Psi(const int nk_in,
                                                            sizeof(T) * nk_in * nbd_in * nbs_in);
 }
 
-// Constructor 3-1: 2D Psi version
+// Constructor 2-1: 2D Psi version
 template <typename T, typename Device>
 Psi<T, Device>::Psi(T* psi_pointer,
                     const int nk_in,
@@ -96,6 +98,9 @@ Psi<T, Device>::Psi(T* psi_pointer,
     this->allocate_inside = false;
 
     this->ngk = nullptr;
+    ngk_vector = std::vector<int>(nk_in, current_nbasis_in);
+
+
     this->psi = psi_pointer;
 
     this->nk = nk_in;
@@ -112,7 +117,7 @@ Psi<T, Device>::Psi(T* psi_pointer,
     base_device::information::print_device_info<Device>(this->ctx, GlobalV::ofs_device);
 }
 
-// Constructor 3-2: 2D Psi version
+// Constructor 2-2: 2D Psi version
 template <typename T, typename Device>
 Psi<T, Device>::Psi(const int nk_in,
                     const int nbd_in,
@@ -127,6 +132,8 @@ Psi<T, Device>::Psi(const int nk_in,
     this->allocate_inside = true;
 
     this->ngk = nullptr;
+    ngk_vector = std::vector<int>(nk_in, current_nbasis_in);
+
     assert(nk_in > 0 && nbd_in >= 0 && nbs_in > 0);
     resize_memory_op()(this->psi, nk_in * static_cast<std::size_t>(nbd_in) * nbs_in, "no_record");
 
@@ -148,12 +155,14 @@ Psi<T, Device>::Psi(const int nk_in,
                                                            sizeof(T) * nk_in * nbd_in * nbs_in);
 }
 
-// Constructor 2-1:
+// Copy Constructor 3-1:
 template <typename T, typename Device>
 Psi<T, Device>::Psi(const Psi& psi_in)
 {
 
     this->ngk = psi_in.ngk;
+    this->ngk_vector = psi_in.ngk_vector;
+
     this->nk = psi_in.get_nk();
     this->nbands = psi_in.get_nbands();
     this->nbasis = psi_in.get_nbasis();
@@ -172,13 +181,15 @@ Psi<T, Device>::Psi(const Psi& psi_in)
 }
 
 
-// Constructor 2-2:
+// Copy Constructor 3-2:
 template <typename T, typename Device>
 template <typename T_in, typename Device_in>
 Psi<T, Device>::Psi(const Psi<T_in, Device_in>& psi_in)
 {
 
     this->ngk = psi_in.get_ngk_pointer();
+    this->ngk_vector = psi_in.get_ngk_vector();
+
     this->nk = psi_in.get_nk();
     this->nbands = psi_in.get_nbands();
     this->nbasis = psi_in.get_nbasis();
@@ -276,6 +287,13 @@ const int* Psi<T, Device>::get_ngk_pointer() const
     return this->ngk;
 }
 
+template <typename T, typename Device>
+const std::vector<int>& Psi<T, Device>::get_ngk_vector() const
+{
+    return this->ngk_vector;
+}
+
+
 template <typename T, typename Device>
 const int& Psi<T, Device>::get_psi_bias() const
 {

source/module_psi/psi.h

@@ -42,26 +42,26 @@ class 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.
+    // Constructor 2-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
+    // Constructor 2-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);
 
+    // Copy Constructor 3-1: initialize a new psi from the given psi_in
+    Psi(const Psi& psi_in);
+    
+    // Copy Constructor 3-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);
+
+
     // Destructor for deleting the psi array manually
     ~Psi();
 
@@ -120,6 +120,8 @@ class Psi
 
     const int* get_ngk_pointer() const;
 
+    const std::vector<int>& get_ngk_vector() const; 
+
     // return k_first
     const bool& get_k_first() const;
 
@@ -156,6 +158,7 @@ class Psi
     mutable int psi_bias = 0;
 
     const int* ngk = nullptr;
+    std::vector<int> ngk_vector;
 
     bool k_first = true;