Skip to content

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

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 7 commits into from
Oct 17, 2024
Merged

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

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
7 commits
Select commit Hold shift + click to select a range
d5358d5
remove dependence of DensityMatrix on K_Vectors
maki49 Oct 13, 2024
8f31794
rename _nks as _nk to avoid ambiguity
maki49 Oct 13, 2024
f15c6e1
[pre-commit.ci lite] apply automatic fixes
pre-commit-ci-lite[bot] Oct 13, 2024
1c2e28c
remove map
maki49 Oct 14, 2024
68158c9
make nk non-optional
maki49 Oct 15, 2024
218e8ad
fix after rebase
maki49 Oct 16, 2024
0897233
Merge branch 'develop' into dm_construct
maki49 Oct 16, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
3 changes: 2 additions & 1 deletion source/module_elecstate/elecstate_lcao.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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 <>
Expand Down
107 changes: 34 additions & 73 deletions source/module_elecstate/module_dm/density_matrix.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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++)
{
Expand Down Expand Up @@ -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());
}
Expand All @@ -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);
}
Expand All @@ -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());
Expand All @@ -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>
Expand Down Expand Up @@ -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();
Expand Down Expand Up @@ -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);
Expand Down Expand Up @@ -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();
Expand Down Expand Up @@ -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);
Expand Down Expand Up @@ -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);
Expand Down Expand Up @@ -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();
Expand Down Expand Up @@ -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);
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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.
Expand All @@ -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();
Expand All @@ -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);
Expand All @@ -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];
}
}

Expand All @@ -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);
Expand All @@ -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();
}
}

Expand Down
27 changes: 15 additions & 12 deletions source/module_elecstate/module_dm/density_matrix.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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"
Expand Down Expand Up @@ -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);

Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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
Expand All @@ -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;


};
Expand Down
Loading
Loading