Skip to content

Refactor: generalize the transition density matrix in module_lr #5852

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 4 commits into from
Jan 13, 2025
Merged

Refactor: generalize the transition density matrix in module_lr #5852

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
6cf277f
generalize dm_trans
maki49 Jan 12, 2025
af8bbf4
pass factor instead of bool renorm_k
maki49 Jan 12, 2025
a7f4331
fix bug
maki49 Jan 12, 2025
2251150
update UT
maki49 Jan 12, 2025
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
9 changes: 6 additions & 3 deletions source/module_lr/ao_to_mo_transformer/ao_to_mo.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -7,15 +7,18 @@
#endif
namespace LR
{
#ifndef MO_TYPE_H
#define MO_TYPE_H
enum MO_TYPE { OO, VO, VV };
#endif
template<typename T>
void ao_to_mo_forloop_serial(
const std::vector<container::Tensor>& mat_ao,
const psi::Psi<T>& coeff,
const int& nocc,
const int& nvirt,
T* const mat_mo,
MO_TYPE type = VO);
const MO_TYPE type = VO);
template<typename T>
void ao_to_mo_blas(
const std::vector<container::Tensor>& mat_ao,
Expand All @@ -24,7 +27,7 @@ namespace LR
const int& nvirt,
T* const mat_mo,
const bool add_on = true,
MO_TYPE type = VO);
const MO_TYPE type = VO);
#ifdef __MPI
template<typename T>
void ao_to_mo_pblas(
Expand All @@ -38,6 +41,6 @@ namespace LR
const Parallel_2D& pmat_mo,
T* const mat_mo,
const bool add_on = true,
MO_TYPE type = VO);
const MO_TYPE type = VO);
#endif
}
4 changes: 2 additions & 2 deletions source/module_lr/ao_to_mo_transformer/ao_to_mo_parallel.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -21,7 +21,7 @@ namespace LR
const Parallel_2D& pmat_mo,
double* mat_mo,
const bool add_on,
MO_TYPE type)
const MO_TYPE type)
{
ModuleBase::TITLE("hamilt_lrtd", "ao_to_mo_pblas");
assert(pmat_ao.comm() == pcoeff.comm() && pmat_ao.comm() == pmat_mo.comm());
Expand Down Expand Up @@ -79,7 +79,7 @@ namespace LR
const Parallel_2D& pmat_mo,
std::complex<double>* const mat_mo,
const bool add_on,
MO_TYPE type)
const MO_TYPE type)
{
ModuleBase::TITLE("hamilt_lrtd", "cal_AX_plas");
assert(pmat_ao.comm() == pcoeff.comm() && pmat_ao.comm() == pmat_mo.comm());
Expand Down
8 changes: 4 additions & 4 deletions source/module_lr/ao_to_mo_transformer/ao_to_mo_serial.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -11,7 +11,7 @@ namespace LR
const int& nocc,
const int& nvirt,
double* mat_mo,
MO_TYPE type)
const MO_TYPE type)
{
ModuleBase::TITLE("hamilt_lrtd", "ao_to_mo_forloop_serial");
const int nks = mat_ao.size();
Expand Down Expand Up @@ -49,7 +49,7 @@ namespace LR
const int& nocc,
const int& nvirt,
std::complex<double>* const mat_mo,
MO_TYPE type)
const MO_TYPE type)
{
ModuleBase::TITLE("hamilt_lrtd", "ao_to_mo_forloop_serial");
const int nks = mat_ao.size();
Expand Down Expand Up @@ -88,7 +88,7 @@ namespace LR
const int& nvirt,
double* mat_mo,
const bool add_on,
MO_TYPE type)
const MO_TYPE type)
{
ModuleBase::TITLE("hamilt_lrtd", "ao_to_mo_blas");
const int nks = mat_ao.size();
Expand Down Expand Up @@ -129,7 +129,7 @@ namespace LR
const int& nvirt,
std::complex<double>* const mat_mo,
const bool add_on,
MO_TYPE type)
const MO_TYPE type)
{
ModuleBase::TITLE("hamilt_lrtd", "ao_to_mo_blas");
const int nks = mat_ao.size();
Expand Down
19 changes: 12 additions & 7 deletions source/module_lr/dm_trans/dm_trans.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -8,7 +8,12 @@
#endif
namespace LR
{
// use templates in the future.

#ifndef MO_TYPE_H
#define MO_TYPE_H
enum MO_TYPE { OO, VO, VV };
#endif

#ifdef __MPI
/// @brief calculate the 2d-block transition density matrix in AO basis using p?gemm
/// \f[ \tilde{\rho}_{\mu_j\mu_b}=\sum_{jb}c_{j,\mu_j}X_{jb}c^*_{b,\mu_b} \f]
Expand All @@ -22,8 +27,8 @@ namespace LR
const int nocc,
const int nvirt,
const Parallel_2D& pmat,
const bool renorm_k = true,
const int nspin = 1);
const T factor = (T)1.0,
const MO_TYPE type = MO_TYPE::VO);
#endif

/// @brief calculate the 2d-block transition density matrix in AO basis using ?gemm
Expand All @@ -32,8 +37,8 @@ namespace LR
const T* const X_istate,
const psi::Psi<T>& c,
const int& nocc, const int& nvirt,
const bool renorm_k = true,
const int nspin = 1);
const T factor = (T)1.0,
const MO_TYPE type = MO_TYPE::VO);

// for test
/// @brief calculate the 2d-block transition density matrix in AO basis using for loop (for test)
Expand All @@ -42,6 +47,6 @@ namespace LR
const T* const X_istate,
const psi::Psi<T>& c,
const int& nocc, const int& nvirt,
const bool renorm_k = true,
const int nspin = 1);
const T factor = (T)1.0,
const MO_TYPE type = MO_TYPE::VO);
}
50 changes: 29 additions & 21 deletions source/module_lr/dm_trans/dm_trans_parallel.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -18,44 +18,49 @@ std::vector<container::Tensor> cal_dm_trans_pblas(const double* const X_istate,
const int nocc,
const int nvirt,
const Parallel_2D& pmat,
const bool renorm_k,
const int nspin)
const double factor,
const MO_TYPE type)
{
ModuleBase::TITLE("hamilt_lrtd", "cal_dm_trans_pblas");
assert(px.comm() == pc.comm() && px.comm() == pmat.comm());
assert(px.blacs_ctxt == pc.blacs_ctxt && px.blacs_ctxt == pmat.blacs_ctxt);
assert(pmat.get_local_size() > 0);

const int nks = c.get_nk();
const int i1 = 1;
const int ivirt = nocc + 1;
const int nmo1 = type == MO_TYPE::VV ? nvirt : nocc;
const int nmo2 = type == MO_TYPE::OO ? nocc : nvirt;
const int imo1 = type == MO_TYPE::VV ? ivirt : i1;
const int imo2 = type == MO_TYPE::OO ? i1 : ivirt;

std::vector<container::Tensor> dm_trans(nks,
container::Tensor(DAT::DT_DOUBLE, DEV::CpuDevice, { pmat.get_col_size(), pmat.get_row_size() }));
for (int isk = 0; isk < nks; ++isk)
{
c.fix_k(isk);
const int x_start = isk * px.get_local_size();
int i1 = 1;
int ivirt = nocc + 1;

char transa = 'N';
char transb = 'T';
const double alpha = 1.0;
const double beta = 0;

// 1. [X*C_occ^T]^T=C_occ*X^T
Parallel_2D pXc; // nvirt*naos
LR_Util::setup_2d_division(pXc, px.get_block_size(), naos, nvirt, px.blacs_ctxt);
LR_Util::setup_2d_division(pXc, px.get_block_size(), naos, nmo2, px.blacs_ctxt);
container::Tensor Xc(DAT::DT_DOUBLE,
DEV::CpuDevice,
{pXc.get_col_size(), pXc.get_row_size()}); // row is "inside"(memory contiguity) for pblas
Xc.zero();
pdgemm_(&transa, &transb, &naos, &nvirt, &nocc,
&alpha, c.get_pointer(), &i1, &i1, pc.desc,
pdgemm_(&transa, &transb, &naos, &nmo2, &nmo1,
&alpha, c.get_pointer(), &i1, &imo1, pc.desc,
X_istate + x_start, &i1, &i1, px.desc,
&beta, Xc.data<double>(), &i1, &i1, pXc.desc);

// 2. C_virt*[X*C_occ^T]
pdgemm_(&transa, &transb, &naos, &naos, &nvirt,
&alpha, c.get_pointer(), &i1, &ivirt, pc.desc,
pdgemm_(&transa, &transb, &naos, &naos, &nmo2,
&factor, c.get_pointer(), &i1, &imo2, pc.desc,
Xc.data<double>(), &i1, &i1, pXc.desc,
&beta, dm_trans[isk].data<double>(), &i1, &i1, pmat.desc);
}
Expand All @@ -70,23 +75,27 @@ std::vector<container::Tensor> cal_dm_trans_pblas(const std::complex<double>* co
const int nocc,
const int nvirt,
const Parallel_2D& pmat,
const bool renorm_k,
const int nspin)
const std::complex<double> factor,
const MO_TYPE type)
{
ModuleBase::TITLE("hamilt_lrtd", "cal_dm_trans_pblas");
assert(px.comm() == pc.comm() && px.comm() == pmat.comm());
assert(px.blacs_ctxt == pc.blacs_ctxt && px.blacs_ctxt == pmat.blacs_ctxt);
assert(pmat.get_local_size() > 0);
const int nks = c.get_nk();
const int i1 = 1;
const int ivirt = nocc + 1;
const int nmo1 = type == MO_TYPE::VV ? nvirt : nocc;
const int nmo2 = type == MO_TYPE::OO ? nocc : nvirt;
const int imo1 = type == MO_TYPE::VV ? ivirt : i1;
const int imo2 = type == MO_TYPE::OO ? i1 : ivirt;

std::vector<container::Tensor> dm_trans(nks,
container::Tensor(DAT::DT_COMPLEX_DOUBLE, DEV::CpuDevice, {pmat.get_col_size(), pmat.get_row_size()}));
for (int isk = 0; isk < nks; ++isk)
{
c.fix_k(isk);
const int x_start = isk * px.get_local_size();
int i1 = 1;
int ivirt = nocc + 1;

// ============== C_virt * X * C_occ^\dagger=============
// char transa = 'N';
Expand Down Expand Up @@ -114,24 +123,23 @@ std::vector<container::Tensor> cal_dm_trans_pblas(const std::complex<double>* co
char transa = 'N';
char transb = 'C';
Parallel_2D pXc;
LR_Util::setup_2d_division(pXc, px.get_block_size(), nvirt, naos, px.blacs_ctxt);
LR_Util::setup_2d_division(pXc, px.get_block_size(), nmo2, naos, px.blacs_ctxt);
container::Tensor Xc(DAT::DT_COMPLEX_DOUBLE,
DEV::CpuDevice,
{pXc.get_col_size(), pXc.get_row_size()}); // row is "inside"(memory contiguity) for pblas
Xc.zero();
std::complex<double> alpha(1.0, 0.0);
const std::complex<double> alpha(1.0, 0.0);
const std::complex<double> beta(0.0, 0.0);
pzgemm_(&transa, &transb, &nvirt, &naos, &nocc, &alpha,
pzgemm_(&transa, &transb, &nmo2, &naos, &nmo1, &alpha,
X_istate + x_start, &i1, &i1, px.desc,
c.get_pointer(), &i1, &i1, pc.desc,
c.get_pointer(), &i1, &imo1, pc.desc,
&beta, Xc.data<std::complex<double>>(), &i1, &i1, pXc.desc);

// 2. [X*C_occ^\dagger]^TC_virt^T
alpha.real(renorm_k ? 1.0 / static_cast<double>(nks) : 1.0);
transa = transb = 'T';
pzgemm_(&transa, &transb, &naos, &naos, &nvirt,
&alpha, Xc.data<std::complex<double>>(), &i1, &i1, pXc.desc,
c.get_pointer(), &i1, &ivirt, pc.desc,
pzgemm_(&transa, &transb, &naos, &naos, &nmo2,
&factor, Xc.data<std::complex<double>>(), &i1, &i1, pXc.desc,
c.get_pointer(), &i1, &imo2, pc.desc,
&beta, dm_trans[isk].data<std::complex<double>>(), &i1, &i1, pmat.desc);
}
return dm_trans;
Expand Down
Loading
Loading