Skip to content

Commit aee5691

Browse files
AsTonyshmentAsTonyshment
authored
Feature: Support outputting partial charge and wave functions under PW basis when kpar > 1 (#6242)
* Fix if_separate_k description in the doc * Use nkstot instead of nks when kpar > 1 * Fix doc: set_wf -> get_wf * Fix kpar nks -> nkstot * Fix write_cube and pgrid reduce logic * Delete cout and nks -> nkstot for out_wfc_re_im * Add parameter reduce_all_pool to control reduce behavior * Delete useless parameters in get_pchg_pw * Fix incorrect pchg when if_separate_k is false due to lack of reduce across diff pools * Add integrate test 113_PW_get_pchg_kpar * Add integrate test 114_PW_get_pchg_sepk * Add integrate test 115_PW_get_wf_kpar * Add MPI compilation macros * Update 115_PW_get_wf_kpar to avoid random behavior in CI/CD test * Change param of test 115_PW_get_wf_kpar
1 parent 98e51b4 commit aee5691

File tree

26 files changed

+1144
-867
lines changed

26 files changed

+1144
-867
lines changed

‎docs/advanced/input_files/input-main.md‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -599,7 +599,7 @@ These variables are used to control general system parameters.
599599
- nao: from numerical atomic orbitals. If they are not enough, other wave functions are initialized with random numbers.
600600
- nao+random: add small random numbers on numerical atomic orbitals
601601

602-
> Only the `file` option is useful for the lcao basis set, which is mostly used when [calculation](#calculation) is set to `set_wf` and `get_pchg`. See more details in [out_wfc_lcao](#out_wfc_lcao).
602+
> Only the `file` option is useful for the lcao basis set, which is mostly used when [calculation](#calculation) is set to `get_wf` and `get_pchg`. See more details in [out_wfc_lcao](#out_wfc_lcao).
603603
- **Default**: atomic
604604

605605
### init_chg
@@ -1921,7 +1921,7 @@ The band (KS orbital) energy for each (k-point, spin, band) will be printed in t
19211921
### if_separate_k
19221922

19231923
- **Type**: Boolean
1924-
- **Availability**: Only for LCAO, used only when `calculation = get_pchg` and `gamma_only` is turned off.
1924+
- **Availability**: For both PW and LCAO. When `basis_type = pw`, used if `out_pchg` is set. When `basis_type = lcao`, used only when `calculation = get_pchg` and `gamma_only` is turned off.
19251925
- **Description**: Specifies whether to write the partial charge densities for all k-points to individual files or merge them. **Warning**: Enabling symmetry may produce incorrect results due to incorrect k-point weights. Therefore, when calculating partial charge densities, it is strongly recommended to set `symmetry = -1`.
19261926
- **Default**: false
19271927

‎source/module_elecstate/module_charge/symmetry_rho.cpp‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -97,7 +97,7 @@ void Symmetry_rho::psymm(double* rho_part,
9797
rhotot.resize(rho_basis->nxyz);
9898
ModuleBase::GlobalFunc::ZEROS(rhotot.data(), rho_basis->nxyz);
9999
}
100-
Pgrid.reduce(rhotot.data(), rho_part);
100+
Pgrid.reduce(rhotot.data(), rho_part, false);
101101

102102
// (2)
103103
if (GlobalV::MY_RANK == 0)

‎source/module_esolver/esolver_ks_pw.cpp‎

Lines changed: 44 additions & 30 deletions
Original file line numberDiff line numberDiff line change
@@ -645,14 +645,24 @@ void ESolver_KS_PW<T, Device>::iter_finish(UnitCell& ucell, const int istep, int
645645
//----------------------------------------------------------
646646
// 3) Print out electronic wavefunctions in pw basis
647647
//----------------------------------------------------------
648-
if (iter % PARAM.inp.out_freq_elec == 0 || iter == PARAM.inp.scf_nmax || conv_esolver)
649-
{
650-
ModuleIO::write_wfc_pw(GlobalV::KPAR, GlobalV::MY_POOL, GlobalV::MY_RANK,
651-
PARAM.inp.nbands, PARAM.inp.nspin, PARAM.globalv.npol,
652-
GlobalV::RANK_IN_POOL, GlobalV::NPROC_IN_POOL,
653-
PARAM.inp.out_wfc_pw, PARAM.inp.ecutwfc, PARAM.globalv.global_out_dir,
654-
this->psi[0], this->kv, this->pw_wfc, GlobalV::ofs_running);
655-
}
648+
if (iter % PARAM.inp.out_freq_elec == 0 || iter == PARAM.inp.scf_nmax || conv_esolver)
649+
{
650+
ModuleIO::write_wfc_pw(GlobalV::KPAR,
651+
GlobalV::MY_POOL,
652+
GlobalV::MY_RANK,
653+
PARAM.inp.nbands,
654+
PARAM.inp.nspin,
655+
PARAM.globalv.npol,
656+
GlobalV::RANK_IN_POOL,
657+
GlobalV::NPROC_IN_POOL,
658+
PARAM.inp.out_wfc_pw,
659+
PARAM.inp.ecutwfc,
660+
PARAM.globalv.global_out_dir,
661+
this->psi[0],
662+
this->kv,
663+
this->pw_wfc,
664+
GlobalV::ofs_running);
665+
}
656666

657667
//----------------------------------------------------------
658668
// 4) check if oscillate for delta_spin method
@@ -718,11 +728,6 @@ void ESolver_KS_PW<T, Device>::after_scf(UnitCell& ucell, const int istep, const
718728
this->pw_rhod->ny,
719729
this->pw_rhod->nz,
720730
this->pw_rhod->nxyz,
721-
this->kv.get_nks(),
722-
this->kv.isk,
723-
this->kv.wk,
724-
this->pw_big->bz,
725-
this->pw_big->nbz,
726731
this->chr.ngmc,
727732
&ucell,
728733
this->psi,
@@ -731,16 +736,29 @@ void ESolver_KS_PW<T, Device>::after_scf(UnitCell& ucell, const int istep, const
731736
this->ctx,
732737
this->Pgrid,
733738
PARAM.globalv.global_out_dir,
734-
PARAM.inp.if_separate_k);
739+
PARAM.inp.if_separate_k,
740+
this->kv,
741+
GlobalV::KPAR,
742+
GlobalV::MY_POOL,
743+
&this->chr);
735744
}
736745

737-
738-
// tmp 2025-05-17, mohan note
739-
ModuleIO::write_wfc_pw(GlobalV::KPAR, GlobalV::MY_POOL, GlobalV::MY_RANK,
740-
PARAM.inp.nbands, PARAM.inp.nspin, PARAM.globalv.npol,
741-
GlobalV::RANK_IN_POOL, GlobalV::NPROC_IN_POOL,
742-
PARAM.inp.out_wfc_pw, PARAM.inp.ecutwfc, PARAM.globalv.global_out_dir,
743-
this->psi[0], this->kv, this->pw_wfc, GlobalV::ofs_running);
746+
// tmp 2025-05-17, mohan note
747+
ModuleIO::write_wfc_pw(GlobalV::KPAR,
748+
GlobalV::MY_POOL,
749+
GlobalV::MY_RANK,
750+
PARAM.inp.nbands,
751+
PARAM.inp.nspin,
752+
PARAM.globalv.npol,
753+
GlobalV::RANK_IN_POOL,
754+
GlobalV::NPROC_IN_POOL,
755+
PARAM.inp.out_wfc_pw,
756+
PARAM.inp.ecutwfc,
757+
PARAM.globalv.global_out_dir,
758+
this->psi[0],
759+
this->kv,
760+
this->pw_wfc,
761+
GlobalV::ofs_running);
744762

745763
//------------------------------------------------------------------
746764
//! 5) calculate Wannier functions in pw basis
@@ -931,7 +949,7 @@ void ESolver_KS_PW<T, Device>::after_all_runners(UnitCell& ucell)
931949
//----------------------------------------------------------
932950

933951
//----------------------------------------------------------
934-
//! The write_psi_r_1 interface will be removed in the very
952+
//! The write_psi_r_1 interface will be removed in the very
935953
//! near future. Don't use it!
936954
//----------------------------------------------------------
937955
// if (PARAM.inp.out_wfc_r == 1) // Peize Lin add 2021.11.21
@@ -951,19 +969,15 @@ void ESolver_KS_PW<T, Device>::after_all_runners(UnitCell& ucell)
951969
this->pw_rhod->ny,
952970
this->pw_rhod->nz,
953971
this->pw_rhod->nxyz,
954-
this->kv.get_nks(),
955-
this->kv.isk,
956-
this->kv.wk,
957-
this->pw_big->bz,
958-
this->pw_big->nbz,
959-
this->chr.ngmc,
960972
&ucell,
961973
this->psi,
962-
this->pw_rhod,
963974
this->pw_wfc,
964975
this->ctx,
965976
this->Pgrid,
966-
PARAM.globalv.global_out_dir);
977+
PARAM.globalv.global_out_dir,
978+
this->kv,
979+
GlobalV::KPAR,
980+
GlobalV::MY_POOL);
967981
}
968982

969983
//----------------------------------------------------------

0 commit comments

Comments
 (0)