diff --git a/source/module_esolver/lcao_others.cpp b/source/module_esolver/lcao_others.cpp index 5587baed1d..2d117aaeed 100644 --- a/source/module_esolver/lcao_others.cpp +++ b/source/module_esolver/lcao_others.cpp @@ -110,13 +110,12 @@ void ESolver_KS_LCAO::others(const int istep) this->pw_big->nbz, PARAM.globalv.gamma_only_local, PARAM.inp.nbands_istate, - PARAM.inp.bands_to_print, + PARAM.inp.out_pchg, PARAM.inp.nbands, PARAM.inp.nelec, PARAM.inp.nspin, PARAM.globalv.nlocal, PARAM.globalv.global_out_dir, - GlobalV::MY_RANK, GlobalV::ofs_warning, &GlobalC::ucell, &GlobalC::GridD, @@ -140,13 +139,12 @@ void ESolver_KS_LCAO::others(const int istep) this->pw_big->nbz, PARAM.globalv.gamma_only_local, PARAM.inp.nbands_istate, - PARAM.inp.bands_to_print, + PARAM.inp.out_pchg, PARAM.inp.nbands, PARAM.inp.nelec, PARAM.inp.nspin, PARAM.globalv.nlocal, PARAM.globalv.global_out_dir, - GlobalV::MY_RANK, GlobalV::ofs_warning, &GlobalC::ucell, &GlobalC::GridD, @@ -174,7 +172,8 @@ void ESolver_KS_LCAO::others(const int istep) this->kv, PARAM.inp.nelec, PARAM.inp.nbands_istate, - PARAM.inp.bands_to_print, + PARAM.inp.out_wfc_norm, + PARAM.inp.out_wfc_re_im, PARAM.inp.nbands, PARAM.inp.nspin, PARAM.globalv.nlocal, @@ -193,7 +192,8 @@ void ESolver_KS_LCAO::others(const int istep) this->kv, PARAM.inp.nelec, PARAM.inp.nbands_istate, - PARAM.inp.bands_to_print, + PARAM.inp.out_wfc_norm, + PARAM.inp.out_wfc_re_im, PARAM.inp.nbands, PARAM.inp.nspin, PARAM.globalv.nlocal, diff --git a/source/module_io/get_pchg_lcao.cpp b/source/module_io/get_pchg_lcao.cpp index 3e8de3d5cf..2858147353 100644 --- a/source/module_io/get_pchg_lcao.cpp +++ b/source/module_io/get_pchg_lcao.cpp @@ -26,7 +26,7 @@ IState_Charge::~IState_Charge() { } -// for gamma only +// For gamma_only void IState_Charge::begin(Gint_Gamma& gg, double** rho, const ModuleBase::matrix& wg, @@ -41,13 +41,12 @@ void IState_Charge::begin(Gint_Gamma& gg, const int bigpw_nbz, const bool gamma_only_local, const int nbands_istate, - const std::vector& out_band_kb, + const std::vector& out_pchg, const int nbands, const double nelec, const int nspin, const int nlocal, const std::string& global_out_dir, - const int my_rank, std::ofstream& ofs_warning, const UnitCell* ucell_in, Grid_Driver* GridD_in, @@ -55,22 +54,22 @@ void IState_Charge::begin(Gint_Gamma& gg, { ModuleBase::TITLE("IState_Charge", "begin"); - std::cout << " Perform |psi(i)|^2 for selected bands (band-decomposed charge densities, gamma only)." << std::endl; + std::cout << " Calculate |psi(i)|^2 for selected bands (band-decomposed charge densities, gamma only)." + << std::endl; + // Determine the mode based on the input parameters int mode = 0; - if (nbands_istate > 0 && static_cast(out_band_kb.size()) == 0) + // mode = 1: select bands below and above the Fermi surface using parameter `nbands_istate` + if (nbands_istate > 0 && static_cast(out_pchg.size()) == 0) { mode = 1; } - else if (static_cast(out_band_kb.size()) > 0) + // mode = 2: select bands directly using parameter `out_pchg` + else if (static_cast(out_pchg.size()) > 0) { - // If out_band_kb (bands_to_print) is not empty, set mode to 2 + // If out_pchg is not empty, set mode to 2 mode = 2; - std::cout << " Notice: INPUT parameter `nbands_istate` overwritten by `bands_to_print`!" << std::endl; - } - else - { - mode = 3; + std::cout << " Notice: INPUT parameter `nbands_istate` overwritten by `out_pchg`!" << std::endl; } // if ucell is odd, it's correct, @@ -81,7 +80,7 @@ void IState_Charge::begin(Gint_Gamma& gg, std::cout << " number of occupied bands = " << fermi_band << std::endl; // Set this->bands_picked_ according to the mode - select_bands(nbands_istate, out_band_kb, nbands, nelec, mode, fermi_band); + select_bands(nbands_istate, out_pchg, nbands, nelec, mode, fermi_band); for (int ib = 0; ib < nbands; ++ib) { @@ -174,13 +173,12 @@ void IState_Charge::begin(Gint_k& gk, const int bigpw_nbz, const bool gamma_only_local, const int nbands_istate, - const std::vector& out_band_kb, + const std::vector& out_pchg, const int nbands, const double nelec, const int nspin, const int nlocal, const std::string& global_out_dir, - const int my_rank, std::ofstream& ofs_warning, UnitCell* ucell_in, Grid_Driver* GridD_in, @@ -191,18 +189,18 @@ void IState_Charge::begin(Gint_k& gk, { ModuleBase::TITLE("IState_Charge", "begin"); - std::cout << " Perform |psi(i)|^2 for selected bands (band-decomposed charge densities, multi-k)." << std::endl; + std::cout << " Calculate |psi(i)|^2 for selected bands (band-decomposed charge densities, multi-k)." << std::endl; int mode = 0; - if (nbands_istate > 0 && static_cast(out_band_kb.size()) == 0) + if (nbands_istate > 0 && static_cast(out_pchg.size()) == 0) { mode = 1; } - else if (static_cast(out_band_kb.size()) > 0) + else if (static_cast(out_pchg.size()) > 0) { - // If out_band_kb (bands_to_print) is not empty, set mode to 2 + // If out_pchg is not empty, set mode to 2 mode = 2; - std::cout << " Notice: INPUT parameter `nbands_istate` overwritten by `bands_to_print`!" << std::endl; + std::cout << " Notice: INPUT parameter `nbands_istate` overwritten by `out_pchg`!" << std::endl; } else { @@ -214,7 +212,7 @@ void IState_Charge::begin(Gint_k& gk, std::cout << " number of occupied bands = " << fermi_band << std::endl; // Set this->bands_picked_ according to the mode - select_bands(nbands_istate, out_band_kb, nbands, nelec, mode, fermi_band); + select_bands(nbands_istate, out_pchg, nbands, nelec, mode, fermi_band); for (int ib = 0; ib < nbands; ++ib) { @@ -322,13 +320,7 @@ void IState_Charge::begin(Gint_k& gk, { rho_save_pointers[i] = rho_save[i].data(); } - srho.begin(is, - rho_save_pointers.data(), - rhog, - ngmc, - nullptr, - rho_pw, - ucell_in->symm); + srho.begin(is, rho_save_pointers.data(), rhog, ngmc, nullptr, rho_pw, ucell_in->symm); } std::cout << " Writing cube files..."; @@ -368,7 +360,7 @@ void IState_Charge::begin(Gint_k& gk, } void IState_Charge::select_bands(const int nbands_istate, - const std::vector& out_band_kb, + const std::vector& out_pchg, const int nbands, const double nelec, const int mode, @@ -382,6 +374,7 @@ void IState_Charge::select_bands(const int nbands_istate, this->bands_picked_.resize(nbands); ModuleBase::GlobalFunc::ZEROS(bands_picked_.data(), nbands); + // mode = 1: select bands below and above the Fermi surface using parameter `nbands_istate` if (mode == 1) { bands_below = nbands_istate; @@ -404,33 +397,28 @@ void IState_Charge::select_bands(const int nbands_istate, } } } + // mode = 2: select bands directly using parameter `out_pchg` else if (mode == 2) { - // Check if length of out_band_kb is valid - if (static_cast(out_band_kb.size()) > nbands) + // Check if length of out_pchg is valid + if (static_cast(out_pchg.size()) > nbands) { - ModuleBase::WARNING_QUIT( - "IState_Charge::select_bands", - "The number of bands specified by `bands_to_print` in the INPUT file exceeds `nbands`!"); + ModuleBase::WARNING_QUIT("IState_Charge::select_bands", + "The number of bands specified by `out_pchg` in the INPUT file exceeds `nbands`!"); } - // Check if all elements in out_band_kb are 0 or 1 - for (int value: out_band_kb) + // Check if all elements in out_pchg are 0 or 1 + for (int value: out_pchg) { if (value != 0 && value != 1) { - ModuleBase::WARNING_QUIT( - "IState_Charge::select_bands", - "The elements of `bands_to_print` must be either 0 or 1. Invalid values found!"); + ModuleBase::WARNING_QUIT("IState_Charge::select_bands", + "The elements of `out_pchg` must be either 0 or 1. Invalid values found!"); } } - // Fill bands_picked_ with values from out_band_kb + // Fill bands_picked_ with values from out_pchg // Remaining bands are already set to 0 - const int length = std::min(static_cast(out_band_kb.size()), nbands); - for (int i = 0; i < length; ++i) - { - // out_band_kb rely on function parse_expression - bands_picked_[i] = out_band_kb[i]; - } + const int length = std::min(static_cast(out_pchg.size()), nbands); + std::copy(out_pchg.begin(), out_pchg.begin() + length, bands_picked_.begin()); // Check if there are selected bands below the Fermi surface bool has_below = false; @@ -485,7 +473,7 @@ void IState_Charge::select_bands(const int nbands_istate, } #ifdef __MPI -// for gamma only +// For gamma_only void IState_Charge::idmatrix(const int& ib, const int nspin, const double& nelec, diff --git a/source/module_io/get_pchg_lcao.h b/source/module_io/get_pchg_lcao.h index 26e15b2579..57617280c8 100644 --- a/source/module_io/get_pchg_lcao.h +++ b/source/module_io/get_pchg_lcao.h @@ -28,7 +28,7 @@ class IState_Charge ~IState_Charge(); - // for gamma only + // For gamma_only void begin(Gint_Gamma& gg, double** rho, const ModuleBase::matrix& wg, @@ -43,13 +43,12 @@ class IState_Charge const int bigpw_nbz, const bool gamma_only_local, const int nbands_istate, - const std::vector& out_band_kb, + const std::vector& out_pchg, const int nbands, const double nelec, const int nspin, const int nlocal, const std::string& global_out_dir, - const int my_rank, std::ofstream& ofs_warning, const UnitCell* ucell_in, Grid_Driver* GridD_in, @@ -72,13 +71,12 @@ class IState_Charge const int bigpw_nbz, const bool gamma_only_local, const int nbands_istate, - const std::vector& out_band_kb, + const std::vector& out_pchg, const int nbands, const double nelec, const int nspin, const int nlocal, const std::string& global_out_dir, - const int my_rank, std::ofstream& ofs_warning, UnitCell* ucell_in, Grid_Driver* GridD_in, @@ -92,14 +90,14 @@ class IState_Charge * @brief Set this->bands_picked_ according to the mode, and process an error if the mode is not recognized. * * @param nbands_istate INPUT parameter nbands_istate. - * @param out_band_kb Calculated from INPUT parameter bands_to_print, vector. + * @param out_pchg INPUT parameter out_pchg, vector. * @param nbands INPUT parameter nbands. * @param nelec Total number of electrons. * @param mode Selected mode. * @param fermi_band Calculated Fermi band. */ void select_bands(const int nbands_istate, - const std::vector& out_band_kb, + const std::vector& out_pchg, const int nbands, const double nelec, const int mode, diff --git a/source/module_io/get_wf_lcao.cpp b/source/module_io/get_wf_lcao.cpp index e6360cb6ae..24a98757d6 100644 --- a/source/module_io/get_wf_lcao.cpp +++ b/source/module_io/get_wf_lcao.cpp @@ -18,10 +18,11 @@ IState_Envelope::~IState_Envelope() { } +// For gamma_only void IState_Envelope::begin(const psi::Psi* psid, - const ModulePW::PW_Basis* rhopw, - const ModulePW::PW_Basis_K* wfcpw, - const ModulePW::PW_Basis_Big* bigpw, + const ModulePW::PW_Basis* pw_rhod, + const ModulePW::PW_Basis_K* pw_wfc, + const ModulePW::PW_Basis_Big* pw_big, const Parallel_Orbitals& para_orb, Gint_Gamma& gg, const int& out_wfc_pw, @@ -29,7 +30,8 @@ void IState_Envelope::begin(const psi::Psi* psid, const K_Vectors& kv, const double nelec, const int nbands_istate, - const std::vector& out_band_kb, + const std::vector& out_wfc_norm, + const std::vector& out_wfc_re_im, const int nbands, const int nspin, const int nlocal, @@ -37,150 +39,134 @@ void IState_Envelope::begin(const psi::Psi* psid, { ModuleBase::TITLE("IState_Envelope", "begin"); - std::cout << " Perform |psi(band, r)| for selected bands." << std::endl; + std::cout << " Calculate |psi(i, r)|, Re[psi(i, r)], Im[psi(i, r)] for selected bands (gamma only)." << std::endl; - int mode = 0; - if (nbands_istate > 0 && static_cast(out_band_kb.size()) == 0) - { - mode = 1; - } - else if (static_cast(out_band_kb.size()) > 0) - { - // If out_band_kb (bands_to_print) is not empty, set mode to 2 - mode = 2; - std::cout << " Notice: INPUT parameter `nbands_istate` overwritten by `bands_to_print`!" << std::endl; - } - - int fermi_band = 0; - int bands_below = 0; - int bands_above = 0; - - this->bands_picked_.resize(nbands); - ModuleBase::GlobalFunc::ZEROS(bands_picked_.data(), nbands); - - // (1) - // mohan update 2011-03-21 // if ucell is odd, it's correct, // if ucell is even, it's also correct. // +1.0e-8 in case like (2.999999999+1)/2 + const int fermi_band = static_cast((nelec + 1) / 2 + 1.0e-8); std::cout << " number of electrons = " << nelec << std::endl; - fermi_band = static_cast((nelec + 1) / 2 + 1.0e-8); std::cout << " number of occupied bands = " << fermi_band << std::endl; - if (mode == 1) + // allocate grid wave functions for gamma_only + std::vector wfc_gamma_grid(nspin); + for (int is = 0; is < nspin; ++is) { - bands_below = nbands_istate; - bands_above = nbands_istate; - - std::cout << " Plot band decomposed charge density below Fermi surface with " << bands_below << " bands." - << std::endl; - - std::cout << " Plot band decomposed charge density above Fermi surface with " << bands_above << " bands." - << std::endl; - - for (int ib = 0; ib < nbands; ib++) + wfc_gamma_grid[is] = new double*[nbands]; + for (int ib = 0; ib < nbands; ++ib) { - if (ib >= fermi_band - bands_below) - { - if (ib < fermi_band + bands_above) - { - bands_picked_[ib] = 1; - } - } + wfc_gamma_grid[is][ib] = new double[gg.gridt->lgd]; } } - else if (mode == 2) + + // for pw_wfc in G space + psi::Psi> psi_g; + if (out_wfc_pw || out_wfc_r) { - // Check if length of out_band_kb is valid - if (static_cast(out_band_kb.size()) > nbands) - { - ModuleBase::WARNING_QUIT( - "IState_Envelope::begin", - "The number of bands specified by `bands_to_print` in the INPUT file exceeds `nbands`!"); - } - // Check if all elements in bands_picked_ are 0 or 1 - for (int value: out_band_kb) - { - if (value != 0 && value != 1) - { - ModuleBase::WARNING_QUIT( - "IState_Envelope::begin", - "The elements of `bands_to_print` must be either 0 or 1. Invalid values found!"); - } - } - // Fill bands_picked_ with values from out_band_kb - // Remaining bands are already set to 0 - int length = std::min(static_cast(out_band_kb.size()), nbands); - for (int i = 0; i < length; ++i) - { - // out_band_kb rely on function parse_expression - bands_picked_[i] = out_band_kb[i]; - } + psi_g.resize(nspin, nbands, kv.ngk[0]); + } - std::cout << " Plot band decomposed charge density below the Fermi surface: band "; - for (int i = 0; i + 1 <= fermi_band; ++i) - { - if (bands_picked_[i] == 1) - { - std::cout << i + 1 << " "; - } - } - std::cout << std::endl; - std::cout << " Plot band decomposed charge density above the Fermi surface: band "; - for (int i = fermi_band; i < nbands; ++i) - { - if (bands_picked_[i] == 1) - { - std::cout << i + 1 << " "; - } - } - std::cout << std::endl; + const double mem_size = sizeof(double) * double(gg.gridt->lgd) * double(nbands) * double(nspin) / 1024.0 / 1024.0; + ModuleBase::Memory::record("IState_Envelope::begin::wfc_gamma_grid", mem_size); + printf(" Estimated on-the-fly memory consuming by IState_Envelope::begin::wfc_gamma_grid: %f MB\n", mem_size); + + int mode_norm = 0; + if (nbands_istate > 0 && static_cast(out_wfc_norm.size()) == 0) + { + mode_norm = 1; } - else + else if (static_cast(out_wfc_norm.size()) > 0) { - ModuleBase::WARNING_QUIT("IState_Envelope::begin", "Invalid mode! Please check the code."); + // If out_wfc_norm is not empty, set mode to 2 + mode_norm = 2; + std::cout << " Notice: INPUT parameter `nbands_istate` overwritten by `out_wfc_norm`!" << std::endl; } - // (2) cicle: + // Set this->bands_picked_ according to the mode + select_bands(nbands_istate, out_wfc_norm, nbands, nelec, mode_norm, fermi_band); - // (2.1) calculate the selected density matrix - // from wave functions. + // Calculate out_wfc_norm + for (int ib = 0; ib < nbands; ++ib) + { + if (bands_picked_[ib]) + { + std::cout << " Performing grid integral over real space grid for band " << ib + 1 << "..." << std::endl; - // (2.2) carry out the grid integration to - // get the charge density. + for (int is = 0; is < nspin; ++is) + { + ModuleBase::GlobalFunc::ZEROS(pes_->charge->rho[is], pw_wfc->nrxx); - // (2.3) output the charge density in .cub format. + psid->fix_k(is); +#ifdef __MPI + wfc_2d_to_grid(psid->get_pointer(), para_orb, wfc_gamma_grid[is], gg.gridt->trace_lo); +#else + // if not MPI enabled, it is the case psid holds a global matrix. use fix_k to switch between different + // spin channels (actually kpoints, because now the same kpoint in different spin channels are treated + // as distinct kpoints) - // allocate grid wavefunction for gamma_only - std::vector wfc_gamma_grid(nspin); - for (int is = 0; is < nspin; ++is) - { - wfc_gamma_grid[is] = new double*[nbands]; - for (int ib = 0; ib < nbands; ++ib) { - wfc_gamma_grid[is][ib] = new double[gg.gridt->lgd]; -} - } + for (int i = 0; i < nbands; ++i) + { + for (int j = 0; j < nlocal; ++j) + { + wfc_gamma_grid[is][i][j] = psid[0](i, j); + } + } +#endif - const double mem_size = sizeof(double) * double(gg.gridt->lgd) * double(nbands) * double(nspin) / 1024.0 / 1024.0; - ModuleBase::Memory::record("IState_Envelope::begin::wfc_gamma_grid", mem_size); - printf(" Estimated on-the-fly memory consuming by IState_Envelope::begin::wfc_gamma_grid: %f MB\n", mem_size); + gg.cal_env(wfc_gamma_grid[is][ib], pes_->charge->rho[is], GlobalC::ucell); - // for pw-wfc in G space - psi::Psi> pw_wfc_g; + pes_->charge->save_rho_before_sum_band(); - if (out_wfc_pw || out_wfc_r) + std::stringstream ss; + ss << global_out_dir << "BAND" << ib + 1 << "_GAMMA" << "_SPIN" << is + 1 << "_ENV.cube"; + + const double ef_tmp = this->pes_->eferm.get_efval(is); + ModuleIO::write_cube( +#ifdef __MPI + pw_big->bz, + pw_big->nbz, + pw_rhod->nplane, + pw_rhod->startz_current, +#endif + pes_->charge->rho_save[is], + is, + nspin, + 0, + ss.str(), + pw_rhod->nx, + pw_rhod->ny, + pw_rhod->nz, + ef_tmp, + &(GlobalC::ucell)); + } + } + } + + int mode_re_im = 0; + if (nbands_istate > 0 && static_cast(out_wfc_re_im.size()) == 0) + { + mode_re_im = 1; + } + else if (static_cast(out_wfc_re_im.size()) > 0) { - pw_wfc_g.resize(1, nbands, kv.ngk[0]); + // If out_wfc_re_im is not empty, set mode to 2 + mode_re_im = 2; + std::cout << " Notice: INPUT parameter `nbands_istate` overwritten by `out_wfc_re_im`!" << std::endl; } - for (int ib = 0; ib < nbands; ib++) + // Set this->bands_picked_ according to the mode + select_bands(nbands_istate, out_wfc_re_im, nbands, nelec, mode_re_im, fermi_band); + + // Calculate out_wfc_re_im + for (int ib = 0; ib < nbands; ++ib) { if (bands_picked_[ib]) { - for (int is = 0; is < nspin; ++is) // loop over spin + std::cout << " Performing grid integral over real space grid for band " << ib + 1 << "..." << std::endl; + + for (int is = 0; is < nspin; ++is) { - std::cout << " Perform envelope function for band " << ib + 1 << std::endl; - ModuleBase::GlobalFunc::ZEROS(pes_->charge->rho[is], wfcpw->nrxx); + ModuleBase::GlobalFunc::ZEROS(pes_->charge->rho[is], pw_wfc->nrxx); psid->fix_k(is); #ifdef __MPI @@ -193,38 +179,77 @@ void IState_Envelope::begin(const psi::Psi* psid, for (int i = 0; i < nbands; ++i) { for (int j = 0; j < nlocal; ++j) + { wfc_gamma_grid[is][i][j] = psid[0](i, j); + } } #endif + gg.cal_env(wfc_gamma_grid[is][ib], pes_->charge->rho[is], GlobalC::ucell); - pes_->charge->save_rho_before_sum_band(); // xiaohui add 2014-12-09 - std::stringstream ss; - ss << global_out_dir << "BAND" << ib + 1 << "_s_" << is + 1 << "_ENV.cube"; + pes_->charge->save_rho_before_sum_band(); + const double ef_tmp = this->pes_->eferm.get_efval(is); + + // only for gamma_only now + psi_g.fix_k(is); + this->set_pw_wfc(pw_wfc, is, ib, nspin, pes_->charge->rho, psi_g); + + // Calculate real-space wave functions + psi_g.fix_k(is); + std::vector> wfc_r(pw_wfc->nrxx); + pw_wfc->recip2real(&psi_g(ib, 0), wfc_r.data(), is); + + // Extract real and imaginary parts + std::vector wfc_real(pw_wfc->nrxx); + std::vector wfc_imag(pw_wfc->nrxx); + for (int ir = 0; ir < pw_wfc->nrxx; ++ir) + { + wfc_real[ir] = wfc_r[ir].real(); + wfc_imag[ir] = wfc_r[ir].imag(); + } + + // Output real part + std::stringstream ss_real; + ss_real << global_out_dir << "BAND" << ib + 1 << "_GAMMA" << "_SPIN" << is + 1 << "_REAL.cube"; ModuleIO::write_cube( #ifdef __MPI - bigpw->bz, - bigpw->nbz, - rhopw->nplane, - rhopw->startz_current, + pw_big->bz, + pw_big->nbz, + pw_rhod->nplane, + pw_rhod->startz_current, #endif - pes_->charge->rho_save[is], + wfc_real.data(), is, nspin, 0, - ss.str(), - rhopw->nx, - rhopw->ny, - rhopw->nz, + ss_real.str(), + pw_rhod->nx, + pw_rhod->ny, + pw_rhod->nz, ef_tmp, - &(GlobalC::ucell), - 3, - 1); + &(GlobalC::ucell)); - if (out_wfc_pw || out_wfc_r) { // only for gamma_only now - this->set_pw_wfc(wfcpw, 0, ib, nspin, pes_->charge->rho_save, pw_wfc_g); -} + // Output imaginary part + std::stringstream ss_imag; + ss_imag << global_out_dir << "BAND" << ib + 1 << "_GAMMA" << "_SPIN" << is + 1 << "_IMAG.cube"; + ModuleIO::write_cube( +#ifdef __MPI + pw_big->bz, + pw_big->nbz, + pw_rhod->nplane, + pw_rhod->startz_current, +#endif + wfc_imag.data(), + is, + nspin, + 0, + ss_imag.str(), + pw_rhod->nx, + pw_rhod->ny, + pw_rhod->nz, + ef_tmp, + &(GlobalC::ucell)); } } } @@ -233,29 +258,32 @@ void IState_Envelope::begin(const psi::Psi* psid, { std::stringstream ssw; ssw << global_out_dir << "WAVEFUNC"; - std::cout << " write G-space wavefunction into \"" << global_out_dir << "/" << ssw.str() << "\" files." + std::cout << " Write G-space wave functions into \"" << global_out_dir << "/" << ssw.str() << "\" files." << std::endl; - ModuleIO::write_wfc_pw(ssw.str(), pw_wfc_g, kv, wfcpw); + ModuleIO::write_wfc_pw(ssw.str(), psi_g, kv, pw_wfc); } + if (out_wfc_r) { - ModuleIO::write_psi_r_1(pw_wfc_g, wfcpw, "wfc_realspace", false, kv); + ModuleIO::write_psi_r_1(psi_g, pw_wfc, "wfc_realspace", false, kv); } for (int is = 0; is < nspin; ++is) { - for (int ib = 0; ib < nbands; ++ib) { + for (int ib = 0; ib < nbands; ++ib) + { delete[] wfc_gamma_grid[is][ib]; -} + } delete[] wfc_gamma_grid[is]; } return; } +// For multi-k void IState_Envelope::begin(const psi::Psi>* psi, - const ModulePW::PW_Basis* rhopw, - const ModulePW::PW_Basis_K* wfcpw, - const ModulePW::PW_Basis_Big* bigpw, + const ModulePW::PW_Basis* pw_rhod, + const ModulePW::PW_Basis_K* pw_wfc, + const ModulePW::PW_Basis_Big* pw_big, const Parallel_Orbitals& para_orb, Gint_k& gk, const int& out_wf, @@ -263,7 +291,8 @@ void IState_Envelope::begin(const psi::Psi>* psi, const K_Vectors& kv, const double nelec, const int nbands_istate, - const std::vector& out_band_kb, + const std::vector& out_wfc_norm, + const std::vector& out_wfc_re_im, const int nbands, const int nspin, const int nlocal, @@ -271,123 +300,17 @@ void IState_Envelope::begin(const psi::Psi>* psi, { ModuleBase::TITLE("IState_Envelope", "begin"); - std::cout << " Perform |psi(band, r)| for selected bands." << std::endl; - - int mode = 0; - if (nbands_istate > 0 && static_cast(out_band_kb.size()) == 0) - { - mode = 1; - } - else if (static_cast(out_band_kb.size()) > 0) - { - // If out_band_kb (bands_to_print) is not empty, set mode to 2 - mode = 2; - std::cout << " Notice: INPUT parameter `nbands_istate` overwritten by `bands_to_print`!" << std::endl; - } + std::cout << " Calculate |psi(i, r)|, Re[psi(i, r)], Im[psi(i, r)] for selected bands (multi-k)." << std::endl; - int fermi_band = 0; - int bands_below = 0; - int bands_above = 0; - - this->bands_picked_.resize(nbands); - ModuleBase::GlobalFunc::ZEROS(bands_picked_.data(), nbands); - - // (1) - // mohan update 2011-03-21 // if ucell is odd, it's correct, // if ucell is even, it's also correct. // +1.0e-8 in case like (2.999999999+1)/2 // if NSPIN=4, each band only one electron, fermi_band should be nelec - + const int fermi_band = nspin < 4 ? static_cast((nelec + 1) / 2 + 1.0e-8) : nelec; std::cout << " number of electrons = " << nelec << std::endl; - fermi_band = nspin < 4 ? static_cast((nelec + 1) / 2 + 1.0e-8) : nelec; std::cout << " number of occupied bands = " << fermi_band << std::endl; - if (mode == 1) - { - bands_below = nbands_istate; - bands_above = nbands_istate; - - std::cout << " Plot band decomposed charge density below Fermi surface with " << bands_below << " bands." - << std::endl; - - std::cout << " Plot band decomposed charge density above Fermi surface with " << bands_above << " bands." - << std::endl; - - for (int ib = 0; ib < nbands; ib++) - { - if (ib >= fermi_band - bands_below) - { - if (ib < fermi_band + bands_above) - { - bands_picked_[ib] = 1; - } - } - } - } - else if (mode == 2) - { - // Check if length of out_band_kb is valid - if (static_cast(out_band_kb.size()) > nbands) - { - ModuleBase::WARNING_QUIT( - "IState_Envelope::begin", - "The number of bands specified by `bands_to_print` in the INPUT file exceeds `nbands`!"); - } - // Check if all elements in bands_picked_ are 0 or 1 - for (int value: out_band_kb) - { - if (value != 0 && value != 1) - { - ModuleBase::WARNING_QUIT( - "IState_Envelope::begin", - "The elements of `bands_to_print` must be either 0 or 1. Invalid values found!"); - } - } - // Fill bands_picked_ with values from out_band_kb - // Remaining bands are already set to 0 - int length = std::min(static_cast(out_band_kb.size()), nbands); - for (int i = 0; i < length; ++i) - { - // out_band_kb rely on function parse_expression - bands_picked_[i] = out_band_kb[i]; - } - - std::cout << " Plot band decomposed charge density below the Fermi surface: band "; - for (int i = 0; i + 1 <= fermi_band; ++i) - { - if (bands_picked_[i] == 1) - { - std::cout << i + 1 << " "; - } - } - std::cout << std::endl; - std::cout << " Plot band decomposed charge density above the Fermi surface: band "; - for (int i = fermi_band; i < nbands; ++i) - { - if (bands_picked_[i] == 1) - { - std::cout << i + 1 << " "; - } - } - std::cout << std::endl; - } - else - { - ModuleBase::WARNING_QUIT("IState_Envelope::begin", "Invalid mode! Please check the code."); - } - - // (2) cicle: - - // (2.1) calculate the selected density matrix - // from wave functions. - - // (2.2) carry out the grid integration to - // get the charge density. - - // (2.3) output the charge density in .cub format. - - // allocate grid wavefunction for gamma_only + // allocate grid wave functions for multi-k const int nks = kv.get_nks(); std::vector**> wfc_k_grid(nks); for (int ik = 0; ik < nks; ++ik) @@ -398,30 +321,48 @@ void IState_Envelope::begin(const psi::Psi>* psi, wfc_k_grid[ik][ib] = new std::complex[gk.gridt->lgd]; } } + const double mem_size = sizeof(std::complex) * double(gk.gridt->lgd) * double(nbands) * double(nks) / 1024.0 / 1024.0; ModuleBase::Memory::record("IState_Envelope::begin::wfc_k_grid", mem_size); printf(" Estimated on-the-fly memory consuming by IState_Envelope::begin::wfc_k_grid: %f MB\n", mem_size); - // for pw-wfc in G space - psi::Psi> pw_wfc_g(kv.ngk.data()); - + // for pw_wfc in G space + psi::Psi> psi_g(kv.ngk.data()); if (out_wf || out_wf_r) { - pw_wfc_g.resize(nks, nbands, wfcpw->npwk_max); + psi_g.resize(nks, nbands, pw_wfc->npwk_max); + } + + int mode_norm = 0; + if (nbands_istate > 0 && static_cast(out_wfc_norm.size()) == 0) + { + mode_norm = 1; } + else if (static_cast(out_wfc_norm.size()) > 0) + { + // If out_wfc_norm is not empty, set mode to 2 + mode_norm = 2; + std::cout << " Notice: INPUT parameter `nbands_istate` overwritten by `out_wfc_norm`!" << std::endl; + } + + // Set this->bands_picked_ according to the mode + select_bands(nbands_istate, out_wfc_norm, nbands, nelec, mode_norm, fermi_band); - for (int ib = 0; ib < nbands; ib++) + // Calculate out_wfc_norm + for (int ib = 0; ib < nbands; ++ib) { if (bands_picked_[ib]) { + std::cout << " Performing grid integral over real space grid for band " << ib + 1 << "..." << std::endl; + const int nspin0 = (nspin == 2) ? 2 : 1; for (int ik = 0; ik < nks; ++ik) // the loop of nspin0 is included { const int ispin = kv.isk[ik]; ModuleBase::GlobalFunc::ZEROS(pes_->charge->rho[ispin], - wfcpw->nrxx); // terrible, you make changes on another instance's data??? - std::cout << " Perform envelope function for kpoint " << ik << ", band" << ib + 1 << std::endl; + pw_wfc->nrxx); // terrible, you make changes on another instance's data??? + std::cout << " Calculate envelope function for kpoint " << ik + 1 << ", band" << ib + 1 << std::endl; // 2d-to-grid conversion is unified into `wfc_2d_to_grid`. psi->fix_k(ik); #ifdef __MPI // need to deal with NSPIN=4 !!!! @@ -437,34 +378,33 @@ void IState_Envelope::begin(const psi::Psi>* psi, gk.cal_env_k(ik, wfc_k_grid[ik][ib], pes_->charge->rho[ispin], kv.kvec_c, kv.kvec_d, GlobalC::ucell); std::stringstream ss; - ss << global_out_dir << "BAND" << ib + 1 << "_k_" << ik / nspin0 + 1 << "_s_" << ispin + 1 - << "_ENV.cube"; + ss << global_out_dir << "BAND" << ib + 1 << "_k_" << ik + 1 << "_s_" << ispin + 1 << "_ENV.cube"; const double ef_tmp = this->pes_->eferm.get_efval(ispin); ModuleIO::write_cube( #ifdef __MPI - bigpw->bz, - bigpw->nbz, - rhopw->nplane, - rhopw->startz_current, + pw_big->bz, + pw_big->nbz, + pw_rhod->nplane, + pw_rhod->startz_current, #endif pes_->charge->rho[ispin], ispin, nspin, 0, ss.str(), - rhopw->nx, - rhopw->ny, - rhopw->nz, + pw_rhod->nx, + pw_rhod->ny, + pw_rhod->nz, ef_tmp, &(GlobalC::ucell), 3, 1); - if (out_wf || out_wf_r) // only for gamma_only now + if (out_wf || out_wf_r) { - pw_wfc_g.fix_k(ik); - this->set_pw_wfc(wfcpw, ik, ib, nspin, pes_->charge->rho, pw_wfc_g); + psi_g.fix_k(ik); + this->set_pw_wfc(pw_wfc, ik, ib, nspin, pes_->charge->rho, psi_g); } } } @@ -476,50 +416,244 @@ void IState_Envelope::begin(const psi::Psi>* psi, { std::stringstream ssw; ssw << global_out_dir << "WAVEFUNC"; - std::cout << " write G-space wavefunction into \"" << global_out_dir << "/" << ssw.str() << "\" files." + std::cout << " write G-space wave functions into \"" << global_out_dir << "/" << ssw.str() << "\" files." << std::endl; - ModuleIO::write_wfc_pw(ssw.str(), pw_wfc_g, kv, wfcpw); + ModuleIO::write_wfc_pw(ssw.str(), psi_g, kv, pw_wfc); } if (out_wf_r) { - ModuleIO::write_psi_r_1(pw_wfc_g, wfcpw, "wfc_realspace", false, kv); + ModuleIO::write_psi_r_1(psi_g, pw_wfc, "wfc_realspace", false, kv); + } + + std::cout << " Outputting real-space wave functions in cube format..." << std::endl; + + for (int ib = 0; ib < nbands; ++ib) + { + if (bands_picked_[ib]) + { + const int nspin0 = (nspin == 2) ? 2 : 1; + for (int ik = 0; ik < nks; ++ik) + { + const int ispin = kv.isk[ik]; + std::cout << " Processing band " << ib + 1 << ", k-point " << ik << ", spin " << ispin + 1 + << std::endl; + + psi_g.fix_k(ik); + + // Calculate real-space wave functions + std::vector> wfc_r(pw_wfc->nrxx); + pw_wfc->recip2real(&psi_g(ib, 0), wfc_r.data(), ik); + + // Extract real and imaginary parts + std::vector wfc_real(pw_wfc->nrxx); + std::vector wfc_imag(pw_wfc->nrxx); + for (int ir = 0; ir < pw_wfc->nrxx; ++ir) + { + wfc_real[ir] = wfc_r[ir].real(); + wfc_imag[ir] = wfc_r[ir].imag(); + } + + // Output real part + std::stringstream ss_real; + ss_real << global_out_dir << "BAND" << ib + 1 << "_k_" << ik + 1 << "_s_" << ispin + 1 + << "_REAL.cube"; + const double ef_tmp = this->pes_->eferm.get_efval(ispin); + ModuleIO::write_cube( +#ifdef __MPI + pw_big->bz, + pw_big->nbz, + pw_rhod->nplane, + pw_rhod->startz_current, +#endif + wfc_real.data(), + ispin, + nspin, + 0, + ss_real.str(), + pw_rhod->nx, + pw_rhod->ny, + pw_rhod->nz, + ef_tmp, + &(GlobalC::ucell)); + + // Output imaginary part + std::stringstream ss_imag; + ss_imag << global_out_dir << "BAND" << ib + 1 << "_k_" << ik + 1 << "_s_" << ispin + 1 + << "_IMAG.cube"; + ModuleIO::write_cube( +#ifdef __MPI + pw_big->bz, + pw_big->nbz, + pw_rhod->nplane, + pw_rhod->startz_current, +#endif + wfc_imag.data(), + ispin, + nspin, + 0, + ss_imag.str(), + pw_rhod->nx, + pw_rhod->ny, + pw_rhod->nz, + ef_tmp, + &(GlobalC::ucell)); + } + } } } for (int ik = 0; ik < nks; ++ik) { - for (int ib = 0; ib < nbands; ++ib) { + for (int ib = 0; ib < nbands; ++ib) + { delete[] wfc_k_grid[ik][ib]; -} + } delete[] wfc_k_grid[ik]; } return; } +void IState_Envelope::select_bands(const int nbands_istate, + const std::vector& out_wfc_kb, + const int nbands, + const double nelec, + const int mode, + const int fermi_band) +{ + ModuleBase::TITLE("IState_Envelope", "select_bands"); + + int bands_below = 0; + int bands_above = 0; + + this->bands_picked_.resize(nbands); + ModuleBase::GlobalFunc::ZEROS(bands_picked_.data(), nbands); + + // mode = 1: select bands below and above the Fermi surface using parameter `nbands_istate` + if (mode == 1) + { + bands_below = nbands_istate; + bands_above = nbands_istate; + + std::cout << " Plot wave functions below the Fermi surface with " << bands_below << " bands." << std::endl; + + std::cout << " Plot wave functions above the Fermi surface with " << bands_above << " bands." << std::endl; + + for (int ib = 0; ib < nbands; ++ib) + { + if (ib >= fermi_band - bands_below) + { + if (ib < fermi_band + bands_above) + { + bands_picked_[ib] = 1; + } + } + } + } + // mode = 2: select bands directly using parameter `out_wfc_norm` or `out_wfc_re_im` + else if (mode == 2) + { + // Check if length of out_wfc_kb is valid + if (static_cast(out_wfc_kb.size()) > nbands) + { + ModuleBase::WARNING_QUIT("IState_Envelope::select_bands", + "The number of bands specified by `out_wfc_norm` or `out_wfc_re_im` in the INPUT " + "file exceeds `nbands`!"); + } + // Check if all elements in out_wfc_kb are 0 or 1 + for (int value: out_wfc_kb) + { + if (value != 0 && value != 1) + { + ModuleBase::WARNING_QUIT( + "IState_Envelope::select_bands", + "The elements of `out_wfc_norm` or `out_wfc_re_im` must be either 0 or 1. Invalid values found!"); + } + } + // Fill bands_picked_ with values from out_wfc_kb + // Remaining bands are already set to 0 + const int length = std::min(static_cast(out_wfc_kb.size()), nbands); + std::copy(out_wfc_kb.begin(), out_wfc_kb.begin() + length, bands_picked_.begin()); + + // Check if there are selected bands below the Fermi surface + bool has_below = false; + for (int i = 0; i + 1 <= fermi_band; ++i) + { + if (bands_picked_[i] == 1) + { + has_below = true; + break; + } + } + if (has_below) + { + std::cout << " Plot wave functions below the Fermi surface: band "; + for (int i = 0; i + 1 <= fermi_band; ++i) + { + if (bands_picked_[i] == 1) + { + std::cout << i + 1 << " "; + } + } + std::cout << std::endl; + } + + // Check if there are selected bands above the Fermi surface + bool has_above = false; + for (int i = fermi_band; i < nbands; ++i) + { + if (bands_picked_[i] == 1) + { + has_above = true; + break; + } + } + if (has_above) + { + std::cout << " Plot wave functions above the Fermi surface: band "; + for (int i = fermi_band; i < nbands; ++i) + { + if (bands_picked_[i] == 1) + { + std::cout << i + 1 << " "; + } + } + std::cout << std::endl; + } + } + else + { + ModuleBase::WARNING_QUIT("IState_Envelope::select_bands", "Invalid mode! Please check the code."); + } +} + // for each band -void IState_Envelope::set_pw_wfc(const ModulePW::PW_Basis_K* wfcpw, +void IState_Envelope::set_pw_wfc(const ModulePW::PW_Basis_K* pw_wfc, const int& ik, const int& ib, const int& nspin, const double* const* const rho, psi::Psi>& wfc_g) { - if (ib == 0) { // once is enough + if (ib == 0) + { + // once is enough ModuleBase::TITLE("IState_Envelope", "set_pw_wfc"); -} + } - std::vector> Porter(wfcpw->nrxx); + std::vector> Porter(pw_wfc->nrxx); // here I refer to v_hartree, but I don't know how to deal with NSPIN=4 const int nspin0 = (nspin == 2) ? 2 : 1; - for (int is = 0; is < nspin0; is++) { - for (int ir = 0; ir < wfcpw->nrxx; ir++) { + for (int is = 0; is < nspin0; ++is) + { + for (int ir = 0; ir < pw_wfc->nrxx; ++ir) + { Porter[ir] += std::complex(rho[is][ir], 0.0); -} -} + } + } // call FFT - wfcpw->real2recip(Porter.data(), &wfc_g(ib, 0), ik); + pw_wfc->real2recip(Porter.data(), &wfc_g(ib, 0), ik); } #ifdef __MPI diff --git a/source/module_io/get_wf_lcao.h b/source/module_io/get_wf_lcao.h index 0b075b8462..b4c8c11d56 100644 --- a/source/module_io/get_wf_lcao.h +++ b/source/module_io/get_wf_lcao.h @@ -16,11 +16,11 @@ class IState_Envelope IState_Envelope(const elecstate::ElecState* pes); ~IState_Envelope(); - /// for gamma_only + /// For gamma_only void begin(const psi::Psi* psid, - const ModulePW::PW_Basis* rhopw, - const ModulePW::PW_Basis_K* wfcpw, - const ModulePW::PW_Basis_Big* bigpw, + const ModulePW::PW_Basis* pw_rhod, + const ModulePW::PW_Basis_K* pw_wfc, + const ModulePW::PW_Basis_Big* pw_big, const Parallel_Orbitals& para_orb, Gint_Gamma& gg, const int& out_wfc_pw, @@ -28,7 +28,8 @@ class IState_Envelope const K_Vectors& kv, const double nelec, const int nbands_istate, - const std::vector& out_band_kb, + const std::vector& out_wfc_norm, + const std::vector& out_wfc_re_im, const int nbands, const int nspin, const int nlocal, @@ -36,9 +37,9 @@ class IState_Envelope /// tmp, delete after Gint is refactored. void begin(const psi::Psi* psid, - const ModulePW::PW_Basis* rhopw, - const ModulePW::PW_Basis_K* wfcpw, - const ModulePW::PW_Basis_Big* bigpw, + const ModulePW::PW_Basis* pw_rhod, + const ModulePW::PW_Basis_K* pw_wfc, + const ModulePW::PW_Basis_Big* pw_big, const Parallel_Orbitals& para_orb, Gint_k& gg, const int& out_wfc_pw, @@ -46,7 +47,8 @@ class IState_Envelope const K_Vectors& kv, const double nelec, const int nbands_istate, - const std::vector& out_band_kb, + const std::vector& out_wfc_norm, + const std::vector& out_wfc_re_im, const int nbands, const int nspin, const int nlocal, @@ -54,11 +56,12 @@ class IState_Envelope { throw std::logic_error("gint_k should use with complex psi."); }; - /// for multi-k + + /// For multi-k void begin(const psi::Psi>* psi, - const ModulePW::PW_Basis* rhopw, - const ModulePW::PW_Basis_K* wfcpw, - const ModulePW::PW_Basis_Big* bigpw, + const ModulePW::PW_Basis* pw_rhod, + const ModulePW::PW_Basis_K* pw_wfc, + const ModulePW::PW_Basis_Big* pw_big, const Parallel_Orbitals& para_orb, Gint_k& gk, const int& out_wfc_pw, @@ -66,7 +69,8 @@ class IState_Envelope const K_Vectors& kv, const double nelec, const int nbands_istate, - const std::vector& out_band_kb, + const std::vector& out_wfc_norm, + const std::vector& out_wfc_re_im, const int nbands, const int nspin, const int nlocal, @@ -74,9 +78,9 @@ class IState_Envelope /// tmp, delete after Gint is refactored. void begin(const psi::Psi>* psi, - const ModulePW::PW_Basis* rhopw, - const ModulePW::PW_Basis_K* wfcpw, - const ModulePW::PW_Basis_Big* bigpw, + const ModulePW::PW_Basis* pw_rhod, + const ModulePW::PW_Basis_K* pw_wfc, + const ModulePW::PW_Basis_Big* pw_big, const Parallel_Orbitals& para_orb, Gint_Gamma& gk, const int& out_wfc_pw, @@ -84,7 +88,8 @@ class IState_Envelope const K_Vectors& kv, const double nelec, const int nbands_istate, - const std::vector& out_band_kb, + const std::vector& out_wfc_norm, + const std::vector& out_wfc_re_im, const int nbands, const int nspin, const int nlocal, @@ -94,16 +99,22 @@ class IState_Envelope }; private: - std::vector bands_picked_; - const elecstate::ElecState* pes_ = nullptr; + void select_bands(const int nbands_istate, + const std::vector& out_wfc_kb, + const int nbands, + const double nelec, + const int mode, + const int fermi_band); - void set_pw_wfc(const ModulePW::PW_Basis_K* wfcpw, + void set_pw_wfc(const ModulePW::PW_Basis_K* pw_wfc, const int& ik, const int& ib, const int& nspin, const double* const* const rho, psi::Psi>& wfc_g); + int globalIndex(int localindex, int nblk, int nprocs, int myproc); + int localIndex(int globalindex, int nblk, int nprocs, int& myproc); #ifdef __MPI @@ -120,5 +131,8 @@ class IState_Envelope template void wfc_2d_to_grid(const T* wfc_2d, const Parallel_Orbitals& pv, T** wfc_grid, const std::vector& trace_lo); #endif + + std::vector bands_picked_; + const elecstate::ElecState* pes_ = nullptr; }; #endif diff --git a/source/module_io/read_input_item_output.cpp b/source/module_io/read_input_item_output.cpp index 7f0651299f..7709f60213 100644 --- a/source/module_io/read_input_item_output.cpp +++ b/source/module_io/read_input_item_output.cpp @@ -42,12 +42,16 @@ void ReadInput::item_output() item.read_value = [](const Input_Item& item, Parameter& para) { size_t count = item.get_size(); std::vector out_chg(count); // create a placeholder vector - std::transform(item.str_values.begin(), item.str_values.end(), out_chg.begin(), [](std::string s) { return std::stoi(s); }); + std::transform(item.str_values.begin(), item.str_values.end(), out_chg.begin(), [](std::string s) { + return std::stoi(s); + }); // assign non-negative values to para.input.out_chg std::copy(out_chg.begin(), out_chg.end(), para.input.out_chg.begin()); }; item.reset_value = [](const Input_Item& item, Parameter& para) { - para.input.out_chg[0] = (para.input.calculation == "get_wf" || para.input.calculation == "get_pchg") ? 1 : para.input.out_chg[0]; + para.input.out_chg[0] = (para.input.calculation == "get_wf" || para.input.calculation == "get_pchg") + ? 1 + : para.input.out_chg[0]; }; sync_intvec(input.out_chg, 2, 0); this->add_item(item); @@ -472,7 +476,7 @@ void ReadInput::item_output() } { Input_Item item("bands_to_print"); - item.annotation = "specify the bands to be calculated in get_wf and get_pchg calculation"; + item.annotation = "specify the bands to be calculated for the partial (band-decomposed) charge densities"; item.read_value = [](const Input_Item& item, Parameter& para) { parse_expression(item.str_values, para.input.bands_to_print); }; @@ -485,6 +489,51 @@ void ReadInput::item_output() add_intvec_bcast(input.bands_to_print, para.input.bands_to_print.size(), 0); this->add_item(item); } + { + Input_Item item("out_pchg"); + item.annotation = "specify the bands to be calculated for the partial (band-decomposed) charge densities"; + item.read_value = [](const Input_Item& item, Parameter& para) { + parse_expression(item.str_values, para.input.out_pchg); + }; + item.get_final_value = [](Input_Item& item, const Parameter& para) { + if (item.is_read()) + { + item.final_value.str(longstring(item.str_values)); + } + }; + add_intvec_bcast(input.out_pchg, para.input.out_pchg.size(), 0); + this->add_item(item); + } + { + Input_Item item("out_wfc_norm"); + item.annotation = "specify the bands to be calculated for the norm of wavefunctions"; + item.read_value = [](const Input_Item& item, Parameter& para) { + parse_expression(item.str_values, para.input.out_wfc_norm); + }; + item.get_final_value = [](Input_Item& item, const Parameter& para) { + if (item.is_read()) + { + item.final_value.str(longstring(item.str_values)); + } + }; + add_intvec_bcast(input.out_wfc_norm, para.input.out_wfc_norm.size(), 0); + this->add_item(item); + } + { + Input_Item item("out_wfc_re_im"); + item.annotation = "specify the bands to be calculated for the real and imaginary parts of wavefunctions"; + item.read_value = [](const Input_Item& item, Parameter& para) { + parse_expression(item.str_values, para.input.out_wfc_re_im); + }; + item.get_final_value = [](Input_Item& item, const Parameter& para) { + if (item.is_read()) + { + item.final_value.str(longstring(item.str_values)); + } + }; + add_intvec_bcast(input.out_wfc_re_im, para.input.out_wfc_re_im.size(), 0); + this->add_item(item); + } { Input_Item item("if_separate_k"); item.annotation = "specify whether to write the partial charge densities for all k-points to individual files " diff --git a/source/module_io/test/read_input_ptest.cpp b/source/module_io/test/read_input_ptest.cpp index 23efe87072..7953002c3a 100644 --- a/source/module_io/test/read_input_ptest.cpp +++ b/source/module_io/test/read_input_ptest.cpp @@ -51,6 +51,9 @@ TEST_F(InputParaTest, ParaRead) EXPECT_EQ(param.inp.nbands_sto, 256); EXPECT_EQ(param.inp.nbands_istate, 5); EXPECT_EQ(param.inp.bands_to_print.size(), 0); + EXPECT_EQ(param.inp.out_pchg.size(), 0); + EXPECT_EQ(param.inp.out_wfc_norm.size(), 0); + EXPECT_EQ(param.inp.out_wfc_re_im.size(), 0); EXPECT_FALSE(param.inp.if_separate_k); EXPECT_EQ(param.inp.pw_seed, 1); EXPECT_EQ(param.inp.emin_sto, 0.0); diff --git a/source/module_io/write_wfc_r.cpp b/source/module_io/write_wfc_r.cpp index f8a15edb4b..b900f8f0be 100644 --- a/source/module_io/write_wfc_r.cpp +++ b/source/module_io/write_wfc_r.cpp @@ -4,21 +4,25 @@ // DATE : 2021-11-21 //====================== +//====================== +// WARNING: These interfaces will be removed in the future! Do not use them! +// Taoni add 2024-10-08 +//====================== + +#include "module_base/timer.h" +#include "module_base/tool_title.h" +#include "module_hamilt_pw/hamilt_pwdft/global.h" #include "write_wfc_r.h" #include #include #include -#include "module_base/timer.h" -#include "module_base/tool_title.h" -#include "module_hamilt_pw/hamilt_pwdft/global.h" - namespace ModuleIO { - // write ||wfc_r|| for all k-points and all bands - // Input: wfc_g(ik, ib, ig) - // loop order is for(z){for(y){for(x)}} +// write ||wfc_r|| for all k-points and all bands +// Input: wfc_g(ik, ib, ig) +// loop order is for(z){for(y){for(x)}} void write_psi_r_1(const psi::Psi>& wfc_g, const ModulePW::PW_Basis_K* wfcpw, const std::string& folder_name, @@ -31,157 +35,174 @@ void write_psi_r_1(const psi::Psi>& wfc_g, const std::string outdir = PARAM.globalv.global_out_dir + folder_name + "/"; ModuleBase::GlobalFunc::MAKE_DIR(outdir); #ifdef __MPI - std::vector mpi_requests; + std::vector mpi_requests; #endif - for(int ik=0; ik> wfc_r = cal_wfc_r(wfcpw, wfc_g, ik, ib); - - std::vector wfc_r2(wfc_r.size()); - std::vector wfc_i2; - if (square) - for (int ir = 0; ir < wfc_r2.size(); ++ir) - wfc_r2[ir] = std::norm(wfc_r[ir]); // "std::norm(z)" returns |z|^2 - else + for (int ik = 0; ik < wfc_g.get_nk(); ++ik) + { + wfc_g.fix_k(ik); + const int ik_out = (PARAM.inp.nspin != 2) + ? ik + GlobalC::Pkpoints.startk_pool[GlobalV::MY_POOL] + : ik - kv.get_nks() / 2 * kv.isk[ik] + kv.get_nkstot() / 2 * kv.isk[ik] + + GlobalC::Pkpoints.startk_pool[GlobalV::MY_POOL]; + for (int ib = 0; ib < wfc_g.get_nbands(); ++ib) + { + const std::vector> wfc_r = cal_wfc_r(wfcpw, wfc_g, ik, ib); + + std::vector wfc_real(wfc_r.size()); + std::vector wfc_imag; + if (square) + { + for (int ir = 0; ir < wfc_real.size(); ++ir) { - wfc_i2.resize(wfc_r.size()); - for (int ir = 0; ir < wfc_r2.size(); ++ir) - { - wfc_r2[ir] = wfc_r[ir].real(); - wfc_i2[ir] = wfc_r[ir].imag(); - } + wfc_real[ir] = std::norm(wfc_r[ir]); // "std::norm(z)" returns |z|^2 } - const std::string file_name = outdir + "wfc_realspace_" - + ModuleBase::GlobalFunc::TO_STRING(ik_out) - + "_" + ModuleBase::GlobalFunc::TO_STRING(ib); + } + else + { + wfc_imag.resize(wfc_r.size()); + for (int ir = 0; ir < wfc_real.size(); ++ir) + { + wfc_real[ir] = wfc_r[ir].real(); + wfc_imag[ir] = wfc_r[ir].imag(); + } + } + const std::string file_name = outdir + "wfc_realspace_" + ModuleBase::GlobalFunc::TO_STRING(ik_out) + "_" + + ModuleBase::GlobalFunc::TO_STRING(ib); #ifdef __MPI - mpi_requests.push_back({}); - write_chg_r_1(wfcpw, wfc_r2, file_name, mpi_requests.back()); - if (!square) - write_chg_r_1(wfcpw, wfc_i2, file_name + "_imag", mpi_requests.back()); + // Use write_chg_r_1 to output the real and imaginary parts of the wave function to file + mpi_requests.push_back({}); + write_chg_r_1(wfcpw, wfc_real, file_name, mpi_requests.back()); + if (!square) + { + write_chg_r_1(wfcpw, wfc_imag, file_name + "_imag", mpi_requests.back()); + } #else - write_chg_r_1(wfcpw, wfc_r2, file_name); - //if (!square) - //write_chg_r_1(wfc_i2, file_name + "_imag", mpi_requests.back()); + write_chg_r_1(wfcpw, wfc_real, file_name); + // if (!square) + // write_chg_r_1(wfc_imag, file_name + "_imag", mpi_requests.back()); #endif - } - } + } + } #ifdef __MPI - MPI_Waitall( mpi_requests.size(), mpi_requests.data(), MPI_STATUSES_IGNORE ); + MPI_Waitall(mpi_requests.size(), mpi_requests.data(), MPI_STATUSES_IGNORE); #endif - ModuleBase::timer::tick("ModuleIO", "write_psi_r_1"); - } - // processes output pipeline: - // - // t0 t1 t2 t3 t4 t5 t6 t7 - // --------------------------------> - // rank0 k0 k1 k2 k3 k4 k5 - // \ \ \ \ \ \ + ModuleBase::timer::tick("ModuleIO", "write_psi_r_1"); +} +// processes output pipeline: +// +// t0 t1 t2 t3 t4 t5 t6 t7 +// --------------------------------> +// rank0 k0 k1 k2 k3 k4 k5 +// \ \ \ \ \ \ // rank1 k0 k1 k2 k3 k4 k5 - // \ \ \ \ \ \ +// \ \ \ \ \ \ // rank2 k0 k1 k2 k3 k4 k5 - // Input: wfc_g(ib,ig) - // Output: wfc_r[ir] - std::vector> cal_wfc_r(const ModulePW::PW_Basis_K* wfcpw, - const psi::Psi>& wfc_g, - const int ik, - const int ib) - { - ModuleBase::timer::tick("ModuleIO", "cal_wfc_r"); +// Input: wfc_g(ib,ig) +// Output: wfc_r[ir] +std::vector> cal_wfc_r(const ModulePW::PW_Basis_K* wfcpw, + const psi::Psi>& wfc_g, + const int ik, + const int ib) +{ + ModuleBase::timer::tick("ModuleIO", "cal_wfc_r"); - std::vector> wfc_r(wfcpw->nrxx); - wfcpw->recip2real(&wfc_g(ib, 0), wfc_r.data(), ik); + std::vector> wfc_r(wfcpw->nrxx); + wfcpw->recip2real(&wfc_g(ib, 0), wfc_r.data(), ik); - ModuleBase::timer::tick("ModuleIO", "cal_wfc_r"); - return wfc_r; - } + ModuleBase::timer::tick("ModuleIO", "cal_wfc_r"); + return wfc_r; +} - // Input: chg_r[ir] +// Input: chg_r[ir] #ifdef __MPI - void write_chg_r_1(const ModulePW::PW_Basis_K* wfcpw, - const std::vector& chg_r, - const std::string& file_name, - MPI_Request& mpi_request) +void write_chg_r_1(const ModulePW::PW_Basis_K* wfcpw, + const std::vector& chg_r, + const std::string& file_name, + MPI_Request& mpi_request) #else - void write_chg_r_1(const ModulePW::PW_Basis_K* wfcpw, - const std::vector& chg_r, - const std::string& file_name) +void write_chg_r_1(const ModulePW::PW_Basis_K* wfcpw, const std::vector& chg_r, const std::string& file_name) #endif - { - ModuleBase::timer::tick("ModuleIO", "write_chg_r_1"); - std::ofstream ofs; - -#ifdef __MPI - constexpr int mpi_tag=100; - if(GlobalV::RANK_IN_POOL==0) - { +{ + ModuleBase::timer::tick("ModuleIO", "write_chg_r_1"); + std::ofstream ofs; + +#ifdef __MPI + constexpr int mpi_tag = 100; + if (GlobalV::RANK_IN_POOL == 0) + { #endif - ofs.open(file_name); - - ofs<<"calculated by ABACUS"<nx<<" "<ny<<" "<nz<nx << " " << wfcpw->ny << " " << wfcpw->nz << std::endl; +#ifdef __MPI + } + else + { + char recv_tmp; + MPI_Recv(&recv_tmp, 1, MPI_CHAR, GlobalV::RANK_IN_POOL - 1, mpi_tag, POOL_WORLD, MPI_STATUS_IGNORE); + + ofs.open(file_name, std::ofstream::app); + } #endif - assert(wfcpw->nx * wfcpw->ny * wfcpw->nplane == chg_r.size()); - for(int iz=0; iznplane; ++iz) - { - for(int iy=0; iyny; ++iy) - { - for(int ix=0; ixnx; ++ix) - { - const int ir = (ix*wfcpw->ny+iy)*wfcpw->nplane+iz; - ofs<nx * wfcpw->ny * wfcpw->nplane == chg_r.size()); + for (int iz = 0; iz < wfcpw->nplane; ++iz) + { + for (int iy = 0; iy < wfcpw->ny; ++iy) + { + for (int ix = 0; ix < wfcpw->nx; ++ix) + { + const int ir = (ix * wfcpw->ny + iy) * wfcpw->nplane + iz; + ofs << chg_r[ir] << " "; + } + ofs << "\n"; + } + } + ofs.close(); + +#ifdef __MPI + if (GlobalV::RANK_IN_POOL < GlobalV::NPROC_IN_POOL - 1) + { + const char send_tmp = 'c'; + MPI_Isend(&send_tmp, 1, MPI_CHAR, GlobalV::RANK_IN_POOL + 1, mpi_tag, POOL_WORLD, &mpi_request); + } + else + { + mpi_request = MPI_REQUEST_NULL; + } #endif - ModuleBase::timer::tick("ModuleIO", "write_chg_r_1"); - } -}; + ModuleBase::timer::tick("ModuleIO", "write_chg_r_1"); +} +}; // namespace ModuleIO diff --git a/source/module_parameter/input_parameter.h b/source/module_parameter/input_parameter.h index f5abc047d9..f870dd7ba0 100644 --- a/source/module_parameter/input_parameter.h +++ b/source/module_parameter/input_parameter.h @@ -57,8 +57,8 @@ struct Input_para std::string stru_file = "STRU"; ///< file contains atomic positions -- ///< xiaohui modify 2015-02-01 std::string kpoint_file = "KPT"; ///< file contains k-points -- xiaohui modify 2015-02-01 - std::string pseudo_dir = ""; ///< directory of pseudopotential - std::string orbital_dir = ""; ///< directory of orbital file + std::string pseudo_dir = ""; ///< directory of pseudopotential + std::string orbital_dir = ""; ///< directory of orbital file std::string read_file_dir = "auto"; ///< directory of files for reading bool restart_load = false; std::string wannier_card = "none"; ///< input card for wannier functions. @@ -110,13 +110,13 @@ struct Input_para bool mixing_dftu = false; ///< whether to mix locale in DFT+U bool mixing_dmr = false; ///< whether to mix real space density matrix - bool gamma_only = false; ///< for plane wave. - int scf_nmax = 100; ///< number of max elec iter - double scf_thr = -1.0; ///< \sum |rhog_out - rhog_in |^2 + bool gamma_only = false; ///< for plane wave. + int scf_nmax = 100; ///< number of max elec iter + double scf_thr = -1.0; ///< \sum |rhog_out - rhog_in |^2 double scf_ene_thr = -1.0; ///< energy threshold for scf convergence, in eV - int scf_thr_type = -1; ///< type of the criterion of scf_thr, 1: reci drho, 2: real drho - bool final_scf= false; ///< whether to do final scf - + int scf_thr_type = -1; ///< type of the criterion of scf_thr, 1: reci drho, 2: real drho + bool final_scf = false; ///< whether to do final scf + bool lspinorb = false; ///< consider the spin-orbit interaction bool noncolin = false; ///< using non-collinear-spin double soc_lambda = 1.0; ///< The fraction of averaged SOC pseudopotential @@ -331,11 +331,11 @@ struct Input_para bool out_mat_hs2 = false; ///< LiuXh add 2019-07-16, output H(R) matrix and ///< S(R) matrix in local basis. bool out_mat_dh = false; - bool out_mat_xc = false; ///< output exchange-correlation matrix in - ///< KS-orbital representation. - bool out_eband_terms = false; ///< output the band energy terms separately - bool out_hr_npz = false; ///< output exchange-correlation matrix in - ///< KS-orbital representation. + bool out_mat_xc = false; ///< output exchange-correlation matrix in + ///< KS-orbital representation. + bool out_eband_terms = false; ///< output the band energy terms separately + bool out_hr_npz = false; ///< output exchange-correlation matrix in + ///< KS-orbital representation. bool out_dm_npz = false; int out_interval = 1; @@ -354,7 +354,10 @@ struct Input_para bool restart_save = false; ///< restart //Peize Lin add 2020-04-04 bool rpa = false; ///< rpa calculation int nbands_istate = 5; ///< number of bands around fermi level for get_pchg calculation. - std::vector bands_to_print = {}; ///< specify the bands to be calculated in the get_pchg + std::vector bands_to_print = {}; ///< specify the bands to be calculated for partial charge + std::vector out_pchg = {}; ///< specify the bands to be calculated for partial charge + std::vector out_wfc_norm = {}; ///< specify the bands to be calculated for norm of wfc + std::vector out_wfc_re_im = {}; ///< specify the bands to be calculated for real and imaginary parts of wfc bool if_separate_k = false; ///< whether to write partial charge for all k-points to individual files or merge them std::vector out_elf = {0, 3}; ///< output the electron localization function (ELF). 0: no; 1: yes diff --git a/tests/integrate/312_NO_GO_wfc_ienvelope/result.ref b/tests/integrate/312_NO_GO_wfc_ienvelope/result.ref index c516128531..2daac8dfa0 100644 --- a/tests/integrate/312_NO_GO_wfc_ienvelope/result.ref +++ b/tests/integrate/312_NO_GO_wfc_ienvelope/result.ref @@ -1,5 +1,5 @@ -BAND1_s_1_ENV.cube 1.06415 -BAND2_s_1_ENV.cube 1.07189 -BAND3_s_1_ENV.cube 1.06892 -BAND4_s_1_ENV.cube 1.06874 +BAND1_GAMMA_SPIN1_ENV.cube 1.06417 +BAND2_GAMMA_SPIN1_ENV.cube 1.0719 +BAND3_GAMMA_SPIN1_ENV.cube 1.06891 +BAND4_GAMMA_SPIN1_ENV.cube 1.06869 totaltimeref 0.16356