Skip to content

Commit

Permalink
Fix: init_chg wfc support npsin = 4 now (#5166)
Browse files Browse the repository at this point in the history 
Fix write_wfc_pw and read_wfc_pw when nspin = 4
  • Loading branch information
@Qianruipku
Qianruipku authored Sep 25, 2024
1 parent fb5beba commit e32341a
Show file tree
Hide file tree
Showing 11 changed files with 371 additions and 167 deletions.
18 changes: 9 additions & 9 deletions source/module_io/read_wfc_pw.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -82,7 +82,7 @@ void ModuleIO::read_wfc_pw(const std::string& filename,
ikstot = ik;
#endif

npwtot *= PARAM.globalv.npol;
int npwtot_npol = npwtot * PARAM.globalv.npol;



Expand Down Expand Up @@ -178,7 +178,7 @@ void ModuleIO::read_wfc_pw(const std::string& filename,
}

// read in miller index
ModuleBase::Vector3<int>* miller = new ModuleBase::Vector3<int>[npwtot_in];
ModuleBase::Vector3<int>* miller = new ModuleBase::Vector3<int>[npwtot];
int* glo_order = nullptr;
if (GlobalV::RANK_IN_POOL == 0)
{
Expand All @@ -188,7 +188,7 @@ void ModuleIO::read_wfc_pw(const std::string& filename,
else if (filetype == "dat")
{
rfs >> size;
for (int i = 0; i < npwtot_in; ++i)
for (int i = 0; i < npwtot; ++i)
{
rfs >> miller[i].x >> miller[i].y >> miller[i].z;
}
Expand All @@ -201,7 +201,7 @@ void ModuleIO::read_wfc_pw(const std::string& filename,
{
glo_order[i] = -1;
}
for (int i = 0; i < npwtot_in / PARAM.globalv.npol; ++i)
for (int i = 0; i < npwtot; ++i)
{
int index = (miller[i].x * ny + miller[i].y) * nz + miller[i].z;
glo_order[index] = i;
Expand All @@ -221,7 +221,7 @@ void ModuleIO::read_wfc_pw(const std::string& filename,
}

// read in wfc
std::complex<double>* wfc_in = new std::complex<double>[npwtot_in];
std::complex<double>* wfc_in = new std::complex<double>[npwtot_npol];
for (int ib = 0; ib < nbands_in; ib++)
{
if (GlobalV::RANK_IN_POOL == 0)
Expand All @@ -232,7 +232,7 @@ void ModuleIO::read_wfc_pw(const std::string& filename,
else if (filetype == "dat")
{
rfs >> size;
for (int i = 0; i < npwtot_in; ++i)
for (int i = 0; i < npwtot_npol; ++i)
{
rfs >> wfc_in[i];
}
Expand Down Expand Up @@ -285,7 +285,7 @@ void ModuleIO::read_wfc_pw(const std::string& filename,
{
for (int i = 0; i < size; i++)
{
wfc_ip[i] = wfc_in[glo_order[ig_ip[i]] + npwtot_in / 2];
wfc_ip[i] = wfc_in[glo_order[ig_ip[i]] + npwtot];
}
MPI_Send(wfc_ip, size, MPI_DOUBLE_COMPLEX, ip, ip + 2 * GlobalV::NPROC_IN_POOL, POOL_WORLD);
}
Expand All @@ -305,7 +305,7 @@ void ModuleIO::read_wfc_pw(const std::string& filename,
{
for (int i = 0; i < pw_wfc->npwk[ik]; ++i)
{
wfc(ib, i + npwk_max) = wfc_in[glo_order[l2g_pw[i]] + npwtot_in / 2];
wfc(ib, i + npwk_max) = wfc_in[glo_order[l2g_pw[i]] + npwtot];
}
}
}
Expand All @@ -321,7 +321,7 @@ void ModuleIO::read_wfc_pw(const std::string& filename,
{
for (int i = 0; i < pw_wfc->npwk[ik]; ++i)
{
wfc(ib, i + npwk_max) = wfc_in[glo_order[l2g_pw[i]] + npwtot_in / 2];
wfc(ib, i + npwk_max) = wfc_in[glo_order[l2g_pw[i]] + npwtot];
}
}
#endif
Expand Down
50 changes: 37 additions & 13 deletions source/module_io/read_wfc_to_rho.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -5,6 +5,7 @@
#include "module_hamilt_pw/hamilt_pwdft/global.h"
#include "module_elecstate/module_charge/symmetry_rho.h"
#include "module_parameter/parameter.h"
#include "module_elecstate/kernels/elecstate_op.h"

void ModuleIO::read_wfc_to_rho(const ModulePW::PW_Basis_K* pw_wfc,
ModuleSymmetry::Symmetry& symm,
Expand All @@ -21,14 +22,14 @@ void ModuleIO::read_wfc_to_rho(const ModulePW::PW_Basis_K* pw_wfc,
const int nbands = GlobalV::NBANDS;
const int nspin = PARAM.inp.nspin;

const int npwk_max = pw_wfc->npwk_max;
const int ng_npol = pw_wfc->npwk_max * PARAM.globalv.npol;
const int nrxx = pw_wfc->nrxx;
for (int is = 0; is < nspin; ++is)
{
ModuleBase::GlobalFunc::ZEROS(chg.rho[is], nrxx);
}

ModuleBase::ComplexMatrix wfc_tmp(nbands, npwk_max);
ModuleBase::ComplexMatrix wfc_tmp(nbands, ng_npol);
std::vector<std::complex<double>> rho_tmp(nrxx);

// read occupation numbers
Expand Down Expand Up @@ -78,21 +79,44 @@ void ModuleIO::read_wfc_to_rho(const ModulePW::PW_Basis_K* pw_wfc,
std::stringstream filename;
filename << PARAM.globalv.global_readin_dir << "WAVEFUNC" << ikstot + 1 << ".dat";
ModuleIO::read_wfc_pw(filename.str(), pw_wfc, ik, nkstot, wfc_tmp);
for (int ib = 0; ib < nbands; ++ib)
if (PARAM.inp.nspin == 4)
{
const std::complex<double>* wfc_ib = wfc_tmp.c + ib * npwk_max;
pw_wfc->recip2real(wfc_ib, rho_tmp.data(), ik);

const double w1 = wg_tmp(ikstot, ib) / pw_wfc->omega;
std::vector<std::complex<double>> rho_tmp2(nrxx);
for (int ib = 0; ib < nbands; ++ib)
{
const std::complex<double>* wfc_ib = wfc_tmp.c + ib * ng_npol;
const std::complex<double>* wfc_ib2 = wfc_tmp.c + ib * ng_npol + ng_npol / 2;
pw_wfc->recip2real(wfc_ib, rho_tmp.data(), ik);
pw_wfc->recip2real(wfc_ib2, rho_tmp2.data(), ik);
const double w1 = wg_tmp(ikstot, ib) / pw_wfc->omega;

if (w1 != 0.0)
if (w1 != 0.0)
{
base_device::DEVICE_CPU* ctx = nullptr;
elecstate::elecstate_pw_op<double, base_device::DEVICE_CPU>()(ctx,
PARAM.globalv.domag,
PARAM.globalv.domag_z,
nrxx,
w1,
chg.rho,
rho_tmp.data(),
rho_tmp2.data());
}
}
}
else
{
for (int ib = 0; ib < nbands; ++ib)
{
#ifdef _OPENMP
#pragma omp parallel for
#endif
for (int ir = 0; ir < nrxx; ir++)
const std::complex<double>* wfc_ib = wfc_tmp.c + ib * ng_npol;
pw_wfc->recip2real(wfc_ib, rho_tmp.data(), ik);

const double w1 = wg_tmp(ikstot, ib) / pw_wfc->omega;

if (w1 != 0.0)
{
chg.rho[is][ir] += w1 * std::norm(rho_tmp[ir]);
base_device::DEVICE_CPU* ctx = nullptr;
elecstate::elecstate_pw_op<double, base_device::DEVICE_CPU>()(ctx, is, nrxx, w1, chg.rho, rho_tmp.data());
}
}
}
Expand Down
Loading

0 comments on commit e32341a

Please sign in to comment.