Skip to content

2026PKUCourseHW5：使用 std::mt19937 替代 rand() 以提高浮点数精度 #7137

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
Tornado1016 wants to merge 1 commit into from
Closed

2026PKUCourseHW5：使用 std::mt19937 替代 rand() 以提高浮点数精度 #7137

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
27 changes: 20 additions & 7 deletions source/source_pw/module_stodft/sto_wf.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -7,6 +7,8 @@
#include <cassert>
#include <ctime>

#include <random>

#include "source_base/global_function.h"

template <typename T, typename Device>
Expand All @@ -19,7 +21,7 @@ Stochastic_WF<T, Device>::~Stochastic_WF()
{
delete chi0_cpu;
Device* ctx = {};
if (base_device::get_device_type(ctx) == base_device::GpuDevice)
if (base_device::get_device_type<Device>(ctx) == base_device::GpuDevice)
{
delete chi0;
}
Expand Down Expand Up @@ -60,18 +62,22 @@ void Stochastic_WF<T, Device>::clean_chiallorder()
delete[] chiallorder;
chiallorder = nullptr;
}

template <typename T, typename Device>
void Stochastic_WF<T, Device>::init_sto_orbitals(const int seed_in)
{
const unsigned int rank_seed_offset = 10000;
unsigned int final_seed;
if (seed_in == 0 || seed_in == -1)
{
srand(static_cast<unsigned int>(time(nullptr)) + GlobalV::MY_RANK * rank_seed_offset); // GlobalV global variables are reserved
final_seed = (unsigned)time(nullptr) + GlobalV::MY_RANK * 10000;
}
else
{
srand(static_cast<unsigned int>(std::abs(seed_in)) + (GlobalV::MY_BNDGROUP * GlobalV::NPROC_IN_BNDGROUP + GlobalV::RANK_IN_BPGROUP) * rank_seed_offset);
final_seed = (unsigned)std::abs(seed_in) + (GlobalV::MY_BNDGROUP * GlobalV::NPROC_IN_BNDGROUP + GlobalV::RANK_IN_BPGROUP) * 10000;
}
Comment on lines +69 to 77
Copy link

Copilot AI Mar 25, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The seed offset value 10000 is now duplicated in both seeding branches. This used to be a named constant, and keeping it as a single constexpr/const improves readability and reduces the chance of future inconsistencies if the offset changes.

Copilot uses AI. Check for mistakes.

// initialize the random number generator with the final seed
this->rng.seed(final_seed);

this->allocate_chi0();
this->update_sto_orbitals(seed_in);
Comment on lines +69 to 83
Copy link

Copilot AI Mar 25, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The PR description says all rand()/srand() usage was replaced, but this file still contains srand(seed)/rand() usage in init_sto_orbitals_Ecut (later in sto_wf.cpp), and module_stodft/sto_tool.cpp also uses std::rand(). Either migrate those remaining call sites to the new std::mt19937 approach as well, or adjust the PR description/scope accordingly to avoid misleading reproducibility/thread-safety expectations.

Copilot uses AI. Check for mistakes.
Expand Down Expand Up @@ -119,7 +125,7 @@ void Stochastic_WF<T, Device>::allocate_chi0()

// allocate chi0
Device* ctx = {};
if (base_device::get_device_type(ctx) == base_device::GpuDevice)
if (base_device::get_device_type<Device>(ctx) == base_device::GpuDevice)
{
this->chi0 = new psi::Psi<T, Device>(nks, this->nchip_max, npwx, this->ngk, true);
}
Expand All @@ -134,19 +140,26 @@ void Stochastic_WF<T, Device>::update_sto_orbitals(const int seed_in)
{
const int nchi = PARAM.inp.nbands_sto;
this->chi0_cpu->fix_k(0);

// Uniform distribution to generate random phases between 0 and 2*pi
std::uniform_real_distribution<double> dist_phi(0.0, 2.0 * ModuleBase::PI);
// Bernoulli distribution to generate +1/sqrt(nchi) or -1/sqrt(nchi) with equal probability
std::bernoulli_distribution dist_coin(0.5);

if (seed_in >= 0)
{
for (int i = 0; i < this->chi0_cpu->size(); ++i)
{
const double phi = 2 * ModuleBase::PI * rand() / double(RAND_MAX);
const double phi = dist_phi(this->rng);
this->chi0_cpu->get_pointer()[i] = std::complex<double>(cos(phi), sin(phi)) / sqrt(double(nchi));
}
}
else
{
for (int i = 0; i < this->chi0_cpu->size(); ++i)
{
if (rand() / double(RAND_MAX) < 0.5)
// use Bernoulli distribution to generate +1/sqrt(nchi) or -1/sqrt(nchi) with equal probability
if (dist_coin(this->rng))
{
this->chi0_cpu->get_pointer()[i] = -1.0 / sqrt(double(nchi));
}
Expand Down
4 changes: 4 additions & 0 deletions source/source_pw/module_stodft/sto_wf.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -59,6 +59,10 @@ class Stochastic_WF
void init_com_orbitals();
// sync chi0 from CPU to GPU
void sync_chi0();

private:
// random number generator
std::mt19937 rng;
Comment on lines +62 to +65
Copy link

Copilot AI Mar 25, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

std::mt19937 is used in the class definition, but sto_wf.h does not include <random>. Any translation unit that includes this header (without separately including <random>) will fail to compile. Add #include <random> to sto_wf.h (or hide the RNG behind a PIMPL so the header doesn’t depend on <random>).

Copilot uses AI. Check for mistakes.

protected:
using setmem_complex_op = base_device::memory::set_memory_op<T, Device>;
Expand Down
Loading