Skip to content

Commit d236416

Browse files
dyzhengclaude
dyzheng
and
claude
committed
feat: extend onsite_projector to support nspin=1/2
- Add isk_in parameter to cal_occupations() for spin channel identification - Add npwx parameter to overlap_proj_psi() and cal_becp() for flexible leading dimension - Add npol=1 branch in occupation calculation for nspin=1/2 systems - Update caller in ctrl_output_pw.cpp to pass isk array Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
1 parent c039e92 commit d236416

5 files changed

Lines changed: 43 additions & 15 deletions

File tree

‎source/source_io/module_ctrl/ctrl_output_pw.cpp‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -237,7 +237,8 @@ void ModuleIO::ctrl_scf_pw(const int istep,
237237
{ // float type has not been implemented
238238
auto* onsite_p = projectors::OnsiteProjector<double, Device>::get_instance();
239239
onsite_p->cal_occupations(reinterpret_cast<psi::Psi<std::complex<double>, Device>*>(stp.psi_t),
240-
pelec->wg);
240+
pelec->wg,
241+
kv.isk.data());
241242
}
242243

243244
ModuleBase::timer::tick("ModuleIO", "ctrl_scf_pw");

‎source/source_pw/module_pwdft/onsite_proj_tools.cpp‎

Lines changed: 4 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -278,12 +278,14 @@ template <typename FPTYPE, typename Device>
278278
void Onsite_Proj_tools<FPTYPE, Device>::cal_becp(int ik,
279279
int npm,
280280
std::complex<FPTYPE>* becp_in,
281-
const std::complex<FPTYPE>* ppsi_in)
281+
const std::complex<FPTYPE>* ppsi_in,
282+
int npwx)
282283
{
283284
ModuleBase::TITLE("Onsite_Proj_tools", "cal_becp");
284285
ModuleBase::timer::tick("Onsite_Proj_tools", "cal_becp");
285286

286287
const int npol = this->ucell_->get_npol();
288+
if(npwx == 0) npwx = this->wfc_basis_->npwk_max;
287289
const std::complex<FPTYPE>* ppsi = ppsi_in == nullptr ? &(this->psi_[0](ik, 0, 0)) : ppsi_in;
288290
const int npw = this->wfc_basis_->npwk[ik];
289291
if (becp_in == nullptr && this->becp == nullptr)
@@ -434,7 +436,7 @@ void Onsite_Proj_tools<FPTYPE, Device>::cal_becp(int ik,
434436
this->ppcell_vkb,
435437
npw,
436438
ppsi,
437-
this->max_npw,
439+
npwx,
438440
&ModuleBase::ZERO,
439441
becp_tmp,
440442
this->nkb);

‎source/source_pw/module_pwdft/onsite_proj_tools.h‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -58,7 +58,7 @@ class Onsite_Proj_tools
5858
/**
5959
* @brief calculate the becp = <psi|beta> for all beta functions
6060
*/
61-
void cal_becp(int ik, int npm, std::complex<FPTYPE>* becp_in = nullptr, const std::complex<FPTYPE>* ppsi_in = nullptr);
61+
void cal_becp(int ik, int npm, std::complex<FPTYPE>* becp_in = nullptr, const std::complex<FPTYPE>* ppsi_in = nullptr, int npwx = 0);
6262
/**
6363
* @brief calculate the dbecp_{ij} = <psi|\partial beta/\partial varepsilon_{ij}> for all beta functions
6464
* stress_{ij} = -1/omega \sum_{n,k}f_{nk} \sum_I \sum_{lm,l'm'}D_{l,l'}^{I} becp * dbecp_{ij} also calculated

‎source/source_pw/module_pwdft/onsite_projector.cpp‎

Lines changed: 33 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -337,9 +337,11 @@ void projectors::OnsiteProjector<T, Device>::tabulate_atomic(const int ik, const
337337
}
338338

339339
template<typename T, typename Device>
340-
void projectors::OnsiteProjector<T, Device>::overlap_proj_psi(
340+
void projectors::OnsiteProjector<T, Device>::overlap_proj_psi(
341341
const int npm,
342-
const std::complex<double>* ppsi)
342+
const std::complex<double>* ppsi,
343+
int npwx
344+
)
343345
{
344346
ModuleBase::timer::tick("OnsiteProj", "overlap");
345347
// STAGE 3 - cal_becp
@@ -384,6 +386,7 @@ void projectors::OnsiteProjector<T, Device>::overlap_proj_psi(
384386
// std::cout << "at " << __FILE__ << ": " << __LINE__ << " output npm: " << npm << std::endl;
385387
// std::cout << "at " << __FILE__ << ": " << __LINE__ << " ik_: " << ik_ << std::endl;
386388
int npol = this->ucell->get_npol();
389+
if(npwx == 0) npwx = this->npwx_;
387390
if(this->becp == nullptr || this->size_becp < npm*this->tot_nproj)
388391
{
389392
this->size_becp = npm*this->tot_nproj;
@@ -397,7 +400,7 @@ void projectors::OnsiteProjector<T, Device>::overlap_proj_psi(
397400
this->h_becp = this->becp;
398401
}
399402
}
400-
this->fs_tools->cal_becp(ik_, npm/npol, this->becp, ppsi); // in cal_becp, npm should be the one not multiplied by npol
403+
this->fs_tools->cal_becp(ik_, npm/npol, this->becp, ppsi, npwx); // in cal_becp, npm should be the one not multiplied by npol
401404
if(this->device == base_device::GpuDevice)
402405
{
403406
syncmem_complex_d2h_op()(h_becp, this->becp, this->size_becp);
@@ -522,8 +525,9 @@ void projectors::OnsiteProjector<T, Device>::read_abacus_orb(std::ifstream& ifs,
522525

523526
template<typename T, typename Device>
524527
void projectors::OnsiteProjector<T, Device>::cal_occupations(
525-
const psi::Psi<std::complex<T>, Device>* psi_in,
526-
const ModuleBase::matrix& wg_in)
528+
const psi::Psi<std::complex<T>, Device>* psi_in,
529+
const ModuleBase::matrix& wg_in,
530+
const int* isk_in)
527531
{
528532
ModuleBase::timer::tick("OnsiteProj", "cal_occupation");
529533
this->tabulate_atomic(0);
@@ -534,6 +538,7 @@ void projectors::OnsiteProjector<T, Device>::cal_occupations(
534538
for(int ik = 0; ik < psi_in->get_nk(); ik++)
535539
{
536540
psi_in->fix_k(ik);
541+
const int sign = isk_in[ik] == 0? 1: -1;
537542
if(ik != 0)
538543
{
539544
this->tabulate_atomic(ik);
@@ -556,6 +561,7 @@ void projectors::OnsiteProjector<T, Device>::cal_occupations(
556561
for(int iat = 0; iat < this->iat_nh.size(); iat++)
557562
{
558563
const int nh = this->get_nh(iat);
564+
if(this->ucell->get_npol() == 2)
559565
for(int ih = 0; ih < nh; ih++)
560566
{
561567
const int occ_index = (begin_ih + ih) * 4;
@@ -565,6 +571,16 @@ void projectors::OnsiteProjector<T, Device>::cal_occupations(
565571
occs[occ_index + 2] += weight * conj(becp_p[index + nkb]) * becp_p[index];
566572
occs[occ_index + 3] += weight * conj(becp_p[index + nkb]) * becp_p[index + nkb];
567573
}
574+
else if(this->ucell->get_npol() == 1)
575+
{
576+
for(int ih = 0; ih < nh; ih++)
577+
{
578+
const int occ_index = (begin_ih + ih) * 4;
579+
const int index = ib*nkb + begin_ih + ih;
580+
occs[occ_index] += weight * conj(becp_p[index]) * becp_p[index];
581+
occs[occ_index + 3] += sign * weight * conj(becp_p[index]) * becp_p[index];
582+
}
583+
}
568584
begin_ih += nh;
569585
}
570586
}
@@ -608,10 +624,18 @@ void projectors::OnsiteProjector<T, Device>::cal_occupations(
608624
std::vector<double> charge_mag(4, 0.0);
609625
for(int ih=0;ih<this->iat_nh[iat];ih++)
610626
{
611-
charge_mag[3] += (occs[occ_index] - occs[occ_index + 3]).real();
612-
charge_mag[1] += (occs[occ_index + 1] + occs[occ_index + 2]).real();
613-
charge_mag[2] += (occs[occ_index + 1] - occs[occ_index + 2]).imag();
614-
charge_mag[0] += (occs[occ_index] + occs[occ_index + 3]).real();
627+
if(this->ucell->get_npol() == 2)
628+
{
629+
charge_mag[3] += (occs[occ_index] - occs[occ_index + 3]).real();
630+
charge_mag[1] += (occs[occ_index + 1] + occs[occ_index + 2]).real();
631+
charge_mag[2] += (occs[occ_index + 1] - occs[occ_index + 2]).imag();
632+
charge_mag[0] += (occs[occ_index] + occs[occ_index + 3]).real();
633+
}
634+
else if (this->ucell->get_npol() == 1)
635+
{
636+
charge_mag[0] += occs[occ_index].real();
637+
charge_mag[3] += occs[occ_index + 3].real();
638+
}
615639
if(ih == current_l * current_l - 1)
616640
{
617641
sum[0] += charge_mag[0];

‎source/source_pw/module_pwdft/onsite_projector.h‎

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -45,7 +45,8 @@ namespace projectors
4545

4646
void overlap_proj_psi(
4747
const int npm,
48-
const std::complex<double>* ppsi
48+
const std::complex<double>* ppsi,
49+
int npwx = 0
4950
);
5051
void read_abacus_orb(std::ifstream& ifs,
5152
std::string& elem,
@@ -70,7 +71,7 @@ namespace projectors
7071
const ModuleBase::matrix& ekb);
7172

7273
/// @brief calculate and print the occupations of all lm orbitals
73-
void cal_occupations(const psi::Psi<std::complex<T>, Device>* psi, const ModuleBase::matrix& wg_in);
74+
void cal_occupations(const psi::Psi<std::complex<T>, Device>* psi, const ModuleBase::matrix& wg_in, const int* isk_in);
7475

7576
int get_size_becp() const { return size_becp; }
7677
std::complex<double>* get_becp() const { return becp; }

0 commit comments

Comments
 (0)