remove dependence of DensityMatrix on K_Vectors

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 = std::map<int, int>({ {1, 1}, {2, 2}, {4, 1} })[nspin];
+    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), _nks((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 = _nks * _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) }), _nks(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++)
     {
@@ -328,10 +290,9 @@ 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() == _nks * _nspin);
 #endif
-    return this->_kv->get_nks();
+    return _nks * _nspin;
 }
 
 template <typename TK, typename TR>
@@ -440,7 +401,7 @@ void DensityMatrix<TK, TR>::cal_DMR_test()
                     // 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);
@@ -515,7 +476,7 @@ void DensityMatrix<std::complex<double>, double>::cal_DMR()
                         // 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);
@@ -561,7 +522,7 @@ void DensityMatrix<std::complex<double>, double>::cal_DMR()
                         // 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);
@@ -678,7 +639,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);
@@ -855,9 +816,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.
@@ -890,7 +851,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);
@@ -927,7 +888,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);

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.
+     * if remains default or large than kvec_d.size(), it will be set to kvec_d.size()
      */
-    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 = -1);
 
     /**
      * @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
@@ -249,7 +252,7 @@ class DensityMatrix
     /**
      * @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
@@ -268,7 +271,7 @@ class DensityMatrix
      * _nks is not equal to _kv->get_nks() when spin-polarization is considered
      * _nks = kv->_nks / nspin
      */
-    int _nks = 0;
+    int _nks = 0;   // comment by lunasea: I want to rename it as nk
 
 
 };

source/module_elecstate/module_dm/test/test_cal_dm_R.cpp

@@ -74,7 +74,7 @@ class DMTest : public testing::Test
         init_parav();
     }
 
-    void TearDown()
+    void TearDown() override
     {
         delete paraV;
         delete[] ucell.atoms[0].tau;
@@ -129,7 +129,7 @@ TEST_F(DMTest, cal_DMR_test)
     kv->set_nks(nks);
     kv->kvec_d.resize(nks);
     // construct DM
-    elecstate::DensityMatrix<double, double> DM(kv, paraV, nspin);
+    elecstate::DensityMatrix<double, double> DM(paraV, nspin, kv->kvec_d, kv->get_nks() / nspin);
     // set this->_DMK
     for (int is = 1; is <= nspin; is++)
     {
@@ -198,7 +198,7 @@ TEST_F(DMTest, cal_DMR_blas_double)
     kv->set_nks(nks);
     kv->kvec_d.resize(nks);
     // construct DM
-    elecstate::DensityMatrix<double, double> DM(kv, paraV, nspin);
+    elecstate::DensityMatrix<double, double> DM(paraV, nspin, kv->kvec_d, kv->get_nks() / nspin);
     // set this->_DMK
     for (int is = 1; is <= nspin; is++)
     {
@@ -269,7 +269,7 @@ TEST_F(DMTest, cal_DMR_blas_complex)
     kv->kvec_d[1].x = 0.5;
     kv->kvec_d[3].x = 0.5;
     // construct DM
-    elecstate::DensityMatrix<std::complex<double>, double> DM(kv, paraV, nspin);
+    elecstate::DensityMatrix<std::complex<double>, double> DM(paraV, nspin, kv->kvec_d, kv->get_nks() / nspin);
     // set this->_DMK
     for (int is = 1; is <= nspin; is++)
     {

source/module_elecstate/module_dm/test/test_dm_R_init.cpp

@@ -74,7 +74,7 @@ class DMTest : public testing::Test
         init_parav();
     }
 
-    void TearDown()
+    void TearDown() override
     {
         delete paraV;
         delete[] ucell.atoms[0].tau;
@@ -118,7 +118,7 @@ TEST_F(DMTest, DMInit1)
     // construct DM
     std::cout << "dim0: " << paraV->dim0 << "    dim1:" << paraV->dim1 << std::endl;
     std::cout << "nrow: " << paraV->nrow << "    ncol:" << paraV->ncol << std::endl;
-    elecstate::DensityMatrix<double, double> DM(kv, paraV, nspin);
+    elecstate::DensityMatrix<double, double> DM(paraV, nspin, kv->kvec_d);
     // initialize this->_DMR
     Grid_Driver gd(0,0);
     DM.init_DMR(&gd, &ucell);
@@ -145,7 +145,7 @@ TEST_F(DMTest, DMInit2)
     // construct DM
     std::cout << "dim0: " << paraV->dim0 << "    dim1:" << paraV->dim1 << std::endl;
     std::cout << "nrow: " << paraV->nrow << "    ncol:" << paraV->ncol << std::endl;
-    elecstate::DensityMatrix<double, double> DM(kv, paraV, nspin);
+    elecstate::DensityMatrix<double, double> DM(paraV, nspin, kv->kvec_d);
     // initialize Record_adj using Grid_Driver
     Grid_Driver gd(0,0);
     Record_adj ra;
@@ -208,12 +208,12 @@ TEST_F(DMTest, DMInit3)
     kv->kvec_d[1].x = 0.5;
     kv->kvec_d[3].x = 0.5;
     // construct a DM
-    elecstate::DensityMatrix<std::complex<double>, double> DM(kv, paraV, nspin);
+    elecstate::DensityMatrix<std::complex<double>, double> DM(paraV, nspin, kv->kvec_d, kv->get_nks() / nspin);
     Grid_Driver gd(0, 0);
     DM.init_DMR(&gd, &ucell);
     std::cout << "dim0: " << paraV->dim0 << "    dim1:" << paraV->dim1 << std::endl;
     // construct another DM
-    elecstate::DensityMatrix<std::complex<double>, double> DM1(kv, paraV, nspin);
+    elecstate::DensityMatrix<std::complex<double>, double> DM1(paraV, nspin, kv->kvec_d, kv->get_nks() / nspin);
     DM1.init_DMR(*DM.get_DMR_pointer(1));
     // compare
     EXPECT_EQ(DM1.get_DMR_pointer(2)->size_atom_pairs(), test_size * test_size);
@@ -266,7 +266,7 @@ TEST_F(DMTest, DMInit4)
         }
     }
     // construct a DM from this HContainer
-    elecstate::DensityMatrix<std::complex<double>, double> DM(kv, paraV, nspin);
+    elecstate::DensityMatrix<std::complex<double>, double> DM(paraV, nspin, kv->kvec_d, kv->get_nks() / nspin);
     DM.init_DMR(*tmp_DMR);
     std::cout << "dim0: " << paraV->dim0 << "    dim1:" << paraV->dim1 << std::endl;
     // compare
@@ -292,11 +292,11 @@ TEST_F(DMTest, saveDMR)
     kv->kvec_d[1].x = 0.5;
     kv->kvec_d[3].x = 0.5;
     // construct a DM
-    elecstate::DensityMatrix<std::complex<double>, double> DM(kv, paraV, nspin);
+    elecstate::DensityMatrix<std::complex<double>, double> DM(paraV, nspin, kv->kvec_d, kv->get_nks() / nspin);
     Grid_Driver gd(0, 0);
     DM.init_DMR(&gd, &ucell);
     // construct another DM
-    elecstate::DensityMatrix<std::complex<double>, double> DM_test(kv, paraV, nspin);
+    elecstate::DensityMatrix<std::complex<double>, double> DM_test(paraV, nspin, kv->kvec_d, kv->get_nks() / nspin);
     DM_test.init_DMR(*DM.get_DMR_pointer(1));
     DM_test.save_DMR();
     EXPECT_EQ(DM_test.get_DMR_pointer(1)->get_nnr(), DM.get_DMR_pointer(1)->get_nnr());

source/module_elecstate/module_dm/test/test_dm_constructor.cpp

@@ -127,7 +127,7 @@ TEST_F(DMTest, DMConstructor_nspin1)
     std::cout << "dim0: " << paraV->dim0 << "    dim1:" << paraV->dim1 << std::endl;
     std::cout << "nrow: " << paraV->nrow << "    ncol:" << paraV->ncol << std::endl;
     int nspin = 1;
-    elecstate::DensityMatrix<double, double> DM(kv, paraV, nspin);
+    elecstate::DensityMatrix<double, double> DM(paraV, nspin, kv->kvec_d);
     // compare
     EXPECT_EQ(DM.get_DMK_nks(), kv->get_nks());
     EXPECT_EQ(DM.get_DMK_nrow(), paraV->nrow);
@@ -196,7 +196,7 @@ TEST_F(DMTest, DMConstructor_nspin2)
     // construct DM
     std::cout << "dim0: " << paraV->dim0 << "    dim1:" << paraV->dim1 << std::endl;
     std::cout << "nrow: " << paraV->nrow << "    ncol:" << paraV->ncol << std::endl;
-    elecstate::DensityMatrix<double, double> DM(kv, paraV, nspin);
+    elecstate::DensityMatrix<double, double> DM(paraV, nspin, kv->kvec_d, kv->get_nks() / nspin);
     // compare
     EXPECT_EQ(DM.get_DMK_nks(), kv->get_nks());
     EXPECT_EQ(DM.get_DMK_nrow(), paraV->nrow);

source/module_elecstate/module_dm/test/test_dm_io.cpp

@@ -142,7 +142,7 @@ TEST_F(DMTest, DMConstructor1)
     int nspin = 1;
     // construct DM
     std::cout << paraV->nrow << paraV->ncol << std::endl;
-    elecstate::DensityMatrix<double, double> DM(kv, paraV, nspin);
+    elecstate::DensityMatrix<double, double> DM(paraV, nspin, kv->kvec_d);
     // read DMK
     std::string directory = "./support/";
     for (int is = 1; is <= nspin; ++is)
@@ -162,7 +162,7 @@ TEST_F(DMTest, DMConstructor1)
         }
     }
     // construct a new DM
-    elecstate::DensityMatrix<double, double> DM1(kv, paraV, nspin);
+    elecstate::DensityMatrix<double, double> DM1(paraV, nspin, kv->kvec_d);
     directory = "./support/output";
     for (int is = 1; is <= nspin; ++is)
     {

source/module_hamilt_lcao/hamilt_lcaodft/fedm_k.cpp

@@ -42,10 +42,11 @@ void Force_LCAO<std::complex<double>>::cal_fedm(const bool isforce,
     ModuleBase::timer::tick("Force_LCAO", "cal_fedm");
 
     const int nspin = PARAM.inp.nspin;
+    const int nspin_dm = std::map<int, int>({ {1,1},{2,2},{4,1} })[nspin];
     const int nbands = PARAM.inp.nbands;
 
     // construct a DensityMatrix object
-    elecstate::DensityMatrix<std::complex<double>, double> edm(kv, &pv, nspin);
+    elecstate::DensityMatrix<std::complex<double>, double> edm(&pv, nspin_dm, kv->kvec_d, kv->get_nks() / nspin_dm);
 
     //--------------------------------------------
     // calculate the energy density matrix here.

source/module_hamilt_lcao/hamilt_lcaodft/hamilt_lcao.h

@@ -2,6 +2,7 @@
 #define W_ABACUS_DEVELOP_ABACUS_DEVELOP_SOURCE_MODULE_HAMILT_LCAO_HAMILT_LCAODFT_HAMILT_LCAO_H
 
 #include "module_basis/module_nao/two_center_bundle.h"
+#include "module_cell/klist.h"
 #include "module_elecstate/module_dm/density_matrix.h"
 #include "module_elecstate/potentials/potential_new.h"
 #include "module_hamilt_general/hamilt.h"

source/module_io/get_pchg_lcao.cpp

@@ -219,7 +219,8 @@ void IState_Charge::begin(Gint_k& gk,
         if (bands_picked_[ib])
         {
             // Using new density matrix inplementation (multi-k)
-            elecstate::DensityMatrix<std::complex<double>, double> DM(&kv, this->ParaV, nspin);
+            const int nspin_dm = std::map<int, int>({ {1,1},{2,2},{4,1} })[nspin];
+            elecstate::DensityMatrix<std::complex<double>, double> DM(this->ParaV, nspin_dm, kv.kvec_d, kv.get_nks() / nspin_dm);
 
 #ifdef __MPI
             this->idmatrix(ib, nspin, nelec, nlocal, wg, DM, kv, if_separate_k);