Skip to content

add generalized eigenvalue test for diago_cg_float_test #7408

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

Open
kevin0x3f wants to merge 4 commits into
base: develop
Choose a base branch
from
Open

add generalized eigenvalue test for diago_cg_float_test #7408

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
2e11533
add generalized eigenvalue test for diago_cg_float_test
May 30, 2026
6951779
Merge branch 'deepmodeling:develop' into diago-generalized-test
kevin0x3f Jun 6, 2026
9f10c7b
small fixes
Jun 6, 2026
215e6dd
add parallel testing of CG algorithm on generalized eigenvalue problems
Jun 6, 2026
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
61 changes: 59 additions & 2 deletions source/source_hsolver/test/diago_cg_float_test.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -58,6 +58,35 @@ void lapackEigen(int &npw, std::vector<std::complex<float>> &hm, float *e, bool
delete[] work2;
}

// LAPACK reference for generalized eigenproblem when S is diagonal (complex<float>)
void lapackGeneralEigen(int &npw, std::vector<std::complex<float>> &hm, const std::vector<std::complex<float>> &sdiag, float *e, bool outtime = false)
{
// build transformed matrix tmp = S^{-1/2} H S^{-1/2}
std::vector<std::complex<float>> tmp(npw * npw);
for (int i = 0; i < npw; ++i) {
std::complex<float> si = std::sqrt(sdiag[i]);
for (int j = 0; j < npw; ++j) {
std::complex<float> sj = std::sqrt(sdiag[j]);
tmp[i * npw + j] = hm[i * npw + j] / (si * sj);
}
}
Comment on lines +62 to +72

// call cheev_ on transformed matrix
clock_t start = clock(), end;
int lwork = 2 * npw;
std::complex<float> *work2 = new std::complex<float>[lwork];
float *rwork = new float[3 * npw - 2];
int info = 0;
char tmp_c1 = 'V', tmp_c2 = 'U';
cheev_(&tmp_c1, &tmp_c2, &npw, tmp.data(), &npw, e, work2, &lwork, rwork, &info);
end = clock();
if (outtime) {
std::cout << "Lapack General Run time: " << (float)(end - start) / CLOCKS_PER_SEC << " S" << std::endl;
}
delete[] rwork;
delete[] work2;
}

class DiagoCGPrepare
{
public:
Expand Down Expand Up @@ -85,8 +114,14 @@ class DiagoCGPrepare
float *e_lapack = new float[npw];
auto ev = DIAGOTEST::hmatrix_f;

if(mypnum == 0) { lapackEigen(npw, ev, e_lapack, false);
}
if(mypnum == 0) {
if (DIAGOTEST::sdiag_f.empty()) {
lapackEigen(npw, ev, e_lapack);
} else {
auto hm_copy = ev; // operate on a copy
lapackGeneralEigen(npw, hm_copy, DIAGOTEST::sdiag_f, e_lapack);
}
}

// initial guess of psi by perturbing lapack psi
std::vector<std::complex<float>> psiguess(nband * npw);
Expand Down Expand Up @@ -135,6 +170,7 @@ class DiagoCGPrepare
precondition_local = new float[DIAGOTEST::npw];
for(int i=0;i<DIAGOTEST::npw;i++) precondition_local[i] = precondition[i];
#endif
DIAGOTEST::sync_sdiag(DIAGOTEST::sdiag_f, DIAGOTEST::sdiag_local_f);
// hsolver::DiagoCG<std::complex<float>> cg(precondition_local);
psi_local.fix_k(0);
// cg.diag(ha,psi_local,en);
Expand Down Expand Up @@ -228,12 +264,33 @@ TEST_P(DiagoCGFloatTest, RandomHamilt)
//std::cout<<"eps "<<hsolver::DiagoIterAssist<std::complex<float>>::PW_DIAG_THR<<std::endl;
HPsi<std::complex<float>> hpsi(dcp.nband, dcp.npw, dcp.sparsity);
DIAGOTEST::hmatrix_f = hpsi.hamilt();
DIAGOTEST::sdiag_f.clear(); // ensure sdiag is empty for this test

DIAGOTEST::npw = dcp.npw;
// ModuleBase::ComplexMatrix psi = hpsi.psi();
dcp.CompareEigen(hpsi.precond());
}

TEST_P(DiagoCGFloatTest, RandomHamiltAndS)
{
DiagoCGPrepare dcp = GetParam();
hsolver::DiagoIterAssist<std::complex<float>>::PW_DIAG_NMAX = dcp.maxiter;
hsolver::DiagoIterAssist<std::complex<float>>::PW_DIAG_THR = dcp.eps;
HPsi<std::complex<float>> hpsi(dcp.nband, dcp.npw, dcp.sparsity);
DIAGOTEST::hmatrix_f = hpsi.hamilt();

DIAGOTEST::npw = dcp.npw;

DIAGOTEST::sdiag_f.resize(dcp.npw);
std::default_random_engine eng(123);
std::uniform_real_distribution<float> ud(0.5f, 1.5f);
for (int i = 0; i < dcp.npw; ++i) {
DIAGOTEST::sdiag_f[i] = std::complex<float>(ud(eng), 0.0f);
}

dcp.CompareEigen(hpsi.precond());
}

INSTANTIATE_TEST_SUITE_P(VerifyCG,
DiagoCGFloatTest,
::testing::Values(
Expand Down
1 change: 1 addition & 0 deletions source/source_hsolver/test/diago_cg_real_test.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -136,6 +136,7 @@ class DiagoCGPrepare
precondition_local = new double[DIAGOTEST::npw];
for (int i = 0;i < DIAGOTEST::npw;i++) precondition_local[i] = precondition[i];
#endif
DIAGOTEST::sync_sdiag(DIAGOTEST::sdiag_d, DIAGOTEST::sdiag_local_d);
// hsolver::DiagoCG<double> cg(precondition_local);
psi_local.fix_k(0);
// double start, end;
Expand Down
65 changes: 62 additions & 3 deletions source/source_hsolver/test/diago_cg_test.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -58,6 +58,36 @@ void lapackEigen(int &npw, std::vector<std::complex<double>> &hm, double *e, boo
delete[] work2;
}

// LAPACK reference for generalized eigenproblem when S is diagonal (complex<double>)
void lapackGeneralEigen(int &npw, std::vector<std::complex<double>> &hm, const std::vector<std::complex<double>> &sdiag, double *e, bool outtime = false)
{
// build transformed matrix tmp = S^{-1/2} H S^{-1/2}
std::vector<std::complex<double>> tmp(npw * npw);
for (int i = 0; i < npw; ++i) {
std::complex<double> si = std::sqrt(sdiag[i]);
for (int j = 0; j < npw; ++j) {
std::complex<double> sj = std::sqrt(sdiag[j]);
tmp[i * npw + j] = hm[i * npw + j] / (si * sj);
}
}

// call cheev_ on transformed matrix
clock_t start = clock(), end;
int lwork = 2 * npw;
std::complex<double> *work2 = new std::complex<double>[lwork];
double *rwork = new double[3 * npw - 2];
int info = 0;
char tmp_c1 = 'V', tmp_c2 = 'U';
zheev_(&tmp_c1, &tmp_c2, &npw, tmp.data(), &npw, e, work2, &lwork, rwork, &info);
end = clock();
if (outtime) {
std::cout << "Lapack General Run time: " << (double)(end - start) / CLOCKS_PER_SEC << " S" << std::endl;
}
delete[] rwork;
delete[] work2;
}


class DiagoCGPrepare
{
public:
Expand All @@ -84,8 +114,14 @@ class DiagoCGPrepare
// calculate eigenvalues by LAPACK;
double *e_lapack = new double[npw];
auto ev = DIAGOTEST::hmatrix;
if(mypnum == 0) { lapackEigen(npw, ev, e_lapack, false);
}
if(mypnum == 0) {
if (DIAGOTEST::sdiag.empty()) {
lapackEigen(npw, ev, e_lapack);
} else {
auto hm_copy = ev; // operate on a copy
lapackGeneralEigen(npw, hm_copy, DIAGOTEST::sdiag, e_lapack);
}
}
// initial guess of psi by perturbing lapack psi
ModuleBase::ComplexMatrix psiguess(nband, npw);
std::default_random_engine p(1);
Expand Down Expand Up @@ -131,6 +167,7 @@ class DiagoCGPrepare
precondition_local = new double[DIAGOTEST::npw];
for(int i=0;i<DIAGOTEST::npw;i++) precondition_local[i] = precondition[i];
#endif
DIAGOTEST::sync_sdiag(DIAGOTEST::sdiag, DIAGOTEST::sdiag_local);

/**************************************************************/
// New interface of cg method
Expand Down Expand Up @@ -201,7 +238,6 @@ class DiagoCGPrepare
{
EXPECT_NEAR(en[i], e_lapack[i], threshold);
}

delete[] en;
delete[] e_lapack;
delete ha;
Expand All @@ -223,12 +259,35 @@ TEST_P(DiagoCGTest, RandomHamilt)
//std::cout<<"eps "<<hsolver::DiagoIterAssist<std::complex<double>>::PW_DIAG_THR<<std::endl;
HPsi<std::complex<double>> hpsi(dcp.nband, dcp.npw, dcp.sparsity);
DIAGOTEST::hmatrix = hpsi.hamilt();
DIAGOTEST::sdiag.clear(); // ensure sdiag is empty for this test

DIAGOTEST::npw = dcp.npw;
// ModuleBase::ComplexMatrix psi = hpsi.psi();
dcp.CompareEigen(hpsi.precond());
}

TEST_P(DiagoCGTest, RandomHamiltAndS)
{
DiagoCGPrepare dcp = GetParam();
hsolver::DiagoIterAssist<std::complex<double>>::PW_DIAG_NMAX = dcp.maxiter;
hsolver::DiagoIterAssist<std::complex<double>>::PW_DIAG_THR = dcp.eps;
HPsi<std::complex<double>> hpsi(dcp.nband, dcp.npw, dcp.sparsity);
DIAGOTEST::hmatrix = hpsi.hamilt();
DIAGOTEST::sdiag.clear(); //ensure sdiag is empty before generating a new one

DIAGOTEST::npw = dcp.npw;

DIAGOTEST::sdiag.resize(dcp.npw);
std::default_random_engine eng(123);
std::uniform_real_distribution<double> ud(0.5f, 1.5f);
for (int i = 0; i < dcp.npw; ++i) {
DIAGOTEST::sdiag[i] = std::complex<double>(ud(eng), 0.0L);
}

dcp.CompareEigen(hpsi.precond());
}


INSTANTIATE_TEST_SUITE_P(VerifyCG,
DiagoCGTest,
::testing::Values(
Expand Down
88 changes: 76 additions & 12 deletions source/source_hsolver/test/diago_mock.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -11,6 +11,15 @@ namespace DIAGOTEST
std::vector<std::complex<double>> hmatrix_local;
std::vector<std::complex<float>> hmatrix_f;
std::vector<std::complex<float>> hmatrix_local_f;

// diagonal representation of overlap (S) for simple mock of generalized eigenproblem
// if empty, sPsi will treat S as identity
std::vector<double> sdiag_d;
std::vector<double> sdiag_local_d;
std::vector<std::complex<float> > sdiag_f;
std::vector<std::complex<float>> sdiag_local_f;
std::vector<std::complex<double>> sdiag;
Comment on lines +15 to +21
std::vector<std::complex<double>> sdiag_local;
int h_nr;
int h_nc;
int npw;
Expand Down Expand Up @@ -121,6 +130,31 @@ namespace DIAGOTEST
template void divide_psi<double>(double* psi, double* psi_local);
template void divide_psi<std::complex<double>>(std::complex<double>* psi, std::complex<double>* psi_local);
template void divide_psi<std::complex<float>>(std::complex<float>* psi, std::complex<float>* psi_local);

template<typename T>
void sync_sdiag(const std::vector<T>& sdiag, std::vector<T>& sdiag_local)
{
if (sdiag.empty())
{
sdiag_local.clear();
return;
}

int nprocs = 1, mypnum = 0;
#ifdef __MPI
MPI_Comm_size(MPI_COMM_WORLD, &nprocs);
MPI_Comm_rank(MPI_COMM_WORLD, &mypnum);
#endif

sdiag_local.resize(npw_local[mypnum]);
divide_psi(const_cast<T*>(sdiag.data()), sdiag_local.data());
}

template void sync_sdiag<double>(const std::vector<double>& sdiag, std::vector<double>& sdiag_local);
template void sync_sdiag<std::complex<double>>(const std::vector<std::complex<double>>& sdiag,
std::vector<std::complex<double>>& sdiag_local);
template void sync_sdiag<std::complex<float>>(const std::vector<std::complex<float>>& sdiag,
std::vector<std::complex<float>>& sdiag_local);
#ifdef __MPI
void cal_division(int &npw)
{
Expand Down Expand Up @@ -409,11 +443,21 @@ void hamilt::HamiltPW<double, base_device::DEVICE_CPU>::sPsi(const double* psi_i
const int npw,
const int nbands) const
{
for (size_t i = 0; i < static_cast<size_t>(nbands * nrow); i++)
{
spsi[i] = psi_in[i];
std::vector<double>& sdiag_local
#ifdef __MPI
= DIAGOTEST::sdiag_local_d;
#else
= DIAGOTEST::sdiag_d;
#endif
if (sdiag_local.size() < static_cast<size_t>(nrow)) {
sdiag_local.assign(nrow, 1.0);
}
for (int v = 0; v < nbands; ++v) {
for (int i = 0; i < nrow; ++i) {
size_t idx = static_cast<size_t>(v) * nrow + i;
spsi[idx] = psi_in[idx] * sdiag_local[i];
}
}
return;
}
template <>
void hamilt::HamiltPW<std::complex<double>, base_device::DEVICE_CPU>::sPsi(const std::complex<double>* psi_in,
Expand All @@ -422,11 +466,21 @@ void hamilt::HamiltPW<std::complex<double>, base_device::DEVICE_CPU>::sPsi(const
const int npw,
const int nbands) const
{
for (size_t i = 0; i < static_cast<size_t>(nbands * nrow); i++)
{
spsi[i] = psi_in[i];
std::vector<std::complex<double>>& sdiag_local
#ifdef __MPI
= DIAGOTEST::sdiag_local;
#else
= DIAGOTEST::sdiag;
#endif
if (sdiag_local.size() < static_cast<size_t>(nrow)) {
sdiag_local.assign(nrow, std::complex<double>(1.0, 0.0));
}
for (int v = 0; v < nbands; ++v) {
for (int i = 0; i < nrow; ++i) {
size_t idx = static_cast<size_t>(v) * nrow + i;
spsi[idx] = psi_in[idx] * sdiag_local[i];
}
}
return;
}
template <>
void hamilt::HamiltPW<std::complex<float>, base_device::DEVICE_CPU>::sPsi(const std::complex<float>* psi_in,
Expand All @@ -435,11 +489,21 @@ void hamilt::HamiltPW<std::complex<float>, base_device::DEVICE_CPU>::sPsi(const
const int npw,
const int nbands) const
{
for (size_t i = 0; i < static_cast<size_t>(nbands * nrow); i++)
{
spsi[i] = psi_in[i];
std::vector<std::complex<float>>& sdiag_local
#ifdef __MPI
= DIAGOTEST::sdiag_local_f;
#else
= DIAGOTEST::sdiag_f;
#endif
if (sdiag_local.size() < static_cast<size_t>(nrow)) {
sdiag_local.assign(nrow, std::complex<float>(1.0f, 0.0f));
}
for (int v = 0; v < nbands; ++v) {
for (int i = 0; i < nrow; ++i) {
size_t idx = static_cast<size_t>(v) * nrow + i;
spsi[idx] = psi_in[idx] * sdiag_local[i];
}
}
return;
}

//Mock function h_psi
Expand Down
Loading