Refactor class DensityMatrix: remove the dependence on K_Vectors and the ambiguity of _nks (#5224)

* remove dependence of DensityMatrix on K_Vectors

* rename _nks as _nk to avoid ambiguity

* [pre-commit.ci lite] apply automatic fixes

* remove map

* make nk non-optional

* fix after rebase

---------

Co-authored-by: pre-commit-ci-lite[bot] <117423508+pre-commit-ci-lite[bot]@users.noreply.github.com>

Author: maki49

source/module_elecstate/elecstate_lcao.cpp

@@ -131,7 +131,8 @@ void ElecStateLCAO<double>::psiToRho(const psi::Psi<double>& psi)
 template <typename TK>
 void ElecStateLCAO<TK>::init_DM(const K_Vectors* kv, const Parallel_Orbitals* paraV, const int nspin)
 {
-    this->DM = new DensityMatrix<TK, double>(kv, paraV, nspin);
+    const int nspin_dm = nspin == 2 ? 2 : 1;
+    this->DM = new DensityMatrix<TK, double>(paraV, nspin_dm, kv->kvec_d, kv->get_nks() / nspin_dm);
 }
 
 template <>

source/module_elecstate/module_dm/density_matrix.cpp

@@ -24,62 +24,24 @@ DensityMatrix<TK, TR>::~DensityMatrix()
     }
 }
 
-// constructor for multi-k
 template <typename TK, typename TR>
-DensityMatrix<TK, TR>::DensityMatrix(const K_Vectors* kv_in, const Parallel_Orbitals* paraV_in, const int nspin)
+DensityMatrix<TK, TR>::DensityMatrix(const Parallel_Orbitals* paraV_in, const int nspin, const std::vector<ModuleBase::Vector3<double>>& kvec_d, const int nk)
+    : _paraV(paraV_in), _nspin(nspin), _kvec_d(kvec_d), _nk((nk > 0 && nk <= _kvec_d.size()) ? nk : _kvec_d.size())
 {
     ModuleBase::TITLE("DensityMatrix", "DensityMatrix-MK");
-    this->_kv = kv_in;
-    this->_paraV = paraV_in;
-    // set this->_nspin
-    if (nspin == 1 || nspin == 4)
-    {
-        this->_nspin = 1;
-    }
-    else if (nspin == 2)
-    {
-        this->_nspin = 2;
-#ifdef __DEBUG
-        assert(kv_in->get_nks() % 2 == 0);
-#endif
-    }
-    else
-    {
-        throw std::string("nspin must be 1, 2 or 4");
-    }
-    // set this->_nks, which is real number of k-points
-    this->_nks = kv_in->get_nks() / this->_nspin;
-    // allocate memory for _DMK
-    this->_DMK.resize(this->_kv->get_nks());
-    for (int ik = 0; ik < this->_kv->get_nks(); ik++)
+    const int nks = _nk * _nspin;
+    this->_DMK.resize(nks);
+    for (int ik = 0; ik < nks; ik++)
     {
         this->_DMK[ik].resize(this->_paraV->get_row_size() * this->_paraV->get_col_size());
     }
     ModuleBase::Memory::record("DensityMatrix::DMK", this->_DMK.size() * this->_DMK[0].size() * sizeof(TK));
 }
 
-// constructor for Gamma-Only
 template <typename TK, typename TR>
-DensityMatrix<TK, TR>::DensityMatrix(const Parallel_Orbitals* paraV_in, const int nspin)
+DensityMatrix<TK, TR>::DensityMatrix(const Parallel_Orbitals* paraV_in, const int nspin) :_paraV(paraV_in), _nspin(nspin), _kvec_d({ ModuleBase::Vector3<double>(0,0,0) }), _nk(1)
 {
     ModuleBase::TITLE("DensityMatrix", "DensityMatrix-GO");
-    this->_paraV = paraV_in;
-    // set this->_nspin
-    if (nspin == 1 || nspin == 4)
-    {
-        this->_nspin = 1;
-    }
-    else if (nspin == 2)
-    {
-        this->_nspin = 2;
-    }
-    else
-    {
-        throw std::string("nspin must be 1, 2 or 4");
-    }
-    // set this->_nks, which is real number of k-points
-    this->_nks = 1;
-    // allocate memory for _DMK
     this->_DMK.resize(_nspin);
     for (int ik = 0; ik < this->_nspin; ik++)
     {
@@ -274,7 +236,7 @@ template <typename TK, typename TR>
 TK* DensityMatrix<TK, TR>::get_DMK_pointer(const int ik) const
 {
 #ifdef __DEBUG
-    assert(ik < this->_nks * this->_nspin);
+    assert(ik < this->_nk * this->_nspin);
 #endif
     return const_cast<TK*>(this->_DMK[ik].data());
 }
@@ -284,7 +246,7 @@ template <typename TK, typename TR>
 void DensityMatrix<TK, TR>::set_DMK_pointer(const int ik, TK* DMK_in)
 {
 #ifdef __DEBUG
-    assert(ik < this->_nks * this->_nspin);
+    assert(ik < this->_nk * this->_nspin);
 #endif
     this->_DMK[ik].assign(DMK_in, DMK_in + this->_paraV->nrow * this->_paraV->ncol);
 }
@@ -295,17 +257,17 @@ void DensityMatrix<TK, TR>::set_DMK(const int ispin, const int ik, const int i,
 {
 #ifdef __DEBUG
     assert(ispin > 0 && ispin <= this->_nspin);
-    assert(ik >= 0 && ik < this->_nks);
+    assert(ik >= 0 && ik < this->_nk);
 #endif
     // consider transpose col=>row
-    this->_DMK[ik + this->_nks * (ispin - 1)][i * this->_paraV->nrow + j] = value;
+    this->_DMK[ik + this->_nk * (ispin - 1)][i * this->_paraV->nrow + j] = value;
 }
 
 // set _DMK element
 template <typename TK, typename TR>
 void DensityMatrix<TK, TR>::set_DMK_zero()
 {
-    for (int ik = 0; ik < _nspin * _nks; ik++)
+    for (int ik = 0; ik < _nspin * _nk; ik++)
     {
         ModuleBase::GlobalFunc::ZEROS(this->_DMK[ik].data(),
                                       this->_paraV->get_row_size() * this->_paraV->get_col_size());
@@ -320,18 +282,17 @@ TK DensityMatrix<TK, TR>::get_DMK(const int ispin, const int ik, const int i, co
     assert(ispin > 0 && ispin <= this->_nspin);
 #endif
     // consider transpose col=>row
-    return this->_DMK[ik + this->_nks * (ispin - 1)][i * this->_paraV->nrow + j];
+    return this->_DMK[ik + this->_nk * (ispin - 1)][i * this->_paraV->nrow + j];
 }
 
 // get _DMK nks, nrow, ncol
 template <typename TK, typename TR>
 int DensityMatrix<TK, TR>::get_DMK_nks() const
 {
 #ifdef __DEBUG
-    assert(this->_DMK.size() != 0);
-    assert(this->_kv != nullptr);
+    assert(this->_DMK.size() == _nk * _nspin);
 #endif
-    return this->_kv->get_nks();
+    return _nk * _nspin;
 }
 
 template <typename TK, typename TR>
@@ -403,7 +364,7 @@ void DensityMatrix<TK, TR>::cal_DMR_test()
 {
     for (int is = 1; is <= this->_nspin; ++is)
     {
-        int ik_begin = this->_nks * (is - 1); // jump this->_nks for spin_down if nspin==2
+        int ik_begin = this->_nk * (is - 1); // jump this->_nk for spin_down if nspin==2
         hamilt::HContainer<TR>* tmp_DMR = this->_DMR[is - 1];
         // set zero since this function is called in every scf step
         tmp_DMR->set_zero();
@@ -435,12 +396,12 @@ void DensityMatrix<TK, TR>::cal_DMR_test()
 #endif
                 std::complex<TR> tmp_res;
                 // loop over k-points
-                for (int ik = 0; ik < this->_nks; ++ik)
+                for (int ik = 0; ik < this->_nk; ++ik)
                 {
                     // cal k_phase
                     // if TK==std::complex<double>, kphase is e^{ikR}
                     const ModuleBase::Vector3<double> dR(r_index[0], r_index[1], r_index[2]);
-                    const double arg = (this->_kv->kvec_d[ik] * dR) * ModuleBase::TWO_PI;
+                    const double arg = (this->_kvec_d[ik] * dR) * ModuleBase::TWO_PI;
                     double sinp, cosp;
                     ModuleBase::libm::sincos(arg, &sinp, &cosp);
                     std::complex<double> kphase = std::complex<double>(cosp, sinp);
@@ -477,7 +438,7 @@ void DensityMatrix<std::complex<double>, double>::cal_DMR()
     int ld_hk2 = 2 * ld_hk;
     for (int is = 1; is <= this->_nspin; ++is)
     {
-        int ik_begin = this->_nks * (is - 1); // jump this->_nks for spin_down if nspin==2
+        int ik_begin = this->_nk * (is - 1); // jump this->_nk for spin_down if nspin==2
         hamilt::HContainer<double>* tmp_DMR = this->_DMR[is - 1];
         // set zero since this function is called in every scf step
         tmp_DMR->set_zero();
@@ -515,12 +476,12 @@ void DensityMatrix<std::complex<double>, double>::cal_DMR()
                 // loop over k-points
                 if (PARAM.inp.nspin != 4)
                 {
-                    for (int ik = 0; ik < this->_nks; ++ik)
+                    for (int ik = 0; ik < this->_nk; ++ik)
                     {
                         // cal k_phase
                         // if TK==std::complex<double>, kphase is e^{ikR}
                         const ModuleBase::Vector3<double> dR(r_index[0], r_index[1], r_index[2]);
-                        const double arg = (this->_kv->kvec_d[ik] * dR) * ModuleBase::TWO_PI;
+                        const double arg = (this->_kvec_d[ik] * dR) * ModuleBase::TWO_PI;
                         double sinp, cosp;
                         ModuleBase::libm::sincos(arg, &sinp, &cosp);
                         std::complex<double> kphase = std::complex<double>(cosp, sinp);
@@ -559,12 +520,12 @@ void DensityMatrix<std::complex<double>, double>::cal_DMR()
                 if (PARAM.inp.nspin == 4)
                 {
                     tmp_DMR.assign(tmp_ap.get_size(), std::complex<double>(0.0, 0.0));
-                    for (int ik = 0; ik < this->_nks; ++ik)
+                    for (int ik = 0; ik < this->_nk; ++ik)
                     {
                         // cal k_phase
                         // if TK==std::complex<double>, kphase is e^{ikR}
                         const ModuleBase::Vector3<double> dR(r_index[0], r_index[1], r_index[2]);
-                        const double arg = (this->_kv->kvec_d[ik] * dR) * ModuleBase::TWO_PI;
+                        const double arg = (this->_kvec_d[ik] * dR) * ModuleBase::TWO_PI;
                         double sinp, cosp;
                         ModuleBase::libm::sincos(arg, &sinp, &cosp);
                         std::complex<double> kphase = std::complex<double>(cosp, sinp);
@@ -644,7 +605,7 @@ void DensityMatrix<std::complex<double>, double>::cal_DMR(const int ik)
     int ld_hk2 = 2 * ld_hk;
     for (int is = 1; is <= this->_nspin; ++is)
     {
-        int ik_begin = this->_nks * (is - 1); // jump this->_nks for spin_down if nspin==2
+        int ik_begin = this->_nk * (is - 1); // jump this->_nk for spin_down if nspin==2
         hamilt::HContainer<double>* tmp_DMR = this->_DMR[is - 1];
         // set zero since this function is called in every scf step
         tmp_DMR->set_zero();
@@ -680,7 +641,7 @@ void DensityMatrix<std::complex<double>, double>::cal_DMR(const int ik)
                     // cal k_phase
                     // if TK==std::complex<double>, kphase is e^{ikR}
                     const ModuleBase::Vector3<double> dR(r_index[0], r_index[1], r_index[2]);
-                    const double arg = (this->_kv->kvec_d[ik] * dR) * ModuleBase::TWO_PI;
+                    const double arg = (this->_kvec_d[ik] * dR) * ModuleBase::TWO_PI;
                     double sinp, cosp;
                     ModuleBase::libm::sincos(arg, &sinp, &cosp);
                     std::complex<double> kphase = std::complex<double>(cosp, sinp);
@@ -734,14 +695,14 @@ void DensityMatrix<double, double>::cal_DMR()
     int ld_hk = this->_paraV->nrow;
     for (int is = 1; is <= this->_nspin; ++is)
     {
-        int ik_begin = this->_nks * (is - 1); // jump this->_nks for spin_down if nspin==2
+        int ik_begin = this->_nk * (is - 1); // jump this->_nk for spin_down if nspin==2
         hamilt::HContainer<double>* tmp_DMR = this->_DMR[is - 1];
         // set zero since this function is called in every scf step
         tmp_DMR->set_zero();
 
 #ifdef __DEBUG
         // assert(tmp_DMR->is_gamma_only() == true);
-        assert(this->_nks == 1);
+        assert(this->_nk == 1);
 #endif
 #ifdef _OPENMP
 #pragma omp parallel for
@@ -857,9 +818,9 @@ void DensityMatrix<TK, TR>::read_DMK(const std::string directory, const int ispi
         // quit the program or not.
         bool quit = false;
 
-        ModuleBase::CHECK_DOUBLE(ifs, this->_kv->kvec_d[ik].x, quit);
-        ModuleBase::CHECK_DOUBLE(ifs, this->_kv->kvec_d[ik].y, quit);
-        ModuleBase::CHECK_DOUBLE(ifs, this->_kv->kvec_d[ik].z, quit);
+        ModuleBase::CHECK_DOUBLE(ifs, this->_kvec_d[ik].x, quit);
+        ModuleBase::CHECK_DOUBLE(ifs, this->_kvec_d[ik].y, quit);
+        ModuleBase::CHECK_DOUBLE(ifs, this->_kvec_d[ik].z, quit);
         ModuleBase::CHECK_INT(ifs, this->_paraV->nrow);
         ModuleBase::CHECK_INT(ifs, this->_paraV->ncol);
     } // If file exist, read in data.
@@ -869,7 +830,7 @@ void DensityMatrix<TK, TR>::read_DMK(const std::string directory, const int ispi
     {
         for (int j = 0; j < this->_paraV->ncol; ++j)
         {
-            ifs >> this->_DMK[ik + this->_nks * (ispin - 1)][i * this->_paraV->ncol + j];
+            ifs >> this->_DMK[ik + this->_nk * (ispin - 1)][i * this->_paraV->ncol + j];
         }
     }
     ifs.close();
@@ -892,7 +853,7 @@ void DensityMatrix<double, double>::write_DMK(const std::string directory, const
     {
         ModuleBase::WARNING("elecstate::write_dmk", "Can't create DENSITY MATRIX File!");
     }
-    ofs << this->_kv->kvec_d[ik].x << " " << this->_kv->kvec_d[ik].y << " " << this->_kv->kvec_d[ik].z << std::endl;
+    ofs << this->_kvec_d[ik].x << " " << this->_kvec_d[ik].y << " " << this->_kvec_d[ik].z << std::endl;
     ofs << "\n  " << this->_paraV->nrow << " " << this->_paraV->ncol << std::endl;
 
     ofs << std::setprecision(3);
@@ -906,7 +867,7 @@ void DensityMatrix<double, double>::write_DMK(const std::string directory, const
             {
                 ofs << "\n";
             }
-            ofs << " " << this->_DMK[ik + this->_nks * (ispin - 1)][i * this->_paraV->ncol + j];
+            ofs << " " << this->_DMK[ik + this->_nk * (ispin - 1)][i * this->_paraV->ncol + j];
         }
     }
 
@@ -929,7 +890,7 @@ void DensityMatrix<std::complex<double>, double>::write_DMK(const std::string di
     {
         ModuleBase::WARNING("elecstate::write_dmk", "Can't create DENSITY MATRIX File!");
     }
-    ofs << this->_kv->kvec_d[ik].x << " " << this->_kv->kvec_d[ik].y << " " << this->_kv->kvec_d[ik].z << std::endl;
+    ofs << this->_kvec_d[ik].x << " " << this->_kvec_d[ik].y << " " << this->_kvec_d[ik].z << std::endl;
     ofs << "\n  " << this->_paraV->nrow << " " << this->_paraV->ncol << std::endl;
 
     ofs << std::setprecision(3);
@@ -943,7 +904,7 @@ void DensityMatrix<std::complex<double>, double>::write_DMK(const std::string di
             {
                 ofs << "\n";
             }
-            ofs << " " << this->_DMK[ik + this->_nks * (ispin - 1)][i * this->_paraV->ncol + j].real();
+            ofs << " " << this->_DMK[ik + this->_nk * (ispin - 1)][i * this->_paraV->ncol + j].real();
         }
     }
 

source/module_elecstate/module_dm/density_matrix.h

@@ -3,7 +3,6 @@
 
 #include <string>
 
-#include "module_cell/klist.h"
 #include "module_cell/module_neighbor/sltk_grid_driver.h"
 #include "module_hamilt_lcao/hamilt_lcaodft/record_adj.h"
 #include "module_hamilt_lcao/module_hcontainer/hcontainer.h"
@@ -45,16 +44,20 @@ class DensityMatrix
 
     /**
      * @brief Constructor of class DensityMatrix for multi-k calculation
-     * @param _kv pointer of K_Vectors object
      * @param _paraV pointer of Parallel_Orbitals object
-     * @param nspin spin setting (1 - none spin; 2 - spin; 4 - SOC)
+     * @param nspin number of spin of the density matrix, set by user according to global nspin
+     *  (usually {nspin_global -> nspin_dm} = {1->1, 2->2, 4->1}, but sometimes 2->1 like in LR-TDDFT)
+     * @param kvec_d direct coordinates of kpoints
+     * @param nk number of k-points, not always equal to K_Vectors::get_nks()/nspin_dm.
+     *               it will be set to kvec_d.size() if the value is invalid
      */
-    DensityMatrix(const K_Vectors* _kv, const Parallel_Orbitals* _paraV, const int nspin);
+    DensityMatrix(const Parallel_Orbitals* _paraV, const int nspin, const std::vector<ModuleBase::Vector3<double>>& kvec_d, const int nk);
 
     /**
      * @brief Constructor of class DensityMatrix for gamma-only calculation, where kvector is not required
      * @param _paraV pointer of Parallel_Orbitals object
-     * @param nspin spin setting (1 - none spin; 2 - spin; 4 - SOC)
+     * @param nspin number of spin of the density matrix, set by user according to global nspin
+     *  (usually {nspin_global -> nspin_dm} = {1->1, 2->2, 4->1}, but sometimes 2->1 like in LR-TDDFT)
      */
     DensityMatrix(const Parallel_Orbitals* _paraV, const int nspin);
 
@@ -169,7 +172,7 @@ class DensityMatrix
      */
     const Parallel_Orbitals* get_paraV_pointer() const {return this->_paraV;}
 
-    const K_Vectors* get_kv_pointer() const {return this->_kv;}
+    const std::vector<ModuleBase::Vector3<double>>& get_kvec_d() const { return this->_kvec_d; }
 
     /**
      * @brief calculate density matrix DMR from dm(k) using blas::axpy
@@ -240,16 +243,16 @@ class DensityMatrix
 
     /**
      * @brief density matrix in k space, which is a vector[ik]
-     * DMK should be a [_nspin][_nks][i][j] matrix,
-     * whose size is _nspin * _nks * _paraV->get_nrow() * _paraV->get_ncol()
+     * DMK should be a [_nspin][_nk][i][j] matrix,
+     * whose size is _nspin * _nk * _paraV->get_nrow() * _paraV->get_ncol()
      */
     // std::vector<ModuleBase::ComplexMatrix> _DMK;
     std::vector<std::vector<TK>> _DMK;
 
     /**
      * @brief K_Vectors object, which is used to get k-point information
      */
-    const K_Vectors* _kv;
+    const std::vector<ModuleBase::Vector3<double>> _kvec_d;
 
     /**
      * @brief Parallel_Orbitals object, which contain all information of 2D block cyclic distribution
@@ -265,10 +268,10 @@ class DensityMatrix
 
     /**
      * @brief real number of k-points
-     * _nks is not equal to _kv->get_nks() when spin-polarization is considered
-     * _nks = kv->_nks / nspin
+     * _nk is not equal to _kv->get_nks() when spin-polarization is considered
+     * _nk = kv->get_nks() / nspin when nspin=2
      */
-    int _nks = 0;
+    int _nk = 0;
 
 
 };