update psi code

Author: haozhihan

source/module_psi/psi.cpp

@@ -31,9 +31,7 @@ Range::Range(const bool k_first_in, const size_t index_1_in, const size_t range_
 // Constructor 0: basic
 template <typename T, typename Device>
 Psi<T, Device>::Psi()
-{
-    this->npol = PARAM.globalv.npol;
-}
+{}
 
 template <typename T, typename Device>
 Psi<T, Device>::~Psi()
@@ -53,7 +51,6 @@ Psi<T, Device>::Psi(const int nk_in, const int nbd_in, const int nbs_in, const i
     assert(nbs_in > 0);
 
     this->k_first = k_first_in;
-    this->npol = PARAM.globalv.npol;
     this->allocate_inside = true;
 
     this->ngk = ngk_in; // modify later
@@ -91,7 +88,6 @@ Psi<T, Device>::Psi(const int nk_in,
     assert(nbs_in > 0);
 
     this->k_first = k_first_in;
-    this->npol = PARAM.globalv.npol;
     this->allocate_inside = true;
 
     this->ngk = ngk_in.data(); // modify later
@@ -129,7 +125,6 @@ Psi<T, Device>::Psi(T* psi_pointer,
     // assert(nk_in == 1); // NOTE because lr/utils/lr_uril.hpp func & get_psi_spin func
 
     this->k_first = k_first_in;
-    this->npol = PARAM.globalv.npol;
     this->allocate_inside = false;
 
     this->ngk = nullptr;
@@ -161,7 +156,6 @@ Psi<T, Device>::Psi(const int nk_in,
     assert(nk_in == 1);
 
     this->k_first = k_first_in;
-    this->npol = PARAM.globalv.npol;
     this->allocate_inside = true;
 
     this->ngk = nullptr;
@@ -191,7 +185,6 @@ template <typename T, typename Device>
 Psi<T, Device>::Psi(const Psi& psi_in)
 {
     this->ngk = psi_in.ngk;
-    this->npol = PARAM.globalv.npol;
     this->nk = psi_in.get_nk();
     this->nbands = psi_in.get_nbands();
     this->nbasis = psi_in.get_nbasis();
@@ -218,7 +211,6 @@ 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->npol = PARAM.globalv.npol;
     this->nk = psi_in.get_nk();
     this->nbands = psi_in.get_nbands();
     this->nbasis = psi_in.get_nbasis();
@@ -331,7 +323,7 @@ const int& Psi<T, Device>::get_psi_bias() const
 template <typename T, typename Device>
 const int& Psi<T, Device>::get_current_ngk() const
 {
-    if (this->npol == 1)
+    if (PARAM.inp.nspin != 4)
     {
         return this->current_nbasis;
     }
@@ -519,13 +511,13 @@ std::tuple<const T*, int> Psi<T, Device>::to_range(const Range& range) const
     else if (i1 < 0) // [r1, r2] is the range of index1 with length m
     {
         const T* p = &this->psi[r1 * (k_first ? this->nbands : this->nk) * this->nbasis];
-        int m = (r2 - r1 + 1) * this->npol;
+        int m = (r2 - r1 + 1) * this->get_npol();
         return std::tuple<const T*, int>(p, m);
     }
     else // [r1, r2] is the range of index2 with length m
     {
         const T* p = &this->psi[(i1 * (k_first ? this->nbands : this->nk) + r1) * this->nbasis];
-        int m = (r2 - r1 + 1) * this->npol;
+        int m = (r2 - r1 + 1) * this->get_npol();
         return std::tuple<const T*, int>(p, m);
     }
 }

source/module_psi/psi.h

@@ -3,6 +3,7 @@
 
 #include "module_base/module_device/memory_op.h"
 #include "module_base/module_device/types.h"
+#include "module_parameter/parameter.h"
 
 #include <tuple>
 #include <vector>
@@ -134,7 +135,17 @@ class Psi
 
     const int& get_current_ngk() const;
 
-    const int& get_npol() const {return this->npol;}
+    const int& get_npol() const 
+    {
+        if (PARAM.inp.nspin == 4)
+        {
+            return 2;
+        }
+        else
+        {
+            return 1;
+        }
+    }
 
     // solve Range: return(pointer of begin, number of bands or k-points)
     std::tuple<const T*, int> to_range(const Range& range) const;
@@ -143,27 +154,22 @@ class Psi
     T* psi = nullptr; // avoid using C++ STL
 
     Device* ctx = {}; // an context identifier for obtaining the device variable
-    int npol = 1;
+    bool allocate_inside = true; ///< whether allocate psi inside Psi class
+    bool k_first = true;
+
+    const int* ngk = nullptr;
 
     // dimensions
     int nk = 1;     // number of k points
     int nbands = 1; // number of bands
     int nbasis = 1; // number of basis
 
+    // mutable values
     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; ///< whether allocate psi inside Psi class
+    mutable T* psi_current = nullptr; // current pointer for getting the psi
+    mutable int psi_bias = 0; // psi_current = psi + psi_bias;
 
 #ifdef __DSP
     using delete_memory_op = base_device::memory::delete_memory_op_mt<T, Device>;