Skip to content

Commit

Permalink
Refactor: remove the use of LOWF in istate_envelope (#4468)
Browse files Browse the repository at this point in the history 
* investigate functions and trying to substitute lowf with paraV

* sort the param list of wfc_2d_to_grid

* temporarily store a commit

* bon week-end

* delete wfc_2d_to_grid in lowf

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

* fix integrated test

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

* delete the redundant delete[] on old member variable

* correct mem_size calculation

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

* roll back

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

---------

Co-authored-by: pre-commit-ci-lite[bot] <117423508+pre-commit-ci-lite[bot]@users.noreply.github.com>
  • Loading branch information
@kirk0830 @pre-commit-ci-lite
kirk0830 and pre-commit-ci-lite[bot] authored Jun 24, 2024
1 parent e733396 commit b55a5f3
Show file tree
Hide file tree
Showing 7 changed files with 422 additions and 299 deletions.
7 changes: 5 additions & 2 deletions source/module_esolver/esolver_ks_lcao_elec.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -453,14 +453,17 @@ void ESolver_KS_LCAO<TK, TR>::others(const int istep)
}
else if (cal_type == "get_wf")
{
// Note: on the removal of LOWF
// here the LOWF is involved by gathering the wave functions from 2D BCD
// to serial. Parameter list of begin() should be updated.
IState_Envelope IEP(this->pelec);
if (GlobalV::GAMMA_ONLY_LOCAL)
{
IEP.begin(this->psi,
this->pw_rho,
this->pw_wfc,
this->pw_big,
this->LOWF,
this->orb_con.ParaV,
this->GG,
INPUT.out_wfc_pw,
this->wf.out_wfc_r,
Expand All @@ -479,7 +482,7 @@ void ESolver_KS_LCAO<TK, TR>::others(const int istep)
this->pw_rho,
this->pw_wfc,
this->pw_big,
this->LOWF,
this->orb_con.ParaV,
this->GK,
INPUT.out_wfc_pw,
this->wf.out_wfc_r,
Expand Down
36 changes: 21 additions & 15 deletions source/module_hamilt_lcao/hamilt_lcaodft/local_orbital_charge.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -3,13 +3,13 @@
#include "module_base/blas_connector.h"
#include "module_base/timer.h"
#include "module_hamilt_pw/hamilt_pwdft/global.h"

#include "module_io/read_wfc_lcao.h"
// Shen Yu add 2019/5/9
extern "C"
{
void Cblacs_gridinfo(int icontxt, int *nprow, int *npcol, int *myprow, int *mypcol);
void Cblacs_pinfo(int *myid, int *nprocs);
void Cblacs_pcoord(int icontxt, int pnum, int *prow, int *pcol);
void Cblacs_gridinfo(int icontxt, int* nprow, int* npcol, int* myprow, int* mypcol);
void Cblacs_pinfo(int* myid, int* nprocs);
void Cblacs_pcoord(int icontxt, int pnum, int* prow, int* pcol);
int Cblacs_pnum(int icontxt, int prow, int pcol);
}

Expand Down Expand Up @@ -58,32 +58,38 @@ Local_Orbital_Charge::~Local_Orbital_Charge()
}

void Local_Orbital_Charge::allocate_dm_wfc(const Grid_Technique& gt,
elecstate::ElecState* pelec,
Local_Orbital_wfc& lowf,
psi::Psi<double>* psi,
const K_Vectors& kv,
const int& istep)
elecstate::ElecState* pelec,
Local_Orbital_wfc& lowf,
psi::Psi<double>* psi,
const K_Vectors& kv,
const int& istep)
{
ModuleBase::TITLE("Local_Orbital_Charge", "allocate_dm_wfc");
this->LOWF = &lowf;
this->LOWF->gridt = &gt;
// here we reset the density matrix dimension.
// it is weird the psi_gamma is allocated inside the function allocate_gamma
// but psi_k is allocated outside the function allocate_DM_k
// and why the allocation of psi and DM are done together???
// lowf.gamma_file(psi, pelec);
// will be replaced by
this->allocate_gamma(gt.lgd, psi, pelec, kv.get_nks(), istep);
return;
}

void Local_Orbital_Charge::allocate_dm_wfc(const Grid_Technique &gt,
elecstate::ElecState* pelec,
Local_Orbital_wfc& lowf,
psi::Psi<std::complex<double>>* psi,
const K_Vectors& kv,
const int& istep)
void Local_Orbital_Charge::allocate_dm_wfc(const Grid_Technique& gt,
elecstate::ElecState* pelec,
Local_Orbital_wfc& lowf,
psi::Psi<std::complex<double>>* psi,
const K_Vectors& kv,
const int& istep)
{
ModuleBase::TITLE("Local_Orbital_Charge", "allocate_dm_wfc");

this->LOWF = &lowf;
this->LOWF->gridt = &gt;
// here we reset the density matrix dimension.
// lgd will be the dimension of global matrix and nbasis of psi to be of the local one
lowf.allocate_k(gt.lgd, psi, pelec, kv.get_nks(), kv.get_nkstot(), kv.kvec_c, istep);
this->allocate_DM_k(kv.get_nks(), gt.nnrg);
return;
Expand Down
93 changes: 0 additions & 93 deletions source/module_hamilt_lcao/hamilt_lcaodft/local_orbital_wfc.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -236,96 +236,3 @@ int Local_Orbital_wfc::localIndex(int globalindex, int nblk, int nprocs, int& my
myproc = int((globalindex % (nblk * nprocs)) / nblk);
return int(globalindex / (nblk * nprocs)) * nblk + globalindex % nblk;
}

#ifdef __MPI
void Local_Orbital_wfc::wfc_2d_to_grid(const double* wfc_2d,
double** wfc_grid,
const int ik,
const ModuleBase::matrix& ekb,
const ModuleBase::matrix& wg)
{
ModuleBase::TITLE(" Local_Orbital_wfc", "wfc_2d_to_grid");
ModuleBase::timer::tick("Local_Orbital_wfc", "wfc_2d_to_grid");

const Parallel_Orbitals* pv = this->ParaV;
const int inc = 1;
int myid = 0;
MPI_Comm_rank(pv->comm_2D, &myid);
int info = 0;

// calculate maxnloc for bcasting 2d-wfc
long maxnloc; // maximum number of elements in local matrix
info = MPI_Reduce(&pv->nloc_wfc, &maxnloc, 1, MPI_LONG, MPI_MAX, 0, pv->comm_2D);
info = MPI_Bcast(&maxnloc, 1, MPI_LONG, 0, pv->comm_2D);
std::vector<double> work(maxnloc); // work/buffer matrix

int naroc[2]; // maximum number of row or column
for (int iprow = 0; iprow < pv->dim0; ++iprow)
{
for (int ipcol = 0; ipcol < pv->dim1; ++ipcol)
{
const int coord[2] = {iprow, ipcol};
int src_rank;
info = MPI_Cart_rank(pv->comm_2D, coord, &src_rank);
if (myid == src_rank)
{
BlasConnector::copy(pv->nloc_wfc, wfc_2d, inc, work.data(), inc);
naroc[0] = pv->nrow;
naroc[1] = pv->ncol_bands;
}
info = MPI_Bcast(naroc, 2, MPI_INT, src_rank, pv->comm_2D);
info = MPI_Bcast(work.data(), maxnloc, MPI_DOUBLE, src_rank, pv->comm_2D);

info = this->set_wfc_grid(naroc, pv->nb, pv->dim0, pv->dim1, iprow, ipcol, work.data(), wfc_grid);

} // loop ipcol
} // loop iprow
ModuleBase::timer::tick("Local_Orbital_wfc", "wfc_2d_to_grid");
}

void Local_Orbital_wfc::wfc_2d_to_grid(const std::complex<double>* wfc_2d,
std::complex<double>** wfc_grid,
int ik,
const ModuleBase::matrix& ekb,
const ModuleBase::matrix& wg,
const std::vector<ModuleBase::Vector3<double>>& kvec_c)
{
ModuleBase::TITLE("Local_Orbital_wfc", "wfc_2d_to_grid");
ModuleBase::timer::tick("Local_Orbital_wfc", "wfc_2d_to_grid");

const Parallel_Orbitals* pv = this->ParaV;
const int inc = 1;
int myid = 0;
MPI_Comm_rank(pv->comm_2D, &myid);
int info = 0;

// calculate maxnloc for bcasting 2d-wfc
long maxnloc = 0; // maximum number of elements in local matrix
info = MPI_Reduce(&pv->nloc_wfc, &maxnloc, 1, MPI_LONG, MPI_MAX, 0, pv->comm_2D);
info = MPI_Bcast(&maxnloc, 1, MPI_LONG, 0, pv->comm_2D);
std::vector<std::complex<double>> work(maxnloc); // work/buffer matrix

int naroc[2] = {0}; // maximum number of row or column
for (int iprow = 0; iprow < pv->dim0; ++iprow)
{
for (int ipcol = 0; ipcol < pv->dim1; ++ipcol)
{
const int coord[2] = {iprow, ipcol};
int src_rank;
info = MPI_Cart_rank(pv->comm_2D, coord, &src_rank);
if (myid == src_rank)
{
BlasConnector::copy(pv->nloc_wfc, wfc_2d, inc, work.data(), inc);
naroc[0] = pv->nrow;
naroc[1] = pv->ncol_bands;
}
info = MPI_Bcast(naroc, 2, MPI_INT, src_rank, pv->comm_2D);
info = MPI_Bcast(work.data(), maxnloc, MPI_DOUBLE_COMPLEX, src_rank, pv->comm_2D);
// mohan update 2021-02-12, delte BFIELD option
info = this->set_wfc_grid(naroc, pv->nb, pv->dim0, pv->dim1, iprow, ipcol, work.data(), wfc_grid);
} // loop ipcol
} // loop iprow

ModuleBase::timer::tick("Local_Orbital_wfc", "wfc_2d_to_grid");
}
#endif
115 changes: 9 additions & 106 deletions source/module_hamilt_lcao/hamilt_lcaodft/local_orbital_wfc.h
View file
Original file line number Diff line number Diff line change
Expand Up @@ -10,22 +10,9 @@

class Local_Orbital_wfc
{
public:

Local_Orbital_wfc();
~Local_Orbital_wfc();
// refactor new implementation: RAII
// a more look-looking name would be LocalOrbitalWfc, I suppose...
Local_Orbital_wfc(const int& nspin,
const int& nks,
const int& nbands,
const int& nlocal,
const int& gamma_only,
const int& nb2d,
const std::string& ks_solver,
const std::string& readin_dir);
//
void initialize();
public:
Local_Orbital_wfc();
~Local_Orbital_wfc();
///=========================================
/// grid wfc
/// used to generate density matrix: LOC.DM_R,
Expand All @@ -45,12 +32,12 @@ class Local_Orbital_wfc
/// read wavefunction coefficients: WFC_NAO_K/GAMMA*.txt
void gamma_file(psi::Psi<double>* psid, elecstate::ElecState* pelec);
void allocate_k(const int& lgd,
psi::Psi<std::complex<double>>* psi,
elecstate::ElecState* pelec,
const int& nks,
const int& nkstot,
const std::vector<ModuleBase::Vector3<double>>& kvec_c,
const int& istep);
psi::Psi<std::complex<double>>* psi,
elecstate::ElecState* pelec,
const int& nks,
const int& nkstot,
const std::vector<ModuleBase::Vector3<double>>& kvec_c,
const int& istep);

//=========================================
// Init Cij, make it satisfy 2 conditions:
Expand All @@ -65,96 +52,12 @@ class Local_Orbital_wfc
static int globalIndex(int localindex, int nblk, int nprocs, int myproc);
static int localIndex(int globalindex, int nblk, int nprocs, int& myproc);

#ifdef __MPI
///=========================================
/// Parallel: convert the distribution of wavefunction from 2D to grid
///=========================================
/// For gamma_only, T = double;
/// For multi-k, T = complex<double>;
/// Set myid and ctot when output is needed;
/// Set wfc as nullptr when 2d-to-grid convertion is not needed.

// Notice: here I reload this function rather than use template
// (in which the implementation should be put in header file )
// because sub-function `write_wfc_nao_complex`contains GlobalC declared in `global.h`
// which will cause lots of "not defined" if included in a header file.
void wfc_2d_to_grid(const double* wfc_2d,
double** wfc_grid,
const int ik,
const ModuleBase::matrix& ekb,
const ModuleBase::matrix& wg);
void wfc_2d_to_grid(const std::complex<double>* wfc_2d,
std::complex<double>** wfc_grid,
int ik,
const ModuleBase::matrix& ekb,
const ModuleBase::matrix& wg,
const std::vector<ModuleBase::Vector3<double>>& kvec_c);
#endif

int error = 0;

private:

template <typename T>
int set_wfc_grid(int naroc[2],
int nb,
int dim0,
int dim1,
int iprow,
int ipcol,
T* work,
T** wfc,
int myid = -1,
T** ctot = nullptr);

bool wfck_flag;
bool complex_flag;
bool allocate_flag;
int nks;
};


// the function should not be defined here!! mohan 2024-03-28
template <typename T>
int Local_Orbital_wfc::set_wfc_grid(int naroc[2],
int nb,
int dim0,
int dim1,
int iprow,
int ipcol,
T* work,
T** wfc,
int myid,
T** ctot)
{
#ifdef __MPI
ModuleBase::TITLE(" Local_Orbital_wfc", "set_wfc_grid");
if (!wfc && !ctot)
{
return 0;
}
for (int j = 0; j < naroc[1]; ++j)
{
int igcol = globalIndex(j, nb, dim1, ipcol);
if (igcol >= GlobalV::NBANDS)
{
continue;
}
for (int i = 0; i < naroc[0]; ++i)
{
int igrow = globalIndex(i, nb, dim0, iprow);
int mu_local = this->gridt->trace_lo[igrow];
if (wfc && mu_local >= 0)
{
wfc[igcol][mu_local] = work[j * naroc[0] + i];
}
if (ctot && myid == 0)
{
ctot[igcol][igrow] = work[j * naroc[0] + i];
}
}
}
#endif
return 0;
}
#endif
Loading

0 comments on commit b55a5f3

Please sign in to comment.