diff --git a/.readthedocs.yaml b/.readthedocs.yaml index 5584c665201..318171496b1 100644 --- a/.readthedocs.yaml +++ b/.readthedocs.yaml @@ -19,7 +19,12 @@ build: - pkg-config jobs: pre_build: - - cmake -Bbuild -GNinja -DCMAKE_INSTALL_PREFIX=install -DENABLE_MPI=OFF -DENABLE_LCAO=OFF + # This executable only exports INPUT metadata; skip optimization and debug + # symbols to keep the documentation build within the RTD time limit. + - >- + cmake -Bbuild -GNinja -DCMAKE_INSTALL_PREFIX=install + -DENABLE_MPI=OFF -DENABLE_LCAO=OFF -DENABLE_NATIVE_OPTIMIZATION=OFF + -DCMAKE_BUILD_TYPE=Debug -DCMAKE_CXX_FLAGS_DEBUG="-O0 -g0" - cmake --build build --target install -j $(nproc) - install/bin/abacus --generate-parameters-yaml > docs/parameters.yaml - python3 docs/generate_input_main.py docs/parameters.yaml diff --git a/docs/advanced/scf/hsolver.md b/docs/advanced/scf/hsolver.md index cc85b4d87be..e394cacef28 100644 --- a/docs/advanced/scf/hsolver.md +++ b/docs/advanced/scf/hsolver.md @@ -4,7 +4,11 @@ Method of explicit solving KS-equation can be chosen by variable "ks_solver" in INPUT file. -When "basis_type = pw", `ks_solver` can be `cg`, `bpcg` or `dav`. The default setting `cg` is recommended, which is band-by-band conjugate gradient diagonalization method. There is a large probability that the use of setting of `dav` , which is block Davidson diagonalization method, can be tried to improve performance. +When "basis_type = pw", `ks_solver` can be `cg`, `bpcg`, `dav`, `dav_subspace`, or `ppcg`. The default setting `cg` is recommended, which is a band-by-band conjugate-gradient diagonalization method. The `dav` and `dav_subspace` settings use Davidson-style subspace diagonalization and can be tried to improve performance. + +The `ppcg` setting uses the projection preconditioned conjugate-gradient method, a restarted block method with bounded workspace. It is optimized and validated for CPU plane-wave calculations; non-CPU devices use a transitional host/device bridge. PPCG is unavailable for `esolver_type = sdft`, including mixed stochastic-deterministic calculations. + +For PPCG, `pw_diag_ndim` controls the block size. `pw_diag_rr_step` sets the interval between full Rayleigh--Ritz diagonalizations; its default of 1 performs one on every iteration. Initialization, startup correction, tentative convergence, final validation and failed overlap orthogonalization can also trigger a full Rayleigh--Ritz step. Larger intervals reduce projected eigensolves but may require more iterations, so energies, forces and stresses should be checked for the intended calculation. Refreshing the cached Hamiltonian and overlap images is independent of this interval: periodic and numerical safeguards recompute them, and convergence is confirmed with freshly applied operators. When "basis_type = lcao", `ks_solver` can be `genelpa` or `scalapack_gvx`. The default setting `genelpa` is recommended, which is based on ELPA (EIGENVALUE SOLVERS FOR PETAFLOP APPLICATIONS) (https://elpa.mpcdf.mpg.de/) and the kernel is auto choosed by GENELPA(https://github.com/pplab/GenELPA), usually faster than the setting of "scalapack_gvx", which is based on ScaLAPACK(Scalable Linear Algebra PACKage) diff --git a/source/Makefile.Objects b/source/Makefile.Objects index 9145af6c2cd..a84f74ffc67 100644 --- a/source/Makefile.Objects +++ b/source/Makefile.Objects @@ -497,6 +497,7 @@ OBJS_HSOLVER=diago_cg.o\ diago_david.o\ diago_dav_subspace.o\ diago_bpcg.o\ + diago_ppcg.o\ para_lin_tf.o\ hsolver.o\ hsolver_pw.o\ diff --git a/source/source_base/parallel_comm.cpp b/source/source_base/parallel_comm.cpp index 5d03447b5aa..99d9cb197b9 100644 --- a/source/source_base/parallel_comm.cpp +++ b/source/source_base/parallel_comm.cpp @@ -3,7 +3,8 @@ #include "mpi.h" #include "parallel_global.h" -MPI_Comm POOL_WORLD; //groups for different plane waves. In this group, only plane waves are different. K-points and bands are the same. +// Plane-wave pool: only the plane waves differ; k-points and bands are shared. +MPI_Comm POOL_WORLD; MPI_Comm KP_WORLD; // groups for differnt k. In this group, only k-points are different. Bands and plane waves are the same. MPI_Comm BP_WORLD; // groups for differnt bands. In this group, only bands are different. K-points and plane waves are the same. MPI_Comm INT_BGROUP; // internal comm groups for same bands. In this group, only bands are the same. K-points and plane waves are different. @@ -52,4 +53,4 @@ void MPICommGroup::divide_group_comm(const int& ngroup, const bool assert_even) } } -#endif \ No newline at end of file +#endif diff --git a/source/source_base/parallel_reduce.cpp b/source/source_base/parallel_reduce.cpp index 741b38f36d5..4a1f113d2c8 100644 --- a/source/source_base/parallel_reduce.cpp +++ b/source/source_base/parallel_reduce.cpp @@ -13,6 +13,31 @@ const MPI_Datatype Parallel_Reduce::MPI_Type>::value = MPI_C const MPI_Datatype Parallel_Reduce::MPI_Type::value = MPI_LONG_LONG; #endif +bool Parallel_Reduce::is_world_root() +{ +#ifdef __MPI + int initialized = 0; + MPI_Initialized(&initialized); + if (!initialized) + { + return true; + } + + int finalized = 0; + MPI_Finalized(&finalized); + if (finalized) + { + return false; + } + + int rank = 0; + MPI_Comm_rank(MPI_COMM_WORLD, &rank); + return rank == 0; +#else + return true; +#endif +} + template void Parallel_Reduce::reduce_all(T& object) { @@ -111,6 +136,7 @@ template void Parallel_Reduce::reduce_pool(double&); template void Parallel_Reduce::reduce_pool>(std::complex&); template void Parallel_Reduce::reduce_pool(int*, const int); +template void Parallel_Reduce::reduce_pool(float*, const int); template void Parallel_Reduce::reduce_pool(double*, const int); template void Parallel_Reduce::reduce_pool>(std::complex*, const int); template void Parallel_Reduce::reduce_pool>(std::complex*, const int); diff --git a/source/source_base/parallel_reduce.h b/source/source_base/parallel_reduce.h index 2e2f842f3a2..9130a002104 100644 --- a/source/source_base/parallel_reduce.h +++ b/source/source_base/parallel_reduce.h @@ -11,6 +11,9 @@ namespace Parallel_Reduce { +/// Return whether this process is the MPI world root (or is running serially). +bool is_world_root(); + #ifdef __MPI template struct MPI_Type; diff --git a/source/source_esolver/esolver_ks.cpp b/source/source_esolver/esolver_ks.cpp index d81f3bf435c..6ac5f07cefa 100644 --- a/source/source_esolver/esolver_ks.cpp +++ b/source/source_esolver/esolver_ks.cpp @@ -157,7 +157,7 @@ void ESolver_KS::hamilt2rho(UnitCell& ucell, const int istep, const int iter, co { diag_ethr = hsolver::reset_diag_ethr(GlobalV::ofs_running, this->inp_->basis_type, policy, this->inp_->precision, hsolver_error, - drho, diag_ethr, this->inp_->nelec); + drho, diag_ethr, this->inp_->nelec, this->inp_->pw_diag_thr_max); this->hamilt2rho_single(ucell, istep, iter, diag_ethr); @@ -287,7 +287,8 @@ void ESolver_KS::iter_init(UnitCell& ucell, const int istep, const int iter) { diag_ethr = hsolver::set_diagethr_ks(this->inp_->basis_type, policy, this->inp_->calculation, this->inp_->init_chg, this->inp_->precision, istep, iter, - drho, this->inp_->pw_diag_thr, diag_ethr, this->inp_->nelec, this->inp_->scf_thr); + drho, this->inp_->pw_diag_thr, diag_ethr, this->inp_->nelec, this->inp_->scf_thr, + this->inp_->pw_diag_thr_max); } else if (policy == "sdft") { diff --git a/source/source_esolver/esolver_ks_pw.cpp b/source/source_esolver/esolver_ks_pw.cpp index c700eefa055..a0388ef39fe 100644 --- a/source/source_esolver/esolver_ks_pw.cpp +++ b/source/source_esolver/esolver_ks_pw.cpp @@ -324,6 +324,7 @@ void ESolver_KS_PW::hamilt2rho_single(UnitCell& ucell, const int iste this->inp_->nbands, this->inp_->diago_smooth_ethr, this->inp_->pw_diag_ndim, + this->inp_->pw_diag_rr_step, this->inp_->diag_subspace, this->inp_->nb2d, this->inp_->use_k_continuity); diff --git a/source/source_esolver/esolver_sdft_pw.cpp b/source/source_esolver/esolver_sdft_pw.cpp index 840dd1175bb..7015156f933 100644 --- a/source/source_esolver/esolver_sdft_pw.cpp +++ b/source/source_esolver/esolver_sdft_pw.cpp @@ -171,6 +171,7 @@ void ESolver_SDFT_PW::hamilt2rho_single(UnitCell& ucell, int istep, i this->inp_->nbands, this->inp_->diago_smooth_ethr, this->inp_->pw_diag_ndim, + this->inp_->pw_diag_rr_step, this->inp_->diag_subspace, this->inp_->nb2d, PARAM.globalv.ks_run, diff --git a/source/source_estate/elecstate_print.cpp b/source/source_estate/elecstate_print.cpp index 4c4a2d9edf5..6acf3c95fa7 100644 --- a/source/source_estate/elecstate_print.cpp +++ b/source/source_estate/elecstate_print.cpp @@ -57,6 +57,7 @@ void print_scf_iterinfo(const std::string& ks_solver, {"scalapack_gvx", "GV"}, {"cusolver", "CU"}, {"bpcg", "BP"}, + {"ppcg", "PP"}, {"pexsi", "PE"}, {"cusolvermp", "CM"}, {"sdft", "CT"}}; // CT = Chebyshev Trace, for pure SDFT (nbands=0) where no H diagonalization is performed diff --git a/source/source_estate/module_charge/chg_mix.cpp b/source/source_estate/module_charge/chg_mix.cpp index 45d93180a4d..361b565f56f 100644 --- a/source/source_estate/module_charge/chg_mix.cpp +++ b/source/source_estate/module_charge/chg_mix.cpp @@ -112,6 +112,18 @@ void Charge_Mixing::init_mixing() "set_mixing must be called before init_mixing"); } + // Real-space charge and kinetic density are stored on the dense grid. + // Reciprocal mixing still separates smooth and high-frequency components. + ModulePW::PW_Basis* real_grid = this->rhopw; + if (this->cfg_.scf_thr_type == 2 && this->cfg_.double_grid) + { + if (this->rhodpw == nullptr) + { + ModuleBase::WARNING_QUIT("Charge_Mixing", "rhodpw is null in real-space init_mixing"); + } + real_grid = this->rhodpw; + } + // (re)construct mixing object if (this->cfg_.mixing_mode == "broyden") { @@ -160,14 +172,13 @@ void Charge_Mixing::init_mixing() } else { - if (this->cfg_.nspin == 4 && this->cfg_.mixing_angle > 0 ) - { - this->mixing->init_mixing_data(this->rho_mdata, this->rhopw->nrxx * 2, sizeof(double)); - } - else + int real_channels = this->cfg_.nspin; + if (this->cfg_.nspin == 4 && this->cfg_.mixing_angle > 0) { - this->mixing->init_mixing_data(this->rho_mdata, this->rhopw->nrxx * this->cfg_.nspin, sizeof(double)); + real_channels = 2; } + const int real_size = real_grid->nrxx * real_channels; + this->mixing->init_mixing_data(this->rho_mdata, real_size, sizeof(double)); } // initailize tau_mdata @@ -181,7 +192,8 @@ void Charge_Mixing::init_mixing() } else { - this->mixing->init_mixing_data(this->tau_mdata, this->rhopw->nrxx * this->cfg_.nspin, sizeof(double)); + const int tau_size = real_grid->nrxx * this->cfg_.nspin; + this->mixing->init_mixing_data(this->tau_mdata, tau_size, sizeof(double)); } } diff --git a/source/source_estate/module_charge/chg_mix_rho.cpp b/source/source_estate/module_charge/chg_mix_rho.cpp index 9ee17da8b86..f3f63ac657d 100644 --- a/source/source_estate/module_charge/chg_mix_rho.cpp +++ b/source/source_estate/module_charge/chg_mix_rho.cpp @@ -258,20 +258,28 @@ void Charge_Mixing::mix_rho_real(Charge* chr) const int nspin = cfg_.nspin; assert(nspin==1 || nspin==2 || nspin==4); + // rho and kin_r use the dense real-space grid when double_grid is set. + // Use that grid for strides, FFT screening, and the history inner product. + ModulePW::PW_Basis* real_grid = this->rhopw; + if (this->cfg_.double_grid) + { + real_grid = this->rhodpw; + } + double* rhor_in=nullptr; double* rhor_out=nullptr; // Kerker screening functor (see mix_rho_recip for the disable flag rationale). - std::function screen = [this](double* p) { + std::function screen = [this, real_grid](double* p) { if (this->kerker_disabled_) { return; } - module_charge::kerker_screen_real(this->cfg_, this->rhopw, *this->tpiba, p); + module_charge::kerker_screen_real(this->cfg_, real_grid, *this->tpiba, p); }; - std::function inner_product = [this](double* rho1, double* rho2) + std::function inner_product = [this, real_grid](double* rho1, double* rho2) { - return module_charge::inner_product_real(rho1, rho2, *this->rhopw, this->cfg_); + return module_charge::inner_product_real(rho1, rho2, *real_grid, this->cfg_); }; if (nspin == 1) @@ -285,7 +293,7 @@ void Charge_Mixing::mix_rho_real(Charge* chr) else if (nspin == 2) { // magnetic density - const int nrxx = this->rhopw->nrxx; + const int nrxx = real_grid->nrxx; std::vector rho_mag(nrxx * nspin); std::vector rho_mag_save(nrxx * nspin); module_charge::detail::pack_rho_mag(rho_mag.data(), chr->rho[0], chr->rho[1], nrxx); @@ -305,7 +313,7 @@ void Charge_Mixing::mix_rho_real(Charge* chr) { rhor_in = chr->rho_save[0]; rhor_out = chr->rho[0]; - const int nrxx = this->rhopw->nrxx; + const int nrxx = real_grid->nrxx; std::function twobeta_mix = module_charge::detail::make_twobeta_mix(4 * nrxx, nrxx, this->cfg_.mixing_beta, this->cfg_.mixing_beta_mag); this->mixing->push_data(this->rho_mdata, rhor_in, rhor_out, screen, twobeta_mix, true); @@ -315,7 +323,7 @@ void Charge_Mixing::mix_rho_real(Charge* chr) else if (nspin == 4 && cfg_.mixing_angle > 0) { // real-space version of the {rho, |m|} broyden mixing - const int nrxx = this->rhopw->nrxx; + const int nrxx = real_grid->nrxx; std::vector rho_magabs(nrxx * 2); std::vector rho_magabs_save(nrxx * 2); for (int ir = 0; ir < nrxx; ir++) diff --git a/source/source_estate/module_charge/unittests/test_chg_mix_rho.cpp b/source/source_estate/module_charge/unittests/test_chg_mix_rho.cpp index ff298ff72cc..adfb58c5452 100644 --- a/source/source_estate/module_charge/unittests/test_chg_mix_rho.cpp +++ b/source/source_estate/module_charge/unittests/test_chg_mix_rho.cpp @@ -9,6 +9,9 @@ #include "source_estate/module_charge/chg_mix_cfg.h" #include "source_io/module_parameter/parameter.h" +#include +#include +#include #include // charge.cpp references Magnetism; provide a lightweight stub. @@ -151,7 +154,6 @@ TEST_F(ChargeMixRhoTest, MixRhoDoubleGridWithoutRhodpwAborts) MixingConfig cfg = make_cfg(1, 2, true, false); // set_mixing with rhodpw == nullptr while double_grid is on cm.set_mixing(cfg, &pw_basis, nullptr, omega, tpiba); - cm.init_mixing(); setup_charge(1); EXPECT_DEATH(cm.mix_rho(&charge), ""); } @@ -190,3 +192,193 @@ TEST_F(ChargeMixRhoTest, MixRhoRealPlainNspin1) EXPECT_NEAR(charge.rho_save[0][ir], rho_new_val, 1e-8); } } + + +namespace +{ +void initialize_dense_charge(Charge& charge, const int nspin, const int nrxx, const bool kinetic) +{ + for (int is = 0; is < nspin; ++is) + { + for (int ir = 0; ir < nrxx; ++ir) + { + const double wave = std::sin(0.11 * ir + 0.7 * is); + charge.rho_save[is][ir] = 1.0 + 0.2 * is + 0.05 * wave; + charge.rho[is][ir] = 1.1 + 0.3 * is + 0.07 * wave; + if (kinetic) + { + charge.kin_r_save[is][ir] = 0.5 + 0.1 * is + 0.03 * wave; + charge.kin_r[is][ir] = 0.8 + 0.15 * is + 0.06 * wave; + } + } + } +} +} + +class ChargeMixDensePlainTest : public ChargeMixRhoTest, + public ::testing::WithParamInterface +{ +}; + +TEST_P(ChargeMixDensePlainTest, MixesAllDensePointsAndSpinChannels) +{ + const int nspin = GetParam(); + const int nrxx = pw_dbasis.nrxx; + ASSERT_GT(nrxx, pw_basis.nrxx); + MixingConfig cfg = make_cfg(nspin, 2, true, true); + cfg.mixing_beta_mag = 0.3; + Charge_Mixing cm; + cm.set_mixing(cfg, &pw_basis, &pw_dbasis, omega, tpiba); + cm.init_mixing(); + charge.set_rhopw(&pw_dbasis); + charge.allocate(nspin, true, true, 0); + initialize_dense_charge(charge, nspin, nrxx, true); + std::vector original(nspin * nrxx); + std::vector original_tau(nspin * nrxx); + std::vector expected(nspin * nrxx); + std::vector expected_tau(nspin * nrxx); + for (int is = 0; is < nspin; ++is) + { + for (int ir = 0; ir < nrxx; ++ir) + { + const int index = is * nrxx + ir; + original[index] = charge.rho[is][ir]; + original_tau[index] = charge.kin_r[is][ir]; + double beta = cfg.mixing_beta; + if (nspin == 4 && is > 0) + { + beta = cfg.mixing_beta_mag; + } + expected[index] = charge.rho_save[is][ir] + + beta * (charge.rho[is][ir] - charge.rho_save[is][ir]); + expected_tau[index] = charge.kin_r_save[is][ir] + + cfg.mixing_beta * (charge.kin_r[is][ir] - charge.kin_r_save[is][ir]); + } + } + if (nspin == 2) + { + for (int ir = 0; ir < nrxx; ++ir) + { + const double sum_in = charge.rho_save[0][ir] + charge.rho_save[1][ir]; + const double sum_out = charge.rho[0][ir] + charge.rho[1][ir]; + const double mag_in = charge.rho_save[0][ir] - charge.rho_save[1][ir]; + const double mag_out = charge.rho[0][ir] - charge.rho[1][ir]; + const double sum = sum_in + cfg.mixing_beta * (sum_out - sum_in); + const double mag = mag_in + cfg.mixing_beta_mag * (mag_out - mag_in); + expected[ir] = 0.5 * (sum + mag); + expected[nrxx + ir] = 0.5 * (sum - mag); + } + } + cm.mix_rho(&charge); + double max_error = 0.0; + bool all_finite = true; + for (int is = 0; is < nspin; ++is) + { + for (int ir = 0; ir < nrxx; ++ir) + { + const int index = is * nrxx + ir; + const double rho_error = std::abs(charge.rho[is][ir] - expected[index]); + const double tau_error = std::abs(charge.kin_r[is][ir] - expected_tau[index]); + const double rho_save_error = std::abs(charge.rho_save[is][ir] - original[index]); + const double tau_save_error = std::abs(charge.kin_r_save[is][ir] - original_tau[index]); + all_finite = all_finite && std::isfinite(rho_error) && std::isfinite(tau_error); + max_error = std::max(max_error, rho_error); + max_error = std::max(max_error, tau_error); + all_finite = all_finite && std::isfinite(rho_save_error) && std::isfinite(tau_save_error); + max_error = std::max(max_error, rho_save_error); + max_error = std::max(max_error, tau_save_error); + } + } + EXPECT_TRUE(all_finite); + EXPECT_LT(max_error, 1.0e-12); +} + +INSTANTIATE_TEST_SUITE_P(SpinChannels, ChargeMixDensePlainTest, ::testing::Values(1, 2, 4)); + +class ChargeMixDenseHistoryTest : public ChargeMixRhoTest, + public ::testing::WithParamInterface> +{ +}; + +TEST_P(ChargeMixDenseHistoryTest, DoubleGridRealKerkerAndHistoryMatchDenseGridOnly) +{ + const int nspin = GetParam().first; + const int nrxx = pw_dbasis.nrxx; + MixingConfig cfg = make_cfg(nspin, 2, true, true); + cfg.mixing_mode = "broyden"; + cfg.mixing_ndim = 3; + cfg.mixing_beta_mag = 0.3; + cfg.mixing_gg0 = 1.0; + cfg.mixing_gg0_mag = 0.8; + if (GetParam().second) + { + cfg.mixing_angle = 0.5; + } + Charge_Mixing actual; + actual.set_mixing(cfg, &pw_basis, &pw_dbasis, omega, tpiba); + actual.init_mixing(); + cfg.double_grid = false; + Charge_Mixing control; + control.set_mixing(cfg, &pw_dbasis, &pw_dbasis, omega, tpiba); + control.init_mixing(); + charge.set_rhopw(&pw_dbasis); + charge.allocate(nspin, true, true, 0); + Charge dense_control; + dense_control.set_rhopw(&pw_dbasis); + dense_control.allocate(nspin, true, true, 0); + initialize_dense_charge(charge, nspin, nrxx, true); + initialize_dense_charge(dense_control, nspin, nrxx, true); + double max_error = 0.0; + bool all_finite = true; + for (int step = 0; step < 3; ++step) + { + if (step > 0) + { + for (int is = 0; is < nspin; ++is) + { + for (int ir = 0; ir < nrxx; ++ir) + { + const double perturbation = 0.03 * std::cos(0.07 * ir + 0.4 * is + 0.9 * step); + charge.rho[is][ir] += perturbation; + dense_control.rho[is][ir] += perturbation; + charge.kin_r[is][ir] += perturbation; + dense_control.kin_r[is][ir] += perturbation; + } + } + } + actual.mix_rho(&charge); + control.mix_rho(&dense_control); + for (int is = 0; is < nspin; ++is) + { + for (int ir = 0; ir < nrxx; ++ir) + { + const double rho_error = std::abs(charge.rho[is][ir] - dense_control.rho[is][ir]); + const double tau_error = std::abs(charge.kin_r[is][ir] - dense_control.kin_r[is][ir]); + all_finite = all_finite && std::isfinite(rho_error) && std::isfinite(tau_error); + max_error = std::max(max_error, rho_error); + max_error = std::max(max_error, tau_error); + } + } + } + EXPECT_TRUE(all_finite); + EXPECT_LT(max_error, 1.0e-12); +} + + +namespace +{ +const std::pair collinear_history(2, false); +const std::pair spinor_history(4, false); +const std::pair angle_history(4, true); +} +INSTANTIATE_TEST_SUITE_P(SpinChannels, + ChargeMixDenseHistoryTest, + ::testing::Values(collinear_history, spinor_history, angle_history)); + +TEST_F(ChargeMixRhoTest, RealDoubleGridInitWithoutDenseGridAborts) +{ + Charge_Mixing cm; + MixingConfig cfg = make_cfg(1, 2, true, false); + cm.set_mixing(cfg, &pw_basis, nullptr, omega, tpiba); + EXPECT_DEATH(cm.init_mixing(), ""); +} diff --git a/source/source_estate/unittests/test_elecstate_print.cpp b/source/source_estate/unittests/test_elecstate_print.cpp index d2a8e0d058f..6a956881818 100644 --- a/source/source_estate/unittests/test_elecstate_print.cpp +++ b/source/source_estate/unittests/test_elecstate_print.cpp @@ -150,7 +150,8 @@ TEST_F(ElecStatePrintTest, PrintEtot) } // iteration of different ks_solver - std::vector ks_solvers = {"cg", "lapack", "genelpa", "dav", "scalapack_gvx", "cusolver"}; + std::vector ks_solvers + = {"cg", "lapack", "genelpa", "dav", "scalapack_gvx", "cusolver", "ppcg"}; for (int i = 0; i < ks_solvers.size(); i++) { inp.ks_solver = ks_solvers[i]; @@ -184,6 +185,10 @@ TEST_F(ElecStatePrintTest, PrintEtot) { EXPECT_THAT(output, testing::HasSubstr("CU")); } + else if (inp.ks_solver == "ppcg") + { + EXPECT_THAT(output, testing::HasSubstr("PP")); + } } GlobalV::ofs_running.close(); ifs.open("test.dat", std::ios::in); diff --git a/source/source_hamilt/module_xc/libxc_pot.cpp b/source/source_hamilt/module_xc/libxc_pot.cpp index eb4275cf5e3..53c7f079ce9 100644 --- a/source/source_hamilt/module_xc/libxc_pot.cpp +++ b/source/source_hamilt/module_xc/libxc_pot.cpp @@ -530,8 +530,12 @@ std::tuple XC_Functional_Li constexpr double rho_th = 1e-8; constexpr double grho_th = 1e-12; constexpr double tau_th = 1e-8; - // sgn for threshold mask + // A vanishing density gradient is a valid uniform-density limit, not a + // reason to discard the local energy or rho/tau/Laplacian derivatives. + // Keep the existing density/tau cutoffs; apply the gradient cutoff only + // to the gradient contribution assembled from vsigma below. std::vector sgn( nrxx * nspin); + std::vector sgn_gradient(nrxx * nspin, 1.0); #ifdef _OPENMP #pragma omp parallel for schedule(static, 1024) #endif @@ -547,10 +551,14 @@ std::tuple XC_Functional_Li #endif for( int ir=0; ir XC_Functional_Li #endif for( int ir=0; ir XC_Functional_Li vsigma[ir] *= (1.0 - XC_Functional::get_hybrid_alpha()); } #endif - h[0][ir] = 2.0 * gdr[0][ir] * vsigma[ir] * 2.0 * sgn[ir]; + h[0][ir] = 2.0 * gdr[0][ir] * vsigma[ir] * 2.0 * sgn_gradient[ir]; } } else @@ -668,10 +680,10 @@ std::tuple XC_Functional_Li vsigma[ir*3+2] *= (1.0 - XC_Functional::get_hybrid_alpha()); } #endif - h[0][ir] = 2.0 * (gdr[0][ir] * vsigma[ir*3 ] * sgn[ir*2 ] * 2.0 - + gdr[1][ir] * vsigma[ir*3+1] * sgn[ir*2] * sgn[ir*2+1]); - h[1][ir] = 2.0 * (gdr[1][ir] * vsigma[ir*3+2] * sgn[ir*2+1] * 2.0 - + gdr[0][ir] * vsigma[ir*3+1] * sgn[ir*2] * sgn[ir*2+1]); + h[0][ir] = 2.0 * (gdr[0][ir] * vsigma[ir*3 ] * sgn_gradient[ir*2 ] * 2.0 + + gdr[1][ir] * vsigma[ir*3+1] * sgn_gradient[ir*2] * sgn_gradient[ir*2+1]); + h[1][ir] = 2.0 * (gdr[1][ir] * vsigma[ir*3+2] * sgn_gradient[ir*2+1] * 2.0 + + gdr[0][ir] * vsigma[ir*3+1] * sgn_gradient[ir*2] * sgn_gradient[ir*2+1]); } } diff --git a/source/source_hamilt/module_xc/test/CMakeLists.txt b/source/source_hamilt/module_xc/test/CMakeLists.txt index 82e8188923c..b983f9b2505 100644 --- a/source/source_hamilt/module_xc/test/CMakeLists.txt +++ b/source/source_hamilt/module_xc/test/CMakeLists.txt @@ -63,7 +63,7 @@ AddTest( AddTest( TARGET MODULE_HAMILT_XCTest_VXC LIBS parameter MPI::MPI_CXX Libxc::xc psi device container - SOURCES test_xc5.cpp test_libxc_tools.cpp ../xc_grad.cpp ../xc_grad_prepare.cpp ../xc_grad_kernel.cpp ../xc_grad_assemble.cpp ../xc_grad_wfc.cpp ../xc_grad_utils.cpp ../xc_functional.cpp + SOURCES test_xc5.cpp test_libxc_tools.cpp test_libxc_pot.cpp ../xc_grad.cpp ../xc_grad_prepare.cpp ../xc_grad_kernel.cpp ../xc_grad_assemble.cpp ../xc_grad_wfc.cpp ../xc_grad_utils.cpp ../xc_functional.cpp ../xc_lda_wrap.cpp ../xc_gga_wrap.cpp ../libxc_setup.cpp ../libxc_lda_wrap.cpp diff --git a/source/source_hamilt/module_xc/test/test_libxc_pot.cpp b/source/source_hamilt/module_xc/test/test_libxc_pot.cpp new file mode 100644 index 00000000000..a3fb8dc72d1 --- /dev/null +++ b/source/source_hamilt/module_xc/test/test_libxc_pot.cpp @@ -0,0 +1,157 @@ +#include "../libxc_abacus.h" +#include "../xc_functional.h" +#include "source_estate/module_charge/charge.h" + +#include + +#include +#include +#include +#include + +namespace +{ +// This target uses the one-grid-point FFT mocks defined in test_xc5.cpp. +// All storage here is owned by the fixture, including the arrays borrowed by +// the mocked Charge and PW_Basis classes (whose destructors are empty). +struct MetaPoint +{ + ModulePW::PW_Basis basis; + Charge charge; + std::array, 1> reciprocal; + std::array density; + std::array kinetic; + std::array density_rows; + std::array kinetic_rows; + std::array core; + + MetaPoint() + { + reciprocal[0] = ModuleBase::Vector3(0.0, 0.0, 0.0); + density = {{0.2, 0.1}}; + kinetic = {{0.6, 0.4}}; + core = {{0.0}}; + density_rows = {{&density[0], &density[1]}}; + kinetic_rows = {{&kinetic[0], &kinetic[1]}}; + basis.nrxx = 1; + basis.npw = 1; + basis.nxyz = 1; + basis.nmaxgr = 1; + basis.gcar = reciprocal.data(); + charge.rhopw = &basis; + charge.rho = density_rows.data(); + charge.kin_r = kinetic_rows.data(); + charge.rho_core = core.data(); + } +}; + +typedef std::tuple MetaResult; + +MetaResult evaluate_meta(MetaPoint& point, const int nspin) +{ + const std::vector functionals = {XC_MGGA_X_SCAN, XC_MGGA_C_SCAN}; + const double omega = 1.0; + const double tpiba = 1.0; + const double hybrid_alpha = 0.0; + const double hse_omega = 0.0; + return XC_Functional_Libxc::v_xc_meta(functionals, point.basis.nrxx, omega, tpiba, + &point.charge, nspin, hybrid_alpha, hse_omega); +} + +void check_uniform_limit(const int nspin) +{ + MetaPoint point; + const MetaResult result = evaluate_meta(point, nspin); + const int polarization = nspin == 1 ? XC_UNPOLARIZED : XC_POLARIZED; + const std::array sigma = {{0.0, 0.0, 0.0}}; + const std::array laplacian = {{0.0, 0.0}}; + const std::array tau = {{0.3, 0.2}}; // ABACUS stores 2 * tau. + double expected_energy = 0.0; + std::array expected_potential = {{0.0, 0.0}}; + std::array expected_tau_potential = {{0.0, 0.0}}; + const std::array ids = {{XC_MGGA_X_SCAN, XC_MGGA_C_SCAN}}; + for (const int id : ids) + { + xc_func_type functional; + ASSERT_EQ(xc_func_init(&functional, id, polarization), 0); + double energy = 0.0; + std::array vrho = {{0.0, 0.0}}; + std::array vsigma = {{0.0, 0.0, 0.0}}; + std::array vlapl = {{0.0, 0.0}}; + std::array vtau = {{0.0, 0.0}}; + xc_mgga_exc_vxc(&functional, 1, point.density.data(), sigma.data(), laplacian.data(), + tau.data(), &energy, vrho.data(), vsigma.data(), vlapl.data(), vtau.data()); + xc_func_end(&functional); + for (int spin = 0; spin < nspin; ++spin) + { + expected_energy += 2.0 * energy * point.density[spin]; + expected_potential[spin] += 2.0 * vrho[spin]; + expected_tau_potential[spin] += vtau[spin]; + } + } + EXPECT_LT(expected_energy, -1e-3); + EXPECT_NEAR(std::get<0>(result), expected_energy, 1e-12); + double expected_vtxc = 0.0; + for (int spin = 0; spin < nspin; ++spin) + { + EXPECT_NEAR(std::get<2>(result)(spin, 0), expected_potential[spin], 1e-12); + EXPECT_NEAR(std::get<3>(result)(spin, 0), expected_tau_potential[spin], 1e-12); + expected_vtxc += expected_potential[spin] * point.density[spin]; + } + EXPECT_NEAR(std::get<1>(result), expected_vtxc, 1e-12); +} + +class LibxcMetaPotentialTest : public testing::Test +{ + protected: + void SetUp() override + { + XC_Functional::set_xc_type("SCAN"); + } +}; +} // namespace + +TEST_F(LibxcMetaPotentialTest, UniformUnpolarizedMatchesLibxc) +{ + check_uniform_limit(1); +} + +TEST_F(LibxcMetaPotentialTest, UniformPolarizedMatchesLibxc) +{ + check_uniform_limit(2); +} + +TEST_F(LibxcMetaPotentialTest, LocalTermsRemainContinuousAcrossGradientCutoff) +{ + for (const int nspin : {1, 2}) + { + MetaPoint point; + const MetaResult zero = evaluate_meta(point, nspin); + point.reciprocal[0].x = 1e-13; + const MetaResult below = evaluate_meta(point, nspin); + point.reciprocal[0].x = 1e-10; + const MetaResult above = evaluate_meta(point, nspin); + EXPECT_NEAR(std::get<0>(zero), std::get<0>(above), 1e-12); + EXPECT_NEAR(std::get<0>(below), std::get<0>(above), 1e-12); + for (int spin = 0; spin < nspin; ++spin) + { + EXPECT_NEAR(std::get<2>(zero)(spin, 0), std::get<2>(above)(spin, 0), 1e-12); + EXPECT_NEAR(std::get<3>(zero)(spin, 0), std::get<3>(above)(spin, 0), 1e-12); + } + } +} + +TEST_F(LibxcMetaPotentialTest, ExistingDensityAndKineticCutoffsArePreserved) +{ + for (const int nspin : {1, 2}) + { + MetaPoint low_density; + low_density.density = {{1e-12, 1e-12}}; + const MetaResult density_result = evaluate_meta(low_density, nspin); + EXPECT_DOUBLE_EQ(std::get<0>(density_result), 0.0); + MetaPoint low_kinetic; + low_kinetic.kinetic = {{1e-12, 1e-12}}; + const MetaResult kinetic_result = evaluate_meta(low_kinetic, nspin); + EXPECT_DOUBLE_EQ(std::get<0>(kinetic_result), 0.0); + } +} diff --git a/source/source_hsolver/CMakeLists.txt b/source/source_hsolver/CMakeLists.txt index adca70290f7..a850bfb3405 100644 --- a/source/source_hsolver/CMakeLists.txt +++ b/source/source_hsolver/CMakeLists.txt @@ -4,6 +4,7 @@ list(APPEND objects diago_david.cpp diago_dav_subspace.cpp diago_bpcg.cpp + diago_ppcg.cpp para_lin_tf.cpp hsolver_pw.cpp hsolver_lcaopw.cpp diff --git a/source/source_hsolver/diago_cg.cpp b/source/source_hsolver/diago_cg.cpp index 95071ca3a78..e259eaf1267 100644 --- a/source/source_hsolver/diago_cg.cpp +++ b/source/source_hsolver/diago_cg.cpp @@ -260,6 +260,16 @@ void DiagoCG::calc_grad(const ct::Tensor& prec, 1.0, pphi.data(), (-lambda)); + + // Remove the residual projection left by cancellation in H*phi-lambda*S*phi. + const Real projection = ModuleBase::dot_real_op()(this->n_basis_, sphi.data(), grad.data()); + const Real correction = -projection / es; + ModuleBase::vector_add_vector_op()(this->n_basis_, + grad.data(), + grad.data(), + 1.0, + pphi.data(), + correction); } template @@ -427,15 +437,30 @@ bool DiagoCG::update_psi(const ct::Tensor& pphi, const Real new_e = (e0 - b0) * cos(2.0 * theta) + a0 * sin(2.0 * theta); - const Real e1 = (e0 + b0 + new_e) / 2.0; - const Real e2 = (e0 + b0 - new_e) / 2.0; - - if (e1 > e2) + // Select the lower root before adding the common energy offset. + if (new_e > 0.0) { theta += ModuleBase::PI_HALF; } - eigen = std::min(e1, e2); + // Evaluate the energy decrease before adding the full Rayleigh quotient. + // Subtracting two rounded full eigenvalues can falsely report zero change. + const Real gap = b0 - e0; + const Real splitting = std::hypot(gap, a0); + Real energy_change = 0.0; + if (gap >= 0.0) + { + const Real denominator = gap + splitting; + if (denominator > 0.0) + { + energy_change = -0.5 * a0 * (a0 / denominator); + } + } + else + { + energy_change = 0.5 * (gap - splitting); + } + eigen = e0 + energy_change; const Real cost = cos(theta); const Real sint_norm = sin(theta) / cg_norm; @@ -453,7 +478,7 @@ bool DiagoCG::update_psi(const ct::Tensor& pphi, cg.data(), sint_norm); - if (std::abs(eigen - e0) < ethreshold) + if (std::abs(energy_change) < ethreshold) { // ModuleBase::timer::start("DiagoCG","update"); return true; diff --git a/source/source_hsolver/diago_dav_subspace.cpp b/source/source_hsolver/diago_dav_subspace.cpp index 9b4764fc9d5..021c3e1ad01 100644 --- a/source/source_hsolver/diago_dav_subspace.cpp +++ b/source/source_hsolver/diago_dav_subspace.cpp @@ -7,6 +7,8 @@ #include "source_base/parallel_cell.h" #include "source_base/timer.h" #include "source_base/tool_quit.h" +#include "source_base/parallel_reduce.h" +#include "source_base/module_container/base/third_party/lapack.h" #include "source_base/kernels/math_kernel_op.h" #include "source_base/kernels/dsp/dsp_connector.h" // #include "source_base/module_container/ATen/kernels/lapack.h" @@ -18,6 +20,8 @@ #include "source_hsolver/kernels/bpcg_kernel_op.h" // normalize_op, precondition_op, apply_eigenvalues_op #include +#include +#include #ifdef __MPI #include @@ -27,6 +31,64 @@ using namespace hsolver; +namespace +{ +template +T dav_conjugate(const T value) +{ + return value; +} + +template +std::complex dav_conjugate(const std::complex value) +{ + return std::conj(value); +} + +// Keep INFO and M by reference: the general connector currently passes M by value. +void dav_hegvx(const int n, const int nband, std::complex* h, std::complex* s, + double* eigenvalues, std::complex* vectors, std::complex* work, + const int lwork, double* rwork, int* iwork, int* ifail, int& found, int& info) +{ + const int itype = 1; + const char job = 'V'; + const char range = 'I'; + const char uplo = 'U'; + const int first = 1; + const double zero = 0; + zhegvx_(&itype, &job, &range, &uplo, &n, h, &n, s, &n, &zero, &zero, &first, &nband, + &zero, &found, eigenvalues, vectors, &n, work, &lwork, rwork, iwork, ifail, &info); +} + +void dav_hegvx(const int n, const int nband, std::complex* h, std::complex* s, + float* eigenvalues, std::complex* vectors, std::complex* work, + const int lwork, float* rwork, int* iwork, int* ifail, int& found, int& info) +{ + const int itype = 1; + const char job = 'V'; + const char range = 'I'; + const char uplo = 'U'; + const int first = 1; + const float zero = 0; + chegvx_(&itype, &job, &range, &uplo, &n, h, &n, s, &n, &zero, &zero, &first, &nband, + &zero, &found, eigenvalues, vectors, &n, work, &lwork, rwork, iwork, ifail, &info); +} + +void dav_hegvx(const int n, const int nband, double* h, double* s, double* eigenvalues, + double* vectors, double* work, const int lwork, double* rwork, int* iwork, + int* ifail, int& found, int& info) +{ + const int itype = 1; + const char job = 'V'; + const char range = 'I'; + const char uplo = 'U'; + const int first = 1; + const double zero = 0; + dsygvx_(&itype, &job, &range, &uplo, &n, h, &n, s, &n, &zero, &zero, &first, &nband, + &zero, &found, eigenvalues, vectors, &n, work, &lwork, iwork, ifail, &info); +} +} // namespace + template Diago_DavSubspace::Diago_DavSubspace(const std::vector& precondition_in, const int& nband_in, @@ -117,6 +179,207 @@ Diago_DavSubspace::~Diago_DavSubspace() #endif } +template +int Diago_DavSubspace::orthonormalize_cpu(const HSOperator& op, + const int first, + const int count) +{ + // Scale before testing rank, so small residuals are not discarded just for + // their magnitude. Test the loss of S norm relative to the scaled input. + const Real rank_tolerance = 64 * std::numeric_limits::epsilon(); + const int matrix_ld = std::max(this->dim, 1); + // Keep valid storage on empty ranks while passing the physical size to H/S. + std::vector candidate(matrix_ld); + std::vector s_candidate(matrix_ld); + std::vector overlaps(this->nbase_x); + int retained = 0; + for (int column = 0; column < count; ++column) + { + const int source_offset = (first + column) * this->dim; + Real scale = 0; + int invalid = 0; + for (int row = 0; row < this->dim; ++row) + { + candidate[row] = this->psi_in_iter[source_offset + row]; + const Real magnitude = std::abs(candidate[row]); + invalid += !std::isfinite(magnitude); + scale = std::max(scale, magnitude); + } + if (this->diag_comm.nproc > 1) + { + Parallel_Reduce::reduce_pool(&invalid, 1); + double global_scale = scale; + Parallel_Reduce::reduce_max_pool(this->diag_comm.nproc, global_scale); + scale = static_cast(global_scale); + } + if (invalid != 0) + { + ModuleBase::WARNING_QUIT("Diago_DavSubspace", "Non-finite CPU basis direction."); + } + if (scale == 0) + { + continue; + } + for (int row = 0; row < this->dim; ++row) + { + candidate[row] /= scale; + } + op.spsi(candidate.data(), s_candidate.data(), this->dim, 1); + T initial_norm = 0; + for (int row = 0; row < this->dim; ++row) + { + initial_norm += dav_conjugate(candidate[row]) * s_candidate[row]; + } + if (this->diag_comm.nproc > 1) + { + Parallel_Reduce::reduce_pool(&initial_norm, 1); + } + const Real initial_norm_squared = std::real(initial_norm); + if (!std::isfinite(initial_norm_squared) || initial_norm_squared <= 0) + { + ModuleBase::WARNING_QUIT("Diago_DavSubspace", "The CPU basis has a non-positive or non-finite S norm."); + } + const int accepted = first + retained; + for (int pass = 0; pass < 2; ++pass) + { + if (accepted > 0) + { + if (this->dim == 0) + { + // The preceding collective stores global coefficients here; + // an empty rank must contribute zero again on the next pass. + const T zero_value = 0; + std::fill_n(overlaps.data(), accepted, zero_value); + } + ModuleBase::gemv_op() + ('C', this->dim, accepted, this->one, this->psi_in_iter, matrix_ld, + s_candidate.data(), 1, this->zero, overlaps.data(), 1); + } + if (this->diag_comm.nproc > 1 && accepted > 0) + { + Parallel_Reduce::reduce_pool(overlaps.data(), accepted); + } + if (accepted > 0) + { + ModuleBase::gemv_op() + ('N', this->dim, accepted, this->neg_one, this->psi_in_iter, matrix_ld, + overlaps.data(), 1, this->one, candidate.data(), 1); + } + // Reapply S to avoid estimating a tiny norm by cancellation of + // large overlap terms, including when S is not the identity. + op.spsi(candidate.data(), s_candidate.data(), this->dim, 1); + } + T final_norm = 0; + for (int row = 0; row < this->dim; ++row) + { + final_norm += dav_conjugate(candidate[row]) * s_candidate[row]; + } + if (this->diag_comm.nproc > 1) + { + Parallel_Reduce::reduce_pool(&final_norm, 1); + } + const Real final_norm_squared = std::real(final_norm); + if (!std::isfinite(final_norm_squared)) + { + ModuleBase::WARNING_QUIT("Diago_DavSubspace", "Non-finite projected CPU S norm."); + } + if (final_norm_squared <= rank_tolerance * initial_norm_squared) + { + continue; + } + const Real norm = std::sqrt(final_norm_squared); + const int target_offset = accepted * this->dim; + for (int row = 0; row < this->dim; ++row) + { + this->psi_in_iter[target_offset + row] = candidate[row] / norm; + this->spsi[target_offset + row] = s_candidate[row] / norm; + } + ++retained; + } + return retained; +} + +template +bool Diago_DavSubspace::restarted_basis_is_orthonormal_cpu() const +{ + // Ritz coefficients are S-orthonormal, so rotation normally preserves the + // basis metric. Check the actual rotated vectors before using that identity. + const int gram_size = this->n_band * this->n_band; + const int matrix_ld = std::max(this->dim, 1); + std::vector gram(gram_size); + ModuleBase::gemm_op() + ('C', 'N', this->n_band, this->n_band, this->dim, this->one, + this->psi_in_iter, matrix_ld, this->spsi, matrix_ld, this->zero, + gram.data(), this->n_band); + if (this->diag_comm.nproc > 1) + { + Parallel_Reduce::reduce_pool(gram.data(), gram_size); + } + const Real tolerance = 64 * std::numeric_limits::epsilon() * this->n_band; + for (int column = 0; column < this->n_band; ++column) + { + for (int row = 0; row < this->n_band; ++row) + { + T error = gram[column * this->n_band + row]; + if (row == column) + { + error -= 1; + } + const Real magnitude = std::abs(error); + if (!std::isfinite(magnitude) || magnitude > tolerance) + { + return false; + } + } + } + return true; +} + +template +bool Diago_DavSubspace::residuals_converged_cpu(const int nbase, + const std::vector& eigenvalues, + const std::vector& ethr_band) const +{ + // This guard is used only when rank filtering leaves no expansion. A + // repeated Ritz value alone is not convergence in that situation. + for (int band = 0; band < this->n_band; ++band) + { + Real norms[3] = {}; + for (int row = 0; row < this->dim; ++row) + { + T hx = 0; + T sx = 0; + for (int basis = 0; basis < nbase; ++basis) + { + const T coefficient = this->vcc[band * this->nbase_x + basis]; + hx += this->hpsi[basis * this->dim + row] * coefficient; + sx += this->spsi[basis * this->dim + row] * coefficient; + } + const T residual = hx - eigenvalues[band] * sx; + const T residual_squared = dav_conjugate(residual) * residual; + const T h_squared = dav_conjugate(hx) * hx; + const T s_squared = dav_conjugate(sx) * sx; + norms[0] += std::real(residual_squared); + norms[1] += std::real(h_squared); + norms[2] += std::real(s_squared); + } + if (this->diag_comm.nproc > 1) + { + Parallel_Reduce::reduce_pool(norms, 3); + } + const Real residual_norm = std::sqrt(norms[0]); + const Real operator_scale = std::sqrt(norms[1]) + std::abs(eigenvalues[band]) * std::sqrt(norms[2]); + const Real roundoff = 32 * std::numeric_limits::epsilon() * nbase * operator_scale; + const Real target = static_cast(ethr_band[band]); + const Real tolerance = std::max(target, roundoff); + if (!std::isfinite(residual_norm) || residual_norm > tolerance) + { + return false; + } + } + return true; +} + template int Diago_DavSubspace::diag_once(const HSOperator& op, T* psi_in, @@ -151,6 +414,15 @@ int Diago_DavSubspace::diag_once(const HSOperator& op, unconv[m] = m; } + if (this->device == base_device::CpuDevice) + { + const int retained = this->orthonormalize_cpu(op, 0, this->n_band); + if (retained != this->n_band) + { + ModuleBase::WARNING_QUIT("Diago_DavSubspace", "Initial CPU trial vectors are linearly dependent in S."); + } + } + // compute h*psi_in_iter // NOTE: bands after the first n_band should yield zero // hpsi[:, 0:nbase_x] = H * psi_in_iter[:, 0:nbase_x] @@ -177,7 +449,7 @@ int Diago_DavSubspace::diag_once(const HSOperator& op, { dav_iter++; - this->cal_grad(op, + const int added = this->cal_grad(op, this->dim, nbase, this->notconv, @@ -188,35 +460,49 @@ int Diago_DavSubspace::diag_once(const HSOperator& op, unconv.data(), &eigenvalue_iter); - this->cal_elem(this->dim, - nbase, - this->notconv, - this->psi_in_iter, - this->spsi, - this->hpsi, - this->hcc, - this->scc); - - this->diag_zhegvx(nbase, this->n_band, this->hcc, this->scc, this->nbase_x, &eigenvalue_iter, this->vcc); - - // check convergence and update eigenvalues - ModuleBase::timer::start("Diago_DavSubspace", "check_update"); - - this->notconv = 0; - for (int m = 0; m < this->n_band; m++) + if (added == 0) { - convflag[m] = (std::abs(eigenvalue_iter[m] - eigenvalue_in_hsolver[m]) < ethr_band[m]); - - if (!convflag[m]) + const bool converged = this->residuals_converged_cpu(nbase, eigenvalue_iter, ethr_band); + if (!converged) { - unconv[this->notconv] = m; - this->notconv++; + ModuleBase::WARNING_QUIT("Diago_DavSubspace", + "No independent CPU expansion direction remains, but the true residual is not converged."); } - - eigenvalue_in_hsolver[m] = eigenvalue_iter[m]; + this->notconv = 0; } + else + { + this->cal_elem(this->dim, + nbase, + added, + this->psi_in_iter, + this->spsi, + this->hpsi, + this->hcc, + this->scc); + + this->diag_zhegvx(nbase, this->n_band, this->hcc, this->scc, this->nbase_x, &eigenvalue_iter, this->vcc); + + // check convergence and update eigenvalues + ModuleBase::timer::start("Diago_DavSubspace", "check_update"); + + this->notconv = 0; + for (int m = 0; m < this->n_band; m++) + { + const Real eigenvalue_change = eigenvalue_iter[m] - eigenvalue_in_hsolver[m]; + convflag[m] = (std::abs(eigenvalue_change) < ethr_band[m]); - ModuleBase::timer::end("Diago_DavSubspace", "check_update"); + if (!convflag[m]) + { + unconv[this->notconv] = m; + this->notconv++; + } + + eigenvalue_in_hsolver[m] = eigenvalue_iter[m]; + } + + ModuleBase::timer::end("Diago_DavSubspace", "check_update"); + } if ((this->notconv == 0) || (nbase + this->notconv + 1 > this->nbase_x) || (dav_iter == this->iter_nmax)) { @@ -261,6 +547,9 @@ int Diago_DavSubspace::diag_once(const HSOperator& op, // update this->psi_in_iter according to psi_in syncmem_complex_2d_op()(this->psi_in_iter, this->dim, psi_in, psi_in_dmax, this->dim, this->n_band); + // Rotate cached H and S images with the same Ritz coefficients. + // This avoids reapplying H and solving an unchanged Ritz problem + // at each restart. Repair the CPU basis only if its metric drifts. this->refresh(this->dim, this->n_band, nbase, @@ -272,6 +561,25 @@ int Diago_DavSubspace::diag_once(const HSOperator& op, this->scc, this->vcc); + if (this->device == base_device::CpuDevice && !this->restarted_basis_is_orthonormal_cpu()) + { + const int retained = this->orthonormalize_cpu(op, 0, this->n_band); + if (retained != this->n_band) + { + ModuleBase::WARNING_QUIT("Diago_DavSubspace", "Restarted CPU Ritz vectors have lost S rank."); + } + op.hpsi(this->psi_in_iter, this->hpsi, this->dim, retained); + nbase = 0; + this->cal_elem(this->dim, nbase, retained, this->psi_in_iter, + this->spsi, this->hpsi, this->hcc, this->scc); + this->diag_zhegvx(nbase, this->n_band, this->hcc, this->scc, + this->nbase_x, &eigenvalue_iter, this->vcc); + for (int band = 0; band < this->n_band; ++band) + { + eigenvalue_in_hsolver[band] = eigenvalue_iter[band]; + } + } + ModuleBase::timer::end("Diago_DavSubspace", "last"); } } @@ -284,7 +592,7 @@ int Diago_DavSubspace::diag_once(const HSOperator& op, } template -void Diago_DavSubspace::cal_grad(const HSOperator& op, +int Diago_DavSubspace::cal_grad(const HSOperator& op, const int& dim, const int& nbase, const int& notconv, @@ -298,6 +606,49 @@ void Diago_DavSubspace::cal_grad(const HSOperator& op, ModuleBase::timer::start("Diago_DavSubspace", "cal_grad"); const int matrix_ld = std::max(this->dim, 1); + if (this->device == base_device::CpuDevice) + { + std::vector expansion_eigenvalues(notconv); + const int coefficient_count = nbase * notconv; + std::vector coefficients(coefficient_count); + for (int column = 0; column < notconv; ++column) + { + const int band = unconv[column]; + const Real eigenvalue = (*eigenvalue_iter)[band]; + expansion_eigenvalues[column] = eigenvalue; + for (int basis = 0; basis < nbase; ++basis) + { + coefficients[column * nbase + basis] = vcc[band * this->nbase_x + basis]; + } + } + T* residuals = psi_iter + nbase * dim; + ModuleBase::gemm_op() + ('N', 'N', dim, notconv, nbase, this->one, hpsi, matrix_ld, + coefficients.data(), nbase, this->zero, residuals, matrix_ld); + for (int column = 0; column < notconv; ++column) + { + const Real negative_eigenvalue = -expansion_eigenvalues[column]; + for (int basis = 0; basis < nbase; ++basis) + { + coefficients[column * nbase + basis] *= negative_eigenvalue; + } + } + ModuleBase::gemm_op() + ('N', 'N', dim, notconv, nbase, this->one, spsi, matrix_ld, + coefficients.data(), nbase, this->one, residuals, matrix_ld); + precondition_op()(dim, psi_iter, nbase, notconv, + this->precondition.data(), expansion_eigenvalues.data()); + const int added = this->orthonormalize_cpu(op, nbase, notconv); + if (added > 0) + { + T* new_basis = psi_iter + nbase * dim; + T* new_h_basis = hpsi + nbase * dim; + op.hpsi(new_basis, new_h_basis, dim, added); + } + ModuleBase::timer::end("Diago_DavSubspace", "cal_grad"); + return added; + } + for (size_t i = 0; i < notconv; i++) { if (unconv[i] != i) @@ -493,7 +844,7 @@ void Diago_DavSubspace::cal_grad(const HSOperator& op, op.spsi(psi_iter + nbase * dim, spsi + nbase * this->dim, this->dim, notconv); ModuleBase::timer::end("Diago_DavSubspace", "cal_grad"); - return; + return notconv; } template @@ -603,7 +954,7 @@ void Diago_DavSubspace::cal_elem(const int& dim, { #ifdef __DSP // Only on dsp hardware need an extra space to reduce data - mtfunc::dsp_dav_subspace_reduce(hcc, scc, nbase, this->nbase_x, this->notconv, this->diag_comm.comm); + mtfunc::dsp_dav_subspace_reduce(hcc, scc, nbase, this->nbase_x, notconv, this->diag_comm.comm); #else assert(this->diag_comm.comm == POOL_WORLD); Parallel_Reduce::reduce_pool(hcc + nbase * this->nbase_x, notconv * this->nbase_x); @@ -646,45 +997,111 @@ void Diago_DavSubspace::diag_zhegvx(const int& nbase, { if (this->diag_subspace == 0) { + const int matrix_size = nbase * nbase; + const int vector_size = nbase * nband; + const T scalar_zero = 0; + const Real real_zero = 0; + std::vector vectors(vector_size, scalar_zero); + std::vector eigenvalues(nbase, real_zero); + // Query INFO, solve INFO, and the actual returned eigenpair count. + int status[3] = {}; if (this->diag_comm.rank == 0) { - std::vector> h_diag(nbase, std::vector(nbase, *this->zero)); - std::vector> s_diag(nbase, std::vector(nbase, *this->zero)); - - for (size_t i = 0; i < nbase; i++) + std::vector h_diag(matrix_size); + std::vector s_diag(matrix_size); + for (int column = 0; column < nbase; ++column) { - for (size_t j = 0; j < nbase; j++) + for (int row = 0; row <= column; ++row) { - h_diag[i][j] = hcc[i * this->nbase_x + j]; - s_diag[i][j] = scc[i * this->nbase_x + j]; + const int source = column * nbase_x + row; + const int target = column * nbase + row; + h_diag[target] = hcc[source]; + s_diag[target] = scc[source]; + const Real h_magnitude = std::abs(h_diag[target]); + const Real s_magnitude = std::abs(s_diag[target]); + if (!std::isfinite(h_magnitude) || !std::isfinite(s_magnitude)) + { + status[0] = -1000; + } } } - hegvx_op()(this->ctx, - nbase, - this->nbase_x, - this->hcc, - this->scc, - nband, - (*eigenvalue_iter).data(), - this->vcc); - // reset: - for (size_t i = 0; i < nbase; i++) + const int real_work_size = 7 * nbase; + const int integer_work_size = 5 * nbase; + std::vector rwork(real_work_size); + std::vector iwork(integer_work_size); + std::vector ifail(nbase); + T work_query = 0; + if (status[0] == 0) { - for (size_t j = 0; j < nbase; j++) + dav_hegvx(nbase, nband, h_diag.data(), s_diag.data(), eigenvalues.data(), vectors.data(), + &work_query, -1, rwork.data(), iwork.data(), ifail.data(), status[2], status[0]); + } + const Real requested_work = std::real(work_query); + const Real maximum_work = static_cast(std::numeric_limits::max()); + if (status[0] == 0 + && (!std::isfinite(requested_work) || requested_work < 1 || requested_work >= maximum_work)) + { + status[0] = -1001; + } + if (status[0] == 0) + { + const int queried_work_size = static_cast(requested_work); + const int minimum_work_size = 8 * nbase; + const int work_size = std::max(queried_work_size, minimum_work_size); + std::vector work(work_size); + status[2] = 0; + dav_hegvx(nbase, nband, h_diag.data(), s_diag.data(), eigenvalues.data(), vectors.data(), + work.data(), work_size, rwork.data(), iwork.data(), ifail.data(), status[2], status[1]); + } + if (status[0] == 0 && status[1] == 0 && status[2] == nband) + { + for (int band = 0; band < nband; ++band) { - hcc[i * this->nbase_x + j] = h_diag[i][j]; - scc[i * this->nbase_x + j] = s_diag[i][j]; + if (!std::isfinite(eigenvalues[band])) + { + status[1] = -1002; + } } - - for (size_t j = nbase; j < this->nbase_x; j++) + for (int index = 0; index < vector_size; ++index) { - hcc[i * this->nbase_x + j] = *this->zero; - hcc[j * this->nbase_x + i] = *this->zero; - scc[i * this->nbase_x + j] = *this->zero; - scc[j * this->nbase_x + i] = *this->zero; + const Real magnitude = std::abs(vectors[index]); + if (!std::isfinite(magnitude)) + { + status[1] = -1002; + } } } } + if (this->diag_comm.nproc > 1) + { + // Only the pool root contributes; publish status before any + // rank decides whether the temporary Ritz result is usable. + Parallel_Reduce::reduce_pool(status, 3); + } + if (status[0] != 0 || status[1] != 0 || status[2] != nband) + { + std::ostringstream message; + message << "CPU reduced solve failed: query INFO=" << status[0] + << ", solve INFO=" << status[1] << ", M=" << status[2] + << ", requested=" << nband << ", dimension=" << nbase; + const std::string diagnostic = message.str(); + ModuleBase::WARNING_QUIT("Diago_DavSubspace", diagnostic); + } + if (this->diag_comm.nproc > 1) + { + Parallel_Reduce::reduce_pool(eigenvalues.data(), nband); + Parallel_Reduce::reduce_pool(vectors.data(), vector_size); + } + for (int band = 0; band < nband; ++band) + { + (*eigenvalue_iter)[band] = eigenvalues[band]; + for (int row = 0; row < nbase; ++row) + { + vcc[band * nbase_x + row] = vectors[band * nbase + row]; + } + } + ModuleBase::timer::end("Diago_DavSubspace", "diag_zhegvx"); + return; } else { diff --git a/source/source_hsolver/diago_dav_subspace.h b/source/source_hsolver/diago_dav_subspace.h index 4e1e97fbbbc..d00e9f4c569 100644 --- a/source/source_hsolver/diago_dav_subspace.h +++ b/source/source_hsolver/diago_dav_subspace.h @@ -95,7 +95,7 @@ class Diago_DavSubspace base_device::DEVICE_CPU* cpu_ctx = {}; base_device::AbacusDevice_t device = {}; - void cal_grad(const HSOperator& op, + int cal_grad(const HSOperator& op, const int& dim, const int& nbase, const int& notconv, @@ -106,6 +106,15 @@ class Diago_DavSubspace const int* unconv, std::vector* eigenvalue_iter); + // CPU basis construction; returns the number of independent retained columns. + int orthonormalize_cpu(const HSOperator& op, const int first, const int count); + + bool restarted_basis_is_orthonormal_cpu() const; + + bool residuals_converged_cpu(const int nbase, + const std::vector& eigenvalues, + const std::vector& ethr_band) const; + void cal_elem(const int& dim, int& nbase, const int& notconv, diff --git a/source/source_hsolver/diago_david.cpp b/source/source_hsolver/diago_david.cpp index 31e6fcfb0c2..ff460a26be0 100644 --- a/source/source_hsolver/diago_david.cpp +++ b/source/source_hsolver/diago_david.cpp @@ -10,6 +10,10 @@ #include "source_base/kernels/math_kernel_op.h" #include "source_base/parallel_comm.h" +#include +#include +#include + using namespace hsolver; @@ -154,33 +158,16 @@ int DiagoDavid::diag_once(const HSOperator& op, // orthogonalise the initial trial psi(0~nband-1) - // plan for SchmidtOrth - std::vector pre_matrix_mm_m(nband, 0); - std::vector pre_matrix_mv_m(nband, 1); - this->planSchmidtOrth(nband, pre_matrix_mm_m, pre_matrix_mv_m); - - for (int m = 0; m < nband; m++) - { - { - // phm_in->sPsi(psi_in + m*ld_psi, &this->spsi[m * dim], dim, dim, 1); - op.spsi(psi_in + m*ld_psi,&this->spsi[m*dim],dim, 1); - } - } // begin SchmidtOrth for (int m = 0; m < nband; m++) { syncmem_complex_op()(basis + dim*m, psi_in + m*ld_psi, dim); - this->SchmidtOrth(dim, - nband, - m, - this->spsi, - &this->lagrange_matrix[m * nband], - pre_matrix_mm_m[m], - pre_matrix_mv_m[m]); + T* coefficients = this->lagrange_matrix + m * nband; + const bool independent = this->SchmidtOrth(op, dim, m, this->spsi, coefficients); + if (!independent) { - // phm_in->sPsi(basis + dim*m, &this->spsi[m * dim], dim, dim, 1); - op.spsi(basis + dim*m, &this->spsi[m * dim], dim, 1); + throw std::runtime_error("DiagoDavid: initial wavefunctions are linearly dependent in the S metric"); } } @@ -207,7 +194,8 @@ int DiagoDavid::diag_once(const HSOperator& op, { dav_iter++; - this->cal_grad(op, + std::vector check_residual(nband, false); + const int added = this->cal_grad(op, dim, nbase, nbase_x, @@ -216,11 +204,15 @@ int DiagoDavid::diag_once(const HSOperator& op, this->spsi, this->vcc, unconv.data(), - this->eigenvalue); - - this->cal_elem(dim, nbase, nbase_x, this->notconv, this->hpsi, this->spsi, this->hcc); + this->eigenvalue, + ethr_band, + check_residual); - this->diag_zhegvx(nbase, nband, this->hcc, nbase_x, this->eigenvalue, this->vcc); + if (added > 0) + { + this->cal_elem(dim, nbase, nbase_x, added, this->hpsi, this->spsi, this->hcc); + this->diag_zhegvx(nbase, nband, this->hcc, nbase_x, this->eigenvalue, this->vcc); + } // check convergence and update eigenvalues ModuleBase::timer::start("DiagoDavid", "check_update"); @@ -229,6 +221,14 @@ int DiagoDavid::diag_once(const HSOperator& op, for (int m = 0; m < nband; m++) { convflag[m] = (std::abs(this->eigenvalue[m] - eigenvalue_in[m]) < ethr_band[m]); + if (check_residual[m]) + { + // A skipped or dependent direction is not evidence of convergence. + // Recheck the current Ritz state after the small solve: it may + // have changed even if its previous residual was already small. + const double residual_norm = this->ritz_residual_norm(nbase, m); + convflag[m] = std::isfinite(residual_norm) && residual_norm <= ethr_band[m]; + } if (!convflag[m]) { unconv[this->notconv] = m; @@ -238,7 +238,14 @@ int DiagoDavid::diag_once(const HSOperator& op, } ModuleBase::timer::end("DiagoDavid", "check_update"); - if (!this->notconv || (nbase + this->notconv > nbase_x) + const bool expansion_stalled = added == 0 && this->notconv > 0; + if (expansion_stalled && diag_comm.rank == 0) + { + std::cout << "DiagoDavid: no independent expansion direction remains for " + << this->notconv << " unconverged bands; retaining the residual failure for retry." + << std::endl; + } + if (!this->notconv || expansion_stalled || (nbase + this->notconv > nbase_x) || (dav_iter == david_maxiter)) { ModuleBase::timer::start("DiagoDavid", "last"); @@ -261,7 +268,7 @@ int DiagoDavid::diag_once(const HSOperator& op, psi_in, // C dim * nband psi_ld); - if (!this->notconv || (dav_iter == david_maxiter)) + if (!this->notconv || expansion_stalled || (dav_iter == david_maxiter)) { // overall convergence or last iteration: exit the iteration ModuleBase::timer::end("DiagoDavid", "last"); @@ -297,7 +304,7 @@ int DiagoDavid::diag_once(const HSOperator& op, template -void DiagoDavid::cal_grad(const HSOperator& op, +int DiagoDavid::cal_grad(const HSOperator& op, const int& dim, const int& nbase, // current dimension of the reduced basis const int nbase_x, // maximum dimension of the reduced basis set @@ -306,14 +313,16 @@ void DiagoDavid::cal_grad(const HSOperator& op, T* spsi, const T* vcc, const int* unconv, - const Real* eigenvalue) + const Real* eigenvalue, + const std::vector& ethr_band, + std::vector& check_residual) { if (test_david == 1) { ModuleBase::TITLE("DiagoDavid", "cal_grad"); } // if all eigenpairs have converged, return if (notconv == 0) { - return; + return 0; } ModuleBase::timer::start("DiagoDavid", "cal_grad"); const int matrix_ld = std::max(dim, 1); @@ -469,116 +478,100 @@ void DiagoDavid::cal_grad(const HSOperator& op, } //<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<< - // Preconditioning - // basis[nbase] = T * basis[nbase] = T * (H - lambda * S) * psi - // where T, the preconditioner, is an approximate inverse of H - // T is a diagonal stored in array `precondition` - // to do preconditioning, divide each column of basis by the corresponding element of precondition + // Keep the number of accepted directions distinct from the number of + // unconverged eigenpairs. Zero residuals and redundant block directions + // must not introduce zero columns into the reduced problem. + T* lagrange = nullptr; + const int coefficient_count = nbase + notconv; + resmem_complex_op()(lagrange, coefficient_count); + int added = 0; for (int m = 0; m < notconv; m++) { - //<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<< + const int band = unconv[m]; + T* candidate = basis + dim * (nbase + m); + const double residual2 = ModuleBase::dot_real_op()(dim, candidate, candidate); + if (!std::isfinite(residual2) || residual2 < 0.0) + { + throw std::runtime_error("DiagoDavid: non-finite residual before preconditioning"); + } + const double residual_norm = std::sqrt(residual2); + if (residual_norm <= ethr_band[band]) + { + check_residual[band] = true; + continue; + } + + const int target = nbase + added; + T* packed_candidate = basis + dim * target; + if (added != m) + { + syncmem_complex_op()(packed_candidate, candidate, dim); + } if (this->device == base_device::GpuDevice) { #if defined(__CUDA) || defined(__ROCM) ModuleBase::vector_div_vector_op()(dim, - basis + dim * (nbase + m), - basis + dim * (nbase + m), - this->d_precondition); + packed_candidate, + packed_candidate, + this->d_precondition); #endif } else { ModuleBase::vector_div_vector_op()(dim, - basis + dim * (nbase + m), - basis + dim * (nbase + m), - this->precondition); + packed_candidate, + packed_candidate, + this->precondition); } - //<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<< - // for (int ig = 0; ig < dim; ig++) - // { - // ppsi[ig] /= this->precondition[ig]; - // } - //<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<< - } - - // there is a nbase to nbase + notconv band orthogonalise - // plan for SchmidtOrth - T* lagrange = nullptr; - resmem_complex_op()(lagrange, notconv * (nbase + notconv)); - setmem_complex_op()(lagrange, 0, notconv * (nbase + notconv)); - - std::vector pre_matrix_mm_m(notconv, 0); - std::vector pre_matrix_mv_m(notconv, 1); - this->planSchmidtOrth(notconv, pre_matrix_mm_m, pre_matrix_mv_m); - for (int m = 0; m < notconv; m++) - { + const bool independent = this->SchmidtOrth(op, dim, target, spsi, lagrange); + if (independent) { - // phm_in->sPsi(basis + dim*(nbase + m), &spsi[(nbase + m) * dim], dim, dim, 1); - op.spsi(basis + dim*(nbase + m), &spsi[(nbase + m) * dim], dim, 1); + ++added; } - } - // first nbase bands psi* dot notconv bands spsi to prepare lagrange_matrix - - // calculate the square matrix for future lagranges - if (notconv == 1){ - //Use gemv for vector case to avoid potential bug using gemm call with n=1 - ModuleBase::gemv_op()('C', - dim, // m: row of A - nbase, // n: col of A - this->one, // alpha - basis, // A dim * nbase - matrix_ld, // LDA: if(N) max(1,m) - &spsi[nbase * dim], // X dim - 1, // incx - this->zero, // beta - lagrange, // Y nbase - 1 - ); - } else - { - ModuleBase::gemm_op()('C', - 'N', - nbase, // m: row of A,C - notconv, // n: col of B,C - dim, // k: col of A, row of B - this->one, // alpha - basis, // A - matrix_ld, // LDA: if(N) max(1,m) if(T) max(1,k) - &spsi[nbase * dim], // B - matrix_ld, // LDB: if(N) max(1,k) if(T) max(1,n) - this->zero, // belta - lagrange, // C - nbase + notconv // LDC: if(N) max(1, m) - ); - } - - for (int m = 0; m < notconv; m++) - { - this->SchmidtOrth(dim, - nbase + notconv, - nbase + m, - spsi, - &lagrange[m * (nbase + notconv)], - pre_matrix_mm_m[m], - pre_matrix_mv_m[m]); + else { - // phm_in->sPsi(basis + dim*(nbase + m), &spsi[(nbase + m) * dim], dim, dim, 1); - op.spsi(basis + dim*(nbase + m), &spsi[(nbase + m) * dim], dim, 1); + check_residual[band] = true; } } - // calculate H|psi> for not convergence bands - // hpsi_info dav_hpsi_in(&basis, - // psi::Range(true, 0, nbase, nbase + notconv - 1), - // &hpsi[nbase * dim]); // &hp(nbase, 0) - // phm_in->ops->hPsi(dav_hpsi_in); - // hpsi[:, nbase:nbase+notcnv] = H basis[:, nbase:nbase+notcnv] - op.hpsi(basis + nbase * dim, hpsi + nbase * dim, dim, notconv); + if (added > 0) + { + T* new_basis = basis + nbase * dim; + T* new_hpsi = hpsi + nbase * dim; + op.hpsi(new_basis, new_hpsi, dim, added); + } delmem_complex_op()(lagrange); delmem_complex_op()(vc_ev_vector); ModuleBase::timer::end("DiagoDavid", "cal_grad"); - return; + return added; +} + + +template +double DiagoDavid::ritz_residual_norm(const int nbase, const int band) const +{ + ModuleBase::timer::start("DiagoDavid", "ritz_residual_norm"); + const int matrix_ld = std::max(dim, 1); + T* residual = nullptr; + T* overlap = nullptr; + resmem_complex_op()(residual, matrix_ld); + resmem_complex_op()(overlap, matrix_ld); + const T* coefficients = vcc + band * nbase_x; + ModuleBase::gemv_op()('N', dim, nbase, this->one, hpsi, matrix_ld, + coefficients, 1, this->zero, residual, 1); + ModuleBase::gemv_op()('N', dim, nbase, this->one, spsi, matrix_ld, + coefficients, 1, this->zero, overlap, 1); + const Real unit = 1.0; + const Real negative_eigenvalue = -eigenvalue[band]; + ModuleBase::vector_add_vector_op()(dim, residual, residual, unit, + overlap, negative_eigenvalue); + const double residual2 = ModuleBase::dot_real_op()(dim, residual, residual); + delmem_complex_op()(residual); + delmem_complex_op()(overlap); + const double residual_norm = std::sqrt(residual2); + ModuleBase::timer::end("DiagoDavid", "ritz_residual_norm"); + return residual_norm; } @@ -775,6 +768,10 @@ void DiagoDavid::refresh(const int& dim, nbase = nband; setmem_complex_op()(hcc, 0, nbase_x * nbase_x); + // The restarted basis already consists of Ritz vectors. Its coefficient + // matrix is the identity, including zero off-diagonal entries. + const int coefficient_matrix_size = nbase_x * nbase_x; + setmem_complex_op()(vcc, 0, coefficient_matrix_size); if (this->device == base_device::GpuDevice) { @@ -820,139 +817,76 @@ void DiagoDavid::refresh(const int& dim, template -void DiagoDavid::SchmidtOrth(const int& dim, - const int nband, - const int m, - const T* spsi, - T* lagrange_m, - const int mm_size, - const int mv_size) +bool DiagoDavid::SchmidtOrth(const HSOperator& op, + const int& dim, + const int m, + T* spsi, + T* lagrange_m) { - // if(test_david == 1) ModuleBase::TITLE("DiagoDavid","SchmidtOrth"); ModuleBase::timer::start("DiagoDavid", "SchmidtOrth"); const int matrix_ld = std::max(dim, 1); - - // orthogonalize starting eigenfunction to those already calculated - // psi_m orthogonalize to psi(0) ~ psi(m-1) - // Attention, the orthogonalize here read as - // psi(m) -> psi(m) - \sum_{i < m} \langle psi(i)|S|psi(m) \rangle psi(i) - // so the orthogonalize is performed about S. - - // assert(basis.get_nbands() >= nband); assert(m >= 0); - assert(m < nband); - - if (dim == 0) + T* candidate = basis + dim * m; + T* s_candidate = spsi + dim * m; + op.spsi(candidate, s_candidate, dim, 1); + const Real initial_squared = ModuleBase::dot_real_op()(dim, candidate, s_candidate); + if (!std::isfinite(initial_squared) || initial_squared < 0) { - // GEMV quick-returns on zero rows without overwriting the projection. - const int projection_size = m + 1; - setmem_complex_op()(lagrange_m, 0, projection_size); + throw std::runtime_error("DiagoDavid: non-positive or non-finite overlap norm"); } - - // psi_m = basis[m] - T* psi_m = basis + dim*m; - - // std::complex *lagrange = new std::complex[m + 1]; - // ModuleBase::GlobalFunc::ZEROS(lagrange, m + 1); - - // calculate the square matrix for future lagranges - if (mm_size != 0) + if (initial_squared == 0) { - // lagrange_m[m - mv_size + 1 - mm_size] - // = basis[m - mv_size + 1 - mm_size]' * spsi[m] - ModuleBase::gemm_op()('C', - 'N', - mm_size, // m: row of A,C - mm_size, // n: col of B,C - dim, // k: col of A, row of B - this->one, // alpha - basis + dim * (m - mv_size + 1 - mm_size), // A - matrix_ld, // LDA: if(N) max(1,m) if(T) max(1,k) - &spsi[m * dim], // B - matrix_ld, // LDB: if(N) max(1,k) if(T) max(1,n) - this->zero, // belta - &lagrange_m[m - mv_size + 1 - mm_size], // C - nband // LDC: if(N) max(1, m) - ); + ModuleBase::timer::end("DiagoDavid", "SchmidtOrth"); + return false; } - // calculate other lagranges for this band - // lagrange_m[m - mv_size + 1] - // = basis[m - mv_size + 1]' * spsi[m] - ModuleBase::gemv_op()('C', - dim, - mv_size, - this->one, - basis + dim * (m - mv_size + 1), - matrix_ld, - &spsi[m * dim], - 1, - this->zero, - &lagrange_m[m - mv_size + 1], - 1); - Parallel_Reduce::reduce_pool(lagrange_m, m + 1); - - T var = *this->zero; - syncmem_d2h_op()(&var, lagrange_m + m, 1); - double psi_norm = get_real(var); - - assert(psi_norm > 0.0); - - // / psi_m = psi_m - \sum_{i < m} \langle psi(i)|S|psi(m) \rangle psi(i) - // psi_m = psi_m - basis * lagrange_m - ModuleBase::gemv_op()('N', - dim, - m, - this->neg_one, - basis, - matrix_ld, - lagrange_m, - 1, - this->one, - psi_m, - 1); - - // psi_norm = psi_norm - lagrange_m · lagrange_m - psi_norm -= ModuleBase::dot_real_op()(m, lagrange_m, lagrange_m, false); - - // for (int j = 0; j < m; j++) - // { - // const std::complex alpha = std::complex(-1, 0) * lagrange_m[j]; - // zaxpy_(&npw, &alpha, &psi(j,0), &inc, psi_m, &inc); - // /*for (int ig = 0; ig < npw; ig++) - // { - // psi_m[ig] -= lagrange[j] * psi(j, ig); - // }*/ - // psi_norm -= (conj(lagrange_m[j]) * lagrange_m[j]).real(); - // } - - assert(psi_norm > 0.0); - - psi_norm = sqrt(psi_norm); + // Normalize the input before testing its independence: the absolute size + // of a preconditioned residual does not measure the rank of its direction. + const Real initial_scale = 1 / std::sqrt(initial_squared); + ModuleBase::vector_mul_real_op()(dim, candidate, candidate, initial_scale); + ModuleBase::vector_mul_real_op()(dim, s_candidate, s_candidate, initial_scale); + for (int pass = 0; pass < 2; ++pass) + { + if (m > 0) + { + if (dim == 0) + { + // GEMV leaves its output untouched on empty local rows. + // Clear each pass before all ranks reduce the projection. + setmem_complex_op()(lagrange_m, 0, m); + } + ModuleBase::gemv_op()('C', dim, m, this->one, basis, matrix_ld, + s_candidate, 1, this->zero, lagrange_m, 1); + // Keep the collective host-side even in a device instantiation. + std::vector coefficients(m); + syncmem_d2h_op()(coefficients.data(), lagrange_m, m); + Parallel_Reduce::reduce_pool(coefficients.data(), m); + syncmem_h2d_op()(lagrange_m, coefficients.data(), m); + ModuleBase::gemv_op()('N', dim, m, this->neg_one, basis, matrix_ld, + lagrange_m, 1, this->one, candidate, 1); + } + op.spsi(candidate, s_candidate, dim, 1); + } - if (psi_norm < 1.0e-12) + // Measure the projected vector itself; subtracting nearly equal squared + // norms loses the information needed to normalize a small direction. + const Real projected_squared = ModuleBase::dot_real_op()(dim, candidate, s_candidate); + if (!std::isfinite(projected_squared)) { - std::cout << "DiagoDavid::SchmidtOrth:aborted for psi_norm <1.0e-12" << std::endl; - std::cout << "This may be due to npwx < nbands: the number of plane waves is less than" << std::endl; - std::cout << "the number of bands, leading to a rank-deficient problem." << std::endl; - std::cout << "Please increase ecutwfc or reduce nbands." << std::endl; - std::cout << "nband = " << nband << std::endl; - std::cout << "m = " << m << std::endl; - exit(0); + throw std::runtime_error("DiagoDavid: non-finite projected overlap norm"); } - else + const Real rank_tolerance = 64 * std::numeric_limits::epsilon(); + const Real rank_tolerance_squared = rank_tolerance * rank_tolerance; + if (projected_squared <= rank_tolerance_squared) { - // psi_m = psi_m / psi_norm - ModuleBase::vector_mul_real_op()(dim, psi_m, psi_m, Real(1.0 / psi_norm)); - // for (int i = 0; i < npw; i++) - // { - // psi_m[i] /= psi_norm; - // } + ModuleBase::timer::end("DiagoDavid", "SchmidtOrth"); + return false; } - - // delete[] lagrange; + const Real projected_scale = 1 / std::sqrt(projected_squared); + ModuleBase::vector_mul_real_op()(dim, candidate, candidate, projected_scale); + ModuleBase::vector_mul_real_op()(dim, s_candidate, s_candidate, projected_scale); ModuleBase::timer::end("DiagoDavid", "SchmidtOrth"); - return; + return true; } diff --git a/source/source_hsolver/diago_david.h b/source/source_hsolver/diago_david.h index 5dd29805378..ce6e4244778 100644 --- a/source/source_hsolver/diago_david.h +++ b/source/source_hsolver/diago_david.h @@ -158,8 +158,11 @@ class DiagoDavid * @param vcc The input array for the eigenvector coefficients. * @param unconv The array of indices for the unconverged eigenpairs. * @param eigenvalue The array of eigenvalues. + * @param ethr_band The requested per-band convergence thresholds. + * @param check_residual Bands whose omitted directions require a fresh residual check. + * @return Number of independent expansion vectors actually added. */ - void cal_grad(const HSOperator& op, + int cal_grad(const HSOperator& op, const int& dim, const int& nbase, const int nbase_x, @@ -168,7 +171,11 @@ class DiagoDavid T* spsi, const T* vcc, const int* unconv, - const Real* eigenvalue); + const Real* eigenvalue, + const std::vector& ethr_band, + std::vector& check_residual); + + double ritz_residual_norm(const int nbase, const int band) const; /** * Calculates the elements of the diagonalization matrix for the DiagoDavid class. @@ -219,24 +226,20 @@ class DiagoDavid /** * SchmidtOrth function performs orthogonalization of the starting eigenfunction to those already calculated. - * It takes the dimension of the basis, number of bands, index of the current band, starting eigenfunction psi_m, - * lagrange_m array, mm_size, and mv_size as input parameters. + * Scales the candidate first, then performs two projections with fresh S images. * + * @param op Hamiltonian and overlap operations. * @param dim The dimension of the basis. - * @param nband The number of bands. * @param m The index of the current band. - * @param spsi Pointer to the starting eigenfunction psi_m. + * @param spsi Overlap images, updated for the accepted candidate. * @param lagrange_m Pointer to the lagrange_m array. - * @param mm_size The size of the square matrix for future lagranges. - * @param mv_size The size of the lagrange_m array. + * @return False if the candidate has no numerically independent direction. */ - void SchmidtOrth(const int& dim, - const int nband, + bool SchmidtOrth(const HSOperator& op, + const int& dim, const int m, - const T* spsi, - T* lagrange_m, - const int mm_size, - const int mv_size); + T* spsi, + T* lagrange_m); /** * @brief Plans the Schmidt orthogonalization for a given number of bands. diff --git a/source/source_hsolver/diago_params.cpp b/source/source_hsolver/diago_params.cpp index 28e1040a974..471e61494a0 100644 --- a/source/source_hsolver/diago_params.cpp +++ b/source/source_hsolver/diago_params.cpp @@ -15,11 +15,8 @@ void setup_diago_params_pw(const int istep, DiagoIterAssist::need_subspace = ((istep == 0 || istep == 1) && iter == 1) ? false : true; DiagoIterAssist::SCF_ITER = iter; DiagoIterAssist::PW_DIAG_THR = ethr; - - if (inp.calculation != "nscf") - { - DiagoIterAssist::PW_DIAG_NMAX = inp.pw_diag_nmax; - } + // NSCF must use this calculation's input, not the initial or previous-run limit. + DiagoIterAssist::PW_DIAG_NMAX = inp.pw_diag_nmax; } template diff --git a/source/source_hsolver/diago_ppcg.cpp b/source/source_hsolver/diago_ppcg.cpp new file mode 100644 index 00000000000..da4578765f1 --- /dev/null +++ b/source/source_hsolver/diago_ppcg.cpp @@ -0,0 +1,2380 @@ +#include "diago_ppcg.h" +#include "source_base/parallel_reduce.h" +#include "source_base/timer.h" +#include +#include +#include +#include +#include +#include "source_base/kernels/math_kernel_op.h" +#include +#include +#include +#include +#include +#include + +namespace hsolver { +namespace { + +const int ppcg_openmp_work_threshold = 4096; +const int ppcg_openmp_column_threshold = 16; +const double ppcg_minimum_diagonalization_threshold = 1.0e-14; +const double ppcg_preconditioner_threshold = 1.0e-12; +const double ppcg_numerical_threshold = 1.0e-30; +const double ppcg_scaling_threshold = 1.0e-15; +const int ppcg_stagnation_restart_interval = 15; +// Operator images are transformed with their vectors. Refresh their accumulated +// roundoff independently of the user-selected Ritz diagonalization interval. +const int ppcg_operator_refresh_interval = 16; +// Relative cutoff for overlap eigenvalues in a local [X, W, P] subspace. +// Directions below this level are numerically dependent; retaining them can +// make the reduced generalized solve stagnate instead of improving a Ritz pair. +const double ppcg_subspace_rank_threshold = 1.0e-8; + +// Diagonal shifts used by the small projected eigensolve fallback. +const double ppcg_subspace_shifts[] = {0.0, 1.0e-10, 1.0e-8, 1.0e-6}; + +// A whole-function timer also closes correctly on early return or exception. +class PPCGFunctionTimer +{ +public: + explicit PPCGFunctionTimer(const char* function) : function_(function) + { + ModuleBase::timer::start("DiagoPPCG", function_); + } + + ~PPCGFunctionTimer() + { + ModuleBase::timer::end("DiagoPPCG", function_); + } + +private: + const char* function_; +}; + +} // namespace +} // namespace hsolver + + + +namespace hsolver { +namespace { + +template +void reduce_pool_values(Value& value) +{ + Parallel_Reduce::reduce_pool(value); +} + +template +void reduce_pool_values(Value* value, const int n) +{ + Parallel_Reduce::reduce_pool(value, n); +} + +bool all_pool_operations_succeeded(const bool local_success) +{ + int failure_count = local_success ? 0 : 1; + reduce_pool_values(&failure_count, 1); + return failure_count == 0; +} + +template +Real max_generalized_residual( + const T* hpsi, + const T* spsi, + const Real* eigenvalue, + int ld, + int n_dim, + int ncol) +{ + Real max_res = 0; + std::vector nrm2_all(ncol, 0.0); +#ifdef _OPENMP +#pragma omp parallel for schedule(static) if (n_dim * ncol > ppcg_openmp_work_threshold) +#endif + for (int j = 0; j < ncol; ++j) + { + double nrm2 = 0.0; + for (int ig = 0; ig < n_dim; ++ig) + { + const T r = hpsi[ig + j * ld] - T(eigenvalue[j]) * spsi[ig + j * ld]; + nrm2 += double(std::norm(r)); + } + nrm2_all[j] = nrm2; + } + reduce_pool_values(nrm2_all.data(), ncol); + for (int j = 0; j < ncol; ++j) + { + const Real residual_norm = std::sqrt(Real(nrm2_all[j])); + if (!std::isfinite(residual_norm)) + { + return std::numeric_limits::infinity(); + } + max_res = std::max(max_res, residual_norm); + } + return max_res; +} + +template +inline void set_zero(std::vector& x) +{ + const int n = int(x.size()); +#ifdef _OPENMP +#pragma omp parallel for schedule(static) if (n > ppcg_openmp_work_threshold) +#endif + for (int i = 0; i < n; ++i) + { + x[i] = T(0); + } +} + +} // anonymous namespace +} // namespace hsolver + + + +namespace hsolver { +namespace { + +template +struct HermitianLapack +{ + using Real = typename container::GetTypeReal::type; + using Device = container::DEVICE_CPU; + + static void sygvd(int n, Scalar* a, Scalar* b, Real* w) + { + std::vector eigenvectors(n * n); + container::kernels::lapack_hegvd()( + n, n, a, b, w, eigenvectors.data()); + std::copy(eigenvectors.begin(), eigenvectors.end(), a); + } +}; + +} // anonymous namespace +} // namespace hsolver + + + +namespace hsolver { +namespace { + +inline bool ppcg_contiguous_cols(const std::vector& cols, int& first) +{ + if (cols.empty()) + { + return false; + } + + first = cols.front(); + for (int j = 0; j < int(cols.size()); ++j) + { + if (cols[j] != first + j) + { + return false; + } + } + return true; +} + +} // anonymous namespace + +bool ppcg_workspace_dimensions_valid(const int ld_psi, + const int nband, + const int dim, + const int sbsize) +{ + if (ld_psi <= 0 || nband <= 0 || dim <= 0 || dim > ld_psi) + { + return false; + } + const int limit = std::numeric_limits::max(); + const int bounded_sbsize = std::min(nband, sbsize); + const int block_size = std::max(1, bounded_sbsize); + if (nband > limit / ld_psi || block_size > limit / 3) + { + return false; + } + // The largest local basis holds [X, W, P]. The global startup correction + // holds up to 2*nband columns only in its projected matrices. + const int local_dim = 3 * block_size; + if (local_dim > limit / ld_psi) + { + return false; + } + int projected_dim = local_dim; + if (block_size < nband) + { + if (nband > limit / 2) + { + return false; + } + const int startup_dim = 2 * nband; + projected_dim = std::max(projected_dim, startup_dim); + } + if (projected_dim > limit / projected_dim) + { + return false; + } + const int matrix_size = projected_dim * projected_dim; + // lapack_heevd/hegvd additionally form 1+6*d+2*d*d in int arithmetic. + // The square bound above also makes the linear term safe to evaluate. + const int linear_workspace = 1 + 6 * projected_dim; + return matrix_size <= (limit - linear_workspace) / 2; +} + +// ============================================================================= +// Constructor +// ============================================================================= +template +DiagoPPCG::DiagoPPCG(const Real& diag_thr, + const int& diag_iter_max, + const int& sbsize, + const int& rr_step) + : maxiter_(diag_iter_max), + sbsize_(sbsize), + rr_step_(rr_step), + diag_thr_(diag_thr) +{ + sbsize_ = std::max(1, sbsize_); + rr_step_ = std::max(1, rr_step_); + diag_thr_ = std::max(diag_thr_, Real(ppcg_minimum_diagonalization_threshold)); +} + +// ============================================================================= +// Input validation +// ============================================================================= +template +void DiagoPPCG::validate_input( + const HPsiFunc& hpsi_func, + const T* psi_in, + const Real* eigenvalue_in, + const std::vector& ethr_band, + const Real* prec) const +{ + std::string local_error; + try + { + if (!hpsi_func) + { + throw std::invalid_argument("PPCG: H operator is empty."); + } + if (psi_in == nullptr || eigenvalue_in == nullptr) + { + throw std::invalid_argument("PPCG: psi/eigenvalue pointer is null."); + } + if (prec == nullptr) + { + throw std::invalid_argument("PPCG: preconditioner pointer is null."); + } + if (ld_psi_ <= 0 || n_band_ <= 0 || n_dim_ <= 0) + { + throw std::invalid_argument("PPCG: invalid dimensions."); + } + if (n_dim_ > ld_psi_) + { + throw std::invalid_argument("PPCG: dim must not exceed ld_psi."); + } + if (!ppcg_workspace_dimensions_valid(ld_psi_, n_band_, n_dim_, sbsize_)) + { + throw std::invalid_argument("PPCG: workspace dimensions exceed integer indexing or LAPACK limits."); + } + if (ethr_band.size() < size_t(n_band_)) + { + throw std::invalid_argument("PPCG: ethr_band size is smaller than nband."); + } + for (int i = 0; i < n_band_; ++i) + { + if (!std::isfinite(ethr_band[i])) + { + throw std::invalid_argument("PPCG: ethr_band contains non-finite value."); + } + } + for (int i = 0; i < n_dim_; ++i) + { + if (!std::isfinite(prec[i])) + { + throw std::invalid_argument("PPCG: preconditioner contains non-finite value."); + } + } + } + catch (const std::invalid_argument& error) + { + local_error = error.what(); + } + // All pool ranks must reject the call before any rank enters an operator + // callback or Gram-matrix collective, even if only one input is invalid. + const bool local_success = local_error.empty(); + if (!all_pool_operations_succeeded(local_success)) + { + if (local_error.empty()) + { + local_error = "PPCG: invalid input on another rank in the MPI pool."; + } + throw std::invalid_argument(local_error); + } +} + +// ============================================================================= +// Operator application +// ============================================================================= +template +void DiagoPPCG::apply_h(const HPsiFunc& hpsi_func, + T* psi_in, T* hpsi_out, + int ncol) const +{ + const PPCGFunctionTimer timer("apply_h"); + hpsi_func(psi_in, hpsi_out, ld_psi_, ncol); +} + +template +void DiagoPPCG::apply_s(const SPsiFunc& spsi_func, + T* psi_in, T* spsi_out, + int ncol) const +{ + const PPCGFunctionTimer timer("apply_s"); + if (spsi_func) + { + spsi_func(psi_in, spsi_out, ld_psi_, ncol); + } + else + { +#ifdef _OPENMP +#pragma omp parallel for schedule(static) if (ld_psi_ * ncol > ppcg_openmp_work_threshold) +#endif + for (int j = 0; j < ncol; ++j) + { + std::copy(psi_in + j * ld_psi_, psi_in + (j + 1) * ld_psi_, + spsi_out + j * ld_psi_); + } + } +} + +template +void DiagoPPCG::apply_s_current(T* psi_in, T* spsi_out, + int ncol) const +{ + apply_s(spsi_func_, psi_in, spsi_out, ncol); +} + +// ============================================================================= +// Inner product (real part only, for Hermitian operators) +// ============================================================================= +template +typename DiagoPPCG::Real +DiagoPPCG::gamma_dot(const T* x, const T* y) const +{ + Real result = ModuleBase::dot_real_op()(n_dim_, x, y, false); + reduce_pool_values(result); + return result; +} + +// ============================================================================= +// Gram matrix: out[i, j] = +// ============================================================================= +template +void DiagoPPCG::gram(const T* mat_a, const T* mat_b, + int ncol_a, int ncol_b, + std::vector& out, + int ld_out) const +{ + out.resize(ld_out * ncol_b); + const T one = T(1); + const T zero = T(0); + ModuleBase::gemm_op()('C', + 'N', + ncol_a, + ncol_b, + n_dim_, + &one, + mat_a, + ld_psi_, + mat_b, + ld_psi_, + &zero, + out.data(), + ld_out); + reduce_pool_values(out.data(), ld_out * ncol_b); +} + +// ============================================================================= +// Column gather: extract selected columns into contiguous storage +// ============================================================================= +template +void DiagoPPCG::copy_cols(const T* src, + const std::vector& cols, + std::vector& dst) const +{ + const int ncols = int(cols.size()); + dst.resize(ld_psi_ * ncols); + if (ncols == 0) + { + return; + } + + int first = 0; + if (ppcg_contiguous_cols(cols, first)) + { + std::copy(src + first * ld_psi_, + src + (first + ncols) * ld_psi_, + dst.begin()); + return; + } + +#ifdef _OPENMP +#pragma omp parallel for schedule(static) if (ld_psi_ * ncols > ppcg_openmp_work_threshold) +#endif + for (int j = 0; j < ncols; ++j) + { + const int c = cols[j]; + std::copy(src + c * ld_psi_, src + c * ld_psi_ + ld_psi_, + dst.begin() + j * ld_psi_); + } +} + +// ============================================================================= +// Column scatter: write contiguous storage back into selected columns +// ============================================================================= +template +void DiagoPPCG::scatter_cols( + T* dst, + const std::vector& cols, + const std::vector& src) const +{ + const int ncols = int(cols.size()); + if (ncols == 0) + { + return; + } + + int first = 0; + if (ppcg_contiguous_cols(cols, first)) + { + std::copy(src.begin(), + src.begin() + ld_psi_ * ncols, + dst + first * ld_psi_); + return; + } + +#ifdef _OPENMP +#pragma omp parallel for schedule(static) if (ld_psi_ * ncols > ppcg_openmp_work_threshold) +#endif + for (int j = 0; j < ncols; ++j) + { + const int c = cols[j]; + std::copy(src.begin() + j * ld_psi_, + src.begin() + (j + 1) * ld_psi_, + dst + c * ld_psi_); + } +} + +// ============================================================================= +// Project x onto vectors orthogonal to S-orthonormal basis +// ============================================================================= +template +void DiagoPPCG::project_against( + const T* basis, const T* sbasis, const T* hbasis, + const std::vector& basis_cols, + std::vector& x, std::vector& sx, std::vector* hx, + const std::vector& x_cols) const +{ + const PPCGFunctionTimer timer("project_against"); + if (basis_cols.empty() || x_cols.empty()) + { + return; + } + + const int nbasis = int(basis_cols.size()); + const int nx = int(x_cols.size()); + + int x_first = 0; + const bool contiguous_x = ppcg_contiguous_cols(x_cols, x_first); + + std::vector x_l; + std::vector sx_l; + std::vector hx_l; + T* x_data = x.data() + x_first * ld_psi_; + T* sx_data = sx.data() + x_first * ld_psi_; + T* hx_data = hx == nullptr ? nullptr : hx->data() + x_first * ld_psi_; + if (!contiguous_x) + { + x_l.reserve(ld_psi_ * nx); + sx_l.reserve(ld_psi_ * nx); + copy_cols(x.data(), x_cols, x_l); + copy_cols(sx.data(), x_cols, sx_l); + x_data = x_l.data(); + sx_data = sx_l.data(); + if (hx != nullptr) + { + copy_cols(hx->data(), x_cols, hx_l); + hx_data = hx_l.data(); + } + } + + int basis_first = 0; + const bool contiguous_basis = + ppcg_contiguous_cols(basis_cols, basis_first); + + std::vector basis_l; + std::vector sbasis_l; + std::vector hbasis_l; + const T* basis_data = basis + basis_first * ld_psi_; + const T* sbasis_data = sbasis + basis_first * ld_psi_; + const T* hbasis_data = hbasis == nullptr ? nullptr : hbasis + basis_first * ld_psi_; + if (!contiguous_basis) + { + basis_l.reserve(ld_psi_ * nbasis); + sbasis_l.reserve(ld_psi_ * nbasis); + copy_cols(basis, basis_cols, basis_l); + copy_cols(sbasis, basis_cols, sbasis_l); + basis_data = basis_l.data(); + sbasis_data = sbasis_l.data(); + if (hx != nullptr) + { + copy_cols(hbasis, basis_cols, hbasis_l); + hbasis_data = hbasis_l.data(); + } + } + + std::vector coeff(nbasis * nx, T(0)); + gram(basis_data, sx_data, nbasis, nx, coeff, nbasis); + + const T minus_one = T(-1); + const T one = T(1); + ModuleBase::gemm_op()('N', + 'N', + n_dim_, + nx, + nbasis, + &minus_one, + basis_data, + ld_psi_, + coeff.data(), + nbasis, + &one, + x_data, + ld_psi_); + ModuleBase::gemm_op()('N', + 'N', + n_dim_, + nx, + nbasis, + &minus_one, + sbasis_data, + ld_psi_, + coeff.data(), + nbasis, + &one, + sx_data, + ld_psi_); + if (hx != nullptr) + { + // H(x-X*C) = Hx-HX*C; the same coefficients update all three images. + ModuleBase::gemm_op()('N', + 'N', + n_dim_, + nx, + nbasis, + &minus_one, + hbasis_data, + ld_psi_, + coeff.data(), + nbasis, + &one, + hx_data, + ld_psi_); + } + + if (!contiguous_x) + { + scatter_cols(x.data(), x_cols, x_l); + scatter_cols(sx.data(), x_cols, sx_l); + if (hx != nullptr) + { + scatter_cols(hx->data(), x_cols, hx_l); + } + } +} + +// ============================================================================= +// Project P and refresh its operator images when subtraction loses precision. +// ============================================================================= +template +void DiagoPPCG::project_search_directions( + const HPsiFunc& hpsi_func, + const T* psi, + const std::vector& all_cols, + const std::vector& active_cols) +{ + const PPCGFunctionTimer timer("project_search_directions"); + const int nactive = int(active_cols.size()); + std::vector norms(2 * nactive, 0.0); + for (int j = 0; j < nactive; ++j) + { + const int column = active_cols[j]; + for (int i = 0; i < n_dim_; ++i) + { + norms[j] += std::norm(p_[idx(i, column, ld_psi_)]); + } + } + project_against(psi, spsi_.data(), hpsi_.data(), all_cols, p_, sp_, &hp_, active_cols); + for (int j = 0; j < nactive; ++j) + { + const int column = active_cols[j]; + for (int i = 0; i < n_dim_; ++i) + { + norms[nactive + j] += std::norm(p_[idx(i, column, ld_psi_)]); + } + } + const int norm_count = int(norms.size()); + reduce_pool_values(norms.data(), norm_count); + std::vector refresh_cols; + for (int j = 0; j < nactive; ++j) + { + if (norms[j] == 0.0 || norms[nactive + j] == 0.0) + { + const int offset = active_cols[j] * ld_psi_; + std::fill(p_.begin() + offset, p_.begin() + offset + ld_psi_, T(0)); + std::fill(sp_.begin() + offset, sp_.begin() + offset + ld_psi_, T(0)); + std::fill(hp_.begin() + offset, hp_.begin() + offset + ld_psi_, T(0)); + continue; + } + // A strongly reduced P suffers cancellation in HP-HX*C and SP-SX*C. + // Normalizing that tiny P would amplify the absolute image error. + // Refresh only these columns; reduction gives every rank one selection. + if (norms[nactive + j] < 0.5 * norms[j]) + { + refresh_cols.push_back(active_cols[j]); + } + } + if (!refresh_cols.empty()) + { + std::vector projected; + copy_cols(p_.data(), refresh_cols, projected); + std::vector hprojected(projected.size(), T(0)); + std::vector sprojected(projected.size(), T(0)); + const int refresh_count = int(refresh_cols.size()); + apply_h(hpsi_func, projected.data(), hprojected.data(), refresh_count); + apply_s_current(projected.data(), sprojected.data(), refresh_count); + scatter_cols(hp_.data(), refresh_cols, hprojected); + scatter_cols(sp_.data(), refresh_cols, sprojected); + } +} + +// ============================================================================= +// Preconditioner: x[c] /= max(prec, eps) for each active column c +// ============================================================================= +template +void DiagoPPCG::divide_by_preconditioner( + const std::vector& active_cols, + const Real* prec, + std::vector& x) const +{ + const int ncols = int(active_cols.size()); +#ifdef _OPENMP +#pragma omp parallel for schedule(static) if (n_dim_ * ncols > ppcg_openmp_work_threshold) +#endif + for (int j = 0; j < ncols; ++j) + { + const int c = active_cols[j]; + for (int ig = 0; ig < n_dim_; ++ig) + { + x[idx(ig, c, ld_psi_)] /= + std::max(prec[ig], Real(ppcg_preconditioner_threshold)); + } + } +} + +} // namespace hsolver + + +namespace hsolver { + +//============================================================================== +// BLOCK_SUBSPACE STRATEGY +//============================================================================== + +// --------------------------------------------------------------------------- +// Lock only a contiguous prefix of converged eigenpairs. A small residual +// proves that a vector is an eigenpair, but it does not prove that it belongs +// to the requested lowest roots. Keeping every band after the first +// unconverged root active lets an evolving lower root replace a later exact +// high-energy state. +// --------------------------------------------------------------------------- +template +void DiagoPPCG::lock_epairs( + const std::vector& residual, + const std::vector& ethr_band, + std::vector& active_cols) +{ + active_cols.clear(); + active_cols.reserve(n_band_); + std::vector nrm2_all(n_band_, 0.0); +#ifdef _OPENMP +#pragma omp parallel for schedule(static) if (n_dim_ * n_band_ > ppcg_openmp_work_threshold) +#endif + for (int j = 0; j < n_band_; ++j) + { + double nrm2 = 0.0; + for (int ig = 0; ig < n_dim_; ++ig) + { + nrm2 += double(std::norm(residual[idx(ig, j, ld_psi_)])); + } + nrm2_all[j] = nrm2; + } + reduce_pool_values(nrm2_all.data(), n_band_); + bool found_unconverged = false; + max_residual_ = Real(0); + for (int j = 0; j < n_band_; ++j) + { + const Real rnrm = std::sqrt(std::max(Real(nrm2_all[j]), Real(0))); + const Real thr = std::max(Real(ethr_band[j]), diag_thr_); + if (std::isfinite(rnrm)) + { + max_residual_ = std::max(max_residual_, rnrm); + } + else + { + max_residual_ = std::numeric_limits::infinity(); + } + if (!std::isfinite(rnrm) || rnrm > thr) + { + found_unconverged = true; + } + if (found_unconverged) + { + active_cols.push_back(j); + } + } +} + +template +void DiagoPPCG::scale_to_unit_snorm(std::vector& x, + std::vector& sx, + std::vector& hx, + const int ncols) const +{ + std::vector inverse_snorm(ncols, 0.0); +#ifdef _OPENMP +#pragma omp parallel for schedule(static) if (n_dim_ * ncols > ppcg_openmp_work_threshold) +#endif + for (int j = 0; j < ncols; ++j) + { + double snorm_squared = 0.0; + for (int ig = 0; ig < n_dim_; ++ig) + { + snorm_squared += double(std::real( + std::conj(x[idx(ig, j, ld_psi_)]) + * sx[idx(ig, j, ld_psi_)])); + } + inverse_snorm[j] = snorm_squared; + } + reduce_pool_values(inverse_snorm.data(), ncols); + for (int j = 0; j < ncols; ++j) + { + const Real snorm = std::sqrt(std::max( + Real(inverse_snorm[j]), Real(ppcg_numerical_threshold))); + inverse_snorm[j] = snorm > Real(ppcg_scaling_threshold) + ? double(Real(1) / snorm) + : 1.0; + } +#ifdef _OPENMP +#pragma omp parallel for collapse(2) schedule(static) if (n_dim_ * ncols > ppcg_openmp_work_threshold) +#endif + for (int j = 0; j < ncols; ++j) + { + for (int ig = 0; ig < n_dim_; ++ig) + { + const Real scale = Real(inverse_snorm[j]); + x[idx(ig, j, ld_psi_)] *= scale; + sx[idx(ig, j, ld_psi_)] *= scale; + hx[idx(ig, j, ld_psi_)] *= scale; + } + } +} + +template +void DiagoPPCG::hermitize_projected(std::vector& matrix, + const int dim) const +{ + for (int j = 0; j < dim; ++j) + { + matrix[j + j * dim] = T(std::real(matrix[j + j * dim]), 0); + for (int i = j + 1; i < dim; ++i) + { + const T value = (matrix[i + j * dim] + + std::conj(matrix[j + i * dim])) + * Real(0.5); + matrix[i + j * dim] = value; + matrix[j + i * dim] = std::conj(value); + } + } +} + +template +void DiagoPPCG::insert_startup_gram_block( + const T* left, + const T* right, + const int row_count, + const int column_count, + const int row_offset, + const int column_offset, + const int matrix_dim, + std::vector& matrix, + std::vector& workspace) const +{ + gram(left, right, row_count, column_count, workspace, row_count); + for (int j = 0; j < column_count; ++j) + { + for (int i = 0; i < row_count; ++i) + { + const int matrix_index = row_offset + i + + (column_offset + j) * matrix_dim; + matrix[matrix_index] = workspace[i + j * row_count]; + } + } +} + +template +void DiagoPPCG::build_startup_global_subspace( + const T* psi, + std::vector& w_active, + std::vector& sw_active, + std::vector& hw_active, + const int active_count, + SmallSubspace& subspace) +{ + scale_to_unit_snorm(w_active, sw_active, hw_active, active_count); + + const int state_count = n_band_; + const int correction_offset = state_count; + const int local_dim = state_count + active_count; + const int matrix_size = local_dim * local_dim; + subspace.k.assign(matrix_size, T(0)); + subspace.m.assign(matrix_size, T(0)); + subspace.eval.resize(local_dim); + std::vector gram_workspace; + + insert_startup_gram_block(psi, + hpsi_.data(), + state_count, + state_count, + 0, + 0, + local_dim, + subspace.k, + gram_workspace); + insert_startup_gram_block(psi, + hw_active.data(), + state_count, + active_count, + 0, + correction_offset, + local_dim, + subspace.k, + gram_workspace); + insert_startup_gram_block(w_active.data(), + hpsi_.data(), + active_count, + state_count, + correction_offset, + 0, + local_dim, + subspace.k, + gram_workspace); + insert_startup_gram_block(w_active.data(), + hw_active.data(), + active_count, + active_count, + correction_offset, + correction_offset, + local_dim, + subspace.k, + gram_workspace); + insert_startup_gram_block(psi, + spsi_.data(), + state_count, + state_count, + 0, + 0, + local_dim, + subspace.m, + gram_workspace); + insert_startup_gram_block(psi, + sw_active.data(), + state_count, + active_count, + 0, + correction_offset, + local_dim, + subspace.m, + gram_workspace); + insert_startup_gram_block(w_active.data(), + spsi_.data(), + active_count, + state_count, + correction_offset, + 0, + local_dim, + subspace.m, + gram_workspace); + insert_startup_gram_block(w_active.data(), + sw_active.data(), + active_count, + active_count, + correction_offset, + correction_offset, + local_dim, + subspace.m, + gram_workspace); + + hermitize_projected(subspace.k, local_dim); + hermitize_projected(subspace.m, local_dim); +} + +template +void DiagoPPCG::update_startup_global_subspace( + T* psi, + const std::vector& w_active, + const std::vector& sw_active, + const std::vector& hw_active, + const int active_count, + const SmallSubspace& subspace) +{ + const int local_dim = n_band_ + active_count; + const T* state_coefficients = subspace.k.data(); + const T* correction_coefficients = state_coefficients + n_band_; + const T one = T(1); + const T zero = T(0); + + set_zero(w_); + set_zero(sw_); + set_zero(hw_); + set_zero(p_); + set_zero(sp_); + set_zero(hp_); + + auto combine_state = [&](const T* state, + const T* correction, + T* output) + { + ModuleBase::gemm_op()('N', + 'N', + n_dim_, + n_band_, + n_band_, + &one, + state, + ld_psi_, + state_coefficients, + local_dim, + &zero, + output, + ld_psi_); + ModuleBase::gemm_op()('N', + 'N', + n_dim_, + n_band_, + active_count, + &one, + correction, + ld_psi_, + correction_coefficients, + local_dim, + &one, + output, + ld_psi_); + }; + auto combine_direction = [&](const T* correction, T* output) + { + ModuleBase::gemm_op()('N', + 'N', + n_dim_, + n_band_, + active_count, + &one, + correction, + ld_psi_, + correction_coefficients, + local_dim, + &zero, + output, + ld_psi_); + }; + + combine_state(psi, w_active.data(), w_.data()); + combine_state(spsi_.data(), sw_active.data(), sw_.data()); + combine_state(hpsi_.data(), hw_active.data(), hw_.data()); + combine_direction(w_active.data(), p_.data()); + combine_direction(sw_active.data(), sp_.data()); + combine_direction(hw_active.data(), hp_.data()); + + std::copy(w_.begin(), w_.end(), psi); + std::copy(sw_.begin(), sw_.end(), spsi_.begin()); + std::copy(hw_.begin(), hw_.end(), hpsi_.begin()); +} + +// --------------------------------------------------------------------------- +// Build K = V^H H V and M = V^H S V where V = [psi, w] +// --------------------------------------------------------------------------- +template +void DiagoPPCG::build_small_subspace( + const T* psi, + const std::vector& cols, + int nblk, + SmallSubspace& subspace) const +{ + const PPCGFunctionTimer timer("build_small_subspace"); + const int l = int(cols.size()); + const int dim = nblk * l; + subspace.k.resize(dim * dim); + subspace.m.resize(dim * dim); + subspace.eval.resize(dim); + + copy_cols(psi, cols, subspace.psi_l); + copy_cols(spsi_.data(), cols, subspace.spsi_l); + copy_cols(hpsi_.data(), cols, subspace.hpsi_l); + copy_cols(w_.data(), cols, subspace.w_l); + copy_cols(sw_.data(), cols, subspace.sw_l); + copy_cols(hw_.data(), cols, subspace.hw_l); + if (nblk >= 3) + { + copy_cols(p_.data(), cols, subspace.p_l); + copy_cols(sp_.data(), cols, subspace.sp_l); + copy_cols(hp_.data(), cols, subspace.hp_l); + } + + // --------------------------------------------------------------------------- + // Normalize w columns to unit S-norm for numerical stability. + // + // The w block of the Gram matrix M has entries O(||w||^2) which become + // tiny when residuals are small, making M nearly singular and causing + // sygvd to produce garbage eigenvectors. + // + // Scaling to unit S-norm keeps M well-conditioned (diagonal ~1) without + // changing the subspace. The same scaled basis is reused in update_one_block. + // --------------------------------------------------------------------------- + scale_to_unit_snorm(subspace.w_l, + subspace.sw_l, + subspace.hw_l, + l); + if (nblk >= 3) + { + scale_to_unit_snorm(subspace.p_l, + subspace.sp_l, + subspace.hp_l, + l); + } + + auto copy_block = [&](const std::vector& src, + const int col0, + std::vector& dst) + { +#ifdef _OPENMP +#pragma omp parallel for schedule(static) if (ld_psi_ * l > ppcg_openmp_work_threshold) +#endif + for (int j = 0; j < l; ++j) + { + std::copy(src.begin() + j * ld_psi_, + src.begin() + (j + 1) * ld_psi_, + dst.begin() + (col0 + j) * ld_psi_); + } + }; + + subspace.basis.resize(ld_psi_ * dim); + subspace.hbasis.resize(ld_psi_ * dim); + subspace.sbasis.resize(ld_psi_ * dim); + copy_block(subspace.psi_l, 0, subspace.basis); + copy_block(subspace.hpsi_l, 0, subspace.hbasis); + copy_block(subspace.spsi_l, 0, subspace.sbasis); + copy_block(subspace.w_l, l, subspace.basis); + copy_block(subspace.hw_l, l, subspace.hbasis); + copy_block(subspace.sw_l, l, subspace.sbasis); + if (nblk >= 3) + { + copy_block(subspace.p_l, 2 * l, subspace.basis); + copy_block(subspace.hp_l, 2 * l, subspace.hbasis); + copy_block(subspace.sp_l, 2 * l, subspace.sbasis); + } + + gram(subspace.basis.data(), subspace.hbasis.data(), dim, dim, subspace.k, dim); + gram(subspace.basis.data(), subspace.sbasis.data(), dim, dim, subspace.m, dim); + hermitize_projected(subspace.k, dim); + hermitize_projected(subspace.m, dim); +} + +// --------------------------------------------------------------------------- +// Solve K v = λ M v (small generalized eigenvalue problem) +// --------------------------------------------------------------------------- +template +int DiagoPPCG::solve_small_generalized( + const int dim, + const int nstates, + SmallSubspace& subspace) const +{ + const PPCGFunctionTimer timer("solve_small_generalized"); + const std::vector k0 = subspace.k; + const std::vector m0 = subspace.m; + + // A nearly dependent local basis makes the generalized solve unstable + // even when LAPACK accepts M as positive definite. Work in the retained + // eigenspace of M and map the Ritz vectors back to the original basis. + std::vector overlap_eigenvectors = m0; + std::vector overlap_eigenvalues(dim); + bool local_rank_truncation_succeeded = false; + int retained_count = 0; + try + { + container::kernels::lapack_heevd()( + dim, + overlap_eigenvectors.data(), + dim, + overlap_eigenvalues.data()); + + const Real largest_overlap_eigenvalue = std::max(overlap_eigenvalues.back(), Real(0)); + const Real precision_floor = Real(10) * std::numeric_limits::epsilon(); + const Real relative_rank_threshold = std::max( + Real(ppcg_subspace_rank_threshold), + precision_floor); + const Real overlap_cutoff = largest_overlap_eigenvalue + * relative_rank_threshold; + std::vector retained_indices; + retained_indices.reserve(dim); + for (int i = 0; i < dim; ++i) + { + if (overlap_eigenvalues[i] > overlap_cutoff) + { + retained_indices.push_back(i); + } + } + + const int retained_dim = int(retained_indices.size()); + if (retained_dim >= nstates && retained_dim < dim) + { + std::vector retained_basis(dim * retained_dim, T(0)); + for (int j = 0; j < retained_dim; ++j) + { + const int source_index = retained_indices[j]; + const Real scale = Real(1) + / std::sqrt(overlap_eigenvalues[source_index]); + for (int row = 0; row < dim; ++row) + { + retained_basis[row + j * dim] = overlap_eigenvectors[row + source_index * dim] + * scale; + } + } + + // Q = U_r Lambda_r^(-1/2) orthonormalizes the retained overlap + // subspace. Use BLAS for Q^H K Q rather than a four-index + // contraction: the latter becomes prohibitively expensive when + // the first global correction contains many nearly dependent + // residual directions. + const T one = T(1); + const T zero = T(0); + std::vector projected_hamiltonian(dim * retained_dim, T(0)); + ModuleBase::gemm_op()('N', + 'N', + dim, + retained_dim, + dim, + &one, + k0.data(), + dim, + retained_basis.data(), + dim, + &zero, + projected_hamiltonian.data(), + dim); + std::vector reduced_hamiltonian(retained_dim * retained_dim, T(0)); + ModuleBase::gemm_op()('C', + 'N', + retained_dim, + retained_dim, + dim, + &one, + retained_basis.data(), + dim, + projected_hamiltonian.data(), + dim, + &zero, + reduced_hamiltonian.data(), + retained_dim); + for (int j = 0; j < retained_dim; ++j) + { + const Real diagonal_value = std::real( + reduced_hamiltonian[j + j * retained_dim]); + reduced_hamiltonian[j + j * retained_dim] = T(diagonal_value, Real(0)); + for (int i = 0; i < j; ++i) + { + const T value = Real(0.5) * ( + reduced_hamiltonian[i + j * retained_dim] + + std::conj(reduced_hamiltonian[j + i * retained_dim])); + reduced_hamiltonian[i + j * retained_dim] = value; + reduced_hamiltonian[j + i * retained_dim] = std::conj(value); + } + } + + std::vector reduced_eigenvalues(retained_dim); + container::kernels::lapack_heevd()( + retained_dim, + reduced_hamiltonian.data(), + retained_dim, + reduced_eigenvalues.data()); + + std::fill(subspace.k.begin(), subspace.k.end(), T(0)); + for (int state = 0; state < retained_dim; ++state) + { + subspace.eval[state] = reduced_eigenvalues[state]; + } + ModuleBase::gemm_op()('N', + 'N', + dim, + retained_dim, + retained_dim, + &one, + retained_basis.data(), + dim, + reduced_hamiltonian.data(), + retained_dim, + &zero, + subspace.k.data(), + dim); + local_rank_truncation_succeeded = true; + retained_count = retained_dim; + } + } + catch (const std::runtime_error&) + { + // The shifted generalized solve below remains the fallback. + } + if (all_pool_operations_succeeded(local_rank_truncation_succeeded)) + { + return retained_count; + } + + // Try with increasing diagonal shifts; fall back to identity (no update) + // if the subspace is too ill-conditioned. sygvd modifies both matrices + // in-place before it may fail. + const Real shifts[] = {Real(ppcg_subspace_shifts[0]), + Real(ppcg_subspace_shifts[1]), + Real(ppcg_subspace_shifts[2]), + Real(ppcg_subspace_shifts[3])}; + for (const Real shift : shifts) + { + subspace.k = k0; + subspace.m = m0; + for (int i = 0; i < dim; ++i) + { + subspace.m[i + i * dim] += T(shift); + } + + bool local_sygvd_succeeded = false; + try + { + HermitianLapack::sygvd(dim, subspace.k.data(), + subspace.m.data(), + subspace.eval.data()); + local_sygvd_succeeded = true; + } + catch (const std::runtime_error&) + { + // Try the next diagonal shift. + } + if (all_pool_operations_succeeded(local_sygvd_succeeded)) + { + return dim; + } + } + // All attempts failed — set eigenvectors to identity (no update). + std::fill(subspace.k.begin(), subspace.k.end(), T(0)); + for (int i = 0; i < dim; ++i) + { + subspace.k[i + i * dim] = T(1); + subspace.eval[i] = Real(std::real(k0[i + i * dim])) + / std::max(Real(std::real(m0[i + i * dim])), + Real(ppcg_numerical_threshold)); + } + return 0; +} + +// Keep valid local Ritz directions that a block-size restriction would discard. +template +void DiagoPPCG::collect_discarded_ritz( + const SmallSubspace& subspace, + const int kept_count, + const int valid_count, + OperatorBlock& candidates) const +{ + const PPCGFunctionTimer timer("collect_discarded_ritz"); + const int count = valid_count - kept_count; + if (count <= 0) + { + return; + } + const int dim = int(subspace.eval.size()); + const int offset = int(candidates.vectors.size()); + const T* coefficients = subspace.k.data() + kept_count * dim; + const T one(1); + const T zero(0); + const auto append = [&](const std::vector& basis, std::vector& output) { + output.resize(offset + ld_psi_ * count, T(0)); + T* destination = output.data() + offset; + ModuleBase::gemm_op()('N', 'N', n_dim_, count, dim, + &one, basis.data(), ld_psi_, coefficients, dim, + &zero, destination, ld_psi_); + }; + append(subspace.basis, candidates.vectors); + append(subspace.hbasis, candidates.hvectors); + append(subspace.sbasis, candidates.svectors); +} + +// Preserve X exactly as a subspace: rank decisions apply only to its complement. +template +void DiagoPPCG::project_candidate_complement( + const HPsiFunc& hpsi_func, const T* psi, OperatorBlock& candidates) const +{ + const PPCGFunctionTimer timer("project_candidate_complement"); + const int count = int(candidates.vectors.size()) / ld_psi_; + std::vector basis_cols(n_band_); + std::iota(basis_cols.begin(), basis_cols.end(), 0); + std::vector candidate_cols(count); + std::iota(candidate_cols.begin(), candidate_cols.end(), 0); + std::vector norms(2 * count, 0.0); + for (int column = 0; column < count; ++column) + { + for (int row = 0; row < n_dim_; ++row) + { + norms[column] += std::norm(candidates.vectors[idx(row, column, ld_psi_)]); + } + } + // Two passes suppress overlap left by cancellation against the retained X. + for (int pass = 0; pass < 2; ++pass) + { + project_against(psi, spsi_.data(), hpsi_.data(), basis_cols, + candidates.vectors, candidates.svectors, &candidates.hvectors, candidate_cols); + } + for (int column = 0; column < count; ++column) + { + for (int row = 0; row < n_dim_; ++row) + { + norms[count + column] += std::norm(candidates.vectors[idx(row, column, ld_psi_)]); + } + } + const int norm_count = int(norms.size()); + reduce_pool_values(norms.data(), norm_count); + const Real precision_floor = Real(10) * std::numeric_limits::epsilon(); + const Real rank_threshold = std::max(Real(ppcg_subspace_rank_threshold), precision_floor); + std::vector retained; + std::vector refresh; + for (int column = 0; column < count; ++column) + { + const double before = norms[column]; + const double after = norms[count + column]; + if (!std::isfinite(before) || !std::isfinite(after) || after <= rank_threshold * before) + { + continue; + } + retained.push_back(column); + if (after < 0.5 * before) + { + refresh.push_back(column); + } + } + // The GPU bridge owns only nband columns, even if the complement is larger. + for (int first = 0; first < int(refresh.size()); first += n_band_) + { + const int last = std::min(first + n_band_, int(refresh.size())); + const std::vector columns(refresh.begin() + first, refresh.begin() + last); + std::vector vectors; + copy_cols(candidates.vectors.data(), columns, vectors); + std::vector hvectors(vectors.size(), T(0)); + std::vector svectors(vectors.size(), T(0)); + const int batch_count = int(columns.size()); + apply_h(hpsi_func, vectors.data(), hvectors.data(), batch_count); + apply_s_current(vectors.data(), svectors.data(), batch_count); + scatter_cols(candidates.hvectors.data(), columns, hvectors); + scatter_cols(candidates.svectors.data(), columns, svectors); + } + OperatorBlock projected; + copy_cols(candidates.vectors.data(), retained, projected.vectors); + copy_cols(candidates.hvectors.data(), retained, projected.hvectors); + copy_cols(candidates.svectors.data(), retained, projected.svectors); + candidates = std::move(projected); +} + +template +bool DiagoPPCG::recover_competing_ritz( + const HPsiFunc& hpsi_func, T* psi, Real* eigenvalue, + const bool ordered_ritz, std::vector& active_cols, + const std::vector& ethr_band, OperatorBlock& candidates) +{ + const PPCGFunctionTimer timer("recover_competing_ritz"); + if (candidates.vectors.empty()) + { + return false; + } + Real upper_bound = eigenvalue[n_band_ - 1]; + if (!ordered_ritz) + { + // Between Ritz steps a column Rayleigh quotient is not an upper bound + // on the largest target Ritz value. Use the Hermitian row-sum bound. + upper_bound = std::numeric_limits::lowest(); + for (int row = 0; row < n_band_; ++row) + { + Real row_bound = std::real(rr_hsub_[row + row * n_band_]); + for (int column = 0; column < n_band_; ++column) + { + if (column != row) + { + row_bound += std::abs(rr_hsub_[row + column * n_band_]); + } + } + upper_bound = std::max(upper_bound, row_bound); + } + } + project_candidate_complement(hpsi_func, psi, candidates); + const int count = int(candidates.vectors.size()) / ld_psi_; + if (count == 0) + { + return false; + } + std::vector moments(2 * count, 0.0); + for (int column = 0; column < count; ++column) + { + for (int row = 0; row < n_dim_; ++row) + { + const int offset = idx(row, column, ld_psi_); + const T conjugate = std::conj(candidates.vectors[offset]); + moments[column] += std::real(conjugate * candidates.svectors[offset]); + moments[count + column] += std::real(conjugate * candidates.hvectors[offset]); + } + } + const int moment_count = int(moments.size()); + reduce_pool_values(moments.data(), moment_count); + const Real scale = std::max(Real(1), std::abs(upper_bound)); + const Real margin = Real(32) * std::numeric_limits::epsilon() * scale; + std::vector competing; + for (int column = 0; column < count; ++column) + { + const double norm = moments[column]; + const double energy = moments[count + column]; + if (!std::isfinite(norm) || !std::isfinite(energy) || norm <= 0.0) + { + continue; + } + // Use the physical quotient, not an eigenvalue of the shifted-M + // fallback; they need not represent the same projected problem. + if (energy / norm < upper_bound - margin) + { + competing.push_back(column); + } + } + if (competing.empty()) + { + return false; + } + OperatorBlock selected; + copy_cols(candidates.vectors.data(), competing, selected.vectors); + copy_cols(candidates.hvectors.data(), competing, selected.hvectors); + copy_cols(candidates.svectors.data(), competing, selected.svectors); + // The selected copies own the surviving directions; release the larger + // candidate buffers before constructing their orthogonal basis. + candidates = OperatorBlock{}; + const int selected_count = int(competing.size()); + scale_to_unit_snorm(selected.vectors, selected.svectors, selected.hvectors, selected_count); + std::vector overlap; + gram(selected.vectors.data(), selected.svectors.data(), selected_count, selected_count, + overlap, selected_count); + hermitize_projected(overlap, selected_count); + std::vector norms(selected_count); + bool local_success = false; + try + { + container::kernels::lapack_heevd()( + selected_count, overlap.data(), selected_count, norms.data()); + local_success = true; + } + catch (const std::runtime_error&) + { + } + if (!all_pool_operations_succeeded(local_success)) + { + throw std::runtime_error("PPCG: competing-direction overlap diagonalization failed."); + } + const Real precision_floor = Real(10) * std::numeric_limits::epsilon(); + const Real rank_threshold = std::max(Real(ppcg_subspace_rank_threshold), precision_floor); + const Real cutoff = norms.back() * rank_threshold; + int retained_count = 0; + for (int column = 0; column < selected_count; ++column) + { + if (norms[column] <= cutoff) + { + continue; + } + const Real inverse_norm = Real(1) / std::sqrt(norms[column]); + for (int row = 0; row < selected_count; ++row) + { + overlap[row + retained_count * selected_count] = overlap[row + column * selected_count] * inverse_norm; + } + ++retained_count; + } + if (retained_count == 0) + { + throw std::runtime_error("PPCG: competing-direction complement has zero numerical rank."); + } + OperatorBlock orthogonal; + const T one(1); + const T zero(0); + const auto transform = [&](const std::vector& input, std::vector& output) { + output.assign(ld_psi_ * retained_count, T(0)); + ModuleBase::gemm_op()('N', 'N', n_dim_, retained_count, selected_count, + &one, input.data(), ld_psi_, overlap.data(), selected_count, + &zero, output.data(), ld_psi_); + }; + transform(selected.vectors, orthogonal.vectors); + transform(selected.hvectors, orthogonal.hvectors); + transform(selected.svectors, orthogonal.svectors); + selected = OperatorBlock{}; + SmallSubspace expanded; + build_startup_global_subspace(psi, orthogonal.vectors, orthogonal.svectors, + orthogonal.hvectors, retained_count, expanded); + const int expanded_dim = n_band_ + retained_count; + local_success = false; + try + { + // No rank truncation of [X,Q]: X is retained intact and Q has already + // been rank-revealed independently in its S-orthogonal complement. + HermitianLapack::sygvd(expanded_dim, expanded.k.data(), expanded.m.data(), expanded.eval.data()); + local_success = true; + } + catch (const std::runtime_error&) + { + } + if (!all_pool_operations_succeeded(local_success)) + { + throw std::runtime_error("PPCG: competing-direction Rayleigh-Ritz solve failed."); + } + update_startup_global_subspace(psi, orthogonal.vectors, orthogonal.svectors, + orthogonal.hvectors, retained_count, expanded); + std::copy(expanded.eval.begin(), expanded.eval.begin() + n_band_, eigenvalue); + set_zero(p_); + set_zero(sp_); + set_zero(hp_); + compute_eigenpair_residual(eigenvalue); + lock_epairs(w_, ethr_band, active_cols); + ++rayleigh_ritz_count_; + ++root_recovery_count_; + return true; +} + +// --------------------------------------------------------------------------- +// Update wavefunctions from small subspace eigenvectors +// --------------------------------------------------------------------------- +template +void DiagoPPCG::update_one_block( + T* psi, + const std::vector& cols, + int l, + int nblk, + SmallSubspace& subspace) +{ + const PPCGFunctionTimer timer("update_one_block"); + const int dim = nblk * l; + const T* eigvec = subspace.k.data(); + + subspace.psi_new.assign(ld_psi_ * l, T(0)); + subspace.spsi_new.assign(ld_psi_ * l, T(0)); + subspace.hpsi_new.assign(ld_psi_ * l, T(0)); + subspace.p_new.assign(ld_psi_ * l, T(0)); + subspace.sp_new.assign(ld_psi_ * l, T(0)); + subspace.hp_new.assign(ld_psi_ * l, T(0)); + + // coeff_state: full Ritz-vector rows [psi, w, p] -> new iterate. + // coeff_p: the [w, p] rows (psi rows zero) -> new search direction. + subspace.coeff_state.resize(dim * l); + subspace.coeff_p.resize(dim * l); +#ifdef _OPENMP +#pragma omp parallel for schedule(static) if (l * l > ppcg_openmp_work_threshold) +#endif + for (int j = 0; j < l; ++j) + { + for (int i = 0; i < l; ++i) + { + const T c_psi = eigvec[i + j * dim]; + const T c_w = eigvec[(l + i) + j * dim]; + subspace.coeff_state[i + j * dim] = c_psi; + subspace.coeff_state[(l + i) + j * dim] = c_w; + subspace.coeff_p[i + j * dim] = T(0); + subspace.coeff_p[(l + i) + j * dim] = c_w; + if (nblk >= 3) + { + const T c_p = eigvec[(2 * l + i) + j * dim]; + subspace.coeff_state[(2 * l + i) + j * dim] = c_p; + subspace.coeff_p[(2 * l + i) + j * dim] = c_p; + } + } + } + + auto fill_basis = [&](const std::vector& a, + const std::vector& b, + const std::vector& c, + std::vector& basis) + { + basis.resize(ld_psi_ * dim); +#ifdef _OPENMP +#pragma omp parallel for schedule(static) if (ld_psi_ * l > ppcg_openmp_work_threshold) +#endif + for (int j = 0; j < l; ++j) + { + std::copy(a.begin() + j * ld_psi_, + a.begin() + (j + 1) * ld_psi_, + basis.begin() + j * ld_psi_); + std::copy(b.begin() + j * ld_psi_, + b.begin() + (j + 1) * ld_psi_, + basis.begin() + (l + j) * ld_psi_); + if (nblk >= 3) + { + std::copy(c.begin() + j * ld_psi_, + c.begin() + (j + 1) * ld_psi_, + basis.begin() + (2 * l + j) * ld_psi_); + } + } + }; + + auto combine = [&](const std::vector& basis, + const std::vector& coeff, + std::vector& out) + { + const T one = T(1); + const T zero = T(0); + ModuleBase::gemm_op()('N', + 'N', + n_dim_, + l, + dim, + &one, + basis.data(), + ld_psi_, + coeff.data(), + dim, + &zero, + out.data(), + ld_psi_); + }; + + fill_basis(subspace.psi_l, subspace.w_l, subspace.p_l, subspace.basis); + fill_basis(subspace.spsi_l, subspace.sw_l, subspace.sp_l, subspace.sbasis); + fill_basis(subspace.hpsi_l, subspace.hw_l, subspace.hp_l, subspace.hbasis); + + combine(subspace.basis, subspace.coeff_state, subspace.psi_new); + combine(subspace.sbasis, subspace.coeff_state, subspace.spsi_new); + combine(subspace.hbasis, subspace.coeff_state, subspace.hpsi_new); + combine(subspace.basis, subspace.coeff_p, subspace.p_new); + combine(subspace.sbasis, subspace.coeff_p, subspace.sp_new); + combine(subspace.hbasis, subspace.coeff_p, subspace.hp_new); + + scatter_cols(psi, cols, subspace.psi_new); + scatter_cols(spsi_.data(), cols, subspace.spsi_new); + scatter_cols(hpsi_.data(), cols, subspace.hpsi_new); + scatter_cols(p_.data(), cols, subspace.p_new); + scatter_cols(sp_.data(), cols, subspace.sp_new); + scatter_cols(hp_.data(), cols, subspace.hp_new); +} + +// --------------------------------------------------------------------------- +// Rayleigh-Ritz: full subspace diagonalization + residual computation +// --------------------------------------------------------------------------- +template +bool DiagoPPCG::rayleigh_ritz( + T* psi, Real* eigenvalue, + std::vector& active_cols, + const std::vector& ethr_band) +{ + const PPCGFunctionTimer timer("rayleigh_ritz"); + ++rayleigh_ritz_count_; + gram(psi, hpsi_.data(), n_band_, n_band_, rr_hsub_, n_band_); + gram(psi, spsi_.data(), n_band_, n_band_, rr_ssub_, n_band_); + + bool local_sygvd_ok = false; + try + { + HermitianLapack::sygvd(n_band_, rr_hsub_.data(), rr_ssub_.data(), + rr_eval_.data()); + local_sygvd_ok = true; + } + catch (const std::runtime_error&) + { + // The collective result below selects one common fallback on all ranks. + } + const bool sygvd_ok = all_pool_operations_succeeded(local_sygvd_ok); + if (!sygvd_ok) + { + // LAPACK may overwrite the projected matrices before failing. Re-form + // them on every rank so all ranks enter the same fallback path. + gram(psi, hpsi_.data(), n_band_, n_band_, rr_hsub_, n_band_); + gram(psi, spsi_.data(), n_band_, n_band_, rr_ssub_, n_band_); + for (int ii = 0; ii < n_band_; ++ii) + { + rr_eval_[ii] = Real(std::real(rr_hsub_[ii + ii * n_band_])) + / std::max(Real( + std::real(rr_ssub_[ii + ii * n_band_])), + Real(ppcg_numerical_threshold)); + } + } + + if (sygvd_ok) + { + const int sz = ld_psi_ * n_band_; + std::copy(psi, psi + sz, w_.begin()); + std::copy(spsi_.begin(), spsi_.end(), sw_.begin()); + std::copy(hpsi_.begin(), hpsi_.end(), hw_.begin()); + + std::fill(psi, psi + ld_psi_ * n_band_, T(0)); + set_zero(spsi_); + set_zero(hpsi_); + + const T one = T(1); + const T zero = T(0); + ModuleBase::gemm_op()('N', + 'N', + n_dim_, + n_band_, + n_band_, + &one, + w_.data(), + ld_psi_, + rr_hsub_.data(), + n_band_, + &zero, + psi, + ld_psi_); + ModuleBase::gemm_op()('N', + 'N', + n_dim_, + n_band_, + n_band_, + &one, + sw_.data(), + ld_psi_, + rr_hsub_.data(), + n_band_, + &zero, + spsi_.data(), + ld_psi_); + ModuleBase::gemm_op()('N', + 'N', + n_dim_, + n_band_, + n_band_, + &one, + hw_.data(), + ld_psi_, + rr_hsub_.data(), + n_band_, + &zero, + hpsi_.data(), + ld_psi_); + + for (int j = 0; j < n_band_; ++j) + { + eigenvalue[j] = rr_eval_[j]; + } + // Keep each conjugate direction attached to the same Ritz column + // when the global solve mixes or permutes the retained states. + const auto rotate_direction = [&](std::vector& values) { + std::copy(values.begin(), values.end(), w_.begin()); + set_zero(values); + ModuleBase::gemm_op()('N', 'N', n_dim_, n_band_, n_band_, + &one, w_.data(), ld_psi_, rr_hsub_.data(), n_band_, + &zero, values.data(), ld_psi_); + }; + rotate_direction(p_); + rotate_direction(sp_); + rotate_direction(hp_); + } + else + { + // No rotation: just update eigenvalues with Rayleigh quotients. + for (int j = 0; j < n_band_; ++j) + { + eigenvalue[j] = rr_eval_[j]; + } + } + + compute_eigenpair_residual(eigenvalue); + + if (sygvd_ok) + { + lock_epairs(w_, ethr_band, active_cols); + } + else + { + // Without a successful global Rayleigh-Ritz solve, the columns are + // not known to be ordered by eigenvalue. Do not hard-lock an + // arbitrary column based on its residual alone. + active_cols.resize(n_band_); + std::iota(active_cols.begin(), active_cols.end(), 0); + max_residual_ = max_generalized_residual(hpsi_.data(), + spsi_.data(), + eigenvalue, + ld_psi_, + n_dim_, + n_band_); + } + return sygvd_ok; +} + +template +void DiagoPPCG::compute_eigenpair_residual(const Real* eigenvalue) +{ + const PPCGFunctionTimer timer("compute_eigenpair_residual"); + set_zero(w_); +#ifdef _OPENMP +#pragma omp parallel for collapse(2) schedule(static) if (n_dim_ * n_band_ > ppcg_openmp_work_threshold) +#endif + for (int j = 0; j < n_band_; ++j) + { + for (int ig = 0; ig < n_dim_; ++ig) + { + w_[idx(ig, j, ld_psi_)] = hpsi_[idx(ig, j, ld_psi_)] + - spsi_[idx(ig, j, ld_psi_)] * eigenvalue[j]; + } + } +} + +template +bool DiagoPPCG::orthonormalize(T* psi) +{ + const PPCGFunctionTimer timer("orthonormalize"); + gram(psi, spsi_.data(), n_band_, n_band_, rr_ssub_, n_band_); + bool local_success = false; + try + { + container::kernels::lapack_potrf()( + 'U', n_band_, rr_ssub_.data(), n_band_); + container::kernels::lapack_trtri()( + 'U', 'N', n_band_, rr_ssub_.data(), n_band_); + local_success = true; + } + catch (const std::runtime_error&) + { + // The caller performs a common Rayleigh-Ritz fallback on all ranks. + } + if (!all_pool_operations_succeeded(local_success)) + { + return false; + } + for (int j = 0; j < n_band_; ++j) + { + for (int i = j + 1; i < n_band_; ++i) + { + rr_ssub_[i + j * n_band_] = T(0); + } + } + // X^H SX = R^H R. Right multiplication by the upper triangular R^{-1} + // keeps a previously locked prefix independent of the active columns. + const int size = ld_psi_ * n_band_; + const T one = T(1); + const T zero = T(0); + const auto rotate = [&](T* values) { + std::copy(values, values + size, w_.begin()); + std::fill(values, values + size, T(0)); + ModuleBase::gemm_op()('N', + 'N', + n_dim_, + n_band_, + n_band_, + &one, + w_.data(), + ld_psi_, + rr_ssub_.data(), + n_band_, + &zero, + values, + ld_psi_); + }; + rotate(psi); + rotate(hpsi_.data()); + rotate(spsi_.data()); + // The non-Ritz sweep changes the column basis too. Transform the + // conjugate directions and both cached images with the same coefficients. + rotate(p_.data()); + rotate(sp_.data()); + rotate(hp_.data()); + return true; +} + +template +bool DiagoPPCG::update_block_gradient( + const T* psi, + Real* eigenvalue, + const std::vector& ethr_band) +{ + const PPCGFunctionTimer timer("update_block_gradient"); + gram(psi, hpsi_.data(), n_band_, n_band_, rr_hsub_, n_band_); + hermitize_projected(rr_hsub_, n_band_); + for (int j = 0; j < n_band_; ++j) + { + // These are Rayleigh quotients for tracing, not ordered Ritz values. + eigenvalue[j] = std::real(rr_hsub_[j + j * n_band_]); + } + std::copy(hpsi_.begin(), hpsi_.end(), w_.begin()); + const T minus_one = T(-1); + const T one = T(1); + ModuleBase::gemm_op()('N', + 'N', + n_dim_, + n_band_, + n_band_, + &minus_one, + spsi_.data(), + ld_psi_, + rr_hsub_.data(), + n_band_, + &one, + w_.data(), + ld_psi_); + std::vector candidates; + lock_epairs(w_, ethr_band, candidates); + // Do not change the actual locked prefix without a successful Ritz solve. + return candidates.empty(); +} + +} // namespace hsolver + + +namespace hsolver { + +//============================================================================== +// MAIN DIAGONALIZATION ROUTINE +//============================================================================== +template +double DiagoPPCG::diag(const HPsiFunc& hpsi_func, + const SPsiFunc& spsi_func, + int ld_psi, + int nband, + int dim, + T* psi_in, + Real* eigenvalue_in, + const std::vector& ethr_band, + const Real* prec) +{ + const PPCGFunctionTimer timer("diag"); + ld_psi_ = ld_psi; + n_band_ = nband; + n_dim_ = dim; + converged_ = false; + active_band_count_ = 0; + iteration_limit_reached_ = false; + max_residual_ = Real(0); + first_active_band_ = -1; + first_active_threshold_ = Real(0); + rayleigh_ritz_count_ = 0; + root_recovery_count_ = 0; + + validate_input(hpsi_func, psi_in, eigenvalue_in, ethr_band, prec); + spsi_func_ = spsi_func; + + // Allocate working storage. + const int ncol = n_band_; + const int sz = ld_psi_ * ncol; + + hpsi_.assign(sz, T(0)); + spsi_.assign(sz, T(0)); + w_.assign(sz, T(0)); + sw_.assign(sz, T(0)); + hw_.assign(sz, T(0)); + p_.assign(sz, T(0)); + sp_.assign(sz, T(0)); + hp_.assign(sz, T(0)); + rr_hsub_.resize(ncol * ncol); + rr_ssub_.resize(ncol * ncol); + rr_eval_.resize(ncol); + + std::vector all_cols(ncol); + std::iota(all_cols.begin(), all_cols.end(), 0); + + apply_h(hpsi_func, psi_in, hpsi_.data(), ncol); + apply_s_current(psi_in, spsi_.data(), ncol); + + double avg_iter = 1.0; + int iter = 1; + std::vector active_cols; + active_cols.reserve(ncol); + + std::ofstream residual_trace; + const char* residual_trace_path = std::getenv("ABACUS_PPCG_RESIDUAL_TRACE"); + if (residual_trace_path != nullptr && Parallel_Reduce::is_world_root()) + { + // A single writer avoids file corruption when several MPI pools solve + // different k-points concurrently. + residual_trace.open(residual_trace_path); + if (residual_trace) + { + residual_trace << std::setprecision(std::numeric_limits::max_digits10); + residual_trace << "iteration,stage,max_residual"; + // Test-side tools compare these Ritz values with analytical or + // LAPACK reference eigenvalues without coupling them to the solver. + for (int ib = 0; ib < ncol; ++ib) + { + residual_trace << ",eigenvalue_" << ib; + } + residual_trace << ",local_overlap_min,local_overlap_max,local_overlap_condition\n"; + } + } + int residual_trace_writer_count = residual_trace ? 1 : 0; + reduce_pool_values(&residual_trace_writer_count, 1); + const bool residual_trace_enabled = residual_trace_writer_count > 0; + const bool write_residual_trace = bool(residual_trace); + Real local_overlap_min = std::numeric_limits::quiet_NaN(); + Real local_overlap_max = std::numeric_limits::quiet_NaN(); + Real local_overlap_condition = std::numeric_limits::quiet_NaN(); + auto record_residual = [&](int iteration, const char* stage) { + if (!residual_trace_enabled) + { + return; + } + const Real max_residual = max_generalized_residual(hpsi_.data(), + spsi_.data(), + eigenvalue_in, + ld_psi_, + n_dim_, + ncol); + if (!write_residual_trace) + { + return; + } + residual_trace << iteration << ',' << stage << ',' << max_residual; + for (int ib = 0; ib < ncol; ++ib) + { + residual_trace << ',' << eigenvalue_in[ib]; + } + residual_trace << ',' << local_overlap_min + << ',' << local_overlap_max + << ',' << local_overlap_condition + << '\n'; + }; + + auto measure_local_overlap = [&](const SmallSubspace& local_subspace, + const int local_dim) { + if (!write_residual_trace) + { + return; + } + std::vector overlap = local_subspace.m; + std::vector overlap_eigenvalues(local_dim); + try + { + container::kernels::lapack_heevd()( + local_dim, + overlap.data(), + local_dim, + overlap_eigenvalues.data()); + } + catch (const std::exception&) + { + // Trace diagnostics must not change the solver's MPI control flow. + local_overlap_min = std::numeric_limits::quiet_NaN(); + local_overlap_max = std::numeric_limits::quiet_NaN(); + local_overlap_condition = std::numeric_limits::infinity(); + return; + } + const Real min_eigenvalue = overlap_eigenvalues.front(); + const Real max_eigenvalue = overlap_eigenvalues.back(); + const Real denominator = std::max(std::abs(min_eigenvalue), + Real(ppcg_numerical_threshold)); + const Real condition = std::abs(max_eigenvalue) / denominator; + if (std::isnan(local_overlap_condition) || condition > local_overlap_condition) + { + local_overlap_min = min_eigenvalue; + local_overlap_max = max_eigenvalue; + local_overlap_condition = condition; + } + }; + + // Initialize with Rayleigh-Ritz. + bool last_rr_succeeded = rayleigh_ritz( + psi_in, eigenvalue_in, active_cols, ethr_band); + const bool needs_startup_global_correction = !active_cols.empty() + && sbsize_ < ncol; + // The Ritz transformation updates X, HX and SX with identical coefficients. + record_residual(0, "initial_rr"); + + std::vector active_vector; + std::vector s_active_vector; + std::vector h_active_vector; + std::vector cols; + cols.reserve(std::min(sbsize_, ncol)); + SmallSubspace subspace; + bool use_p = false; // previous search direction becomes available after + // the first block update. + Real best_res = std::numeric_limits::max(); // best residual seen + int no_improve = 0; // iterations since the residual last decreased + + while (!active_cols.empty() && iter <= maxiter_) + { + OperatorBlock discarded; + // Near the floating-point cancellation scale, form a reliable + // residual before using it as a new search direction. This safeguard + // also covers cancellation in the W/P recurrence, not only in the + // explicit projection of P. The final fresh residual remains the + // convergence certificate; this scale is not an error bound. + Real spectral_scale = Real(1); + for (int ib = 0; ib < ncol; ++ib) + { + spectral_scale = std::max(spectral_scale, std::abs(eigenvalue_in[ib])); + } + const Real reliable_update_threshold = std::sqrt(std::numeric_limits::epsilon()) + * spectral_scale; + const bool near_precision = max_residual_ <= reliable_update_threshold; + if (near_precision) + { + apply_h(hpsi_func, psi_in, hpsi_.data(), ncol); + apply_s_current(psi_in, spsi_.data(), ncol); + apply_h(hpsi_func, p_.data(), hp_.data(), ncol); + apply_s_current(p_.data(), sp_.data(), ncol); + if (last_rr_succeeded) + { + compute_eigenpair_residual(eigenvalue_in); + } + else + { + update_block_gradient(psi_in, eigenvalue_in, ethr_band); + } + } + local_overlap_min = std::numeric_limits::quiet_NaN(); + local_overlap_max = std::numeric_limits::quiet_NaN(); + local_overlap_condition = std::numeric_limits::quiet_NaN(); + const int nact = int(active_cols.size()); + const int nsb = 1 + (nact - 1) / sbsize_; + // Keep the same number of local blocks and the same size limit, but + // distribute the remainder evenly. When the limit permits it, this + // avoids a trailing one-column block that cannot exchange search + // directions with a neighboring degenerate Ritz vector. + const int base_block_size = nact / nsb; + const int larger_block_count = nact % nsb; + + // Precondition the residual. + divide_by_preconditioner(active_cols, prec, w_); + copy_cols(w_.data(), active_cols, active_vector); + s_active_vector.assign(ld_psi_ * nact, T(0)); + apply_s_current(active_vector.data(), s_active_vector.data(), nact); + scatter_cols(sw_.data(), active_cols, s_active_vector); + project_against(psi_in, spsi_.data(), nullptr, all_cols, w_, sw_, nullptr, active_cols); + + // Apply H to the search direction. + copy_cols(w_.data(), active_cols, active_vector); + h_active_vector.assign(ld_psi_ * nact, T(0)); + copy_cols(sw_.data(), active_cols, s_active_vector); + apply_h(hpsi_func, active_vector.data(), h_active_vector.data(), nact); + scatter_cols(hw_.data(), active_cols, h_active_vector); + + // Project the cached P, SP and HP together. Linearity preserves their + // consistency without another Hamiltonian/overlap application except + // when strong cancellation requires a fresh image of selected columns. + if (use_p) + { + project_search_directions(hpsi_func, psi_in, all_cols, active_cols); + } + + avg_iter += double(nact) / double(ncol); + + // LOBPCG-style block subspace. On the first sweep only [psi, w] is + // available; afterwards the previous search direction p is added as a + // third block, which restores the conjugate-gradient acceleration. + // The w/p blocks are normalized to unit S-norm before building the + // Gram matrix (see build_small_subspace), keeping M well-conditioned. + + // The startup correction needs a global [X, W_active] comparison to + // retain the requested lowest roots. All current Ritz vectors remain + // available, while a converged prefix contributes no correction + // vectors and stays outside the active iteration count. + const bool startup_correction = iter == 1 && needs_startup_global_correction; + if (startup_correction) + { + const int local_dim = ncol + nact; + build_startup_global_subspace(psi_in, + active_vector, + s_active_vector, + h_active_vector, + nact, + subspace); + measure_local_overlap(subspace, local_dim); + solve_small_generalized(local_dim, ncol, subspace); + update_startup_global_subspace(psi_in, + active_vector, + s_active_vector, + h_active_vector, + nact, + subspace); + } + else + { + const int nblk = use_p ? 3 : 2; + int i0 = 0; + for (int isb = 0; isb < nsb; ++isb) + { + int l = base_block_size; + if (isb < larger_block_count) + { + ++l; + } + cols.assign(active_cols.begin() + i0, + active_cols.begin() + i0 + l); + + build_small_subspace(psi_in, cols, nblk, subspace); + const int local_dim = nblk * l; + measure_local_overlap(subspace, local_dim); + const int valid_ritz_count = solve_small_generalized(local_dim, l, subspace); + collect_discarded_ritz(subspace, l, valid_ritz_count, discarded); + update_one_block(psi_in, cols, l, nblk, subspace); + i0 += l; + } + } + use_p = true; + + // Refresh cached images independently of the full Ritz cadence, so + // rr_step=1 does not add two full H/S applications on every sweep. + // Intermediate sweeps maintain S-orthonormality without diagonalizing + // the full projected Hamiltonian. + const bool scheduled_rr = (iter % rr_step_) == 0; + const bool refresh_operators = (iter % ppcg_operator_refresh_interval) == 0 + && !near_precision; + if (refresh_operators) + { + apply_h(hpsi_func, psi_in, hpsi_.data(), ncol); + apply_s_current(psi_in, spsi_.data(), ncol); + apply_h(hpsi_func, p_.data(), hp_.data(), ncol); + apply_s_current(p_.data(), sp_.data(), ncol); + } + // The startup subspace can insert a newly discovered lower root into + // an old locked prefix. Reclassify it before freezing any column. + bool perform_rr = scheduled_rr || startup_correction || iter == maxiter_; + if (!perform_rr) + { + perform_rr = !orthonormalize(psi_in); + } + if (!perform_rr) + { + perform_rr = update_block_gradient(psi_in, eigenvalue_in, ethr_band); + } + last_rr_succeeded = false; + if (perform_rr) + { + last_rr_succeeded = rayleigh_ritz( + psi_in, eigenvalue_in, active_cols, ethr_band); + } + // A residual only certifies the retained eigenpairs, not that no lower + // state was discarded by another local block. Inspect those available + // directions before accepting a newly converged set of columns. + if (!perform_rr || last_rr_succeeded) + { + const bool recovered = recover_competing_ritz(hpsi_func, psi_in, eigenvalue_in, + last_rr_succeeded, active_cols, + ethr_band, discarded); + if (recovered) + { + last_rr_succeeded = true; + perform_rr = true; + use_p = false; + } + } + if (last_rr_succeeded && active_cols.empty()) + { + // A small cached residual is only tentative. Refresh the + // operator images before stopping, and keep iterating if + // roundoff had made the cached criterion overly optimistic. + apply_h(hpsi_func, psi_in, hpsi_.data(), ncol); + apply_s_current(psi_in, spsi_.data(), ncol); + compute_eigenpair_residual(eigenvalue_in); + lock_epairs(w_, ethr_band, active_cols); + if (!active_cols.empty()) + { + set_zero(p_); + set_zero(sp_); + set_zero(hp_); + use_p = false; + } + } + // Restart the search direction if the residual keeps rising. With a + // poor preconditioner the LOBPCG recurrence can stagnate (or slowly + // diverge) instead of reducing the residual. Requiring several + // consecutive rises avoids resetting on a transient bump, and a reset + // falls back to a steepest-descent step to recover the low eigenpairs. + { + const Real cur_res = max_residual_; + if (cur_res < best_res) + { + best_res = cur_res; + no_improve = 0; + } + else + { + ++no_improve; + } + if (no_improve >= ppcg_stagnation_restart_interval) + { + std::fill(p_.begin(), p_.end(), T(0)); + std::fill(sp_.begin(), sp_.end(), T(0)); + std::fill(hp_.begin(), hp_.end(), T(0)); + use_p = false; + no_improve = 0; + best_res = cur_res; + } + } + const char* stage = perform_rr ? "rayleigh_ritz" : "block_sweep"; + record_residual(iter, stage); + + ++iter; + } + + // Final consistency: ensure H|psi> and S|psi> match the returned vectors, + // then classify convergence from those final residuals rather than from a + // cached Rayleigh-Ritz state. + apply_h(hpsi_func, psi_in, hpsi_.data(), ncol); + apply_s_current(psi_in, spsi_.data(), ncol); + compute_eigenpair_residual(eigenvalue_in); + if (last_rr_succeeded) + { + lock_epairs(w_, ethr_band, active_cols); + } + else + { + active_cols = all_cols; + max_residual_ = max_generalized_residual(hpsi_.data(), + spsi_.data(), + eigenvalue_in, + ld_psi_, + n_dim_, + ncol); + } + active_band_count_ = int(active_cols.size()); + if (!active_cols.empty()) + { + first_active_band_ = active_cols.front(); + first_active_threshold_ = std::max(Real(ethr_band[first_active_band_]), + diag_thr_); + } + converged_ = last_rr_succeeded && active_cols.empty(); + iteration_limit_reached_ = !converged_ && iter > maxiter_; + if (!active_cols.empty()) + { + const char* failure_stage = iteration_limit_reached_ + ? "max_iterations" + : "final_residual_failure"; + record_residual(iter - 1, failure_stage); + } + record_residual(iter - 1, "final"); + return avg_iter; +} + +} // namespace hsolver + +namespace hsolver { + +template class DiagoPPCG, base_device::DEVICE_CPU>; +template class DiagoPPCG, base_device::DEVICE_CPU>; + +} // namespace hsolver diff --git a/source/source_hsolver/diago_ppcg.h b/source/source_hsolver/diago_ppcg.h new file mode 100644 index 00000000000..0958b56e9e2 --- /dev/null +++ b/source/source_hsolver/diago_ppcg.h @@ -0,0 +1,319 @@ +#ifndef DIAGO_PPCG_H +#define DIAGO_PPCG_H + +#include "source_base/module_device/types.h" + +#include +#include +#include +#include + +namespace hsolver { + +// ----------------------------------------------------------------------------- +// DiagoPPCG: Projection Preconditioned solver +// ----------------------------------------------------------------------------- +// +// Implements the block-subspace diagonalization strategy used by +// ks_solver=ppcg. The solver updates groups of active Ritz pairs in a bounded +// local subspace with configurable global Rayleigh-Ritz steps. Operator images +// are cached, with independent periodic and cancellation-triggered refreshes. +// ----------------------------------------------------------------------------- + +namespace base_device = ::base_device; + +// Check all possible CPU workspaces before allocating or copying through the +// GPU bridge. Dense kernels and pool reductions use int element counts. +bool ppcg_workspace_dimensions_valid(int ld_psi, int nband, int dim, int sbsize); + +template +class DiagoPPCG +{ +public: + // ------------------------------------------------------------------------- + // Type aliases + // ------------------------------------------------------------------------- + using Real = typename std::conditional< + std::is_same>::value, double, + float>::type; + using HPsiFunc = std::function; + using SPsiFunc = std::function; + + // ------------------------------------------------------------------------- + // Constructor + // rr_step is the explicit Ritz interval; PW INPUT defaults are supplied by + // the caller. The internal H/S cache-refresh interval is independent of it. + // ------------------------------------------------------------------------- + DiagoPPCG(const Real& diag_thr, + const int& diag_iter_max, + const int& sbsize, + const int& rr_step); + + // ------------------------------------------------------------------------- + // Main entry point + // + // Returns average number of subspace iterations per band. + // Optional ABACUS_PPCG_RESIDUAL_TRACE CSV: initial_rr, rayleigh_ritz and + // final rows contain ordered Ritz values after a successful Ritz solve. + // block_sweep rows instead contain the current column Rayleigh quotients; + // exclude those rows when plotting ordered eigenvalue errors. + // ------------------------------------------------------------------------- + double diag(const HPsiFunc& hpsi_func, + const SPsiFunc& spsi_func, + int ld_psi, + int nband, + int dim, + T* psi_in, + Real* eigenvalue_in, + const std::vector& ethr_band, + const Real* prec); + + bool converged() const + { + return converged_; + } + + int active_band_count() const + { + return active_band_count_; + } + + bool iteration_limit_reached() const + { + return iteration_limit_reached_; + } + + Real max_residual() const + { + return max_residual_; + } + + int first_active_band() const + { + return first_active_band_; + } + + Real first_active_threshold() const + { + return first_active_threshold_; + } + + int rayleigh_ritz_count() const + { + return rayleigh_ritz_count_; + } + + int root_recovery_count() const + { + return root_recovery_count_; + } + +private: + friend class DiagoPPCGProjectionTest; + // ------------------------------------------------------------------------- + // Data members + // ------------------------------------------------------------------------- + int maxiter_; + int sbsize_; + int rr_step_; + Real diag_thr_; + + // Problem dimensions (set in diag()) + int ld_psi_ = 0; + int n_band_ = 0; + int n_dim_ = 0; + bool converged_ = false; + int active_band_count_ = 0; + bool iteration_limit_reached_ = false; + Real max_residual_ = Real(0); + int first_active_band_ = -1; + Real first_active_threshold_ = Real(0); + int rayleigh_ritz_count_ = 0; + int root_recovery_count_ = 0; + + // Cached S-operator (null if identity). + SPsiFunc spsi_func_; + + // Working storage (column-major: ld_psi_ rows, n_band_ columns). + std::vector hpsi_; + std::vector spsi_; + std::vector w_; // residual / preconditioned residual + std::vector sw_; // S * w + std::vector hw_; // H * w + std::vector p_; // previous search direction (LOBPCG) + std::vector sp_; // S * p + std::vector hp_; // H * p + std::vector rr_hsub_; + std::vector rr_ssub_; + std::vector rr_eval_; + + // ------------------------------------------------------------------------- + // Internal helpers + // ------------------------------------------------------------------------- + static inline int idx(int row, int col, int ld) + { + return row + col * ld; + } + + void validate_input(const HPsiFunc& hpsi_func, + const T* psi_in, const Real* eigenvalue_in, + const std::vector& ethr_band, + const Real* prec) const; + + // S-application (identity fallback if spsi_func is null). + void apply_h(const HPsiFunc& hpsi_func, T* psi_in, T* hpsi_out, + int ncol) const; + void apply_s(const SPsiFunc& spsi_func, T* psi_in, T* spsi_out, + int ncol) const; + void apply_s_current(T* psi_in, T* spsi_out, int ncol) const; + + // Inner product (real part only). + Real gamma_dot(const T* x, const T* y) const; + + // Gram matrix: out[i, j] = . + void gram(const T* mat_a, const T* mat_b, + int ncol_a, int ncol_b, + std::vector& out, int ld_out) const; + + // Gather / scatter columns. + void copy_cols(const T* src, const std::vector& cols, + std::vector& dst) const; + void scatter_cols(T* dst, const std::vector& cols, + const std::vector& src) const; + + // Project x onto vectors orthogonal to the S-orthonormal basis. + void project_against(const T* basis, const T* sbasis, const T* hbasis, + const std::vector& basis_cols, + std::vector& x, std::vector& sx, std::vector* hx, + const std::vector& x_cols) const; + + void project_search_directions(const HPsiFunc& hpsi_func, const T* psi, + const std::vector& all_cols, + const std::vector& active_cols); + + // x[c] /= max(prec, eps) for each active column c. + void divide_by_preconditioner(const std::vector& active_cols, + const Real* prec, + std::vector& x) const; + + // ------------------------------------------------------------------------- + // Block-subspace strategy helpers (File 1 style) + // ------------------------------------------------------------------------- + struct SmallSubspace + { + std::vector k; // K matrix (projected H) + std::vector m; // M matrix (projected S) + std::vector eval; // eigenvalues + std::vector psi_l; + std::vector spsi_l; + std::vector hpsi_l; + std::vector w_l; + std::vector sw_l; + std::vector hw_l; + std::vector p_l; + std::vector sp_l; + std::vector hp_l; + std::vector basis; + std::vector hbasis; + std::vector sbasis; + std::vector coeff_state; + std::vector coeff_p; + std::vector psi_new; + std::vector spsi_new; + std::vector hpsi_new; + std::vector p_new; + std::vector sp_new; + std::vector hp_new; + }; + + struct OperatorBlock + { + std::vector vectors; + std::vector hvectors; + std::vector svectors; + }; + + void scale_to_unit_snorm(std::vector& x, + std::vector& sx, + std::vector& hx, + int ncols) const; + + void hermitize_projected(std::vector& matrix, int dim) const; + + void insert_startup_gram_block(const T* left, + const T* right, + int row_count, + int column_count, + int row_offset, + int column_offset, + int matrix_dim, + std::vector& matrix, + std::vector& workspace) const; + + void build_startup_global_subspace(const T* psi, + std::vector& w_active, + std::vector& sw_active, + std::vector& hw_active, + int active_count, + SmallSubspace& subspace); + + void update_startup_global_subspace(T* psi, + const std::vector& w_active, + const std::vector& sw_active, + const std::vector& hw_active, + int active_count, + const SmallSubspace& subspace); + + void lock_epairs(const std::vector& residual, + const std::vector& ethr_band, + std::vector& active_cols); + + void build_small_subspace(const T* psi, + const std::vector& cols, + int nblk, + SmallSubspace& subspace) const; + + // Returns the number of valid Ritz vectors, including discarded states. + // An identity fallback reports zero rather than pretending it is a solve. + int solve_small_generalized(const int dim, + const int nstates, + SmallSubspace& subspace) const; + + void collect_discarded_ritz(const SmallSubspace& subspace, + int kept_count, int valid_count, + OperatorBlock& candidates) const; + + void project_candidate_complement(const HPsiFunc& hpsi_func, + const T* psi, OperatorBlock& candidates) const; + + bool recover_competing_ritz(const HPsiFunc& hpsi_func, + T* psi, Real* eigenvalue, + bool ordered_ritz, + std::vector& active_cols, + const std::vector& ethr_band, + OperatorBlock& candidates); + + void update_one_block(T* psi, + const std::vector& cols, + int l, + int nblk, + SmallSubspace& subspace); + + bool rayleigh_ritz(T* psi, Real* eigenvalue, + std::vector& active_cols, + const std::vector& ethr_band); + + void compute_eigenpair_residual(const Real* eigenvalue); + + bool orthonormalize(T* psi); + + // Between Ritz solves the columns need not be individual eigenvectors. + // Use the block gradient, and request a Ritz solve before declaring them + // converged or changing the locked prefix. + bool update_block_gradient(const T* psi, Real* eigenvalue, + const std::vector& ethr_band); +}; + +} // namespace hsolver + +#endif // DIAGO_PPCG_H diff --git a/source/source_hsolver/hsolver.cpp b/source/source_hsolver/hsolver.cpp index 66e3db65c25..bf0f28b84c7 100644 --- a/source/source_hsolver/hsolver.cpp +++ b/source/source_hsolver/hsolver.cpp @@ -19,7 +19,8 @@ double set_diagethr_ks(const std::string basis_type, const double pw_diag_thr_init, const double diag_ethr_in, const double nelec_in, - const double scf_thr_in) + const double scf_thr_in, + const double pw_diag_thr_max) { double res_diag_ethr = diag_ethr_in; @@ -72,6 +73,13 @@ double set_diagethr_ks(const std::string basis_type, static_cast(0.1) * drho / std::max(static_cast(1.0), static_cast(nelec_in))); } + // Apply the optional bound to every electronic cycle, including ionic relaxation. + const bool supports_bound = calculation_in == "scf" || calculation_in == "relax"; + if (supports_bound && pw_diag_thr_max > 0.0) + { + res_diag_ethr = std::min(res_diag_ethr, pw_diag_thr_max); + } + // Keep the single-precision diagonalization threshold at or above the floor // to avoid excessively strict convergence requirements. if (precision_flag_in == "single") @@ -151,7 +159,8 @@ double reset_diag_ethr(std::ofstream& ofs_running, const double hsover_error, const double drho_in, const double diag_ethr_in, - const double nelec_in) + const double nelec_in, + const double pw_diag_thr_max) { double new_diag_ethr = 0.0; @@ -167,6 +176,12 @@ double reset_diag_ethr(std::ofstream& ofs_running, new_diag_ethr = 0.1 * drho_in / nelec_in; + // INPUT validation restricts a positive bound to PW KSDFT SCF calculations. + if (pw_diag_thr_max > 0.0) + { + new_diag_ethr = std::min(new_diag_ethr, pw_diag_thr_max); + } + // Keep the single-precision diagonalization threshold at or above the floor // to avoid excessively strict convergence requirements. if (precision_flag_in == "single") diff --git a/source/source_hsolver/hsolver.h b/source/source_hsolver/hsolver.h index e2dc3c7fbcf..97904a8d8bc 100644 --- a/source/source_hsolver/hsolver.h +++ b/source/source_hsolver/hsolver.h @@ -18,7 +18,8 @@ double set_diagethr_ks(const std::string basis_type, const double pw_diag_thr_init, const double diag_ethr_in, const double nelec_in, - const double scf_thr_in); + const double scf_thr_in, + const double pw_diag_thr_max); double set_diagethr_sdft(const std::string basis_type, const std::string esolver_type, @@ -43,7 +44,8 @@ double reset_diag_ethr(std::ofstream& ofs_running, const double hsover_error, const double drho_in, const double diag_ethr_in, - const double nelec_in); + const double nelec_in, + const double pw_diag_thr_max); // calculate hsolver_error (for sdft, lcao and lcao-in-pw, we suppose the error is zero) double cal_hsolve_error(const std::string basis_type, @@ -52,4 +54,4 @@ double cal_hsolve_error(const std::string basis_type, const double nelec_in); } // namespace hsolver -#endif \ No newline at end of file +#endif diff --git a/source/source_hsolver/hsolver_pw.cpp b/source/source_hsolver/hsolver_pw.cpp index 1b316b567e0..7c881daff6c 100644 --- a/source/source_hsolver/hsolver_pw.cpp +++ b/source/source_hsolver/hsolver_pw.cpp @@ -1,6 +1,8 @@ #include "hsolver_pw.h" #include "source_base/parallel_comm.h" +#include "source_base/parallel_reduce.h" +#include "source_base/module_device/memory_op.h" #include "source_base/timer.h" #include "source_base/tool_quit.h" #include "source_estate/elecstate_pw.h" @@ -10,16 +12,192 @@ #include "source_hsolver/diago_cg.h" #include "source_hsolver/diago_dav_subspace.h" #include "source_hsolver/diago_david.h" +#include "source_hsolver/diago_ppcg.h" #include "source_hsolver/diago_iter_assist.h" #include "source_psi/psi.h" #include +#include +#include +#include #include +#include +#include +#include #include namespace hsolver { +namespace +{ +struct PPCGRunResult +{ + double average_iterations; + bool converged; + int active_band_count; + bool iteration_limit_reached; + double max_residual; +}; + +template +PPCGRunResult run_ppcg_pw(const HPsiFunc& hpsi_func, + const SPsiFunc& spsi_func, + const int ld_psi, + const int nband, + const int dim, + T* psi, + Real* eigenvalue, + const std::vector& ethr_band, + const Real* pre_condition, + const double diag_thr, + const int diag_iter_max, + const int pw_diag_ndim, + const int rr_step, + std::true_type) +{ + const int sbsize = std::max(1, std::min(nband, pw_diag_ndim)); + const int rr_step_safe = std::max(1, rr_step); + const Real solver_diag_thr = static_cast(diag_thr); + + DiagoPPCG ppcg(solver_diag_thr, + diag_iter_max, + sbsize, + rr_step_safe); + + PPCGRunResult result; + result.average_iterations = ppcg.diag(hpsi_func, + spsi_func, + ld_psi, + nband, + dim, + psi, + eigenvalue, + ethr_band, + pre_condition); + result.converged = ppcg.converged(); + result.active_band_count = ppcg.active_band_count(); + result.iteration_limit_reached = ppcg.iteration_limit_reached(); + result.max_residual = double(ppcg.max_residual()); + return result; +} + +template +PPCGRunResult run_ppcg_pw(const HPsiFunc& hpsi_func, + const SPsiFunc& spsi_func, + const int ld_psi, + const int nband, + const int dim, + T* psi, + Real* eigenvalue, + const std::vector& ethr_band, + const Real* pre_condition, + const double diag_thr, + const int diag_iter_max, + const int pw_diag_ndim, + const int rr_step, + std::false_type) +{ + const int sbsize = std::max(1, std::min(nband, pw_diag_ndim)); + const int rr_step_safe = std::max(1, rr_step); + // Check the host solver's full workspace bounds before bridge allocation + // or copying, and reject unsupported shapes on every pool rank together. + const bool valid_dimensions = ppcg_workspace_dimensions_valid(ld_psi, nband, dim, sbsize); + int invalid_dimensions = valid_dimensions ? 0 : 1; + Parallel_Reduce::reduce_pool(&invalid_dimensions, 1); + if (invalid_dimensions != 0) + { + ModuleBase::WARNING_QUIT("PPCG GPU bridge", "Invalid or oversized workspace dimensions on a pool rank."); + } + const int nelem = ld_psi * nband; + const Real solver_diag_thr = static_cast(diag_thr); + + // Transitional GPU path: keep PPCG's control logic and small dense solves + // on host, while applying H/S through the device operators. + struct DeviceBuffer + { + T* ptr = nullptr; + explicit DeviceBuffer(const int size) + { + base_device::memory::resize_memory_op()(ptr, size, "PPCG device bridge"); + } + ~DeviceBuffer() + { + if (ptr != nullptr) + base_device::memory::delete_memory_op()(ptr); + } + DeviceBuffer(const DeviceBuffer&) = delete; + DeviceBuffer& operator=(const DeviceBuffer&) = delete; + }; + + std::vector psi_host(nelem, T(0)); + T* psi_host_ptr = psi_host.data(); + base_device::memory::synchronize_memory_op()( + psi_host_ptr, psi, nelem); + + DeviceBuffer psi_dev(nelem); + DeviceBuffer out_dev(nelem); + // The bridge buffers are allocated once per diagonalization and reused for + // every H/S application. Keep the leading dimension explicit: operators + // may receive padded wavefunction columns, so copying only ``dim`` would + // corrupt the column stride expected by the device implementation. + const auto copy_to_device = [&](T* host_ptr, const int ld, const int nvec) { + int invalid_block = ld <= 0 || nvec <= 0 || nvec > nelem / ld ? 1 : 0; + Parallel_Reduce::reduce_pool(&invalid_block, 1); + if (invalid_block != 0) + { + ModuleBase::WARNING_QUIT("PPCG GPU bridge", "Input block exceeds allocated workspace on a pool rank."); + } + const int count = ld * nvec; + base_device::memory::synchronize_memory_op()( + psi_dev.ptr, host_ptr, count); + }; + const auto copy_from_device = [&](T* host_ptr, const int ld, const int nvec) { + int invalid_block = ld <= 0 || nvec <= 0 || nvec > nelem / ld ? 1 : 0; + Parallel_Reduce::reduce_pool(&invalid_block, 1); + if (invalid_block != 0) + { + ModuleBase::WARNING_QUIT("PPCG GPU bridge", "Output block exceeds allocated workspace on a pool rank."); + } + const int count = ld * nvec; + base_device::memory::synchronize_memory_op()( + host_ptr, out_dev.ptr, count); + }; + auto bridge_hpsi = [&](T* psi_in, T* hpsi_out, const int ld, const int nvec) { + copy_to_device(psi_in, ld, nvec); + hpsi_func(psi_dev.ptr, out_dev.ptr, ld, nvec); + copy_from_device(hpsi_out, ld, nvec); + }; + auto bridge_spsi = [&](T* psi_in, T* spsi_out, const int ld, const int nvec) { + copy_to_device(psi_in, ld, nvec); + spsi_func(psi_dev.ptr, out_dev.ptr, ld, nvec); + copy_from_device(spsi_out, ld, nvec); + }; + + DiagoPPCG ppcg(solver_diag_thr, + diag_iter_max, + sbsize, + rr_step_safe); + PPCGRunResult result; + result.average_iterations = ppcg.diag(bridge_hpsi, + bridge_spsi, + ld_psi, + nband, + dim, + psi_host_ptr, + eigenvalue, + ethr_band, + pre_condition); + result.converged = ppcg.converged(); + result.active_band_count = ppcg.active_band_count(); + result.iteration_limit_reached = ppcg.iteration_limit_reached(); + result.max_residual = double(ppcg.max_residual()); + base_device::memory::synchronize_memory_op()( + psi, psi_host_ptr, nelem); + return result; +} +} // namespace + template void HSolverPW::cal_smooth_ethr(const double& wk, const double* wg, @@ -79,7 +257,7 @@ void HSolverPW::solve(HSOperator& op, this->nproc_in_pool = nproc_in_pool_in; // report if the specified diagonalization method is not supported - const std::initializer_list _methods = {"cg", "dav", "dav_subspace", "bpcg"}; + const std::initializer_list _methods = {"cg", "dav", "dav_subspace", "bpcg", "ppcg"}; if (std::find(std::begin(_methods), std::end(_methods), this->method) == std::end(_methods)) { ModuleBase::WARNING_QUIT("HSolverPW::solve", "This type of eigensolver is not supported!"); @@ -90,6 +268,14 @@ void HSolverPW::solve(HSOperator& op, std::vector eigenvalues(this->wfc_basis->nks * psi.get_nbands(), 0.0); ethr_band.resize(psi.get_nbands(), this->diag_thr); + std::vector diagonalization_status; + const auto solve_kpoint = [&](const int ik) { + Real* eigenvalues_k = eigenvalues.data() + ik * psi.get_nbands(); + const DiagonalizationStatus status = this->hamiltSolvePsiK( + op, psi, precondition, eigenvalues_k, this->wfc_basis->nks); + diagonalization_status.push_back(status); + }; + // Initialize k-point continuity if enabled static int count = 0; if (use_k_continuity) { @@ -131,7 +317,7 @@ void HSolverPW::solve(HSOperator& op, // solve eigenvector and eigenvalue for H(k) - this->hamiltSolvePsiK(op, psi, precondition, eigenvalues.data() + ik * psi.get_nbands(), this->wfc_basis->nks); + solve_kpoint(ik); if (skip_charge) { @@ -169,7 +355,7 @@ void HSolverPW::solve(HSOperator& op, // solve eigenvector and eigenvalue for H(k) - this->hamiltSolvePsiK(op, psi, precondition, eigenvalues.data() + ik * psi.get_nbands(), this->wfc_basis->nks); + solve_kpoint(ik); // output iteration information and reset avg_iter if (skip_charge) @@ -184,6 +370,11 @@ void HSolverPW::solve(HSOperator& op, } } // else (use_k_continuity) + if (this->method == "ppcg") + { + this->report_ppcg_convergence(diagonalization_status, log); + } + // output average iteration information and reset avg_iter this->output_iterInfo(log); @@ -224,38 +415,40 @@ void HSolverPW::solve(HSOperator& op, } template -void HSolverPW::hamiltSolvePsiK(const HSOperator& op, +typename HSolverPW::DiagonalizationStatus HSolverPW::hamiltSolvePsiK(const HSOperator& op, psi::Psi& psi, std::vector& pre_condition, Real* eigenvalue, const int& nk_nums) { ModuleBase::timer::start("HSolverPW", "solve_psik"); + DiagonalizationStatus status; #ifdef __MPI const diag_comm_info comm_info = {POOL_WORLD, this->rank_in_pool, this->nproc_in_pool}; #else const diag_comm_info comm_info = {this->rank_in_pool, this->nproc_in_pool}; #endif - const int cur_nbasis = psi.get_current_nbas(); + const int local_basis_rank = psi.get_current_nbas() * psi.get_npol(); - // Check for rank deficiency: the total number of plane waves (summed across - // all MPI processes) must be >= nbands. When npw_total < nbands, the basis is - // rank-deficient, leading to psi_norm <= 0 during Schmidt orthogonalization. - // Note: we sum cur_nbasis (local npw for this k-point) across the pool because - // psi.get_nbasis() gives the local storage dimension, not the total. + // Check for rank deficiency: the active wavefunction dimension (summed + // across all MPI processes) must be at least nbands. Otherwise Schmidt + // orthogonalization encounters a non-positive norm. + // Count both spinor components for nspin=4. psi.get_nbasis() is unsuitable + // here because it is the padded storage stride rather than the active rank. const int nbands = psi.get_nbands(); - int npw_total = cur_nbasis; + int basis_rank_total = local_basis_rank; #ifdef __MPI if (this->nproc_in_pool > 1) { - MPI_Allreduce(&cur_nbasis, &npw_total, 1, MPI_INT, MPI_SUM, POOL_WORLD); + Parallel_Reduce::reduce_pool(&basis_rank_total, 1); } #endif - if (npw_total < nbands) + if (basis_rank_total < nbands) { - const std::string msg = "npw_total < nbands (" + std::to_string(npw_total) + " < " + std::to_string(nbands) - + "): the total number of plane waves across all MPI processes " + const std::string msg = "basis_rank_total < nbands (" + std::to_string(basis_rank_total) + " < " + + std::to_string(nbands) + + "): the total active wavefunction dimension across all MPI processes " + "is less than the number of bands, " + "which leads to a rank-deficient problem. " + "Please increase ecutwfc or reduce nbands."; @@ -347,8 +540,99 @@ void HSolverPW::hamiltSolvePsiK(const HSOperator& op, ntry_max, notconv_max)); } + else if (this->method == "ppcg") + { + const auto hpsi_func = [&op](T* psi_in, T* hpsi_out, const int ld, const int nvec) { + op.hpsi(psi_in, hpsi_out, ld, nvec); + }; + const auto spsi_func = [&op](T* psi_in, T* spsi_out, const int ld, const int nvec) { + op.spsi(psi_in, spsi_out, ld, nvec); + }; + const std::is_same cpu_path; + const int ld_psi = psi.get_nbasis(); + const int nband = psi.get_nbands(); + const int dim = psi.get_current_ngk(); + T* psi_pointer = psi.get_pointer(); + const Real* precondition_pointer = pre_condition.data(); + // PPCG uses a residual-norm stopping criterion. Preserve the requested + // norm tolerance: taking its square root can leave forces and stresses + // inaccurate even when the total energy appears converged. + try + { + const PPCGRunResult ppcg_result = run_ppcg_pw( + hpsi_func, + spsi_func, + ld_psi, + nband, + dim, + psi_pointer, + eigenvalue, + this->ethr_band, + precondition_pointer, + this->diag_thr, + this->diag_iter_max, + this->pw_diag_ndim, + this->pw_diag_rr_step, + cpu_path); + DiagoIterAssist::avg_iter += ppcg_result.average_iterations; + status.converged = ppcg_result.converged; + status.iteration_limit_reached = ppcg_result.iteration_limit_reached; + status.active_band_count = ppcg_result.active_band_count; + status.max_residual = ppcg_result.max_residual; + } + catch (const std::exception& error) + { + const std::string message = std::string("PPCG diagonalization failed: ") + error.what(); + ModuleBase::WARNING_QUIT("HSolverPW::hamiltSolvePsiK", message); + } + } ModuleBase::timer::end("HSolverPW", "solve_psik"); - return; + return status; +} + +template +void HSolverPW::report_ppcg_convergence(const std::vector& status, + std::ostream& log) const +{ + int unconverged_kpoints = 0; + int iteration_limits = 0; + int active_bands = 0; + double max_residual = 0.0; + // Each pool contributes each k point once, including pools not owning + // world rank zero. All ranks enter the same final reductions. + if (this->rank_in_pool == 0) + { + for (const DiagonalizationStatus& item : status) + { + if (!item.converged) + { + ++unconverged_kpoints; + if (item.iteration_limit_reached) + { + ++iteration_limits; + } + active_bands += item.active_band_count; + max_residual = std::max(max_residual, item.max_residual); + } + } + } + Parallel_Reduce::reduce_all(unconverged_kpoints); + Parallel_Reduce::reduce_all(iteration_limits); + Parallel_Reduce::reduce_all(active_bands); + Parallel_Reduce::reduce_max(max_residual); + if (unconverged_kpoints > 0) + { + std::ostringstream warning_stream; + warning_stream << "PPCG did not reach the requested residual thresholds at " + << unconverged_kpoints << " k point(s) in electronic iteration " << this->scf_iter + << "; " << iteration_limits << " reached the iteration limit " << this->diag_iter_max + << "; active bands across those k points = " << active_bands + << "; final max residual = " << std::scientific << std::setprecision(6) + << max_residual << '.'; + const std::string warning = warning_stream.str(); + log << warning << '\n'; + ModuleBase::WARNING("HSolverPW::solve", warning); + } } template diff --git a/source/source_hsolver/hsolver_pw.h b/source/source_hsolver/hsolver_pw.h index db6764154c8..c502373f99a 100644 --- a/source/source_hsolver/hsolver_pw.h +++ b/source/source_hsolver/hsolver_pw.h @@ -44,6 +44,7 @@ class HSolverPW const int nbands_in, const bool diago_smooth_ethr_in, const int pw_diag_ndim_in, + const int pw_diag_rr_step_in, const int diag_subspace_in, const int nb2d_in, const bool use_k_continuity_in = false) @@ -51,6 +52,7 @@ class HSolverPW use_uspp(use_uspp_in), nspin(nspin_in), scf_iter(scf_iter_in), diag_iter_max(diag_iter_max_in), diag_thr(diag_thr_in), need_subspace(need_subspace_in), nbands(nbands_in), diago_smooth_ethr(diago_smooth_ethr_in), pw_diag_ndim(pw_diag_ndim_in), + pw_diag_rr_step(pw_diag_rr_step_in), diag_subspace(diag_subspace_in), nb2d(nb2d_in), use_k_continuity(use_k_continuity_in) {}; @@ -69,13 +71,24 @@ class HSolverPW const bool skip_charge); protected: + struct DiagonalizationStatus + { + bool converged = true; + bool iteration_limit_reached = false; + int active_band_count = 0; + double max_residual = 0.0; + }; + // diago caller - void hamiltSolvePsiK(const HSOperator& op, + DiagonalizationStatus hamiltSolvePsiK(const HSOperator& op, psi::Psi& psi, std::vector& pre_condition, Real* eigenvalue, const int& nk_nums); + void report_ppcg_convergence(const std::vector& status, + std::ostream& log) const; + // calculate the precondition array for diagonalization in PW base void update_precondition(std::vector& h_diag, const int ik, const int npw, const Real vl_of_0); @@ -98,6 +111,7 @@ class HSolverPW const int nbands; // global number of bands, may differ from psi.get_nbands() under band parallelism const bool diago_smooth_ethr; // use a band-wise smoothed threshold for all iter methods const int pw_diag_ndim; // dimension of the workspace for Davidson-type methods + const int pw_diag_rr_step; // interval between full PPCG Rayleigh-Ritz steps const int diag_subspace; // subspace eigensolver for dav_subspace: 0 Lapack, 1 elpa, 2 scalapack const int nb2d; // 2d block size used by the dav_subspace scalapack path diff --git a/source/source_hsolver/unittests/CMakeLists.txt b/source/source_hsolver/unittests/CMakeLists.txt index d013a06ed42..414c5166678 100644 --- a/source/source_hsolver/unittests/CMakeLists.txt +++ b/source/source_hsolver/unittests/CMakeLists.txt @@ -2,6 +2,59 @@ abacus_disable_feature_definitions(__CUDA) abacus_disable_feature_definitions(__ROCM) abacus_disable_feature_definitions(__EXX) +# Each adversarial case must run in its own process. Requiring the completion +# marker also catches historical solver failure paths that called exit(0). +add_executable(MODULE_HSOLVER_davidson_stability + test_diago_david.cpp + ../diago_david.cpp + ../diago_dav_subspace.cpp +) +target_link_libraries(MODULE_HSOLVER_davidson_stability PRIVATE + parameter base psi device container abacus::linalg_libs +) +target_compile_definitions(MODULE_HSOLVER_davidson_stability PRIVATE + ABACUS_DAVIDSON_STABILITY_HARNESS +) +if(ENABLE_MPI) + target_sources(MODULE_HSOLVER_davidson_stability PRIVATE + ../diag_hs_para.cpp + ../diago_pxxxgvx.cpp + ) + if(TARGET ELPA::ELPA) + target_link_libraries(MODULE_HSOLVER_davidson_stability PRIVATE ELPA::ELPA genelpa) + endif() +endif() +if(ENABLE_COVERAGE) + add_coverage(MODULE_HSOLVER_davidson_stability) +endif() +install(TARGETS MODULE_HSOLVER_davidson_stability DESTINATION ${CMAKE_BINARY_DIR}/tests) +foreach(solver dav dav_subspace) + foreach(case exact partial_exact collinear near_collinear degenerate generalized_spd warm_start small_direction restart_spd) + set(test_name MODULE_HSOLVER_davidson_stability_${solver}_${case}) + add_test(NAME ${test_name} + COMMAND MODULE_HSOLVER_davidson_stability ${solver} ${case} + ) + set_tests_properties(${test_name} PROPERTIES + PASS_REGULAR_EXPRESSION "STABILITY_RESULT solver=${solver} case=${case} accepted=1 " + FAIL_REGULAR_EXPRESSION "STABILITY_EXCEPTION;accepted=0" + ENVIRONMENT "OMP_NUM_THREADS=1;OPENBLAS_NUM_THREADS=1;MKL_NUM_THREADS=1" + LABELS "hsolver;davidson;stability" + TIMEOUT 60 + ) + endforeach() +endforeach() + +AddTest( + TARGET MODULE_HSOLVER_cg_accuracy + LIBS parameter base psi device container + SOURCES test_diago_cg_accuracy.cpp ../diago_cg.cpp ../diago_iter_assist.cpp ../diag_const_nums.cpp +) +set_tests_properties(MODULE_HSOLVER_cg_accuracy PROPERTIES + PASS_REGULAR_EXPRESSION "CG_ACCURACY_COMPLETE status=0" + ENVIRONMENT "OMP_NUM_THREADS=1;OPENBLAS_NUM_THREADS=1;MKL_NUM_THREADS=1" + TIMEOUT 60 +) + if (ENABLE_MPI) AddTest( @@ -63,9 +116,11 @@ if (ENABLE_MPI) AddTest( TARGET MODULE_HSOLVER_pw LIBS parameter psi device base container symmetry - SOURCES test_hsolver_pw.cpp ../hsolver_pw.cpp ../hsolver_lcaopw.cpp ../diago_bpcg.cpp ../diago_dav_subspace.cpp ../diag_const_nums.cpp ../diago_iter_assist.cpp ../para_lin_tf.cpp - ../../source_estate/elecstate_tools.cpp ../../source_estate/occupy.cpp ../../source_base/module_fft/fft_bundle.cpp ../../source_base/module_fft/fft_cpu.cpp - ../../source_cell/klist.cpp ../../source_cell/klist_io.cpp ../../source_cell/parallel_kpoints.cpp ../../source_cell/reciprocal_grid.cpp + SOURCES test_hsolver_pw.cpp ../hsolver_pw.cpp ../hsolver_lcaopw.cpp ../diago_bpcg.cpp ../diago_dav_subspace.cpp ../diago_ppcg.cpp ../diag_const_nums.cpp ../diago_iter_assist.cpp ../para_lin_tf.cpp + ../../source_estate/elecstate_tools.cpp ../../source_estate/occupy.cpp ../../source_base/module_fft/fft_bundle.cpp ../../source_base/module_fft/fft_cpu.cpp + ../../source_hamilt/module_xc/exx_info.cpp + ../../source_cell/klist.cpp ../../source_cell/klist_io.cpp ../../source_cell/parallel_kpoints.cpp ../../source_cell/reciprocal_grid.cpp + ../../source_basis/module_pw/test/test_tool.cpp ) if(ENABLE_LCAO) @@ -102,6 +157,39 @@ if (ENABLE_MPI) target_compile_definitions(MODULE_HSOLVER_LCAO_cusolver PRIVATE __CUDA) endif() endif() +AddTest( + TARGET MODULE_HSOLVER_ppcg + LIBS parameter base psi device container + SOURCES test_diago_ppcg.cpp ../diago_ppcg.cpp ../../source_basis/module_pw/test/test_tool.cpp +) +AddTest( + TARGET MODULE_HSOLVER_diago_params + LIBS parameter base psi device + SOURCES test_diago_params.cpp ../diago_params.cpp +) +AddTest( + TARGET MODULE_HSOLVER_ppcg_float + LIBS parameter base psi device container + SOURCES test_diago_ppcg_float.cpp ../diago_ppcg.cpp ../../source_basis/module_pw/test/test_tool.cpp +) + +if (ENABLE_MPI) +AddTest( + TARGET MODULE_HSOLVER_compare + LIBS parameter base psi device container + SOURCES test_diago_compare.cpp ../diago_cg.cpp ../diago_bpcg.cpp ../diago_david.cpp ../diago_ppcg.cpp ../diago_iter_assist.cpp ../diag_const_nums.cpp ../para_lin_tf.cpp ../../source_basis/module_pw/test/test_tool.cpp +) +set_tests_properties(MODULE_HSOLVER_compare PROPERTIES + DISABLED TRUE + LABELS benchmark +) +AddTest( + TARGET MODULE_HSOLVER_ppcg_parallel + LIBS parameter base psi device container + SOURCES test_diago_ppcg_parallel.cpp ../diago_ppcg.cpp ../../source_basis/module_pw/test/test_tool.cpp +) +endif() + abacus_copy_test_files( H-KPoints-Si2.dat H-GammaOnly-Si2.dat @@ -118,6 +206,7 @@ abacus_copy_test_files( diago_cg_parallel_test.sh diago_david_parallel_test.sh diago_lcao_parallel_test.sh + diago_ppcg_parallel_test.sh PEXSI-H-GammaOnly-Si2.dat PEXSI-S-GammaOnly-Si2.dat PEXSI-DM-GammaOnly-Si2.dat @@ -179,6 +268,11 @@ if (ENABLE_MPI) "$" WORKING_DIRECTORY ${CMAKE_CURRENT_BINARY_DIR} ) + add_test(NAME MODULE_HSOLVER_ppcg_parallel_test + COMMAND ${BASH} ${CMAKE_CURRENT_SOURCE_DIR}/diago_ppcg_parallel_test.sh + "$" + WORKING_DIRECTORY ${CMAKE_CURRENT_BINARY_DIR} + ) if(ENABLE_LCAO) add_test(NAME MODULE_HSOLVER_LCAO_parallel COMMAND ${BASH} diago_lcao_parallel_test.sh diff --git a/source/source_hsolver/unittests/diago_ppcg_parallel_test.sh b/source/source_hsolver/unittests/diago_ppcg_parallel_test.sh new file mode 100644 index 00000000000..6c2cd2625a0 --- /dev/null +++ b/source/source_hsolver/unittests/diago_ppcg_parallel_test.sh @@ -0,0 +1,44 @@ +#!/bin/bash + +ppcg_binary=${1:-./MODULE_HSOLVER_ppcg_parallel} +np=$(nproc --all) +echo "nprocs in this machine is $np" + +runs=0 +for i in 2 3 6; do + if [[ $i -gt $np ]]; then + continue + fi + echo "TEST DIAGO PPCG in parallel, nprocs=$i" + OMP_NUM_THREADS=1 mpirun -np $i "$ppcg_binary" + e=$? + if [[ $e -ne 0 ]]; then + echo -e "\e[1;33m [ FAILED ] \e[0m"\ + "execute UT with $i cores error." + exit 1 + fi + runs=$((runs + 1)) +done + +if [[ $np -ge 4 ]]; then + echo "TEST DIAGO PPCG with two pools, nprocs=4" + OMP_NUM_THREADS=1 mpirun -np 4 "$ppcg_binary" 2 + e=$? + if [[ $e -ne 0 ]]; then + echo -e "\e[1;33m [ FAILED ] \e[0m"\ + "execute two-pool UT with 4 cores error." + exit 1 + fi + runs=$((runs + 1)) +fi + +if [[ $runs -eq 0 ]]; then + echo "No 2/3/6-process case is supported; running the 1-process fallback." + OMP_NUM_THREADS=1 mpirun -np 1 "$ppcg_binary" + e=$? + if [[ $e -ne 0 ]]; then + echo -e "\e[1;33m [ FAILED ] \e[0m"\ + "execute UT with 1 core error." + exit 1 + fi +fi diff --git a/source/source_hsolver/unittests/test_diago_cg_accuracy.cpp b/source/source_hsolver/unittests/test_diago_cg_accuracy.cpp new file mode 100644 index 00000000000..bb8e0f8d868 --- /dev/null +++ b/source/source_hsolver/unittests/test_diago_cg_accuracy.cpp @@ -0,0 +1,315 @@ +#include "source_hsolver/diago_cg.h" +#include "source_base/parallel_global.h" +#include "source_hsolver/hs_operator.h" +#include +#include +#include +#include +#include +#include +#include + +namespace +{ +using Complex = std::complex; +template +class DiagonalCGOperator : public hsolver::HSOperator> +{ + public: + using Value = std::complex; + explicit DiagonalCGOperator(const std::vector& diagonal) : diagonal_(diagonal) {} + void update_k(const int) override {} + void hpsi(const Value* x, Value* hx, const int ld, const int columns) const override + { + for (int column = 0; column < columns; ++column) + { + for (int row = 0; row < static_cast(diagonal_.size()); ++row) + { + hx[column * ld + row] = diagonal_[row] * x[column * ld + row]; + } + } + } + void spsi(const Value* x, Value* sx, const int ld, const int columns) const override + { + const Value* end = x + ld * columns; + std::copy(x, end, sx); + } + private: + const std::vector& diagonal_; +}; + +class ShiftedTwoByTwoCGOperator : public hsolver::HSOperator +{ + public: + ShiftedTwoByTwoCGOperator(const double shift, const double coupling) + : shift_(shift), coupling_(coupling) {} + void update_k(const int) override {} + void hpsi(const Complex* x, Complex* hx, const int ld, const int columns) const override + { + for (int column = 0; column < columns; ++column) + { + const int first = column * ld; + const int second = first + 1; + hx[first] = shift_ * x[first] + coupling_ * x[second]; + hx[second] = coupling_ * x[first] + shift_ * x[second]; + } + } + void spsi(const Complex* x, Complex* sx, const int ld, const int columns) const override + { + const Complex* end = x + ld * columns; + std::copy(x, end, sx); + } + private: + const double shift_; + const double coupling_; +}; + +TEST(DiagoCGAccuracy, LargeCommonEnergyChoosesLowerRoot) +{ + const int dimension = 2; + const int bands = 1; + const double shift = std::ldexp(1.0, 30); + const double coupling = std::ldexp(1.0, -30); + ShiftedTwoByTwoCGOperator op(shift, coupling); + std::vector psi{Complex(1.0, 0.0), Complex(0.0, 0.0)}; + std::vector eigen(bands, 0.0); + std::vector threshold(bands, 1.0e-20); + std::vector prec(dimension, 1.0); +#ifdef __MPI + const hsolver::diag_comm_info comm(MPI_COMM_SELF, 0, 1); +#else + const hsolver::diag_comm_info comm(0, 1); +#endif + hsolver::DiagoCG solver("pw", "scf", false, comm, 1.0e-20, 100); + solver.diag(op, dimension, bands, dimension, psi.data(), eigen.data(), threshold, prec.data()); + // Both eigenvalues round to shift; the exact lower mode still has opposite components. + const Complex upper_component = psi[0] + psi[1]; + const double error = std::abs(upper_component); + ASSERT_TRUE(std::isfinite(error)); + const double norm = std::norm(psi[0]) + std::norm(psi[1]); + ASSERT_TRUE(std::isfinite(norm)); + EXPECT_NEAR(norm, 1.0, 1.0e-12); + EXPECT_DOUBLE_EQ(eigen[0], shift); + + EXPECT_LT(error, 1.0e-12); +} + +class GeneralizedCGOperator : public DiagonalCGOperator +{ + public: + GeneralizedCGOperator(const std::vector& h, const std::vector& metric) + : DiagonalCGOperator(h), metric_(metric) {} + void spsi(const Complex* x, Complex* sx, const int ld, const int columns) const override + { + for (int column = 0; column < columns; ++column) + { + for (int row = 0; row < static_cast(metric_.size()); ++row) + { + sx[column * ld + row] = metric_[row] * x[column * ld + row]; + } + } + } + private: + const std::vector& metric_; +}; + +TEST(DiagoCGAccuracy, GeneralizedPreconditionedNegativeRoot) +{ + const int dimension = 2; + const int bands = 1; + std::vector h{-6.0, -3.0}; + std::vector metric{2.0, 3.0}; + GeneralizedCGOperator op(h, metric); + const double first = 1.0 / std::sqrt(2.0); + const double second = 1.0e-9 / std::sqrt(6.0); + std::vector psi{Complex(first, 0.0), Complex(second, second)}; + std::vector eigen(bands, 0.0); + std::vector threshold(bands, 1.0e-20); + std::vector prec{2.0, 10.0}; +#ifdef __MPI + const hsolver::diag_comm_info comm(MPI_COMM_SELF, 0, 1); +#else + const hsolver::diag_comm_info comm(0, 1); +#endif + hsolver::DiagoCG solver("pw", "scf", false, comm, 1.0e-20, 100); + solver.diag(op, dimension, bands, dimension, psi.data(), eigen.data(), threshold, prec.data()); + double residual_squared = 0.0; + double norm = 0.0; + for (int row = 0; row < dimension; ++row) + { + const Complex residual = (h[row] - eigen[0] * metric[row]) * psi[row]; + residual_squared += std::norm(residual); + norm += metric[row] * std::norm(psi[row]); + } + const double residual = std::sqrt(residual_squared); + ASSERT_TRUE(std::isfinite(residual)); + ASSERT_TRUE(std::isfinite(norm)); + EXPECT_NEAR(eigen[0], -3.0, 1.0e-14); + EXPECT_NEAR(norm, 1.0, 1.0e-14); + EXPECT_LT(residual, 1.0e-12); +} + +TEST(DiagoCGAccuracy, FloatNearlyConvergedDirection) +{ + using FloatComplex = std::complex; + const int dimension = 2; + const int bands = 1; + std::vector diagonal{1.0f, 2.0f}; + DiagonalCGOperator op(diagonal); + std::vector psi{FloatComplex(1.0f, 0.0f), FloatComplex(1.0e-4f, 0.0f)}; + std::vector eigen(bands, 0.0f); + std::vector threshold(bands, 1.0e-12); + std::vector prec(dimension, 1.0f); +#ifdef __MPI + const hsolver::diag_comm_info comm(MPI_COMM_SELF, 0, 1); +#else + const hsolver::diag_comm_info comm(0, 1); +#endif + hsolver::DiagoCG solver("pw", "scf", false, comm, 1.0e-12f, 100); + solver.diag(op, dimension, bands, dimension, psi.data(), eigen.data(), threshold, prec.data()); + double residual_squared = 0.0; + for (int row = 0; row < dimension; ++row) + { + const FloatComplex residual = (diagonal[row] - eigen[0]) * psi[row]; + residual_squared += std::norm(residual); + } + const double residual = std::sqrt(residual_squared); + ASSERT_TRUE(std::isfinite(residual)); + const double norm = std::norm(psi[0]) + std::norm(psi[1]); + ASSERT_TRUE(std::isfinite(norm)); + EXPECT_NEAR(norm, 1.0, 1.0e-6); + EXPECT_NEAR(eigen[0], 1.0, 1.0e-6); + + EXPECT_LT(residual, 1.0e-6); +} + +TEST(DiagoCGAccuracy, NearlyConvergedDirection) +{ + const int dimension = 2; + const int bands = 1; + std::vector diagonal{1.0, 2.0}; + DiagonalCGOperator op(diagonal); + std::vector psi{Complex(1.0, 0.0), Complex(1.0e-9, 0.0)}; + std::vector eigen(bands, 0.0); + std::vector threshold(bands, 1.0e-20); + std::vector prec(dimension, 1.0); +#ifdef __MPI + const hsolver::diag_comm_info comm(MPI_COMM_SELF, 0, 1); +#else + const hsolver::diag_comm_info comm(0, 1); +#endif + hsolver::DiagoCG solver("pw", "scf", false, comm, 1.0e-20, 100); + solver.diag(op, dimension, bands, dimension, psi.data(), eigen.data(), threshold, prec.data()); + double residual_squared = 0.0; + for (int row = 0; row < dimension; ++row) + { + const Complex residual = (diagonal[row] - eigen[0]) * psi[row]; + residual_squared += std::norm(residual); + } + const double residual = std::sqrt(residual_squared); + std::cout << std::setprecision(17) << "CG_PROBE eigen=" << eigen[0] << " residual=" << residual << std::endl; + ASSERT_TRUE(std::isfinite(residual)); + const double norm = std::norm(psi[0]) + std::norm(psi[1]); + ASSERT_TRUE(std::isfinite(norm)); + EXPECT_NEAR(norm, 1.0, 1.0e-12); + EXPECT_NEAR(eigen[0], 1.0, 1.0e-14); + + EXPECT_LT(residual, 1.0e-12); +} +TEST(DiagoCGAccuracy, TightMultiBandOrthogonality) +{ + const int dimension = 48; + const int bands = 6; + std::vector diagonal(dimension); + for (int row = 0; row < dimension; ++row) + { + diagonal[row] = 1.0 + 0.1 * row; + } + DiagonalCGOperator op(diagonal); + std::vector psi(dimension * bands); + for (int column = 0; column < bands; ++column) + { + for (int row = 0; row < dimension; ++row) + { + const double real_angle = (row + 1) * (column + 2) * 0.75; + const double imag_angle = (row + 3) * (column + 1) * 0.37; + const double a = std::sin(real_angle); + const double b = std::cos(imag_angle); + psi[row + dimension * column] = Complex(a, b); + } + } + std::vector eigen(bands, 0.0); + std::vector threshold(bands, 1.0e-20); + std::vector prec(dimension, 1.0); +#ifdef __MPI + const hsolver::diag_comm_info comm(MPI_COMM_SELF, 0, 1); +#else + const hsolver::diag_comm_info comm(0, 1); +#endif + hsolver::DiagoCG solver("pw", "scf", true, comm, 1.0e-20, 500); + solver.diag(op, dimension, bands, dimension, psi.data(), eigen.data(), threshold, prec.data()); + double max_residual = 0.0; + double max_gram = 0.0; + for (int column = 0; column < bands; ++column) + { + double residual_squared = 0.0; + for (int row = 0; row < dimension; ++row) + { + const Complex residual = (diagonal[row] - eigen[column]) * psi[row + column * dimension]; + residual_squared += std::norm(residual); + } + const double residual = std::sqrt(residual_squared); + ASSERT_TRUE(std::isfinite(residual)); + max_residual = std::max(max_residual, residual); + const double eigen_error = std::abs(eigen[column] - diagonal[column]); + EXPECT_LT(eigen_error, 1.0e-8); + for (int other = 0; other < bands; ++other) + { + Complex gram(0.0, 0.0); + for (int row = 0; row < dimension; ++row) + { + gram += std::conj(psi[row + column * dimension]) * psi[row + other * dimension]; + } + if (column == other) { gram -= Complex(1.0, 0.0); } + const double error = std::abs(gram); + ASSERT_TRUE(std::isfinite(error)); + max_gram = std::max(max_gram, error); + } + } + std::cout << std::setprecision(17) << "CG_MULTI residual=" << max_residual << " gram=" << max_gram << std::endl; + EXPECT_LT(max_residual, 1.0e-8); + EXPECT_LT(max_gram, 1.0e-10); +} + +} + +int main(int argc, char** argv) +{ + int nproc = 1; + int threads = 1; + int rank = 0; + Parallel_Global::read_pal_param(argc, argv, nproc, threads, rank); +#ifdef __MPI + POOL_WORLD = MPI_COMM_WORLD; + KP_WORLD = MPI_COMM_NULL; + BP_WORLD = MPI_COMM_NULL; + INT_BGROUP = MPI_COMM_NULL; + GRID_WORLD = MPI_COMM_NULL; + DIAG_WORLD = MPI_COMM_NULL; +#endif + if (nproc != 1) + { + std::cerr << "CG accuracy tests require one process." << std::endl; +#ifdef __MPI + Parallel_Global::finalize_mpi(); +#endif + return 2; + } + testing::InitGoogleTest(&argc, argv); + const int status = RUN_ALL_TESTS(); + std::cout << "CG_ACCURACY_COMPLETE status=" << status << std::endl; +#ifdef __MPI + Parallel_Global::finalize_mpi(); +#endif + return status; +} diff --git a/source/source_hsolver/unittests/test_diago_compare.cpp b/source/source_hsolver/unittests/test_diago_compare.cpp new file mode 100644 index 00000000000..68949ad300e --- /dev/null +++ b/source/source_hsolver/unittests/test_diago_compare.cpp @@ -0,0 +1,545 @@ +/** + * diago_compare_test.cpp — head-to-head comparison of the iterative + * diagonalization solvers available in source_hsolver, on identical + * random Hermitian matrices: + * - PPCG (DiagoPPCG, BLOCK_SUBSPACE) + * - CG (DiagoCG, band-by-band Polak-Ribiere) + * - BPCG (DiagoBPCG, block PCG) + * - Davidson (DiagoDavid) + * + * Every solver is fed the SAME Hamiltonian, the SAME initial guess and the + * SAME numerical per-band threshold. Each result is checked against the SAME + * LAPACK-reference accuracy (max_eval_err < err_target). The solvers retain + * their own stopping rules and call budgets, so timings must be read together + * with the reported accuracy. + * + * This is a benchmark/audit aid, not a correctness unit test: it is DISABLED + * by default and must be run explicitly. + */ + +#include "../diago_ppcg.h" +#include "../diago_cg.h" +#include "../diago_bpcg.h" +#include "../diago_david.h" + +#include "source_base/module_external/lapack_connector.h" +#include "source_base/parallel_comm.h" +#include "source_basis/module_pw/test/test_tool.h" + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#if defined(__GLIBC__) && (__GLIBC__ > 2 || (__GLIBC__ == 2 && __GLIBC_MINOR__ >= 33)) +#include +#endif + +using T = std::complex; +using Real = double; + +// Optional PPCG parameter overrides (set from argv) for exploring the block +// size (sbsize) and Rayleigh-Ritz frequency (rr_step). A negative value keeps +// the default used by the comparison benchmark (sbsize = nband, rr_step = 16). +static int g_sbsize = -1; +static int g_rr_step = -1; + +// --------------------------------------------------------------------------- +// Common acceptance criterion: max_eval_err < err_target against LAPACK. +// Identical ethr values have different internal meanings: PPCG checks residual +// norms, CG and Davidson check eigenvalue changes, and BPCG checks its gradient +// error. These quantities do not imply the same eigenvalue accuracy. +// +// PPCG, CG and BPCG are re-driven for up to max_outer_passes calls to diag(), +// checking the common target after each call. Davidson receives one call with +// its own iteration budget. The final error remains visible when a solver +// misses the target; the run then returns a nonzero status. This fixed target +// is a benchmark choice, not a guarantee relating ethr to eigenvalue error. +// --------------------------------------------------------------------------- +const double err_target = 1e-5; // common LAPACK-reference acceptance threshold +const int max_outer_passes = 20; // outer diag() re-drives before giving up + +// max over bands of |eval_i - ref_i|: the reference-based accuracy that makes +// wall-clock times comparable across solvers. +static double max_eval_err(const Real* eval, const Real* ref, int nband) +{ + double err = 0.0; + for (int i = 0; i < nband; ++i) + { + const double band_err = std::abs(eval[i] - ref[i]); + if (!std::isfinite(band_err)) + { + return std::numeric_limits::infinity(); + } + err = std::max(err, band_err); + } + return err; +} + +// Allocator-reported process heap in bytes. Samples before solver construction +// and after diag() measure retained heap changes while the solver is alive. +// Temporary workspaces freed inside diag() are absent; allocator caching and +// other process allocations can affect the result. This is not peak memory. +// mallinfo2 is optional and available only with glibc >= 2.33. +static double heap_bytes() +{ +#if defined(__GLIBC__) && (__GLIBC__ > 2 || (__GLIBC__ == 2 && __GLIBC_MINOR__ >= 33)) + const struct mallinfo2 mi = mallinfo2(); + return static_cast(mi.uordblks) + static_cast(mi.hblkhd); +#else + return std::numeric_limits::quiet_NaN(); +#endif +} + +static std::string format_heap_delta(const double retained_heap_delta_bytes) +{ + if (!std::isfinite(retained_heap_delta_bytes)) + { + return "unavailable"; + } + const double retained_heap_delta_mib = retained_heap_delta_bytes / 1048576.0; + char text[32]; + std::snprintf(text, sizeof(text), "%.2f", retained_heap_delta_mib); + return text; +} + +// Sparse symmetric band of half-bandwidth `bw`, stored in LAPACK upper-band +// format: H[i, j] (i <= j, d = j - i) lives at band[bd*j + (bw - d)], +// with bd = bw + 1. This matvec is O(n * bw), modeling the real plane-wave H +// application (kinetic energy is a diagonal/tridiagonal discretization of +// -Laplacian/2 plus a local potential) instead of an O(n^2) dense zgemm. +static void banded_h_multiply(const Real* band, int n, int bw, int bd, const T* in, T* out, int ld, int ncol) +{ + for (int j = 0; j < ncol; ++j) + { + for (int i = 0; i < n; ++i) + { + Real acc = 0.0; + const int hi = std::min(n - 1, i + bw); + for (int k = std::max(0, i - bw); k <= hi; ++k) + { + // H[i, k]: reuse the upper-triangular stored entry (k may be < i) + const int r = std::min(i, k); + const int c = std::max(i, k); + const int d = c - r; + acc += band[bd * c + (bw - d)] * std::real(in[k + j * ld]); + } + out[i + j * ld] = T(acc, 0.0); + } + } +} + +static void identity_s(const T* in, T* out, int ld, int ncol) +{ + for (int j = 0; j < ncol; ++j) + { + for (int i = 0; i < ld; ++i) + { + out[i + j * ld] = in[i + j * ld]; + } + } +} + +extern "C" void dsbev_(const char* jobz, const char* uplo, const int* n, const int* kd, + double* ab, const int* ldab, double* w, double* z, const int* ldz, + double* work, int* info); + +// Reference eigenvalues via LAPACK dsbev on the symmetric band matrix. Only +// the lowest `nband_req` are kept. dsbev is O(n * bw^2), avoiding the O(n^3) +// tridiagonalization that dense zheev/zheevx would require for large n. +static void ref_eigen(const Real* band, int n, int bw, int bd, int nband_req, Real* e) +{ + std::vector ab(band, band + size_t(bd) * n); + std::vector w(n); + const char jobz = 'N'; + const char uplo = 'U'; + std::vector work(3 * n); + int info = 0; + dsbev_(&jobz, &uplo, &n, &bw, ab.data(), &bd, w.data(), nullptr, &n, work.data(), &info); + if (info != 0) + { + std::fprintf(stderr, "[ref_eigen] dsbev info=%d\n", info); + } + for (int i = 0; i < nband_req; ++i) + { + e[i] = w[i]; + } +} + +// Diagonally-dominant symmetric band matrix: a local potential on the diagonal +// plus random couplings within a half-bandwidth `bw`. Stored in LAPACK upper +// band format `band[bd*j + (bw - (j - i))] = H[i, j]` for i <= j, bd = bw + 1. +// This is the discrete analogue of H = -Laplacian/2 + V(r) that plane-wave +// solvers actually apply. +static void make_H(int n, int bw, std::vector& band, int& bd, std::vector& prec) +{ + bd = bw + 1; + band.assign(size_t(bd) * n, 0.0); + std::mt19937 rng(unsigned(n * 100 + bw)); + std::uniform_real_distribution dist(-1.0, 1.0); + for (int i = 0; i < n; ++i) + { + band[bd * i + bw] = std::abs(dist(rng)) * n + 1.0; // diagonal (d=0) + } + for (int i = 0; i < n; ++i) + { + for (int d = 1; d <= bw; ++d) + { + if (i + d < n) + { + band[bd * (i + d) + (bw - d)] = dist(rng) * 0.5 / double(d); + } + } + } + prec.resize(n); + for (int i = 0; i < n; ++i) + { + prec[i] = std::max(band[bd * i + bw], 1e-6); + } +} + +// Random orthonormalized initial guess (identical for every solver). +static void make_psi(int n, int nband, std::vector& psi) +{ + int ld = n; + psi.assign(ld * nband, T(0)); + std::mt19937 rng(42); + std::uniform_real_distribution dist(-1.0, 1.0); + for (int j = 0; j < nband; ++j) + { + for (int i = 0; i < n; ++i) + { + psi[i + j * ld] = T(dist(rng), 0.0); + } + } + for (int j = 0; j < nband; ++j) + { + for (int k = 0; k < j; ++k) + { + T d = 0; + for (int i = 0; i < n; ++i) + { + d += std::conj(psi[i + k * ld]) * psi[i + j * ld]; + } + for (int i = 0; i < n; ++i) + { + psi[i + j * ld] -= d * psi[i + k * ld]; + } + } + Real nr = 0; + for (int i = 0; i < n; ++i) + { + nr += std::norm(psi[i + j * ld]); + } + nr = std::sqrt(nr); + for (int i = 0; i < n; ++i) + { + psi[i + j * ld] /= nr; + } + } +} + +// Rayleigh-Ritz subspace diagonalization used as CG's subspace_func. +static void rr_subspace(const Real* band, int n, int bw, int bd, T* psi_in, T* psi_out, int ld, int nband) +{ + std::vector hpsi(size_t(n) * nband, T(0)); + banded_h_multiply(band, n, bw, bd, psi_in, hpsi.data(), n, nband); + + // S_sub = Psi^H Psi (S = I), H_sub = Psi^H H Psi + std::vector s_sub(nband * nband, T(0)); + std::vector h_sub(nband * nband, T(0)); + for (int i = 0; i < nband; ++i) + { + for (int j = 0; j < nband; ++j) + { + T s = 0; + T h = 0; + for (int k = 0; k < n; ++k) + { + T pk = psi_in[k + i * ld]; + s += std::conj(pk) * psi_in[k + j * ld]; + h += std::conj(pk) * hpsi[k + j * n]; + } + s_sub[i + j * nband] = s; + h_sub[i + j * nband] = h; + } + } + + // Generalized Hermitian eigenproblem: H_sub C = S_sub C Lambda + int lwork = 2 * nband; + std::vector work(lwork); + std::vector rwork(3 * nband - 2); + std::vector w(nband); + int info = 0; + int itype = 1; + int nn = nband; + char jobz = 'V'; + char uplo = 'U'; + zhegv_(&itype, &jobz, &uplo, &nn, h_sub.data(), &nn, s_sub.data(), &nn, w.data(), work.data(), &lwork, rwork.data(), + &info); + + // psi_out = psi_in * C (C now holds the eigenvectors in h_sub) + for (int j = 0; j < nband; ++j) + { + for (int i = 0; i < n; ++i) + { + T acc = 0; + for (int c = 0; c < nband; ++c) + { + acc += psi_in[i + c * ld] * h_sub[c + j * nband]; + } + psi_out[i + j * ld] = acc; + } + } +} + +struct Result +{ + double wall_s = 0.0; + double max_err = 0.0; // max |eval_i - ref_i| over the requested bands + double retained_heap_delta_bytes = std::numeric_limits::quiet_NaN(); + bool ok = false; +}; + +class CompareOperator : public hsolver::HSOperator +{ + public: + using Apply = std::function; + + CompareOperator(Apply h, Apply s) : h_(std::move(h)), s_(std::move(s)) + { + } + + void update_k(const int) override + { + } + + void hpsi(const T* x, T* hx, const int ld, const int nvec) const override + { + h_(x, hx, ld, nvec); + } + + void spsi(const T* x, T* sx, const int ld, const int nvec) const override + { + s_(x, sx, ld, nvec); + } + + private: + Apply h_; + Apply s_; +}; + +static Result run_ppcg(const std::vector& band, int n, int bw, int bd, int nband, const std::vector& prec, + const std::vector& psi0, const std::vector& ethr, const Real* ref) +{ + Result r; + std::vector psi = psi0; + std::vector eval(nband, 0.0); + const double heap_before = heap_bytes(); + const int sbsize = (g_sbsize > 0) ? g_sbsize : nband; + const int rr_step = (g_rr_step > 0) ? g_rr_step : 16; + hsolver::DiagoPPCG solver(1e-8, 500, sbsize, rr_step); + auto h_op = [&band, n, bw, bd](T* in, T* out, int ld, int nc) { banded_h_multiply(band.data(), n, bw, bd, in, out, ld, nc); }; + auto t0 = std::chrono::high_resolution_clock::now(); + int pass = 0; + for (; pass < max_outer_passes; ++pass) + { + solver.diag(h_op, nullptr, n, nband, n, psi.data(), eval.data(), ethr, prec.data()); + if (max_eval_err(eval.data(), ref, nband) < err_target) + { + break; + } + } + auto t1 = std::chrono::high_resolution_clock::now(); + r.wall_s = std::chrono::duration(t1 - t0).count(); + r.retained_heap_delta_bytes = heap_bytes() - heap_before; + r.max_err = max_eval_err(eval.data(), ref, nband); + r.ok = std::isfinite(r.max_err) && r.max_err < err_target; + return r; +} + +static Result run_cg(const std::vector& band, int n, int bw, int bd, int nband, const std::vector& prec, + const std::vector& psi0, const std::vector& ethr, const Real* ref) +{ + Result r; + std::vector psi = psi0; + std::vector eval(nband, 0.0); + const double heap_before = heap_bytes(); + hsolver::diag_comm_info comm(MPI_COMM_WORLD, 0, 1); + hsolver::DiagoCG cg("pw", "scf", true, comm, 1e-8, 500); + CompareOperator op( + [&band, n, bw, bd](const T* in, T* out, int ld, int nc) { + banded_h_multiply(band.data(), n, bw, bd, in, out, ld, nc); + }, + [](const T* in, T* out, int ld, int nc) { identity_s(in, out, ld, nc); }); + auto t0 = std::chrono::high_resolution_clock::now(); + int pass = 0; + for (; pass < max_outer_passes; ++pass) + { + cg.diag(op, n, nband, n, psi.data(), eval.data(), ethr, prec.data()); + if (max_eval_err(eval.data(), ref, nband) < err_target) + { + break; + } + } + + auto t1 = std::chrono::high_resolution_clock::now(); + r.wall_s = std::chrono::duration(t1 - t0).count(); + r.retained_heap_delta_bytes = heap_bytes() - heap_before; + r.max_err = max_eval_err(eval.data(), ref, nband); + r.ok = std::isfinite(r.max_err) && r.max_err < err_target; + return r; +} + +static Result run_bpcg(const std::vector& band, int n, int bw, int bd, int nband, const std::vector& prec, + const std::vector& psi0, const std::vector& ethr, const Real* ref) +{ + Result r; + std::vector psi = psi0; + std::vector eval(nband, 0.0); + const double heap_before = heap_bytes(); + hsolver::DiagoBPCG bpcg(prec.data()); + bpcg.init_iter(nband, nband, n, n); + CompareOperator op( + [&band, n, bw, bd](const T* in, T* out, int ld, int nc) { + banded_h_multiply(band.data(), n, bw, bd, in, out, ld, nc); + }, + [](const T* in, T* out, int ld, int nc) { identity_s(in, out, ld, nc); }); + // Each BPCG call has an internal sweep budget; re-drive for the common target. + int it = 0; + auto t0 = std::chrono::high_resolution_clock::now(); + for (; it < max_outer_passes; ++it) + { + bpcg.diag(op, psi.data(), eval.data(), ethr); + if (max_eval_err(eval.data(), ref, nband) < err_target) + { + break; + } + } + auto t1 = std::chrono::high_resolution_clock::now(); + r.wall_s = std::chrono::duration(t1 - t0).count(); + r.retained_heap_delta_bytes = heap_bytes() - heap_before; + r.max_err = max_eval_err(eval.data(), ref, nband); + r.ok = std::isfinite(r.max_err) && r.max_err < err_target; + return r; +} + +static Result run_dav(const std::vector& band, int n, int bw, int bd, int nband, const std::vector& prec, + const std::vector& psi0, const std::vector& ethr, const Real* ref) +{ + Result r; + std::vector psi = psi0; + std::vector eval(nband, 0.0); + hsolver::diag_comm_info comm(MPI_COMM_WORLD, 0, 1); + const double heap_before = heap_bytes(); + hsolver::DiagoDavid dav(prec.data(), nband, n, 4, comm); + CompareOperator op( + [&band, n, bw, bd](const T* in, T* out, int ld, int nc) { + banded_h_multiply(band.data(), n, bw, bd, in, out, ld, nc); + }, + [](const T* in, T* out, int ld, int nc) { identity_s(in, out, ld, nc); }); + auto t0 = std::chrono::high_resolution_clock::now(); + // Davidson expands its basis up to david_ndim * nband and restarts at the + // cap. This benchmark uses one call with a 500-iteration budget. + dav.diag(op, n, psi.data(), eval.data(), ethr, 500); + auto t1 = std::chrono::high_resolution_clock::now(); + r.wall_s = std::chrono::duration(t1 - t0).count(); + r.retained_heap_delta_bytes = heap_bytes() - heap_before; + r.max_err = max_eval_err(eval.data(), ref, nband); + r.ok = std::isfinite(r.max_err) && r.max_err < err_target; + return r; +} + +int main(int argc, char** argv) +{ + int nproc = 1; + int myrank = 0; + int nproc_in_pool = 0; + int kpar = 1; + int mypool = 0; + int rank_in_pool = 0; + setupmpi(argc, argv, nproc, myrank); + divide_pools(nproc, myrank, nproc_in_pool, kpar, mypool, rank_in_pool); + BP_WORLD = MPI_COMM_WORLD; + + struct Case + { + int n; + int nband; + int bw; + }; + // Without arguments a small default grid is used. To benchmark a single + // (possibly large) problem, pass: [sbsize] [rr_step] + std::vector cases; + if (argc >= 4) + { + cases.push_back({std::atoi(argv[1]), std::atoi(argv[2]), std::atoi(argv[3])}); + } + else + { + cases = { + {50, 10, 1}, {50, 10, 3}, {100, 10, 3}, {200, 10, 5}, {500, 10, 5}, + }; + } + if (argc >= 5) + { + g_sbsize = std::atoi(argv[4]); + } + if (argc >= 6) + { + g_rr_step = std::atoi(argv[5]); + } + + std::printf("\n=== Solver comparison (identical H, psi0, ethr) ===\n"); + std::printf("Heap values are retained process-heap deltas after diag(), not peak memory.\n"); + std::printf("The optional allocator metric is unavailable without glibc >= 2.33.\n"); + std::printf("%-5s %-5s %-6s %-10s %-14s %-10s %-20s\n", "n", "nband", "bw", "solver", "wall_time(s)", + "max_err", "retained_heap(MiB)"); + std::printf("-----------------------------------------------------------------\n"); + + bool all_ok = true; + for (const auto& c : cases) + { + std::vector band; + int bd = 0; + std::vector prec; + make_H(c.n, c.bw, band, bd, prec); + std::vector ref(c.n, 0.0); + ref_eigen(band.data(), c.n, c.bw, bd, c.nband, ref.data()); + std::vector psi0; + make_psi(c.n, c.nband, psi0); + std::vector ethr(c.nband, 1e-6); + + Result r_ppcg = run_ppcg(band, c.n, c.bw, bd, c.nband, prec, psi0, ethr, ref.data()); + Result r_cg = run_cg(band, c.n, c.bw, bd, c.nband, prec, psi0, ethr, ref.data()); + Result r_bpcg = run_bpcg(band, c.n, c.bw, bd, c.nband, prec, psi0, ethr, ref.data()); + Result r_dav = run_dav(band, c.n, c.bw, bd, c.nband, prec, psi0, ethr, ref.data()); + all_ok = all_ok && r_ppcg.ok && r_cg.ok && r_bpcg.ok && r_dav.ok; + + const std::string ppcg_heap = format_heap_delta(r_ppcg.retained_heap_delta_bytes); + const std::string cg_heap = format_heap_delta(r_cg.retained_heap_delta_bytes); + const std::string bpcg_heap = format_heap_delta(r_bpcg.retained_heap_delta_bytes); + const std::string davidson_heap = format_heap_delta(r_dav.retained_heap_delta_bytes); + std::printf("%-5d %-5d %-6d %-10s %-14.5f %-10.2e %-20s\n", c.n, c.nband, c.bw, "PPCG", r_ppcg.wall_s, + r_ppcg.max_err, ppcg_heap.c_str()); + std::printf("%-5s %-5s %-6s %-10s %-14.5f %-10.2e %-20s\n", "", "", "", "CG", r_cg.wall_s, + r_cg.max_err, cg_heap.c_str()); + std::printf("%-5s %-5s %-6s %-10s %-14.5f %-10.2e %-20s\n", "", "", "", "BPCG", r_bpcg.wall_s, + r_bpcg.max_err, bpcg_heap.c_str()); + std::printf("%-5s %-5s %-6s %-10s %-14.5f %-10.2e %-20s\n", "", "", "", "Davidson", r_dav.wall_s, + r_dav.max_err, davidson_heap.c_str()); + std::printf("-----------------------------------------------------------------\n"); + } + + const int exit_code = all_ok ? 0 : 1; + finishmpi(); + return exit_code; +} diff --git a/source/source_hsolver/unittests/test_diago_david.cpp b/source/source_hsolver/unittests/test_diago_david.cpp index b2556513944..70bf2ee3a1f 100644 --- a/source/source_hsolver/unittests/test_diago_david.cpp +++ b/source/source_hsolver/unittests/test_diago_david.cpp @@ -1,3 +1,599 @@ +// The standalone stability harness isolates adversarial cases in separate +// processes; the default GoogleTest target exercises distributed ownership. +// Both test executables retain their original entry points and checks. +#ifdef ABACUS_DAVIDSON_STABILITY_HARNESS +// Public Davidson and Davidson-subspace stability regression tests. +// CTest runs one case per process and requires the final success marker: some +// historical failure paths returned exit(0) before producing valid eigenpairs. +#include "source_hsolver/diago_david.h" +#include "source_hsolver/diago_dav_subspace.h" +#include "source_hsolver/hs_operator.h" +#include "source_base/parallel_global.h" +#include "source_base/module_container/base/third_party/lapack.h" + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +namespace +{ +typedef std::complex Complex; + +const char* const case_names[] = { + "exact", "partial_exact", "collinear", "near_collinear", + "degenerate", "generalized_spd", "warm_start", "small_direction", "restart_spd" +}; + +struct Problem +{ + int dimension = 12; + int bands = 2; + int leading_dimension = 15; + int subspace_multiplier = 2; + int max_iterations = 200; + int solves = 1; + double solver_threshold = 1.0e-12; + double eigenvalue_tolerance = 1.0e-10; + double residual_tolerance = 1.0e-10; + double gram_tolerance = 1.0e-10; + std::vector h; + std::vector s; + std::vector initial; + std::vector precondition; +}; + +void multiply(const std::vector& matrix, const int dimension, + const Complex* in, Complex* out, const int ld, const int columns) +{ + for (int column = 0; column < columns; ++column) + { + for (int row = 0; row < dimension; ++row) + { + Complex value(0.0, 0.0); + for (int index = 0; index < dimension; ++index) + { + value += matrix[row + index * dimension] * in[index + column * ld]; + } + out[row + column * ld] = value; + } + } +} + +class DenseOperator : public hsolver::HSOperator +{ + public: + DenseOperator(const std::vector& h, const std::vector& s, const int dimension) + : h_(h), s_(s), dimension_(dimension) + { + } + + void update_k(const int ik) override + { + if (ik != 0) + { + throw std::runtime_error("Only one algebraic k-point is available."); + } + } + + void hpsi(const Complex* in, Complex* out, const int ld, const int columns) const override + { + ++h_calls_; + multiply(h_, dimension_, in, out, ld, columns); + } + + int h_calls() const + { + return h_calls_; + } + + void spsi(const Complex* in, Complex* out, const int ld, const int columns) const override + { + multiply(s_, dimension_, in, out, ld, columns); + } + + private: + const std::vector& h_; + const std::vector& s_; + const int dimension_; + mutable int h_calls_ = 0; +}; + +void set_hermitian(Problem& problem, const int row, const int column, const Complex value) +{ + problem.h[row + column * problem.dimension] = value; + problem.h[column + row * problem.dimension] = std::conj(value); +} + +Problem make_problem(const std::string& name) +{ + Problem problem; + if (name == "small_direction") + { + problem.bands = 1; + problem.max_iterations = 50; + problem.residual_tolerance = 1.0e-12; + problem.eigenvalue_tolerance = 1.0e-12; + } + else if (name == "restart_spd") + { + problem.dimension = 24; + problem.leading_dimension = 27; + problem.bands = 3; + // Eigenvalue-change stopping needs a tighter request to satisfy the + // independent residual check below across LAPACK implementations. + problem.solver_threshold = 1.0e-15; + problem.residual_tolerance = 2.0e-7; + } + const int matrix_size = problem.dimension * problem.dimension; + const int block_size = problem.leading_dimension * problem.bands; + const Complex zero(0.0, 0.0); + problem.h.assign(matrix_size, zero); + problem.s.assign(matrix_size, zero); + problem.initial.assign(block_size, zero); + problem.precondition.assign(problem.dimension, 1.0); + for (int row = 0; row < problem.dimension; ++row) + { + problem.h[row + row * problem.dimension] = 10.0 + row; + problem.s[row + row * problem.dimension] = 1.0; + } + problem.h[0] = 1.0; + problem.h[1 + problem.dimension] = 2.0; + problem.h[2 + 2 * problem.dimension] = 5.0; + problem.initial[0] = 1.0; + if (problem.bands > 1) + { + problem.initial[1 + problem.leading_dimension] = 1.0; + } + if (problem.bands > 2) + { + problem.initial[2 + 2 * problem.leading_dimension] = 1.0; + } + + if (name == "exact") + { + // Both initial columns are exact lowest eigenvectors. Expanding a zero + // residual must not cause normalization by zero or an abort. + } + else if (name == "partial_exact") + { + const Complex coupling(0.3, 0.0); + set_hermitian(problem, 1, 2, coupling); + } + else if (name == "collinear" || name == "near_collinear" || name == "warm_start") + { + const Complex first_coupling(0.2, 0.0); + const Complex second_coupling(0.3, 0.0); + set_hermitian(problem, 0, 2, first_coupling); + set_hermitian(problem, 1, 2, second_coupling); + if (name == "near_collinear") + { + const Complex perturbation(1.0e-12, 0.0); + set_hermitian(problem, 1, 3, perturbation); + } + if (name == "warm_start") + { + problem.solves = 2; + } + } + else if (name == "degenerate") + { + problem.h[1 + problem.dimension] = 1.0; + problem.h[3 + 3 * problem.dimension] = 6.0; + const double first_norm = std::sqrt(1.04); + const double second_norm = std::sqrt(1.09); + problem.initial[0] = 1.0 / first_norm; + problem.initial[2] = 0.2 / first_norm; + problem.initial[1 + problem.leading_dimension] = 1.0 / second_norm; + problem.initial[3 + problem.leading_dimension] = 0.3 / second_norm; + } + else if (name == "generalized_spd" || name == "restart_spd") + { + // H = sqrt(S) A sqrt(S), with non-real Hermitian couplings in A. + // The two initial correction directions remain exactly dependent. + const Complex first_coupling(0.0, 0.2); + const Complex second_coupling(0.3, 0.1); + set_hermitian(problem, 0, 2, first_coupling); + set_hermitian(problem, 1, 2, second_coupling); + if (name == "restart_spd") + { + // A dense Hermitian problem and a two-block budget require repeated + // restarts. Independent LAPACK residual and S-Gram checks verify + // that reuse of cached H/S products preserves the generalized solve. + for (int column = 1; column < problem.dimension; ++column) + { + for (int row = 0; row < column; ++row) + { + const double real_part = 0.15 / (column - row); + const double imaginary_part = 0.03 * (row + 1) / column; + const Complex coupling(real_part, imaginary_part); + set_hermitian(problem, row, column, coupling); + } + } + } + std::vector scales(problem.dimension, 1.0); + for (int row = 0; row < problem.dimension; ++row) + { + const double diagonal = 1.0 + 0.2 * row; + scales[row] = std::sqrt(diagonal); + problem.s[row + row * problem.dimension] = diagonal; + } + for (int column = 0; column < problem.dimension; ++column) + { + for (int row = 0; row < problem.dimension; ++row) + { + problem.h[row + column * problem.dimension] *= scales[row] * scales[column]; + } + } + for (int column = 0; column < problem.bands; ++column) + { + for (int row = 0; row < problem.dimension; ++row) + { + problem.initial[row + column * problem.leading_dimension] /= scales[row]; + } + } + } + else if (name == "small_direction") + { + // True residual is about 2e-11, above ethr; preconditioning shrinks + // the useful direction to about 2e-13. Absolute step-size rejection + // must not accept the inaccurate input or abort before returning. + for (int row = 0; row < problem.dimension; ++row) + { + problem.h[row + row * problem.dimension] = row; + } + problem.h[0] = -1.0; + const double norm_squared = 1.0 + 1.0e-22; + const double norm = std::sqrt(norm_squared); + problem.initial[0] = 1.0 / norm; + problem.initial[1] = 1.0e-11 / norm; + problem.precondition.assign(problem.dimension, 100.0); + } + else + { + const std::string message = "Unknown case: " + name; + throw std::runtime_error(message); + } + // Padding is deliberately nonzero: the operator contract uses only n rows. + for (int column = 0; column < problem.bands; ++column) + { + for (int row = problem.dimension; row < problem.leading_dimension; ++row) + { + problem.initial[row + column * problem.leading_dimension] = Complex(17.0, -9.0); + } + } + return problem; +} + +struct Reference +{ + std::vector eigenvalues; + std::vector vectors; + int info = -999; +}; + +Reference reference_solve(const Problem& problem) +{ + Reference reference; + const int n = problem.dimension; + reference.eigenvalues.assign(n, 0.0); + reference.vectors = problem.h; + std::vector overlap(problem.s); + const int itype = 1; + const char jobz = 'V'; + const char uplo = 'U'; + int lwork = -1; + int lrwork = -1; + int liwork = -1; + Complex work_query(0.0, 0.0); + double rwork_query = 0.0; + int iwork_query = 0; + zhegvd_(&itype, &jobz, &uplo, &n, reference.vectors.data(), &n, overlap.data(), &n, + reference.eigenvalues.data(), &work_query, &lwork, &rwork_query, &lrwork, + &iwork_query, &liwork, &reference.info); + if (reference.info != 0) + { + return reference; + } + lwork = static_cast(work_query.real()); + lrwork = static_cast(rwork_query); + liwork = iwork_query; + if (lwork < 1 || lrwork < 1 || liwork < 1) + { + throw std::runtime_error("Invalid LAPACK workspace query result."); + } + std::vector work(lwork); + std::vector rwork(lrwork); + std::vector iwork(liwork); + reference.vectors = problem.h; + overlap = problem.s; + zhegvd_(&itype, &jobz, &uplo, &n, reference.vectors.data(), &n, overlap.data(), &n, + reference.eigenvalues.data(), work.data(), &lwork, rwork.data(), &lrwork, + iwork.data(), &liwork, &reference.info); + return reference; +} + +struct Metrics +{ + bool finite = true; + double eigenvalue_error = 0.0; + double residual = 0.0; + double normalized_residual = 0.0; + double gram_error = 0.0; +}; + +Metrics measure(const Problem& problem, const Complex* vectors, const int ld, const int columns, + const std::vector& eigenvalues, const std::vector& reference) +{ + const int length = ld * columns; + const Complex zero(0.0, 0.0); + std::vector hv(length, zero); + std::vector sv(length, zero); + multiply(problem.h, problem.dimension, vectors, hv.data(), ld, columns); + multiply(problem.s, problem.dimension, vectors, sv.data(), ld, columns); + Metrics result; + for (int column = 0; column < columns; ++column) + { + const double eigenvalue_difference = eigenvalues[column] - reference[column]; + const double delta = std::abs(eigenvalue_difference); + result.finite = result.finite && std::isfinite(delta); + result.eigenvalue_error = std::max(result.eigenvalue_error, delta); + double residual2 = 0.0; + double norm2 = 0.0; + for (int row = 0; row < problem.dimension; ++row) + { + const int index = row + column * ld; + const Complex difference = hv[index] - eigenvalues[column] * sv[index]; + residual2 += std::norm(difference); + norm2 += std::norm(vectors[index]); + } + const double residual = std::sqrt(residual2); + const double normalized_squared = residual2 / norm2; + const double normalized = std::sqrt(normalized_squared); + result.finite = result.finite && norm2 > 0.0 && std::isfinite(residual) && std::isfinite(normalized); + result.residual = std::max(result.residual, residual); + result.normalized_residual = std::max(result.normalized_residual, normalized); + } + for (int right = 0; right < columns; ++right) + { + for (int left = 0; left < columns; ++left) + { + Complex product(0.0, 0.0); + for (int row = 0; row < problem.dimension; ++row) + { + const Complex left_value = vectors[row + left * ld]; + product += std::conj(left_value) * sv[row + right * ld]; + } + if (left == right) + { + product -= 1.0; + } + const double error = std::abs(product); + result.finite = result.finite && std::isfinite(error); + result.gram_error = std::max(result.gram_error, error); + } + } + return result; +} + +void print_metrics(const std::string& marker, const std::string& solver, const std::string& name, + const int pass, const Metrics& metrics) +{ + std::cout << marker << " solver=" << solver << " case=" << name << " pass=" << pass + << " finite=" << metrics.finite << " eigenvalue_error=" << metrics.eigenvalue_error + << " residual=" << metrics.residual << " normalized_residual=" << metrics.normalized_residual + << " SGram_error=" << metrics.gram_error << std::endl; +} + +bool validate_return(const Problem& problem, const Reference& reference, + const std::vector& guarded, const std::vector& eigenvalues, + const std::string& solver, const std::string& name, const int pass, + const int iterations) +{ + const int guard = 3; + const int block_size = problem.leading_dimension * problem.bands; + const Complex sentinel(123.0, -456.0); + bool canaries = true; + for (int index = 0; index < guard; ++index) + { + canaries = canaries && guarded[index] == sentinel; + canaries = canaries && guarded[guard + block_size + index] == sentinel; + } + const Complex* vectors = guarded.data() + guard; + // Both public diag implementations explicitly clear the entire output + // block before writing active rows. Their padding contract is zero, not + // preservation of the nonzero input padding used by this test. + bool padding = true; + const Complex zero_padding(0.0, 0.0); + for (int column = 0; column < problem.bands; ++column) + { + for (int row = problem.dimension; row < problem.leading_dimension; ++row) + { + const int index = row + column * problem.leading_dimension; + padding = padding && vectors[index] == zero_padding; + } + } + const Metrics metrics = measure(problem, vectors, problem.leading_dimension, problem.bands, + eigenvalues, reference.eigenvalues); + print_metrics("STABILITY_METRIC", solver, name, pass, metrics); + bool ordered = true; + for (int band = 0; band < problem.bands; ++band) + { + if (band > 0) + { + ordered = ordered && eigenvalues[band] >= eigenvalues[band - 1]; + } + std::cout << "STABILITY_EIGENVALUE solver=" << solver << " case=" << name + << " pass=" << pass << " band=" << band + 1 << " value=" << eigenvalues[band] + << " reference=" << reference.eigenvalues[band] << std::endl; + } + const bool accepted = metrics.finite && canaries && padding && ordered && iterations >= 0 + && metrics.eigenvalue_error <= problem.eigenvalue_tolerance + && metrics.residual <= problem.residual_tolerance + && metrics.gram_error <= problem.gram_tolerance; + std::cout << "STABILITY_CHECK solver=" << solver << " case=" << name << " pass=" << pass + << " accepted=" << accepted << " canaries=" << canaries + << " padding=" << padding << " ordered=" << ordered + << " iterations=" << iterations << std::endl; + return accepted; +} + +int run_case(const std::string& solver_name, const std::string& name, + const hsolver::diag_comm_info& communicator) +{ + const Problem problem = make_problem(name); + const Reference reference = reference_solve(problem); + std::cout << "STABILITY_INPUT solver=" << solver_name << " case=" << name + << " n=" << problem.dimension << " bands=" << problem.bands + << " ld=" << problem.leading_dimension << " ndim=" << problem.subspace_multiplier + << " ethr=" << problem.solver_threshold << " maxiter=" << problem.max_iterations + << " solves=" << problem.solves << " eigenvalue_tolerance=" << problem.eigenvalue_tolerance + << " residual_tolerance=" << problem.residual_tolerance + << " gram_tolerance=" << problem.gram_tolerance << " reference_info=" << reference.info + << " matrix_energy_unit=algebraic" << std::endl; + if (reference.info != 0) + { + throw std::runtime_error("Independent direct LAPACK reference solve failed."); + } + const Metrics reference_metrics = measure(problem, reference.vectors.data(), problem.dimension, + problem.dimension, reference.eigenvalues, reference.eigenvalues); + print_metrics("STABILITY_REFERENCE", solver_name, name, 0, reference_metrics); + if (!reference_metrics.finite || reference_metrics.residual > 1.0e-11 || reference_metrics.gram_error > 1.0e-11) + { + throw std::runtime_error("Independent direct LAPACK reference validation failed."); + } + const int guard = 3; + const int block_size = problem.leading_dimension * problem.bands; + const int guarded_size = block_size + 2 * guard; + const Complex sentinel(123.0, -456.0); + std::vector guarded(guarded_size, sentinel); + Complex* vectors = guarded.data() + guard; + std::copy(problem.initial.begin(), problem.initial.end(), vectors); + std::vector eigenvalues(problem.bands, 0.0); + const std::vector thresholds(problem.bands, problem.solver_threshold); + DenseOperator op(problem.h, problem.s, problem.dimension); + op.update_k(0); + bool accepted = true; + if (solver_name == "dav") + { + hsolver::DiagoDavid solver(problem.precondition.data(), problem.bands, + problem.dimension, problem.subspace_multiplier, communicator); + const int restart_budget = 5; + const int allowed_unconverged = 0; + for (int pass = 1; pass <= problem.solves; ++pass) + { + const int iterations = solver.diag(op, problem.leading_dimension, vectors, eigenvalues.data(), + thresholds, problem.max_iterations, restart_budget, allowed_unconverged); + std::cout << "STABILITY_RETURNED solver=" << solver_name << " case=" << name + << " pass=" << pass << std::endl; + const bool pass_ok = validate_return(problem, reference, guarded, eigenvalues, + solver_name, name, pass, iterations); + accepted = accepted && pass_ok; + } + } + else if (solver_name == "dav_subspace") + { + const int lapack_path = 0; + const int block_size_2d = 0; + const bool scf_type = false; + hsolver::Diago_DavSubspace solver(problem.precondition, problem.bands, + problem.dimension, problem.subspace_multiplier, problem.solver_threshold, problem.max_iterations, + communicator, lapack_path, block_size_2d); + for (int pass = 1; pass <= problem.solves; ++pass) + { + const int h_calls_before = op.h_calls(); + const int iterations = solver.diag(op, vectors, problem.leading_dimension, eigenvalues.data(), + thresholds, scf_type); + std::cout << "STABILITY_RETURNED solver=" << solver_name << " case=" << name + << " pass=" << pass << std::endl; + const bool pass_ok = validate_return(problem, reference, guarded, eigenvalues, + solver_name, name, pass, iterations); + accepted = accepted && pass_ok; + if (name == "restart_spd") + { + const int h_calls = op.h_calls() - h_calls_before; + // One initial H application plus at most one per expansion. + // A restart of an accurate Ritz basis needs no extra H call. + const bool reused_images = iterations >= 3 && h_calls <= iterations + 1; + std::cout << "STABILITY_RESTART h_calls=" << h_calls << " iterations=" << iterations + << " reused_images=" << reused_images << std::endl; + accepted = accepted && reused_images; + } + } + } + else + { + const std::string message = "Unknown solver: " + solver_name; + throw std::runtime_error(message); + } + std::cout << "STABILITY_RESULT solver=" << solver_name << " case=" << name + << " accepted=" << accepted << " completed_solves=" << problem.solves << std::endl; + return accepted ? 0 : 1; +} +} // namespace + +int main(int argc, char** argv) +{ + if (argc == 2 && std::string(argv[1]) == "--list") + { + for (const char* name : case_names) + { + std::cout << name << '\n'; + } + return 0; + } + if (argc != 3) + { + std::cerr << "Usage: MODULE_HSOLVER_davidson_stability {dav|dav_subspace} CASE\n"; + return 2; + } + int nproc = 1; + int rank = 0; + int threads = 1; + Parallel_Global::read_pal_param(argc, argv, nproc, threads, rank); +#ifdef __MPI + // Existing reductions consume POOL_WORLD; the serial harness shares WORLD. + // Other communicator handles are unused, and must be valid at finalize. + POOL_WORLD = MPI_COMM_WORLD; + KP_WORLD = MPI_COMM_NULL; + BP_WORLD = MPI_COMM_NULL; + INT_BGROUP = MPI_COMM_NULL; + GRID_WORLD = MPI_COMM_NULL; + DIAG_WORLD = MPI_COMM_NULL; + const hsolver::diag_comm_info communicator(POOL_WORLD, rank, nproc); +#else + const hsolver::diag_comm_info communicator(rank, nproc); +#endif + std::cout << std::scientific << std::setprecision(17); + int status = 2; + try + { + if (nproc != 1) + { + throw std::runtime_error("This algebraic harness requires exactly one MPI process."); + } + const std::string solver_name(argv[1]); + const std::string case_name(argv[2]); + status = run_case(solver_name, case_name, communicator); + } + catch (const std::exception& error) + { + std::cerr << "STABILITY_EXCEPTION " << error.what() << std::endl; + status = 2; + } +#ifdef __MPI + Parallel_Global::finalize_mpi(); +#endif + return status; +} + +#else #include "source_hsolver/diago_david.h" #include "source_hsolver/diago_dav_subspace.h" #include "source_hsolver/diago_cg.h" @@ -190,3 +786,5 @@ int main(int argc, char** argv) #endif return result; } + +#endif diff --git a/source/source_hsolver/unittests/test_diago_params.cpp b/source/source_hsolver/unittests/test_diago_params.cpp new file mode 100644 index 00000000000..33bb19055fd --- /dev/null +++ b/source/source_hsolver/unittests/test_diago_params.cpp @@ -0,0 +1,109 @@ +#include "../diago_iter_assist.h" +#include "../diago_params.h" + +#include + +#include + +template +class DiagoParamsPWTest : public ::testing::Test +{ + protected: + using Assist = hsolver::DiagoIterAssist; + using Real = typename GetTypeReal::type; + + void SetUp() override + { + saved_nmax = Assist::PW_DIAG_NMAX; + saved_thr = Assist::PW_DIAG_THR; + saved_scf_iter = Assist::SCF_ITER; + saved_need_subspace = Assist::need_subspace; + } + + void TearDown() override + { + Assist::PW_DIAG_NMAX = saved_nmax; + Assist::PW_DIAG_THR = saved_thr; + Assist::SCF_ITER = saved_scf_iter; + Assist::need_subspace = saved_need_subspace; + } + + void setup(const Input_para& input) + { + hsolver::setup_diago_params_pw(0, 1, 1.0e-8, input); + } + + int saved_nmax; + Real saved_thr; + int saved_scf_iter; + bool saved_need_subspace; +}; + +using DiagoParamsTypes = ::testing::Types, std::complex, double>; +TYPED_TEST_SUITE(DiagoParamsPWTest, DiagoParamsTypes); + +TYPED_TEST(DiagoParamsPWTest, AppliesPwDiagNmaxForScf) +{ + using Assist = typename TestFixture::Assist; + Input_para input; + input.calculation = "scf"; + input.pw_diag_nmax = 73; + + const char* solvers[] = {"cg", "dav", "dav_subspace", "ppcg"}; + for (const char* solver : solvers) + { + input.ks_solver = solver; + Assist::PW_DIAG_NMAX = 19; + this->setup(input); + EXPECT_EQ(Assist::PW_DIAG_NMAX, 73) << solver; + } +} + +TYPED_TEST(DiagoParamsPWTest, AppliesPwDiagNmaxForNscf) +{ + using Assist = typename TestFixture::Assist; + Input_para input; + input.calculation = "nscf"; + input.pw_diag_nmax = 73; + + // These solvers consume the shared limit; BPCG has its own update budget. + const char* solvers[] = {"cg", "dav", "dav_subspace", "ppcg"}; + for (const char* solver : solvers) + { + input.ks_solver = solver; + Assist::PW_DIAG_NMAX = 19; + this->setup(input); + EXPECT_EQ(Assist::PW_DIAG_NMAX, 73) << solver; + } +} + +TYPED_TEST(DiagoParamsPWTest, NscfBudgetDoesNotDependOnPreviousCalculation) +{ + using Assist = typename TestFixture::Assist; + Input_para input; + + const char* solvers[] = {"cg", "dav", "dav_subspace", "ppcg"}; + for (const char* solver : solvers) + { + input.ks_solver = solver; + input.calculation = "scf"; + input.pw_diag_nmax = 91; + this->setup(input); + EXPECT_EQ(Assist::PW_DIAG_NMAX, 91) << solver; + + input.calculation = "nscf"; + input.pw_diag_nmax = 7; + this->setup(input); + EXPECT_EQ(Assist::PW_DIAG_NMAX, 7) << solver; + + input.pw_diag_nmax = 200; + this->setup(input); + EXPECT_EQ(Assist::PW_DIAG_NMAX, 200) << solver; + + Input_para default_input; + default_input.calculation = "nscf"; + default_input.ks_solver = solver; + this->setup(default_input); + EXPECT_EQ(Assist::PW_DIAG_NMAX, 50) << solver; + } +} diff --git a/source/source_hsolver/unittests/test_diago_ppcg.cpp b/source/source_hsolver/unittests/test_diago_ppcg.cpp new file mode 100644 index 00000000000..1e4e9e65d42 --- /dev/null +++ b/source/source_hsolver/unittests/test_diago_ppcg.cpp @@ -0,0 +1,4332 @@ +/** + * diago_ppcg_test.cpp — unit test for DiagoPPCG solver + * + * Test matrices include both S = I and non-trivial overlap operators: + * 1. Tridiagonal Laplacian (1D particle-in-a-box): H[i,i]=2, H[i,i±1]=-1 + * Exact λ_k = 2 - 2·cos(k·π/(n+1)). Realistic but sparse. + * 2. Diagonal matrix: H = diag(1, 2, 3, 4, 5) + * Exact eigenvalues are the diagonal entries. Simplest possible + * smoke test — should converge in very few iterations. + * + * Tests exercise the production BLOCK_SUBSPACE strategy (the band-by-band + * CONJUGATE_GRADIENT variant was removed from this PR). + */ + +#include "../diago_ppcg.h" +#ifdef __MPI +#include "source_basis/module_pw/test/test_tool.h" +#endif + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#ifndef M_PI +#define M_PI 3.14159265358979323846 +#endif + +using T = std::complex; +using Real = double; + +// ----------------------------------------------------------------------------- +// Helper: dense H-matrix times a set of column vectors +// H is stored column-major: H(row, col) = H_data[row + col * n_dim] +// ----------------------------------------------------------------------------- +static void dense_h_multiply(const T* H_data, int n_dim, const T* in, T* out, int ld, int ncol) +{ + for (int j = 0; j < ncol; ++j) + { + for (int i = 0; i < n_dim; ++i) + { + T sum = 0; + for (int k = 0; k < n_dim; ++k) + { + sum += H_data[i + k * n_dim] * in[k + j * ld]; + } + out[i + j * ld] = sum; + } + } +} + +namespace hsolver +{ + +// Narrow friend access exercises cancellation in the projection helper without +// relying on a particular sequence of LAPACK rotations in an end-to-end solve. +class DiagoPPCGProjectionTest : public ::testing::Test +{ +protected: + void check_direction_rotation(const bool ritz) + { + const int dimension = 3; + const int leading_dimension = 4; + const int bands = 2; + const int size = leading_dimension * bands; + const std::vector h = {T(1), T(0, -0.4), T(0), + T(0, 0.4), T(4), T(0), T(0), T(0), T(10)}; + const std::vector s = {T(1), T(0), T(0), T(0), T(2), T(0), T(0), T(0), T(3)}; + std::vector x(size, T(0)); + x[0] = T(1); + x[leading_dimension] = T(0.3); + x[leading_dimension + 1] = T(1); + DiagoPPCG solver(1e-12, 20, 1, 4); + solver.ld_psi_ = leading_dimension; + solver.n_dim_ = dimension; + solver.n_band_ = bands; + solver.hpsi_.assign(size, T(0)); + solver.spsi_.assign(size, T(0)); + solver.w_.assign(size, T(0)); + solver.sw_.assign(size, T(0)); + solver.hw_.assign(size, T(0)); + solver.p_.assign(size, T(0)); + solver.p_[2] = T(0.4, 0.3); + solver.p_[leading_dimension + 2] = T(-0.2, 0.6); + solver.hp_.assign(size, T(0)); + solver.sp_.assign(size, T(0)); + solver.rr_eval_.resize(bands); + dense_h_multiply(h.data(), dimension, x.data(), solver.hpsi_.data(), leading_dimension, bands); + dense_h_multiply(s.data(), dimension, x.data(), solver.spsi_.data(), leading_dimension, bands); + dense_h_multiply(h.data(), dimension, solver.p_.data(), solver.hp_.data(), leading_dimension, bands); + dense_h_multiply(s.data(), dimension, solver.p_.data(), solver.sp_.data(), leading_dimension, bands); + const std::vector original_p = solver.p_; + if (ritz) + { + std::vector eigenvalues(bands); + const std::vector thresholds(bands, 1e-12); + std::vector active; + ASSERT_TRUE(solver.rayleigh_ritz(x.data(), eigenvalues.data(), active, thresholds)); + } + else + { + ASSERT_TRUE(solver.orthonormalize(x.data())); + } + double change = 0.0; + for (int column = 0; column < bands; ++column) + { + // Recover C from X_new = X_old C using the known triangular + // leading 2-by-2 block, independently of the solver's workspace. + const int first = column * leading_dimension; + const T c0 = x[first] - 0.3 * x[first + 1]; + const T c1 = x[first + 1]; + for (int row = 0; row < dimension; ++row) + { + const int index = first + row; + const T expected = original_p[row] * c0 + original_p[row + leading_dimension] * c1; + const T direction_difference = solver.p_[index] - expected; + const double error = std::abs(direction_difference); + EXPECT_LT(error, 1e-13); + const T direction_change = solver.p_[index] - original_p[index]; + change += std::norm(direction_change); + } + EXPECT_EQ(solver.p_[first + dimension], T(0)); + } + EXPECT_GT(change, 1e-3); + std::vector hp(size, T(0)); + std::vector sp(size, T(0)); + dense_h_multiply(h.data(), dimension, solver.p_.data(), hp.data(), leading_dimension, bands); + dense_h_multiply(s.data(), dimension, solver.p_.data(), sp.data(), leading_dimension, bands); + for (int index = 0; index < size; ++index) + { + const T h_difference = hp[index] - solver.hp_[index]; + const T s_difference = sp[index] - solver.sp_[index]; + const double h_error = std::abs(h_difference); + const double s_error = std::abs(s_difference); + EXPECT_LT(h_error, 1e-13); + EXPECT_LT(s_error, 1e-13); + } + } + + void check_rank_truncated_outputs(const int kept_count) + { + const int dimension = 8; + const int rank = 7; + const int matrix_size = dimension * dimension; + DiagoPPCG solver(1.0e-12, 20, kept_count, 1); + DiagoPPCG::SmallSubspace subspace; + subspace.k.assign(matrix_size, T(0)); + subspace.m.assign(matrix_size, T(0)); + subspace.eval.assign(dimension, 0.0); + for (int column = 0; column < rank; ++column) + { + const double weight = column + 1.0; + subspace.m[column + column * dimension] = T(weight); + subspace.k[column + column * dimension] = T(weight * weight); + if (column + 1 < rank) + { + const double scale = std::sqrt(weight * (weight + 1.0)); + const T coupling = T(-0.3, 0.2) * scale; + subspace.k[column + (column + 1) * dimension] = coupling; + subspace.k[column + 1 + column * dimension] = std::conj(coupling); + } + } + const std::vector hamiltonian = subspace.k; + const std::vector overlap = subspace.m; + const int valid_count = solver.solve_small_generalized(dimension, kept_count, subspace); + ASSERT_EQ(valid_count, rank); + // Check every exported vector, including those not used for the local + // update: the recovery path must receive valid, normalized Ritz pairs. + for (int column = 0; column < valid_count; ++column) + { + double residual_squared = 0.0; + for (int row = 0; row < dimension; ++row) + { + T residual(0); + for (int inner = 0; inner < dimension; ++inner) + { + const int offset = row + inner * dimension; + residual += (hamiltonian[offset] - subspace.eval[column] * overlap[offset]) + * subspace.k[inner + column * dimension]; + } + residual_squared += std::norm(residual); + } + const double residual_norm = std::sqrt(residual_squared); + EXPECT_LT(residual_norm, 1.0e-12); + for (int other = 0; other < valid_count; ++other) + { + T inner_product(0); + for (int row = 0; row < rank; ++row) + { + inner_product += std::conj(subspace.k[row + column * dimension]) + * (row + 1.0) * subspace.k[row + other * dimension]; + } + const T expected = column == other ? T(1) : T(0); + const double error = std::abs(inner_product - expected); + EXPECT_LT(error, 1.0e-12); + } + } + } + + void check_projection(const bool nearly_parallel, const bool zero_direction) + { + const int dimension = 2; + const int bands = 1; + const std::vector h_diagonal = {1.0, 1000.0}; + const std::vector s_diagonal = {1.0, 3.0}; + const std::vector x = {T(0.5), T(0.5)}; + const double scale = 1.0 / std::sqrt(12.0); + const std::vector z = {T(3.0 * scale), T(-scale)}; + DiagoPPCG solver(1.0e-12, 20, bands, 1); + solver.ld_psi_ = dimension; + solver.n_dim_ = dimension; + solver.n_band_ = bands; + solver.hpsi_.resize(dimension); + solver.spsi_.resize(dimension); + solver.p_.resize(dimension); + solver.hp_.resize(dimension); + solver.sp_.resize(dimension); + for (int i = 0; i < dimension; ++i) + { + solver.hpsi_[i] = h_diagonal[i] * x[i]; + solver.spsi_[i] = s_diagonal[i] * x[i]; + solver.p_[i] = z[i]; + if (nearly_parallel) + { + solver.p_[i] = x[i] + 1.0e-10 * z[i]; + } + solver.hp_[i] = h_diagonal[i] * solver.p_[i]; + solver.sp_[i] = s_diagonal[i] * solver.p_[i]; + if (zero_direction) + { + solver.p_[i] = T(0); + solver.hp_[i] = T(1.0e-15); + solver.sp_[i] = T(1.0e-15); + } + } + int h_calls = 0; + int s_calls = 0; + const auto h_op = [&h_calls, &h_diagonal, dimension](T* in, T* out, int ld, int ncol) { + ++h_calls; + for (int j = 0; j < ncol; ++j) + { + for (int i = 0; i < dimension; ++i) + { + out[i + j * ld] = h_diagonal[i] * in[i + j * ld]; + } + } + }; + solver.spsi_func_ = [&s_calls, &s_diagonal, dimension](T* in, T* out, int ld, int ncol) { + ++s_calls; + for (int j = 0; j < ncol; ++j) + { + for (int i = 0; i < dimension; ++i) + { + out[i + j * ld] = s_diagonal[i] * in[i + j * ld]; + } + } + }; + const std::vector columns = {0}; + solver.project_search_directions(h_op, x.data(), columns, columns); + double direction_norm_squared = 0.0; + double h_error_squared = 0.0; + double s_error_squared = 0.0; + for (int i = 0; i < dimension; ++i) + { + direction_norm_squared += std::norm(solver.p_[i]); + h_error_squared += std::norm(solver.hp_[i] - h_diagonal[i] * solver.p_[i]); + s_error_squared += std::norm(solver.sp_[i] - s_diagonal[i] * solver.p_[i]); + } + if (zero_direction) + { + EXPECT_DOUBLE_EQ(direction_norm_squared, 0.0); + EXPECT_DOUBLE_EQ(h_error_squared, 0.0); + EXPECT_DOUBLE_EQ(s_error_squared, 0.0); + EXPECT_EQ(h_calls, 0); + EXPECT_EQ(s_calls, 0); + return; + } + ASSERT_GT(direction_norm_squared, 0.0); + const double h_error = std::sqrt(h_error_squared / direction_norm_squared); + const double s_error = std::sqrt(s_error_squared / direction_norm_squared); + EXPECT_LT(h_error, 1.0e-12); + EXPECT_LT(s_error, 1.0e-12); + const int expected_calls = nearly_parallel ? 1 : 0; + EXPECT_EQ(h_calls, expected_calls); + EXPECT_EQ(s_calls, expected_calls); + } +}; + +TEST_F(DiagoPPCGProjectionTest, RefreshesImagesAfterSevereProjectionCancellation) +{ + check_projection(true, false); +} + +TEST_F(DiagoPPCGProjectionTest, RitzRotationPreservesDirectionColumnAssociation) +{ + check_direction_rotation(true); +} + +TEST_F(DiagoPPCGProjectionTest, OrthonormalizationRotatesDirectionImagesTogether) +{ + check_direction_rotation(false); +} + +TEST_F(DiagoPPCGProjectionTest, ReusesImagesForWellSeparatedDirection) +{ + check_projection(false, false); +} + +TEST_F(DiagoPPCGProjectionTest, ClearsStaleImagesForZeroDirection) +{ + check_projection(false, true); +} + +TEST_F(DiagoPPCGProjectionTest, RankTruncationExportsKeptAndDiscardedRitzVectors) +{ + const int kept_counts[] = {2, 3, 7}; + for (const int kept_count : kept_counts) + { + SCOPED_TRACE(kept_count); + check_rank_truncated_outputs(kept_count); + } +} + +} // namespace hsolver + +// ============================================================================= +// Test fixture: 1D particle-in-a-box (tridiagonal Laplacian) +// ============================================================================= +class DiagoPPCGTridiagTest : public ::testing::Test +{ + protected: + void SetUp() override + { + n_dim = 10; + nband = 3; + ld = n_dim; + + // Build tridiagonal H: H[i,i] = 2, H[i,i±1] = -1 + H_mat.assign(n_dim * n_dim, T(0)); + for (int i = 0; i < n_dim; ++i) + { + H_mat[i + i * n_dim] = T(2.0, 0); + if (i > 0) + { + H_mat[i + (i - 1) * n_dim] = T(-1.0, 0); + } + if (i < n_dim - 1) + { + H_mat[i + (i + 1) * n_dim] = T(-1.0, 0); + } + } + + // Preconditioner — diagonal of H (all 2.0) + prec.assign(n_dim, 2.0); + + // Exact reference eigenvalues + exact.resize(nband); + for (int k = 0; k < nband; ++k) + { + exact[k] = 2.0 - 2.0 * std::cos(Real(k + 1) * M_PI / Real(n_dim + 1)); + } + + // Convergence thresholds + ethr.assign(nband, 1e-10); + + // Random initial guess (fixed seed for reproducibility) + std::mt19937 rng(42); + std::uniform_real_distribution dist(-1.0, 1.0); + + psi.assign(ld * nband, T(0)); + for (int j = 0; j < nband; ++j) + { + for (int i = 0; i < n_dim; ++i) + { + psi[i + j * ld] = T(dist(rng), 0.0); + } + } + + // Gram-Schmidt orthonormalisation (S = I) + for (int j = 0; j < nband; ++j) + { + for (int k = 0; k < j; ++k) + { + T dot = 0; + for (int i = 0; i < n_dim; ++i) + { + dot += std::conj(psi[i + k * ld]) * psi[i + j * ld]; + } + for (int i = 0; i < n_dim; ++i) + { + psi[i + j * ld] -= dot * psi[i + k * ld]; + } + } + Real nrm = 0; + for (int i = 0; i < n_dim; ++i) + { + nrm += std::norm(psi[i + j * ld]); + } + nrm = std::sqrt(nrm); + for (int i = 0; i < n_dim; ++i) + { + psi[i + j * ld] /= nrm; + } + } + } + + int n_dim; + int nband; + int ld; + std::vector H_mat; + std::vector prec; + std::vector exact; + std::vector ethr; + std::vector psi; +}; + +TEST_F(DiagoPPCGTridiagTest, BlockSubspace) +{ + std::vector psi_run = psi; + std::vector eval(nband, 0.0); + + hsolver::DiagoPPCG solver( + /* diag_thr = */ 1e-12, + /* max_iter = */ 100, + /* sbsize = */ 4, + /* rr_step = */ 4); + + auto h_op = [this](T* in, T* out, int ld_in, int ncol) { + dense_h_multiply(H_mat.data(), n_dim, in, out, ld_in, ncol); + }; + + double avg_iter = solver.diag(h_op, nullptr, ld, nband, n_dim, psi_run.data(), eval.data(), ethr, prec.data()); + + for (int i = 0; i < nband; ++i) + { + EXPECT_NEAR(eval[i], exact[i], 1e-8) << "Tridiag BLOCK: eigenvalue[" << i << "] mismatch"; + } + EXPECT_LE(avg_iter, double(100)) << "Tridiag BLOCK: too many iterations"; +} + +TEST_F(DiagoPPCGTridiagTest, ReportsIterationLimit) +{ + std::vector psi_run = psi; + std::vector eval(nband, 0.0); + std::vector strict_ethr(nband, 1.0e-14); + hsolver::DiagoPPCG solver( + /* diag_thr = */ 1.0e-14, + /* max_iter = */ 1, + /* sbsize = */ 3, + /* rr_step = */ 4); + + auto h_op = [this](T* in, T* out, int ld_in, int ncol) { + dense_h_multiply(H_mat.data(), n_dim, in, out, ld_in, ncol); + }; + + solver.diag(h_op, nullptr, ld, nband, n_dim, psi_run.data(), eval.data(), strict_ethr, prec.data()); + + EXPECT_FALSE(solver.converged()); + EXPECT_GT(solver.active_band_count(), 0); + EXPECT_TRUE(solver.iteration_limit_reached()); + EXPECT_GE(solver.first_active_band(), 0); + EXPECT_GT(solver.max_residual(), Real(1.0e-14)); +} + +TEST_F(DiagoPPCGTridiagTest, ReusesProjectedOperatorImages) +{ + std::vector psi_run = psi; + std::vector eval(nband, 0.0); + const std::vector strict_ethr(nband, 1.0e-14); + int h_calls = 0; + int s_calls = 0; + const auto h_op = [this, &h_calls](T* in, T* out, int ld_in, int ncol) { + ++h_calls; + dense_h_multiply(H_mat.data(), n_dim, in, out, ld_in, ncol); + }; + const auto s_op = [this, &s_calls](T* in, T* out, int ld_in, int ncol) { + ++s_calls; + for (int j = 0; j < ncol; ++j) + { + for (int i = 0; i < n_dim; ++i) + { + out[i + j * ld_in] = (1.0 + 0.1 * i) * in[i + j * ld_in]; + } + } + }; + hsolver::DiagoPPCG solver(1.0e-14, 2, 3, 16); + solver.diag(h_op, s_op, ld, nband, n_dim, + psi_run.data(), eval.data(), strict_ethr, prec.data()); + + ASSERT_TRUE(solver.iteration_limit_reached()); + // Initial/final X applications, one W application per sweep, and one + // cancellation-triggered P refresh. No unconditional HP or projected SW + // recomputation is needed, including when S is not I. + EXPECT_EQ(h_calls, 5); + EXPECT_EQ(s_calls, 5); + double max_residual = 0.0; + for (int j = 0; j < nband; ++j) + { + double residual_squared = 0.0; + for (int i = 0; i < n_dim; ++i) + { + T hpsi = T(0); + for (int k = 0; k < n_dim; ++k) + { + hpsi += H_mat[i + k * n_dim] * psi_run[k + j * ld]; + } + const T spsi = (1.0 + 0.1 * i) * psi_run[i + j * ld]; + residual_squared += std::norm(hpsi - eval[j] * spsi); + } + max_residual = std::max(max_residual, std::sqrt(residual_squared)); + } + EXPECT_NEAR(solver.max_residual(), max_residual, 1.0e-13); +} + +TEST(DiagoPPCGReliableUpdateTest, RefreshesBeforeUsingNearPrecisionResidual) +{ + const int dimension = 2; + const std::vector prec(dimension, 1.0); + const std::vector thresholds(1, 1.0e-14); + const double perturbations[] = {1.0e-5, 1.0e-10}; + for (const double perturbation : perturbations) + { + SCOPED_TRACE(perturbation); + const double normalization = std::sqrt(1.0 + perturbation * perturbation); + std::vector psi = {T(1.0 / normalization), T(perturbation / normalization)}; + std::vector eigenvalues(1, 0.0); + int h_calls = 0; + const auto h_op = [&h_calls](T* in, T* out, int ld, int ncol) { + ++h_calls; + for (int j = 0; j < ncol; ++j) + { + out[j * ld] = in[j * ld]; + out[1 + j * ld] = 3.0 * in[1 + j * ld]; + } + }; + hsolver::DiagoPPCG solver( + 1.0e-14, 1, 1, 1); + solver.diag(h_op, nullptr, dimension, 1, dimension, + psi.data(), eigenvalues.data(), thresholds, prec.data()); + EXPECT_TRUE(solver.converged()); + EXPECT_LT(solver.max_residual(), 1.0e-14); + EXPECT_NEAR(eigenvalues[0], 1.0, 1.0e-14); + const T residual0 = (1.0 - eigenvalues[0]) * psi[0]; + const T residual1 = (3.0 - eigenvalues[0]) * psi[1]; + const double residual = std::sqrt(std::norm(residual0) + std::norm(residual1)); + EXPECT_LT(residual, 1.0e-14); + // Initial X, W, tentative-convergence validation, and final validation. + // Only the near-precision solve adds a reliable X/P refresh before W. + const int expected_calls = perturbation < 1.0e-8 ? 6 : 4; + EXPECT_EQ(h_calls, expected_calls); + } +} + +TEST(DiagoPPCGPeriodicRitzTest, HonorsIntervalAndChecksFinalRitzResidual) +{ + const int dimension = 48; + const int bands = 8; + std::mt19937 random(184); + std::uniform_real_distribution uniform(-1.0, 1.0); + std::vector initial(dimension * bands); + for (T& value : initial) + { + const double real = uniform(random); + const double imaginary = uniform(random); + value = T(real, imaginary); + } + const std::vector prec(dimension, 2.0); + const std::vector thresholds(bands, 1.0e-14); + struct Configuration + { + int sweeps; + int interval; + }; + // Ten sweeps need no periodic refresh; eighteen exercise the one at 16. + // Each configuration additionally encounters one projection cancellation. + // Each run encounters one competing discarded direction and + // therefore add one conditional expanded Ritz solve, without another H call. + const Configuration configurations[] = {{10, 1}, {10, 4}, {10, 16}, + {18, 1}, {18, 4}, {18, 16}}; + for (const Configuration& configuration : configurations) + { + const int sweeps = configuration.sweeps; + const int interval = configuration.interval; + SCOPED_TRACE(sweeps); + SCOPED_TRACE(interval); + std::vector psi = initial; + std::vector eigenvalues(bands, 0.0); + int applications = 0; + const auto h_op = [&applications, dimension](T* in, T* out, int ld, int ncol) { + ++applications; + for (int j = 0; j < ncol; ++j) + { + for (int i = 0; i < dimension; ++i) + { + T value = 2.0 * in[i + j * ld]; + if (i > 0) + { + value -= in[i - 1 + j * ld]; + } + if (i + 1 < dimension) + { + value -= in[i + 1 + j * ld]; + } + out[i + j * ld] = value; + } + } + }; + hsolver::DiagoPPCG solver( + 1.0e-14, sweeps, 2, interval); + solver.diag(h_op, nullptr, dimension, bands, dimension, + psi.data(), eigenvalues.data(), thresholds, prec.data()); + ASSERT_TRUE(solver.iteration_limit_reached()); + const int scheduled = sweeps / interval; + const int final_extra = sweeps % interval == 0 ? 0 : 1; + const int startup_extra = interval == 1 ? 0 : 1; + const int competing_extra = 1; + const int expected_ritz = 1 + scheduled + final_extra + startup_extra + competing_extra; + const int operator_refreshes = sweeps / 16; + const int cancellation_refreshes = 1; + const int expected_applications = 2 + sweeps + 2 * operator_refreshes + + cancellation_refreshes; + EXPECT_EQ(solver.root_recovery_count(), competing_extra); + EXPECT_EQ(solver.rayleigh_ritz_count(), expected_ritz); + EXPECT_EQ(applications, expected_applications); + EXPECT_TRUE(std::is_sorted(eigenvalues.begin(), eigenvalues.end())); + std::vector hpsi(psi.size()); + h_op(psi.data(), hpsi.data(), dimension, bands); + double residual = 0.0; + for (int j = 0; j < bands; ++j) + { + double norm_squared = 0.0; + for (int i = 0; i < dimension; ++i) + { + const T value = hpsi[i + j * dimension] - eigenvalues[j] * psi[i + j * dimension]; + norm_squared += std::norm(value); + } + residual = std::max(residual, std::sqrt(norm_squared)); + for (int k = 0; k < bands; ++k) + { + T overlap = T(0); + for (int i = 0; i < dimension; ++i) + { + overlap += std::conj(psi[i + k * dimension]) * psi[i + j * dimension]; + } + const T expected = j == k ? T(1) : T(0); + const double error = std::abs(overlap - expected); + EXPECT_LT(error, 1.0e-11); + } + } + EXPECT_NEAR(solver.max_residual(), residual, 1.0e-13); + } +} + +TEST(DiagoPPCGPeriodicRitzTest, StartupReclassifiesDisplacedLockedPrefix) +{ + const int dimension = 5; + const int bands = 2; + std::vector psi(dimension * bands, T(0)); + // The initial exact eigenvalue 2 appears converged, but it is not the + // lowest root. The startup correction discovers a lower approximate + // root, which must replace and unlock this column before the next sweep. + psi[1] = T(1); + const double scale = 1.0 / std::sqrt(3.0); + psi[dimension] = T(scale); + psi[2 + dimension] = T(scale); + psi[4 + dimension] = T(scale); + std::vector eigenvalues(bands, 0.0); + const std::vector prec(dimension, 1.0); + const std::vector thresholds(bands, 1.0e-10); + const auto h_op = [dimension](T* in, T* out, int ld, int ncol) { + for (int j = 0; j < ncol; ++j) + { + for (int i = 0; i < dimension; ++i) + { + out[i + j * ld] = double(i + 1) * in[i + j * ld]; + } + } + }; + hsolver::DiagoPPCG solver(1.0e-12, 5, 1, 16); + solver.diag(h_op, nullptr, dimension, bands, dimension, + psi.data(), eigenvalues.data(), thresholds, prec.data()); + EXPECT_TRUE(solver.converged()); + EXPECT_NEAR(eigenvalues[0], 1.0, 1.0e-10); + EXPECT_NEAR(eigenvalues[1], 2.0, 1.0e-10); + EXPECT_LE(solver.max_residual(), 1.0e-10); +} + +TEST(DiagoPPCGPeriodicRitzTest, ClusteredGeneralizedSpectrumWithLockedPrefix) +{ + const int dimension = 16; + const int bands = 4; + std::vector spectrum(dimension); + spectrum[0] = -1.0; + spectrum[1] = 0.0; + spectrum[2] = 1.0e-6; + spectrum[3] = 1.0e-6; + spectrum[4] = 2.0e-6; + std::vector overlap(dimension); + std::vector prec(dimension, 1.0); + for (int i = 0; i < dimension; ++i) + { + overlap[i] = std::pow(10.0, 3.0 * i / (dimension - 1)); + if (i >= 5) + { + spectrum[i] = 1.0 + i; + } + prec[i] = overlap[i]; + } + std::mt19937 random(392); + std::uniform_real_distribution uniform(-1.0, 1.0); + std::vector initial(dimension * bands, T(0)); + initial[0] = T(1.0 / std::sqrt(overlap[0])); + for (int j = 1; j < bands; ++j) + { + for (int i = 1; i < dimension; ++i) + { + const double real = uniform(random); + const double imaginary = uniform(random); + initial[i + j * dimension] = T(real, imaginary) / std::sqrt(overlap[i]); + } + } + const auto h_op = [&spectrum, &overlap, dimension](T* in, T* out, int ld, int ncol) { + for (int j = 0; j < ncol; ++j) + { + for (int i = 0; i < dimension; ++i) + { + out[i + j * ld] = spectrum[i] * overlap[i] * in[i + j * ld]; + } + } + }; + const auto s_op = [&overlap, dimension](T* in, T* out, int ld, int ncol) { + for (int j = 0; j < ncol; ++j) + { + for (int i = 0; i < dimension; ++i) + { + out[i + j * ld] = overlap[i] * in[i + j * ld]; + } + } + }; + const std::vector thresholds(bands, 1.0e-10); + const int intervals[] = {1, 16}; + for (const int interval : intervals) + { + SCOPED_TRACE(interval); + std::vector psi = initial; + std::vector eigenvalues(bands, 0.0); + hsolver::DiagoPPCG solver( + 1.0e-12, 600, 2, interval); + solver.diag(h_op, s_op, dimension, bands, dimension, + psi.data(), eigenvalues.data(), thresholds, prec.data()); + EXPECT_TRUE(solver.converged()); + EXPECT_LE(solver.max_residual(), 1.0e-10); + for (int j = 0; j < bands; ++j) + { + EXPECT_NEAR(eigenvalues[j], spectrum[j], 1.0e-10); + for (int k = 0; k < bands; ++k) + { + T product = T(0); + for (int i = 0; i < dimension; ++i) + { + product += std::conj(psi[i + k * dimension]) * overlap[i] * psi[i + j * dimension]; + } + const T expected = j == k ? T(1) : T(0); + const double error = std::abs(product - expected); + EXPECT_LT(error, 1.0e-10); + } + } + const double prefix_weight = std::norm(psi[0]) * overlap[0]; + EXPECT_NEAR(prefix_weight, 1.0, 1.0e-12); + } +} + +TEST(DiagoPPCGStatusTest, ReportsActiveTailStart) +{ + const int n_dim = 4; + const int nband = 3; + const int ld = n_dim; + std::vector h_mat(n_dim * n_dim, T(0)); + for (int i = 0; i < n_dim; ++i) + { + h_mat[i + i * ld] = T(Real(i + 1), 0); + } + const Real inverse_sqrt_two = Real(1) / std::sqrt(Real(2)); + std::vector psi(ld * nband, T(0)); + psi[0] = T(1, 0); + psi[1 + ld] = T(1, 0); + psi[2 + 2 * ld] = T(inverse_sqrt_two, 0); + psi[3 + 2 * ld] = T(inverse_sqrt_two, 0); + std::vector eval(nband, Real(0)); + const std::vector ethr(nband, 1.0e-12); + const std::vector prec = {1, 2, 3, 4}; + const auto h_op = [&h_mat, n_dim](T* in, T* out, int ld_in, int ncol) { + dense_h_multiply(h_mat.data(), n_dim, in, out, ld_in, ncol); + }; + hsolver::DiagoPPCG solver( + /* diag_thr = */ 1.0e-12, + /* max_iter = */ 0, + /* sbsize = */ 3, + /* rr_step = */ 1); + + solver.diag(h_op, nullptr, ld, nband, n_dim, + psi.data(), eval.data(), ethr, prec.data()); + + EXPECT_FALSE(solver.converged()); + EXPECT_EQ(solver.first_active_band(), 2); + EXPECT_EQ(solver.active_band_count(), 1); + EXPECT_DOUBLE_EQ(solver.first_active_threshold(), Real(1.0e-12)); + EXPECT_GT(solver.max_residual(), solver.first_active_threshold()); +} + +TEST(DiagoPPCGStatusTest, NonFiniteFinalResidualIsNotConverged) +{ + const int n_dim = 2; + const int nband = 2; + const int ld = n_dim; + std::vector psi(ld * nband, T(0)); + psi[0] = T(1, 0); + psi[1 + ld] = T(1, 0); + std::vector eval(nband, Real(0)); + const std::vector ethr(nband, 1e-12); + const std::vector prec = {1, 2}; + int h_apply_count = 0; + const auto h_op = [&h_apply_count](T* in, T* out, int ld_in, int ncol) { + ++h_apply_count; + for (int j = 0; j < ncol; ++j) + { + for (int i = 0; i < ld_in; ++i) + { + const Real eigenvalue = Real(i + 1); + out[i + j * ld_in] = eigenvalue * in[i + j * ld_in]; + } + } + if (h_apply_count >= 2) + { + out[0] = T(std::numeric_limits::quiet_NaN(), 0); + } + }; + hsolver::DiagoPPCG solver( + /* diag_thr = */ 1e-12, + /* max_iter = */ 10, + /* sbsize = */ 2, + /* rr_step = */ 2); + + solver.diag(h_op, nullptr, ld, nband, n_dim, + psi.data(), eval.data(), ethr, prec.data()); + + EXPECT_FALSE(solver.converged()); + EXPECT_EQ(solver.active_band_count(), nband); + EXPECT_FALSE(solver.iteration_limit_reached()); + EXPECT_EQ(solver.first_active_band(), 0); + EXPECT_TRUE(std::isinf(solver.max_residual())); +} + +TEST(DiagoPPCGStatusTest, FailedRayleighRitzCannotCertifyLowestRoots) +{ + const int n_dim = 3; + const int nband = 2; + const int ld = n_dim; + std::vector h_mat(n_dim * n_dim, T(0)); + h_mat[0] = T(1, 0); + h_mat[1 + ld] = T(2, 0); + h_mat[2 + 2 * ld] = T(3, 0); + std::vector psi(ld * nband, T(0)); + psi[2] = T(1, 0); + psi[2 + ld] = T(1, 0); + std::vector eval(nband, Real(0)); + const std::vector ethr(nband, 1e-12); + const std::vector prec = {1, 2, 3}; + const auto h_op = [&h_mat, n_dim](T* in, T* out, int ld_in, int ncol) { + dense_h_multiply(h_mat.data(), n_dim, in, out, ld_in, ncol); + }; + hsolver::DiagoPPCG solver( + /* diag_thr = */ 1e-12, + /* max_iter = */ 0, + /* sbsize = */ 2, + /* rr_step = */ 2); + + solver.diag(h_op, nullptr, ld, nband, n_dim, + psi.data(), eval.data(), ethr, prec.data()); + + EXPECT_FALSE(solver.converged()); + EXPECT_EQ(solver.active_band_count(), nband); + EXPECT_TRUE(solver.iteration_limit_reached()); +} + +TEST_F(DiagoPPCGTridiagTest, ResidualTraceWritesCsv) +{ + const char* env_name = "ABACUS_PPCG_RESIDUAL_TRACE"; + const char* old_env = std::getenv(env_name); + const bool had_old_env = old_env != nullptr; + const std::string old_env_value = had_old_env ? old_env : ""; + const std::string trace_path = "ppcg_residual_trace_test.csv"; + std::remove(trace_path.c_str()); + ASSERT_EQ(::setenv(env_name, trace_path.c_str(), 1), 0); + + std::vector psi_run = psi; + std::vector eval(nband, 0.0); + hsolver::DiagoPPCG solver( + /* diag_thr = */ 1e-12, + /* max_iter = */ 100, + /* sbsize = */ 4, + /* rr_step = */ 4); + + auto h_op = [this](T* in, T* out, int ld_in, int ncol) { + dense_h_multiply(H_mat.data(), n_dim, in, out, ld_in, ncol); + }; + + solver.diag(h_op, nullptr, ld, nband, n_dim, psi_run.data(), eval.data(), ethr, prec.data()); + + if (had_old_env) + { + ASSERT_EQ(::setenv(env_name, old_env_value.c_str(), 1), 0); + } + else + { + ASSERT_EQ(::unsetenv(env_name), 0); + } + + std::ifstream trace(trace_path); + ASSERT_TRUE(trace.good()); + std::string header; + std::string first_record; + std::getline(trace, header); + std::getline(trace, first_record); + EXPECT_EQ(header, + "iteration,stage,max_residual,eigenvalue_0,eigenvalue_1,eigenvalue_2," + "local_overlap_min,local_overlap_max,local_overlap_condition"); + EXPECT_NE(first_record.find("initial_rr"), std::string::npos); + EXPECT_EQ(std::count(first_record.begin(), first_record.end(), ','), nband + 5); + + double previous_max_error = std::numeric_limits::max(); + std::string record = first_record; + do + { + std::stringstream record_stream(record); + std::string field; + for (int field_index = 0; field_index < 3; ++field_index) + { + std::getline(record_stream, field, ','); + } + double max_error = 0.0; + for (int ib = 0; ib < nband; ++ib) + { + std::getline(record_stream, field, ','); + const double eigenvalue = std::stod(field); + const double difference = eigenvalue - exact[ib]; + const double error = std::abs(difference); + max_error = std::max(max_error, error); + } + const double allowed_error = previous_max_error + 1.0e-10; + EXPECT_LE(max_error, allowed_error); + previous_max_error = max_error; + } while (std::getline(trace, record)); + + trace.close(); + std::remove(trace_path.c_str()); +} + +// ============================================================================= +// Test fixture: diagonal matrix (simplest possible Hamiltonian) +// ============================================================================= +class DiagoPPCGDiagonalTest : public ::testing::Test +{ + protected: + void SetUp() override + { + n_dim = 5; + nband = 3; + ld = n_dim; + + // Build diagonal H: H[i,i] = i+1 + H_mat.assign(n_dim * n_dim, T(0)); + for (int i = 0; i < n_dim; ++i) + { + H_mat[i + i * n_dim] = T(Real(i + 1), 0); + } + + // Preconditioner — diagonal of H + prec.resize(n_dim); + for (int i = 0; i < n_dim; ++i) + { + prec[i] = Real(i + 1); + } + + // Lowest 3 eigenvalues: 1, 2, 3 + exact = {1.0, 2.0, 3.0}; + + // Convergence thresholds + ethr.assign(nband, 1e-10); + + // Random initial guess (fixed seed) + std::mt19937 rng(42); + std::uniform_real_distribution dist(-1.0, 1.0); + + psi.assign(ld * nband, T(0)); + for (int j = 0; j < nband; ++j) + { + for (int i = 0; i < n_dim; ++i) + { + psi[i + j * ld] = T(dist(rng), 0.0); + } + } + + // Gram-Schmidt orthonormalisation (S = I) + for (int j = 0; j < nband; ++j) + { + for (int k = 0; k < j; ++k) + { + T dot = 0; + for (int i = 0; i < n_dim; ++i) + { + dot += std::conj(psi[i + k * ld]) * psi[i + j * ld]; + } + for (int i = 0; i < n_dim; ++i) + { + psi[i + j * ld] -= dot * psi[i + k * ld]; + } + } + Real nrm = 0; + for (int i = 0; i < n_dim; ++i) + { + nrm += std::norm(psi[i + j * ld]); + } + nrm = std::sqrt(nrm); + for (int i = 0; i < n_dim; ++i) + { + psi[i + j * ld] /= nrm; + } + } + } + + int n_dim; + int nband; + int ld; + std::vector H_mat; + std::vector prec; + std::vector exact; + std::vector ethr; + std::vector psi; +}; + +TEST_F(DiagoPPCGDiagonalTest, BlockSubspace) +{ + std::vector psi_run = psi; + std::vector eval(nband, 0.0); + + hsolver::DiagoPPCG solver( + /* diag_thr = */ 1e-12, + /* max_iter = */ 50, + /* sbsize = */ 3, + /* rr_step = */ 3); + + auto h_op = [this](T* in, T* out, int ld_in, int ncol) { + dense_h_multiply(H_mat.data(), n_dim, in, out, ld_in, ncol); + }; + + double avg_iter = solver.diag(h_op, nullptr, ld, nband, n_dim, psi_run.data(), eval.data(), ethr, prec.data()); + + for (int i = 0; i < nband; ++i) + { + EXPECT_NEAR(eval[i], exact[i], 1e-8) << "Diagonal BLOCK: eigenvalue[" << i << "] mismatch"; + } + EXPECT_LE(avg_iter, double(50)) << "Diagonal BLOCK: too many iterations"; +} + +TEST(DiagoPPCGRootSelectionTest, StartupCorrectionKeepsLowestRoots) +{ + const int n_dim = 8; + const int nband = 4; + const int ld = n_dim; + const Real inv_sqrt_two = Real(1) / std::sqrt(Real(2)); + + std::vector h_mat(n_dim * n_dim, T(0)); + std::vector prec(n_dim); + for (int i = 0; i < n_dim; ++i) + { + const Real diagonal = Real(i + 1); + h_mat[i + i * n_dim] = T(diagonal, 0); + prec[i] = diagonal; + } + + // The first two columns contain the missing third and fourth roots. The + // last two are exact high-energy eigenvectors and would otherwise lock + // before their lower replacements can enter the active subspace. + std::vector psi(ld * nband, T(0)); + psi[0] = T(inv_sqrt_two, 0); + psi[2] = T(inv_sqrt_two, 0); + psi[1 + ld] = T(inv_sqrt_two, 0); + psi[3 + ld] = T(inv_sqrt_two, 0); + psi[6 + 2 * ld] = T(1, 0); + psi[7 + 3 * ld] = T(1, 0); + + std::vector eval(nband, Real(0)); + std::vector ethr(nband, 1e-12); + hsolver::DiagoPPCG solver( + /* diag_thr = */ 1e-12, + /* max_iter = */ 50, + /* sbsize = */ 2, + /* rr_step = */ 4); + + auto h_op = [&h_mat, n_dim](T* in, T* out, int ld_in, int ncol) { + dense_h_multiply(h_mat.data(), n_dim, in, out, ld_in, ncol); + }; + + const double avg_iter = solver.diag( + h_op, nullptr, ld, nband, n_dim, psi.data(), eval.data(), ethr, prec.data()); + for (int i = 0; i < nband; ++i) + { + EXPECT_NEAR(eval[i], Real(i + 1), 1e-10) << "root[" << i << "] mismatch"; + } + EXPECT_TRUE(solver.converged()); + EXPECT_LE(avg_iter, 50.0); +} + +TEST(DiagoPPCGRootSelectionTest, StartupCorrectionKeepsLowestRootsWithoutInitialGap) +{ + const int n_dim = 8; + const int nband = 4; + const int ld = n_dim; + const Real inv_sqrt_two = Real(1) / std::sqrt(Real(2)); + + std::vector h_mat(n_dim * n_dim, T(0)); + std::vector prec(n_dim); + for (int i = 0; i < n_dim; ++i) + { + const Real diagonal = Real(i + 1); + h_mat[i + i * n_dim] = T(diagonal, 0); + prec[i] = diagonal; + } + + // Every initial vector has a nonzero residual. Independent two-column + // blocks would nevertheless retain roots 1, 2, 5, and 6, discarding the + // third and fourth roots before a global comparison is possible. + std::vector psi(ld * nband, T(0)); + psi[0] = T(inv_sqrt_two, 0); + psi[2] = T(inv_sqrt_two, 0); + psi[1 + ld] = T(inv_sqrt_two, 0); + psi[3 + ld] = T(inv_sqrt_two, 0); + psi[4 + 2 * ld] = T(inv_sqrt_two, 0); + psi[6 + 2 * ld] = T(inv_sqrt_two, 0); + psi[5 + 3 * ld] = T(inv_sqrt_two, 0); + psi[7 + 3 * ld] = T(inv_sqrt_two, 0); + + std::vector eval(nband, Real(0)); + std::vector ethr(nband, 1e-12); + hsolver::DiagoPPCG solver( + /* diag_thr = */ 1e-12, + /* max_iter = */ 50, + /* sbsize = */ 2, + /* rr_step = */ 4); + + auto h_op = [&h_mat, n_dim](T* in, T* out, int ld_in, int ncol) { + dense_h_multiply(h_mat.data(), n_dim, in, out, ld_in, ncol); + }; + + const double avg_iter = solver.diag( + h_op, nullptr, ld, nband, n_dim, psi.data(), eval.data(), ethr, prec.data()); + for (int i = 0; i < nband; ++i) + { + EXPECT_NEAR(eval[i], Real(i + 1), 1e-10) << "root[" << i << "] mismatch"; + } + EXPECT_TRUE(solver.converged()); + EXPECT_LE(avg_iter, 50.0); +} + +TEST(DiagoPPCGRootSelectionTest, StartupCorrectionKeepsConvergedPrefixInactive) +{ + const int n_dim = 10; + const int nband = 6; + const int ld = n_dim; + const Real inv_sqrt_two = Real(1) / std::sqrt(Real(2)); + + std::vector h_mat(n_dim * n_dim, T(0)); + std::vector prec(n_dim); + for (int i = 0; i < n_dim; ++i) + { + const Real diagonal = Real(i + 1); + h_mat[i + i * n_dim] = T(diagonal, 0); + prec[i] = diagonal; + } + + // The first two roots are already converged. The active tail still needs + // one global comparison: its first two vectors contain roots 3--6, while + // its last two vectors are exact high-energy states. + std::vector psi(ld * nband, T(0)); + psi[0] = T(1, 0); + psi[1 + ld] = T(1, 0); + psi[2 + 2 * ld] = T(inv_sqrt_two, 0); + psi[4 + 2 * ld] = T(inv_sqrt_two, 0); + psi[3 + 3 * ld] = T(inv_sqrt_two, 0); + psi[5 + 3 * ld] = T(inv_sqrt_two, 0); + psi[8 + 4 * ld] = T(1, 0); + psi[9 + 5 * ld] = T(1, 0); + + std::vector eval(nband, Real(0)); + const std::vector ethr(nband, 1e-12); + hsolver::DiagoPPCG solver( + /* diag_thr = */ 1e-12, + /* max_iter = */ 1, + /* sbsize = */ 2, + /* rr_step = */ 4); + const auto h_op = [&h_mat, n_dim](T* in, T* out, int ld_in, int ncol) { + dense_h_multiply(h_mat.data(), n_dim, in, out, ld_in, ncol); + }; + + const double avg_iter = solver.diag( + h_op, nullptr, ld, nband, n_dim, psi.data(), eval.data(), ethr, prec.data()); + + for (int i = 0; i < nband; ++i) + { + EXPECT_NEAR(eval[i], Real(i + 1), 1e-10) << "root[" << i << "] mismatch"; + } + const double expected_avg_iter = 1.0 + 4.0 / 6.0; + EXPECT_TRUE(solver.converged()); + EXPECT_NEAR(avg_iter, expected_avg_iter, 1e-12); +} + +TEST(DiagoPPCGRootSelectionTest, WarmStartAfterHamiltonianChangeKeepsLowestRoots) +{ + const int n_dim = 8; + const int nband = 4; + const int ld = n_dim; + const Real inv_sqrt_two = Real(1) / std::sqrt(Real(2)); + + std::vector h_mat(n_dim * n_dim, T(0)); + std::vector prec(n_dim, Real(1)); + std::vector psi(ld * nband, T(0)); + const int first_indices[nband] = {0, 1, 4, 5}; + const int second_indices[nband] = {2, 3, 6, 7}; + for (int state = 0; state < nband; ++state) + { + const int first = first_indices[state]; + const int second = second_indices[state]; + const Real low_eigenvalue = Real(state + 1); + const Real high_eigenvalue = Real(state + 5); + const Real diagonal = Real(0.5) * (low_eigenvalue + high_eigenvalue); + const Real off_diagonal = Real(0.5) * (low_eigenvalue - high_eigenvalue); + h_mat[first + first * n_dim] = T(diagonal, 0); + h_mat[second + second * n_dim] = T(diagonal, 0); + h_mat[first + second * n_dim] = T(off_diagonal, 0); + h_mat[second + first * n_dim] = T(off_diagonal, 0); + psi[first + state * ld] = T(inv_sqrt_two, 0); + psi[second + state * ld] = T(inv_sqrt_two, 0); + prec[first] = diagonal; + prec[second] = diagonal; + } + + std::vector eval(nband, Real(0)); + std::vector ethr(nband, 1e-12); + hsolver::DiagoPPCG solver( + /* diag_thr = */ 1e-12, + /* max_iter = */ 50, + /* sbsize = */ 2, + /* rr_step = */ 4); + auto h_op = [&h_mat, n_dim](T* in, T* out, int ld_in, int ncol) { + dense_h_multiply(h_mat.data(), n_dim, in, out, ld_in, ncol); + }; + + solver.diag( + h_op, nullptr, ld, nband, n_dim, psi.data(), eval.data(), ethr, prec.data()); + ASSERT_TRUE(solver.converged()); + for (int i = 0; i < nband; ++i) + { + ASSERT_NEAR(eval[i], Real(i + 1), 1e-10); + } + + // Model a later SCF call: the previous eigenvectors are a valid warm start, + // but the changed Hamiltonian makes each one mix a low and a high root. + std::fill(h_mat.begin(), h_mat.end(), T(0)); + for (int i = 0; i < n_dim; ++i) + { + const Real diagonal = Real(i + 1); + h_mat[i + i * n_dim] = T(diagonal, 0); + prec[i] = diagonal; + } + + solver.diag( + h_op, nullptr, ld, nband, n_dim, psi.data(), eval.data(), ethr, prec.data()); + EXPECT_TRUE(solver.converged()); + for (int i = 0; i < nband; ++i) + { + EXPECT_NEAR(eval[i], Real(i + 1), 1e-10) << "root[" << i << "] mismatch"; + } +} + +TEST(DiagoPPCGRootSelectionTest, StartupCorrectionKeepsLowestGeneralizedRoots) +{ + const int n_dim = 8; + const int nband = 4; + const int ld = n_dim; + const Real inv_sqrt_two = Real(1) / std::sqrt(Real(2)); + + std::vector h_mat(n_dim * n_dim, T(0)); + std::vector s_mat(n_dim * n_dim, T(0)); + std::vector s_diagonal(n_dim); + std::vector prec(n_dim); + for (int i = 0; i < n_dim; ++i) + { + const Real eigenvalue = Real(i + 1); + const Real overlap_value = Real(1) + Real(0.1) * Real(i % 3); + const Real hamiltonian_value = eigenvalue * overlap_value; + h_mat[i + i * n_dim] = T(hamiltonian_value, 0); + s_mat[i + i * n_dim] = T(overlap_value, 0); + s_diagonal[i] = overlap_value; + prec[i] = hamiltonian_value; + } + + // H = S diag(1,...,8), so the generalized eigenvectors are the + // coordinate vectors. The initial columns create the same cross-block + // root-selection hazard as the identity-overlap regression above. + std::vector psi(ld * nband, T(0)); + std::vector normalized_coefficients(n_dim); + for (int i = 0; i < n_dim; ++i) + { + normalized_coefficients[i] = inv_sqrt_two / std::sqrt(s_diagonal[i]); + } + psi[0] = T(normalized_coefficients[0], 0); + psi[2] = T(normalized_coefficients[2], 0); + psi[1 + ld] = T(normalized_coefficients[1], 0); + psi[3 + ld] = T(normalized_coefficients[3], 0); + psi[4 + 2 * ld] = T(normalized_coefficients[4], 0); + psi[6 + 2 * ld] = T(normalized_coefficients[6], 0); + psi[5 + 3 * ld] = T(normalized_coefficients[5], 0); + psi[7 + 3 * ld] = T(normalized_coefficients[7], 0); + + std::vector eval(nband, Real(0)); + std::vector ethr(nband, 1e-12); + hsolver::DiagoPPCG solver( + /* diag_thr = */ 1e-12, + /* max_iter = */ 50, + /* sbsize = */ 2, + /* rr_step = */ 4); + + auto h_op = [&h_mat, n_dim](T* in, T* out, int ld_in, int ncol) { + dense_h_multiply(h_mat.data(), n_dim, in, out, ld_in, ncol); + }; + auto s_op = [&s_mat, n_dim](T* in, T* out, int ld_in, int ncol) { + dense_h_multiply(s_mat.data(), n_dim, in, out, ld_in, ncol); + }; + + const double avg_iter = solver.diag( + h_op, s_op, ld, nband, n_dim, psi.data(), eval.data(), ethr, prec.data()); + for (int i = 0; i < nband; ++i) + { + EXPECT_NEAR(eval[i], Real(i + 1), 1e-10) << "root[" << i << "] mismatch"; + } + EXPECT_TRUE(solver.converged()); + EXPECT_LE(avg_iter, 50.0); +} + +TEST(DiagoPPCGRootSelectionTest, LockingKeepsRequestedLowRoots) +{ + const int n_dim = 4; + const int nband = 2; + const int ld = n_dim; + std::vector h_mat(n_dim * n_dim, T(0)); + h_mat[0] = T(1, 0); + h_mat[1 + ld] = T(2, 0); + h_mat[2 + 2 * ld] = T(4, 0); + h_mat[3 + 3 * ld] = T(3, 0); + const std::vector prec = {1, 2, 4, 3}; + const std::vector ethr(nband, 1e-12); + const Real epsilons[] = {1e-4, 1e-3, 1e-2, 1e-1}; + + const auto h_op = [&h_mat, n_dim](T* in, T* out, int ld_in, int ncol) { + dense_h_multiply(h_mat.data(), n_dim, in, out, ld_in, ncol); + }; + + for (const Real epsilon : epsilons) + { + SCOPED_TRACE("epsilon = " + std::to_string(epsilon)); + std::vector psi(ld * nband, T(0)); + const Real norm = std::sqrt(Real(2) + epsilon * epsilon); + psi[0] = T(Real(1) / norm, 0); + psi[1] = T(epsilon / norm, 0); + psi[2] = T(Real(1) / norm, 0); + psi[3 + ld] = T(1, 0); + std::vector eval(nband, Real(0)); + + hsolver::DiagoPPCG solver( + /* diag_thr = */ 1e-12, + /* max_iter = */ 80, + /* sbsize = */ 2, + /* rr_step = */ 1); + const double avg_iter = solver.diag( + h_op, nullptr, ld, nband, n_dim, psi.data(), eval.data(), ethr, prec.data()); + EXPECT_NEAR(eval[0], 1.0, 1e-8); + EXPECT_NEAR(eval[1], 2.0, 1e-8); + EXPECT_TRUE(solver.converged()); + EXPECT_LE(avg_iter, 80.0); + } +} + +TEST(DiagoPPCGRootSelectionTest, LaterCorrectionKeepsLowestRootsAcrossBlocks) +{ + const int n_dim = 4; + const int nband = 2; + const int ld = n_dim; + const std::vector diagonal = {1.0, 2.0, 3.0, 100.0}; + const std::vector prec(n_dim, 1.0); + const std::vector ethr(nband, 1.0e-10); + const Real coefficient = 1.0 / std::sqrt(3.0); + const int block_sizes[] = {1, 2}; + const auto h_op = [&diagonal, n_dim](T* in, T* out, int ld_in, int ncol) { + for (int column = 0; column < ncol; ++column) + { + for (int row = 0; row < n_dim; ++row) + { + out[row + column * ld_in] = diagonal[row] * in[row + column * ld_in]; + } + } + }; + + for (const int block_size : block_sizes) + { + SCOPED_TRACE("block_size = " + std::to_string(block_size)); + // The first expanded subspace keeps a low mixture and the exact + // third root. The missing second root needs a later correction; + // convergence of the retained vectors alone does not certify it. + std::vector psi(ld * nband, T(0)); + psi[0] = T(coefficient); + psi[1] = T(coefficient); + psi[3] = T(coefficient); + psi[2 + ld] = T(1); + std::vector eval(nband, 0.0); + hsolver::DiagoPPCG solver( + 1.0e-10, 100, block_size, 1); + solver.diag(h_op, nullptr, ld, nband, n_dim, + psi.data(), eval.data(), ethr, prec.data()); + + double max_residual = 0.0; + for (int column = 0; column < nband; ++column) + { + double residual_squared = 0.0; + for (int row = 0; row < n_dim; ++row) + { + const T residual = (diagonal[row] - eval[column]) * psi[row + column * ld]; + residual_squared += std::norm(residual); + } + const double residual_norm = std::sqrt(residual_squared); + max_residual = std::max(max_residual, residual_norm); + for (int other = 0; other < nband; ++other) + { + T overlap(0); + for (int row = 0; row < n_dim; ++row) + { + overlap += std::conj(psi[row + column * ld]) * psi[row + other * ld]; + } + const T expected_overlap = column == other ? T(1) : T(0); + const double overlap_error = std::abs(overlap - expected_overlap); + EXPECT_LE(overlap_error, 1.0e-12); + } + } + EXPECT_TRUE(solver.converged()); + EXPECT_LE(max_residual, 1.0e-10); + EXPECT_NEAR(solver.max_residual(), max_residual, 1.0e-13); + EXPECT_NEAR(eval[0], 1.0, 1.0e-9); + EXPECT_NEAR(eval[1], 2.0, 1.0e-9) + << "returned roots = " << eval[0] << ", " << eval[1] + << "; fresh max residual = " << max_residual; + if (block_size == 1) + { + EXPECT_GT(solver.root_recovery_count(), 0); + } + else + { + EXPECT_EQ(solver.root_recovery_count(), 0); + } + } +} + +TEST(DiagoPPCGRootSelectionTest, LaterCorrectionsRecoverMultipleComplexGeneralizedRoots) +{ + const int dimension = 6; + const int bands = 3; + const Real spectrum[] = {1.0, 2.0, 3.0, 4.0, 5.0, 100.0}; + const Real overlap_spectrum[] = {1.0, 2.0, 4.0, 7.0, 3.0, 5.0}; + std::vector unitary(dimension * dimension); + const Real normalization = 1.0 / std::sqrt(Real(dimension)); + for (int column = 0; column < dimension; ++column) + { + for (int row = 0; row < dimension; ++row) + { + const Real angle = 2.0 * M_PI * row * column / dimension; + const T phase(std::cos(angle), std::sin(angle)); + unitary[row + column * dimension] = normalization * phase; + } + } + std::vector hamiltonian(dimension * dimension, T(0)); + std::vector overlap(dimension * dimension, T(0)); + for (int column = 0; column < dimension; ++column) + { + for (int row = 0; row < dimension; ++row) + { + for (int state = 0; state < dimension; ++state) + { + const T projector = unitary[row + state * dimension] + * std::conj(unitary[column + state * dimension]); + hamiltonian[row + column * dimension] += projector * overlap_spectrum[state] * spectrum[state]; + overlap[row + column * dimension] += projector * overlap_spectrum[state]; + } + } + } + const auto h_op = [&hamiltonian, dimension](T* in, T* out, int ld, int count) { + dense_h_multiply(hamiltonian.data(), dimension, in, out, ld, count); + }; + const auto s_op = [&overlap, dimension](T* in, T* out, int ld, int count) { + dense_h_multiply(overlap.data(), dimension, in, out, ld, count); + }; + const std::vector preconditioner(dimension, 1.0); + const std::vector thresholds(bands, 1.0e-9); + const int block_sizes[] = {1, 2, 3}; + const int ritz_intervals[] = {1, 4}; + for (const int block_size : block_sizes) + { + for (const int ritz_interval : ritz_intervals) + { + SCOPED_TRACE("block_size = " + std::to_string(block_size) + + ", ritz_interval = " + std::to_string(ritz_interval)); + std::vector psi(dimension * bands, T(0)); + // Four spectral components in the first vector require more + // than the first [X,W] expansion to resolve all three low roots. + const int mixed_states[] = {0, 1, 2, 5}; + for (int row = 0; row < dimension; ++row) + { + for (const int state : mixed_states) + { + psi[row] += 0.5 * unitary[row + state * dimension] + / std::sqrt(overlap_spectrum[state]); + } + psi[row + dimension] = unitary[row + 3 * dimension] / std::sqrt(overlap_spectrum[3]); + psi[row + 2 * dimension] = unitary[row + 4 * dimension] / std::sqrt(overlap_spectrum[4]); + } + std::vector eigenvalues(bands, 0.0); + hsolver::DiagoPPCG solver( + 1.0e-9, 120, block_size, ritz_interval); + solver.diag(h_op, s_op, dimension, bands, dimension, + psi.data(), eigenvalues.data(), thresholds, preconditioner.data()); + std::vector hpsi(psi.size(), T(0)); + std::vector spsi(psi.size(), T(0)); + h_op(psi.data(), hpsi.data(), dimension, bands); + s_op(psi.data(), spsi.data(), dimension, bands); + EXPECT_TRUE(solver.converged()); + for (int column = 0; column < bands; ++column) + { + EXPECT_NEAR(eigenvalues[column], spectrum[column], 1.0e-8); + double residual_squared = 0.0; + for (int row = 0; row < dimension; ++row) + { + const int offset = row + column * dimension; + const T residual = hpsi[offset] - eigenvalues[column] * spsi[offset]; + residual_squared += std::norm(residual); + } + const double residual_norm = std::sqrt(residual_squared); + EXPECT_LE(residual_norm, 1.0e-9); + for (int other = 0; other < bands; ++other) + { + T inner_product(0); + for (int row = 0; row < dimension; ++row) + { + inner_product += std::conj(psi[row + column * dimension]) * spsi[row + other * dimension]; + } + const T expected = column == other ? T(1) : T(0); + const double orthogonality_error = std::abs(inner_product - expected); + EXPECT_LE(orthogonality_error, 1.0e-11); + } + } + } + } +} + +TEST(DiagoPPCGRankTruncationTest, ComplexHermitianStartupSubspace) +{ + const int n_dim = 6; + const int nband = 3; + const int ld = n_dim; + const Real inv_sqrt_two = Real(1) / std::sqrt(Real(2)); + + // Build H = U diag(1,...,6) U^H. The first two columns of U form a + // complex rotation, so the reduced Hamiltonian needs a conjugate + // transpose when it is projected into the retained overlap subspace. + std::vector unitary(n_dim * n_dim, T(0)); + unitary[0] = T(inv_sqrt_two, 0); + unitary[1] = T(0, inv_sqrt_two); + unitary[n_dim] = T(0, inv_sqrt_two); + unitary[1 + n_dim] = T(inv_sqrt_two, 0); + for (int i = 2; i < n_dim; ++i) + { + unitary[i + i * n_dim] = T(1, 0); + } + + std::vector h_mat(n_dim * n_dim, T(0)); + std::vector eigenvalues(n_dim); + for (int state = 0; state < n_dim; ++state) + { + eigenvalues[state] = Real(state + 1); + for (int col = 0; col < n_dim; ++col) + { + for (int row = 0; row < n_dim; ++row) + { + h_mat[row + col * n_dim] += unitary[row + state * n_dim] + * eigenvalues[state] + * std::conj(unitary[col + state * n_dim]); + } + } + } + EXPECT_GT(std::abs(std::imag(h_mat[n_dim])), Real(0.1)); + + std::vector prec(n_dim); + for (int i = 0; i < n_dim; ++i) + { + prec[i] = std::max(Real(1), std::real(h_mat[i + i * n_dim])); + } + + // The third column is an exact high-energy state. Its zero residual + // makes [psi, w] rank deficient during the first global correction, + // while the first two columns and their residuals span the three lowest + // roots. This exercises Q^H K Q for a complex Hermitian matrix. + std::vector psi(ld * nband, T(0)); + for (int i = 0; i < n_dim; ++i) + { + psi[i] = (unitary[i] + unitary[i + 3 * n_dim]) * inv_sqrt_two; + psi[i + ld] = (unitary[i + n_dim] + unitary[i + 2 * n_dim]) * inv_sqrt_two; + psi[i + 2 * ld] = unitary[i + 5 * n_dim]; + } + + std::vector eval(nband, Real(0)); + std::vector ethr(nband, 1e-12); + hsolver::DiagoPPCG solver( + /* diag_thr = */ 1e-12, + /* max_iter = */ 30, + /* sbsize = */ 2, + /* rr_step = */ 4); + + auto h_op = [&h_mat, n_dim](T* in, T* out, int ld_in, int ncol) { + dense_h_multiply(h_mat.data(), n_dim, in, out, ld_in, ncol); + }; + + const double avg_iter = solver.diag( + h_op, nullptr, ld, nband, n_dim, psi.data(), eval.data(), ethr, prec.data()); + for (int i = 0; i < nband; ++i) + { + EXPECT_NEAR(eval[i], Real(i + 1), 1e-10) << "root[" << i << "] mismatch"; + } + EXPECT_TRUE(solver.converged()); + EXPECT_LE(avg_iter, 30.0); +} + +TEST_F(DiagoPPCGDiagonalTest, EmptyHOperatorThrows) +{ + std::vector psi_run = psi; + std::vector eval(nband, 0.0); + + hsolver::DiagoPPCG solver( + /* diag_thr = */ 1e-12, + /* max_iter = */ 50, + /* sbsize = */ 3, + /* rr_step = */ 3); + + hsolver::DiagoPPCG::HPsiFunc h_op; + EXPECT_THROW(solver.diag(h_op, nullptr, ld, nband, n_dim, psi_run.data(), eval.data(), ethr, prec.data()), + std::invalid_argument); +} + +TEST_F(DiagoPPCGDiagonalTest, NonFiniteInputThrows) +{ + std::vector psi_run = psi; + std::vector eval(nband, 0.0); + + hsolver::DiagoPPCG solver( + /* diag_thr = */ 1e-12, + /* max_iter = */ 50, + /* sbsize = */ 3, + /* rr_step = */ 3); + + auto h_op = [this](T* in, T* out, int ld_in, int ncol) { + dense_h_multiply(H_mat.data(), n_dim, in, out, ld_in, ncol); + }; + + std::vector bad_ethr = ethr; + bad_ethr[0] = std::numeric_limits::quiet_NaN(); + EXPECT_THROW(solver.diag(h_op, nullptr, ld, nband, n_dim, psi_run.data(), eval.data(), bad_ethr, prec.data()), + std::invalid_argument); + + std::vector bad_prec = prec; + bad_prec[0] = std::numeric_limits::infinity(); + EXPECT_THROW(solver.diag(h_op, nullptr, ld, nband, n_dim, psi_run.data(), eval.data(), ethr, bad_prec.data()), + std::invalid_argument); +} + +TEST(DiagoPPCGWorkspaceTest, RejectsOversizedWorkspacesBeforeUsingInputBuffers) +{ + struct Shape + { + const char* description; + int leading_dimension; + int bands; + }; + const int int_max = std::numeric_limits::max(); + const Shape shapes[] = { + {"wavefunction storage", int_max / 2 + 1, 2}, + {"local basis storage", int_max / 3 + 1, 2}, + {"projected matrix storage", 1, 32768}, + {"LAPACK workspace", 1, 20000}, + }; + + for (const Shape& shape : shapes) + { + SCOPED_TRACE(shape.description); + T psi_value(1.0, 0.0); + Real eigenvalue = -1.0; + const Real preconditioner = 1.0; + // The size diagnostic must precede threshold/buffer validation. An + // empty threshold vector also keeps this regression bounded if that + // diagnostic is accidentally removed: no huge allocations are needed. + const std::vector thresholds; + int h_calls = 0; + int s_calls = 0; + const auto h_op = [&h_calls](T*, T*, int, int) { ++h_calls; }; + const auto s_op = [&s_calls](T*, T*, int, int) { ++s_calls; }; + hsolver::DiagoPPCG solver(1e-12, 1, 1, 1); + + try + { + solver.diag(h_op, + s_op, + shape.leading_dimension, + shape.bands, + 1, + &psi_value, + &eigenvalue, + thresholds, + &preconditioner); + FAIL() << "Oversized PPCG workspace was accepted"; + } + catch (const std::invalid_argument& error) + { + EXPECT_STREQ(error.what(), "PPCG: workspace dimensions exceed integer indexing or LAPACK limits."); + } + EXPECT_EQ(h_calls, 0); + EXPECT_EQ(s_calls, 0); + EXPECT_EQ(psi_value, T(1.0, 0.0)); + EXPECT_DOUBLE_EQ(eigenvalue, -1.0); + } +} + +TEST_F(DiagoPPCGDiagonalTest, HandlesSubblockSizeAboveBandCount) +{ + std::vector psi_run = psi; + std::vector eval(nband, 0.0); + const int subblock_size = std::numeric_limits::max(); + hsolver::DiagoPPCG solver(1e-12, 100, subblock_size, 3); + const auto h_op = [this](T* in, T* out, int ld_in, int ncol) { + dense_h_multiply(H_mat.data(), n_dim, in, out, ld_in, ncol); + }; + + solver.diag(h_op, nullptr, ld, nband, n_dim, psi_run.data(), eval.data(), ethr, prec.data()); + + EXPECT_TRUE(solver.converged()); + for (int band = 0; band < nband; ++band) + { + EXPECT_NEAR(eval[band], exact[band], ethr[band]); + } +} + +TEST(DiagoPPCGLeadingDimensionTest, BlockSubspaceWithPadding) +{ + const int n_dim = 5; + const int nband = 3; + const int ld = 8; + + std::vector H_mat(n_dim * n_dim, T(0)); + for (int i = 0; i < n_dim; ++i) + { + H_mat[i + i * n_dim] = T(Real(i + 1), 0); + } + + std::vector prec(n_dim); + for (int i = 0; i < n_dim; ++i) + { + prec[i] = Real(i + 1); + } + + std::vector psi(ld * nband, T(17.0, -3.0)); + std::mt19937 rng(7); + std::uniform_real_distribution dist(-1.0, 1.0); + for (int j = 0; j < nband; ++j) + { + for (int i = 0; i < n_dim; ++i) + { + psi[i + j * ld] = T(dist(rng), 0.0); + } + } + + for (int j = 0; j < nband; ++j) + { + for (int k = 0; k < j; ++k) + { + T dot = 0; + for (int i = 0; i < n_dim; ++i) + { + dot += std::conj(psi[i + k * ld]) * psi[i + j * ld]; + } + for (int i = 0; i < n_dim; ++i) + { + psi[i + j * ld] -= dot * psi[i + k * ld]; + } + } + Real nrm = 0; + for (int i = 0; i < n_dim; ++i) + { + nrm += std::norm(psi[i + j * ld]); + } + nrm = std::sqrt(nrm); + for (int i = 0; i < n_dim; ++i) + { + psi[i + j * ld] /= nrm; + } + } + + std::vector eval(nband, 0.0); + std::vector ethr(nband, 1e-10); + hsolver::DiagoPPCG solver( + /* diag_thr = */ 1e-12, + /* max_iter = */ 80, + /* sbsize = */ 2, + /* rr_step = */ 3); + + auto h_op = [&H_mat, n_dim](T* in, T* out, int ld_in, int ncol) { + dense_h_multiply(H_mat.data(), n_dim, in, out, ld_in, ncol); + }; + + double avg_iter = solver.diag(h_op, nullptr, ld, nband, n_dim, psi.data(), eval.data(), ethr, prec.data()); + + const Real exact[] = {1.0, 2.0, 3.0}; + for (int i = 0; i < nband; ++i) + { + EXPECT_NEAR(eval[i], exact[i], 1e-8) << "Padded ld BLOCK: eigenvalue[" << i << "] mismatch"; + } + EXPECT_LE(avg_iter, double(80)) << "Padded ld BLOCK: too many iterations"; +} + +// ============================================================================= +// Test fixture: 2×2 matrix — smallest non-trivial case +// H = [[2, 1], [1, 2]], eigenvalues: 1, 3 +// ============================================================================= +class DiagoPPCG2x2Test : public ::testing::Test +{ + protected: + void SetUp() override + { + n_dim = 2; + nband = 2; + ld = n_dim; + + // H = [[2, 1], [1, 2]] + H_mat.assign(n_dim * n_dim, T(0)); + H_mat[0 + 0 * n_dim] = T(2.0, 0); + H_mat[1 + 1 * n_dim] = T(2.0, 0); + H_mat[0 + 1 * n_dim] = T(1.0, 0); + H_mat[1 + 0 * n_dim] = T(1.0, 0); + + prec.assign(n_dim, 2.0); + + // λ₁ = 1, λ₂ = 3 + exact = {1.0, 3.0}; + + ethr.assign(nband, 1e-10); + + std::mt19937 rng(123); + std::uniform_real_distribution dist(-1.0, 1.0); + + psi.assign(ld * nband, T(0)); + for (int j = 0; j < nband; ++j) + { + for (int i = 0; i < n_dim; ++i) + { + psi[i + j * ld] = T(dist(rng), 0.0); + } + } + + // Gram-Schmidt orthonormalisation (S = I) + for (int j = 0; j < nband; ++j) + { + for (int k = 0; k < j; ++k) + { + T dot = 0; + for (int i = 0; i < n_dim; ++i) + { + dot += std::conj(psi[i + k * ld]) * psi[i + j * ld]; + } + for (int i = 0; i < n_dim; ++i) + { + psi[i + j * ld] -= dot * psi[i + k * ld]; + } + } + Real nrm = 0; + for (int i = 0; i < n_dim; ++i) + { + nrm += std::norm(psi[i + j * ld]); + } + nrm = std::sqrt(nrm); + for (int i = 0; i < n_dim; ++i) + { + psi[i + j * ld] /= nrm; + } + } + } + + int n_dim; + int nband; + int ld; + std::vector H_mat; + std::vector prec; + std::vector exact; + std::vector ethr; + std::vector psi; +}; + +TEST_F(DiagoPPCG2x2Test, BlockSubspace) +{ + std::vector psi_run = psi; + std::vector eval(nband, 0.0); + + hsolver::DiagoPPCG solver( + /* diag_thr = */ 1e-12, + /* max_iter = */ 50, + /* sbsize = */ 2, + /* rr_step = */ 2); + + auto h_op = [this](T* in, T* out, int ld_in, int ncol) { + dense_h_multiply(H_mat.data(), n_dim, in, out, ld_in, ncol); + }; + + double avg_iter = solver.diag(h_op, nullptr, ld, nband, n_dim, psi_run.data(), eval.data(), ethr, prec.data()); + + for (int i = 0; i < nband; ++i) + { + EXPECT_NEAR(eval[i], exact[i], 1e-8) << "2x2 BLOCK: eigenvalue[" << i << "] mismatch"; + } + EXPECT_LE(avg_iter, double(50)) << "2x2 BLOCK: too many iterations"; +} + +TEST(DiagoPPCGComplexHermitianTest, DefaultKeepsImaginaryProjection) +{ + const int n_dim = 2; + const int nband = 2; + const int ld = n_dim; + + // H = [[2, i], [-i, 3]]. Dropping Im() would incorrectly + // diagonalize diag(2, 3); the Hermitian eigenvalues are 2.5 +/- sqrt(1.25). + std::vector H_mat(n_dim * n_dim, T(0)); + H_mat[0 + 0 * n_dim] = T(2.0, 0.0); + H_mat[1 + 1 * n_dim] = T(3.0, 0.0); + H_mat[0 + 1 * n_dim] = T(0.0, 1.0); + H_mat[1 + 0 * n_dim] = T(0.0, -1.0); + + std::vector psi(ld * nband, T(0)); + psi[0 + 0 * ld] = T(1.0, 0.0); + psi[1 + 1 * ld] = T(1.0, 0.0); + + std::vector prec(n_dim, 2.0); + std::vector ethr(nband, 1e-12); + std::vector eval(nband, 0.0); + + hsolver::DiagoPPCG solver( + /* diag_thr = */ 1e-12, + /* max_iter = */ 10, + /* sbsize = */ 2, + /* rr_step = */ 1); + + auto h_op = [&H_mat, n_dim](T* in, T* out, int ld_in, int ncol) { + dense_h_multiply(H_mat.data(), n_dim, in, out, ld_in, ncol); + }; + + solver.diag(h_op, nullptr, ld, nband, n_dim, psi.data(), eval.data(), ethr, prec.data()); + + const Real delta = std::sqrt(1.25); + EXPECT_NEAR(eval[0], 2.5 - delta, 1e-10); + EXPECT_NEAR(eval[1], 2.5 + delta, 1e-10); +} + +TEST(DiagoPPCGComplexHermitianTest, ExactImaginaryLeadingCoefficientNeedsNoProjection) +{ + const int n_dim = 3; + const int nband = 1; + const int ld = n_dim; + const int matrix_size = n_dim * n_dim; + const Real diagonal[n_dim] = {Real(1), Real(4), Real(9)}; + std::vector h_mat(matrix_size, T(0)); + for (int i = 0; i < n_dim; ++i) + { + h_mat[i + i * n_dim] = T(diagonal[i], 0); + } + + // A valid full-complex eigenvector may have an imaginary first + // coefficient. Clearing that coefficient would destroy its norm and + // rank. A Ritz solve is free to choose another overall phase, so check + // the eigenpair and normalization rather than a particular phase. + std::vector psi(ld, T(0)); + psi[0] = T(0, 1); + std::vector eval(nband, Real(0)); + const std::vector ethr(nband, 1e-12); + const std::vector prec(n_dim, Real(1)); + const auto h_op = [&h_mat, n_dim](T* in, T* out, int ld_in, int ncol) { + dense_h_multiply(h_mat.data(), n_dim, in, out, ld_in, ncol); + }; + hsolver::DiagoPPCG solver( + /* diag_thr = */ 1e-12, + /* max_iter = */ 0, + /* sbsize = */ 1, + /* rr_step = */ 4); + + solver.diag( + h_op, nullptr, ld, nband, n_dim, psi.data(), eval.data(), ethr, prec.data()); + + Real norm_squared = Real(0); + Real residual_squared = Real(0); + for (int i = 0; i < n_dim; ++i) + { + const T residual = diagonal[i] * psi[i] - eval[0] * psi[i]; + norm_squared += std::norm(psi[i]); + residual_squared += std::norm(residual); + } + const Real residual_norm = std::sqrt(residual_squared); + EXPECT_TRUE(solver.converged()); + EXPECT_FALSE(solver.iteration_limit_reached()); + EXPECT_NEAR(eval[0], diagonal[0], 1e-12); + EXPECT_NEAR(norm_squared, Real(1), 1e-12); + EXPECT_LE(residual_norm, Real(1e-12)); +} + +TEST(DiagoPPCGComplexHermitianTest, BlockSubspaceSmokeNoNaN) +{ + const int n_dim = 2; + const int nband = 2; + const int ld = n_dim; + + std::vector H_mat(n_dim * n_dim, T(0)); + H_mat[0 + 0 * n_dim] = T(2.0, 0.0); + H_mat[1 + 1 * n_dim] = T(3.0, 0.0); + H_mat[0 + 1 * n_dim] = T(0.0, 1.0); + H_mat[1 + 0 * n_dim] = T(0.0, -1.0); + + std::vector psi(ld * nband, T(0)); + psi[0 + 0 * ld] = T(1.0, 0.0); + psi[1 + 1 * ld] = T(1.0, 0.0); + + std::vector prec(n_dim, 2.0); + std::vector ethr(nband, 1e-10); + std::vector eval(nband, 0.0); + + hsolver::DiagoPPCG solver( + /* diag_thr = */ 1e-10, + /* max_iter = */ 8, + /* sbsize = */ 2, + /* rr_step = */ 1); + + auto h_op = [&H_mat, n_dim](T* in, T* out, int ld_in, int ncol) { + dense_h_multiply(H_mat.data(), n_dim, in, out, ld_in, ncol); + }; + + solver.diag(h_op, nullptr, ld, nband, n_dim, psi.data(), eval.data(), ethr, prec.data()); + + const Real delta = std::sqrt(1.25); + for (int i = 0; i < nband; ++i) + { + EXPECT_TRUE(std::isfinite(eval[i])) << "BLOCK_SUBSPACE produced NaN/Inf"; + } + EXPECT_NEAR(eval[0], 2.5 - delta, 1e-8); + EXPECT_NEAR(eval[1], 2.5 + delta, 1e-8); +} + +// ============================================================================= +// Test fixture: degenerate eigenvalues +// H = I + J (identity plus all-ones), 4×4. +// J has eigenvector [1,1,1,1]^T with eigenvalue 4. +// All vectors orthogonal to [1,1,1,1]^T are eigenvectors with eigenvalue 0. +// So H = I + J has: λ₁ = 1 (multiplicity 3), λ₄ = 5. +// ============================================================================= +class DiagoPPCGDegenerateTest : public ::testing::Test +{ + protected: + void SetUp() override + { + n_dim = 4; + nband = 4; + ld = n_dim; + + // H = I + J where J is the all-ones matrix + H_mat.assign(n_dim * n_dim, T(0)); + for (int i = 0; i < n_dim; ++i) + { + for (int j = 0; j < n_dim; ++j) + { + H_mat[i + j * n_dim] = T(1.0, 0); // all-ones J + } + } + for (int i = 0; i < n_dim; ++i) + { + H_mat[i + i * n_dim] += T(1.0, 0); // J → I+J + } + + // Preconditioner: diagonal = 2 + prec.assign(n_dim, 2.0); + + // λ = {1, 1, 1, 5} + exact = {1.0, 1.0, 1.0, 5.0}; + + ethr.assign(nband, 1e-10); + + std::mt19937 rng(456); + std::uniform_real_distribution dist(-1.0, 1.0); + + psi.assign(ld * nband, T(0)); + for (int j = 0; j < nband; ++j) + { + for (int i = 0; i < n_dim; ++i) + { + psi[i + j * ld] = T(dist(rng), 0.0); + } + } + + // Gram-Schmidt (S = I) + for (int j = 0; j < nband; ++j) + { + for (int k = 0; k < j; ++k) + { + T dot = 0; + for (int i = 0; i < n_dim; ++i) + { + dot += std::conj(psi[i + k * ld]) * psi[i + j * ld]; + } + for (int i = 0; i < n_dim; ++i) + { + psi[i + j * ld] -= dot * psi[i + k * ld]; + } + } + Real nrm = 0; + for (int i = 0; i < n_dim; ++i) + { + nrm += std::norm(psi[i + j * ld]); + } + nrm = std::sqrt(nrm); + for (int i = 0; i < n_dim; ++i) + { + psi[i + j * ld] /= nrm; + } + } + } + + int n_dim; + int nband; + int ld; + std::vector H_mat; + std::vector prec; + std::vector exact; + std::vector ethr; + std::vector psi; +}; + +TEST_F(DiagoPPCGDegenerateTest, BlockSubspace) +{ + std::vector psi_run = psi; + std::vector eval(nband, 0.0); + + hsolver::DiagoPPCG solver( + /* diag_thr = */ 1e-12, + /* max_iter = */ 100, + /* sbsize = */ 4, + /* rr_step = */ 4); + + auto h_op = [this](T* in, T* out, int ld_in, int ncol) { + dense_h_multiply(H_mat.data(), n_dim, in, out, ld_in, ncol); + }; + + double avg_iter = solver.diag(h_op, nullptr, ld, nband, n_dim, psi_run.data(), eval.data(), ethr, prec.data()); + + for (int i = 0; i < nband; ++i) + { + EXPECT_NEAR(eval[i], exact[i], 1e-8) << "Degenerate BLOCK: eigenvalue[" << i << "] mismatch"; + } + EXPECT_LE(avg_iter, double(100)) << "Degenerate BLOCK: too many iterations"; +} + +// ============================================================================= +// Test fixture: larger tridiagonal, more bands +// 20×20 tridiagonal Laplacian, nband=5. +// ============================================================================= +class DiagoPPCGLargeTridiagTest : public ::testing::Test +{ + protected: + void SetUp() override + { + n_dim = 20; + nband = 5; + ld = n_dim; + + H_mat.assign(n_dim * n_dim, T(0)); + for (int i = 0; i < n_dim; ++i) + { + H_mat[i + i * n_dim] = T(2.0, 0); + if (i > 0) + { + H_mat[i + (i - 1) * n_dim] = T(-1.0, 0); + } + if (i < n_dim - 1) + { + H_mat[i + (i + 1) * n_dim] = T(-1.0, 0); + } + } + + prec.assign(n_dim, 2.0); + + exact.resize(nband); + for (int k = 0; k < nband; ++k) + { + exact[k] = 2.0 - 2.0 * std::cos(Real(k + 1) * M_PI / Real(n_dim + 1)); + } + + ethr.assign(nband, 1e-10); + + std::mt19937 rng(789); + std::uniform_real_distribution dist(-1.0, 1.0); + + psi.assign(ld * nband, T(0)); + for (int j = 0; j < nband; ++j) + { + for (int i = 0; i < n_dim; ++i) + { + psi[i + j * ld] = T(dist(rng), 0.0); + } + } + + for (int j = 0; j < nband; ++j) + { + for (int k = 0; k < j; ++k) + { + T dot = 0; + for (int i = 0; i < n_dim; ++i) + { + dot += std::conj(psi[i + k * ld]) * psi[i + j * ld]; + } + for (int i = 0; i < n_dim; ++i) + { + psi[i + j * ld] -= dot * psi[i + k * ld]; + } + } + Real nrm = 0; + for (int i = 0; i < n_dim; ++i) + { + nrm += std::norm(psi[i + j * ld]); + } + nrm = std::sqrt(nrm); + for (int i = 0; i < n_dim; ++i) + { + psi[i + j * ld] /= nrm; + } + } + } + + int n_dim; + int nband; + int ld; + std::vector H_mat; + std::vector prec; + std::vector exact; + std::vector ethr; + std::vector psi; +}; + +TEST_F(DiagoPPCGLargeTridiagTest, BlockSubspace) +{ + std::vector psi_run = psi; + std::vector eval(nband, 0.0); + + hsolver::DiagoPPCG solver( + /* diag_thr = */ 1e-12, + /* max_iter = */ 150, + /* sbsize = */ 5, + /* rr_step = */ 5); + + auto h_op = [this](T* in, T* out, int ld_in, int ncol) { + dense_h_multiply(H_mat.data(), n_dim, in, out, ld_in, ncol); + }; + + double avg_iter = solver.diag(h_op, nullptr, ld, nband, n_dim, psi_run.data(), eval.data(), ethr, prec.data()); + + for (int i = 0; i < nband; ++i) + { + EXPECT_NEAR(eval[i], exact[i], 1e-8) << "Large Tridiag BLOCK: eigenvalue[" << i << "] mismatch"; + } + EXPECT_LE(avg_iter, double(150)) << "Large Tridiag BLOCK: too many iterations"; +} + +// ============================================================================= +// Test fixture: dense matrix with known eigenvalues +// H = Q^T * D * Q where Q is a known orthogonal matrix (a Givens rotation +// repeated on different index pairs) and D is diagonal. +// For an 8×8 case: D = diag(1, 2, 3, 4, 5, 6, 7, 8), then apply several +// Givens rotations to mix all rows/cols. The exact eigenvalues remain 1..8. +// +// This addresses the "full/dense matrix" test that was originally missing. +// ============================================================================= +class DiagoPPCGDenseTest : public ::testing::Test +{ + protected: + void SetUp() override + { + n_dim = 8; + nband = 4; + ld = n_dim; + + // Start with diagonal matrix + std::vector dense(n_dim * n_dim, Real(0)); + for (int i = 0; i < n_dim; ++i) + { + dense[i + i * n_dim] = Real(i + 1); + } + + // Apply several Givens rotations to make it dense while preserving + // eigenvalues. Each rotation: A' = G(i,j,θ)^T * A * G(i,j,θ) + auto apply_givens = [&](int p, int q, Real theta) { + Real c = std::cos(theta); + Real s = std::sin(theta); + // Apply to columns + for (int i = 0; i < n_dim; ++i) + { + Real aip = dense[i + p * n_dim]; + Real aiq = dense[i + q * n_dim]; + dense[i + p * n_dim] = c * aip + s * aiq; + dense[i + q * n_dim] = -s * aip + c * aiq; + } + // Apply to rows + for (int j = 0; j < n_dim; ++j) + { + Real apj = dense[p + j * n_dim]; + Real aqj = dense[q + j * n_dim]; + dense[p + j * n_dim] = c * apj + s * aqj; + dense[q + j * n_dim] = -s * apj + c * aqj; + } + }; + + // Several rotations with different angles to create a genuinely + // dense matrix (all off-diagonals become non-zero) + std::mt19937 rng_dense(111); + std::uniform_real_distribution angle_dist(Real(0.2), Real(1.3)); + for (int k = 0; k < 20; ++k) + { + int p = k % (n_dim - 1); + int q = p + 1 + (k / (n_dim - 1)) % (n_dim - 1 - p); + if (q >= n_dim) + { + q = n_dim - 1; + } + if (p == q) + { + continue; + } + apply_givens(p, q, angle_dist(rng_dense)); + } + + // Copy to complex H_mat + H_mat.assign(n_dim * n_dim, T(0)); + for (int i = 0; i < n_dim * n_dim; ++i) + { + H_mat[i] = T(dense[i], 0); + } + + // Preconditioner: use diagonal of the rotated H + prec.resize(n_dim); + for (int i = 0; i < n_dim; ++i) + { + prec[i] = std::real(H_mat[i + i * n_dim]); + } + + // Lowest 4 eigenvalues: 1, 2, 3, 4 + exact = {1.0, 2.0, 3.0, 4.0}; + + ethr.assign(nband, 1e-10); + + std::mt19937 rng_psi(222); + std::uniform_real_distribution dist(-1.0, 1.0); + + psi.assign(ld * nband, T(0)); + for (int j = 0; j < nband; ++j) + { + for (int i = 0; i < n_dim; ++i) + { + psi[i + j * ld] = T(dist(rng_psi), 0.0); + } + } + + for (int j = 0; j < nband; ++j) + { + for (int k = 0; k < j; ++k) + { + T dot = 0; + for (int i = 0; i < n_dim; ++i) + { + dot += std::conj(psi[i + k * ld]) * psi[i + j * ld]; + } + for (int i = 0; i < n_dim; ++i) + { + psi[i + j * ld] -= dot * psi[i + k * ld]; + } + } + Real nrm = 0; + for (int i = 0; i < n_dim; ++i) + { + nrm += std::norm(psi[i + j * ld]); + } + nrm = std::sqrt(nrm); + for (int i = 0; i < n_dim; ++i) + { + psi[i + j * ld] /= nrm; + } + } + } + + int n_dim; + int nband; + int ld; + std::vector H_mat; + std::vector prec; + std::vector exact; + std::vector ethr; + std::vector psi; +}; + +TEST_F(DiagoPPCGDenseTest, BlockSubspace) +{ + std::vector psi_run = psi; + std::vector eval(nband, 0.0); + + hsolver::DiagoPPCG solver( + /* diag_thr = */ 1e-12, + /* max_iter = */ 200, + /* sbsize = */ 4, + /* rr_step = */ 4); + + auto h_op = [this](T* in, T* out, int ld_in, int ncol) { + dense_h_multiply(H_mat.data(), n_dim, in, out, ld_in, ncol); + }; + + double avg_iter = solver.diag(h_op, nullptr, ld, nband, n_dim, psi_run.data(), eval.data(), ethr, prec.data()); + + for (int i = 0; i < nband; ++i) + { + EXPECT_NEAR(eval[i], exact[i], 1e-8) << "Dense BLOCK: eigenvalue[" << i << "] mismatch"; + } + EXPECT_LE(avg_iter, double(200)) << "Dense BLOCK: too many iterations"; +} + +// ============================================================================= +// Helper: compute Hψ for eigenvector residual check +// ============================================================================= +static void compute_residual(const T* H_data, int n_dim, const T* psi, const Real eval, int ld, T* residual) +{ + // residual = H*psi - eval*psi + dense_h_multiply(H_data, n_dim, psi, residual, ld, 1); + for (int i = 0; i < n_dim; ++i) + { + residual[i] -= eval * psi[i]; + } +} + +static Real column_norm(const T* x, int n_dim) +{ + Real nrm = 0; + for (int i = 0; i < n_dim; ++i) + { + nrm += std::norm(x[i]); + } + return std::sqrt(nrm); +} + +// ============================================================================= +// Test fixture: non-trivial S matrix (diagonal overlap, S ≠ I) +// H = tridiag 6×6 Laplacian, S = diag(1.1, 1.0, 0.9, 1.0, 1.1, 1.0) +// Tests that the solver correctly handles a non-identity overlap matrix. +// ============================================================================= +class DiagoPPCGWithSTest : public ::testing::Test +{ + protected: + void SetUp() override + { + n_dim = 6; + nband = 3; + ld = n_dim + 2; // exercise custom S with padded leading dimension + + // Tridiagonal H + H_mat.assign(n_dim * n_dim, T(0)); + for (int i = 0; i < n_dim; ++i) + { + H_mat[i + i * n_dim] = T(2.0, 0); + if (i > 0) + { + H_mat[i + (i - 1) * n_dim] = T(-1.0, 0); + } + if (i < n_dim - 1) + { + H_mat[i + (i + 1) * n_dim] = T(-1.0, 0); + } + } + + // S = diag(1.1, 1.0, 0.9, 1.0, 1.1, 1.0) + s_diag = {1.1, 1.0, 0.9, 1.0, 1.1, 1.0}; + S_mat.assign(n_dim * n_dim, T(0)); + for (int i = 0; i < n_dim; ++i) + { + S_mat[i + i * n_dim] = T(s_diag[i], 0); + } + + prec.assign(n_dim, 2.0); + + ethr.assign(nband, 1e-8); + + std::mt19937 rng(333); + std::uniform_real_distribution dist(-1.0, 1.0); + + psi.assign(ld * nband, T(0)); + for (int j = 0; j < nband; ++j) + { + for (int i = 0; i < n_dim; ++i) + { + psi[i + j * ld] = T(dist(rng), 0.0); + } + } + + // S-orthonormalize initial guess + for (int j = 0; j < nband; ++j) + { + for (int k = 0; k < j; ++k) + { + T dot = 0; + for (int i = 0; i < n_dim; ++i) + { + dot += std::conj(psi[i + k * ld]) * T(s_diag[i], 0) * psi[i + j * ld]; + } + for (int i = 0; i < n_dim; ++i) + { + psi[i + j * ld] -= dot * psi[i + k * ld]; + } + } + Real nrm = 0; + for (int i = 0; i < n_dim; ++i) + { + nrm += s_diag[i] * std::norm(psi[i + j * ld]); + } + nrm = std::sqrt(nrm); + for (int i = 0; i < n_dim; ++i) + { + psi[i + j * ld] /= nrm; + } + } + } + + int n_dim; + int nband; + int ld; + std::vector H_mat; + std::vector S_mat; + std::vector s_diag; + std::vector prec; + std::vector ethr; + std::vector psi; +}; + +TEST_F(DiagoPPCGWithSTest, BlockSubspace) +{ + std::vector psi_run = psi; + std::vector eval(nband, 0.0); + + // S-apply function: S * psi = diag(s_diag) * psi (element-wise) + auto spsi_func = [this](T* in, T* out, int ld_in, int ncol) { + for (int j = 0; j < ncol; ++j) + { + for (int i = 0; i < n_dim; ++i) + { + out[i + j * ld_in] = T(s_diag[i], 0) * in[i + j * ld_in]; + } + } + }; + + hsolver::DiagoPPCG solver( + /* diag_thr = */ 1e-10, + /* max_iter = */ 100, + /* sbsize = */ 3, + /* rr_step = */ 3); + + auto h_op = [this](T* in, T* out, int ld_in, int ncol) { + dense_h_multiply(H_mat.data(), n_dim, in, out, ld_in, ncol); + }; + + double avg_iter = solver.diag(h_op, spsi_func, ld, nband, n_dim, psi_run.data(), eval.data(), ethr, prec.data()); + + // Eigenvalue check: skip absolute comparison (exact values not + // analytically known for non-trivial S). Instead verify via residual. + // Just check eigenvalues are reasonable (not NaN, not negative for + // this positive-definite problem). + for (int i = 0; i < nband; ++i) + { + EXPECT_GT(eval[i], 0.0) << "WithS BLOCK: eigenvalue[" << i << "] should be positive"; + EXPECT_LT(eval[i], 10.0) << "WithS BLOCK: eigenvalue[" << i << "] unreasonably large"; + } + + // Residual check: ||Hψ_i - ε_i S ψ_i|| / |ε_i| < ethr + std::vector hpsi(n_dim); + std::vector spsi(n_dim); + std::vector res(n_dim); + for (int i = 0; i < nband; ++i) + { + dense_h_multiply(H_mat.data(), n_dim, psi_run.data() + i * ld, hpsi.data(), n_dim, 1); + spsi_func(psi_run.data() + i * ld, spsi.data(), n_dim, 1); + for (int j = 0; j < n_dim; ++j) + { + res[j] = hpsi[j] - T(eval[i], 0) * spsi[j]; + } + Real res_nrm = column_norm(res.data(), n_dim); + EXPECT_LE(res_nrm, std::max(1e-6, 1e2 * ethr[i])) + << "WithS BLOCK: residual[" << i << "] too large, r=" << res_nrm; + } + + EXPECT_LE(avg_iter, double(100)) << "WithS BLOCK: too many iterations"; +} + +// ============================================================================= +// Test fixture: single-band (nband = 1) +// Minimal test — extract only the lowest eigenvalue of a 5×5 tridiagonal +// Laplacian. This exercises the degenerate code paths for a single band. +// ============================================================================= +class DiagoPPCGSingleBandTest : public ::testing::Test +{ + protected: + void SetUp() override + { + n_dim = 5; + nband = 1; + ld = n_dim; + + H_mat.assign(n_dim * n_dim, T(0)); + for (int i = 0; i < n_dim; ++i) + { + H_mat[i + i * n_dim] = T(2.0, 0); + if (i > 0) + { + H_mat[i + (i - 1) * n_dim] = T(-1.0, 0); + } + if (i < n_dim - 1) + { + H_mat[i + (i + 1) * n_dim] = T(-1.0, 0); + } + } + + prec.assign(n_dim, 2.0); + + // Lowest eigenvalue of 5×5 tridiagonal Laplacian + exact = {2.0 - 2.0 * std::cos(M_PI / 6.0)}; + + ethr.assign(nband, 1e-10); + + std::mt19937 rng(42); + std::uniform_real_distribution dist(-1.0, 1.0); + + psi.assign(ld * nband, T(0)); + for (int i = 0; i < n_dim; ++i) + { + psi[i] = T(dist(rng), 0.0); + } + + Real nrm = 0; + for (int i = 0; i < n_dim; ++i) + { + nrm += std::norm(psi[i]); + } + nrm = std::sqrt(nrm); + for (int i = 0; i < n_dim; ++i) + { + psi[i] /= nrm; + } + } + + int n_dim; + int nband; + int ld; + std::vector H_mat; + std::vector prec; + std::vector exact; + std::vector ethr; + std::vector psi; +}; + +TEST_F(DiagoPPCGSingleBandTest, BlockSubspace) +{ + std::vector psi_run = psi; + std::vector eval(nband, 0.0); + + hsolver::DiagoPPCG solver( + /* diag_thr = */ 1e-12, + /* max_iter = */ 50, + /* sbsize = */ 1, + /* rr_step = */ 1); + + auto h_op = [this](T* in, T* out, int ld_in, int ncol) { + dense_h_multiply(H_mat.data(), n_dim, in, out, ld_in, ncol); + }; + + double avg_iter = solver.diag(h_op, nullptr, ld, nband, n_dim, psi_run.data(), eval.data(), ethr, prec.data()); + + EXPECT_NEAR(eval[0], exact[0], 1e-8) << "SingleBand BLOCK: eigenvalue mismatch"; + EXPECT_LE(avg_iter, double(50)) << "SingleBand BLOCK: too many iterations"; +} + +// ============================================================================= +// Test fixture: eigenvector quality — verify Hψ ≈ εψ and ψ^H ψ = I +// Uses the 10×10 tridiagonal Laplacian. After convergence, check: +// 1. ||Hψ_i - ε_i ψ_i|| < tol for each band +// 2. |ψ_i^H ψ_j - δ_ij| < tol for all i,j +// ============================================================================= +class DiagoPPCGEigenvectorTest : public ::testing::Test +{ + protected: + void SetUp() override + { + n_dim = 10; + nband = 3; + ld = n_dim; + + H_mat.assign(n_dim * n_dim, T(0)); + for (int i = 0; i < n_dim; ++i) + { + H_mat[i + i * n_dim] = T(2.0, 0); + if (i > 0) + { + H_mat[i + (i - 1) * n_dim] = T(-1.0, 0); + } + if (i < n_dim - 1) + { + H_mat[i + (i + 1) * n_dim] = T(-1.0, 0); + } + } + + prec.assign(n_dim, 2.0); + + exact.resize(nband); + for (int k = 0; k < nband; ++k) + { + exact[k] = 2.0 - 2.0 * std::cos(Real(k + 1) * M_PI / Real(n_dim + 1)); + } + + ethr.assign(nband, 1e-8); + + std::mt19937 rng(888); + std::uniform_real_distribution dist(-1.0, 1.0); + + psi.assign(ld * nband, T(0)); + for (int j = 0; j < nband; ++j) + { + for (int i = 0; i < n_dim; ++i) + { + psi[i + j * ld] = T(dist(rng), 0.0); + } + } + + for (int j = 0; j < nband; ++j) + { + for (int k = 0; k < j; ++k) + { + T dot = 0; + for (int i = 0; i < n_dim; ++i) + { + dot += std::conj(psi[i + k * ld]) * psi[i + j * ld]; + } + for (int i = 0; i < n_dim; ++i) + { + psi[i + j * ld] -= dot * psi[i + k * ld]; + } + } + Real nrm = 0; + for (int i = 0; i < n_dim; ++i) + { + nrm += std::norm(psi[i + j * ld]); + } + nrm = std::sqrt(nrm); + for (int i = 0; i < n_dim; ++i) + { + psi[i + j * ld] /= nrm; + } + } + } + + int n_dim; + int nband; + int ld; + std::vector H_mat; + std::vector prec; + std::vector exact; + std::vector ethr; + std::vector psi; +}; + +TEST_F(DiagoPPCGEigenvectorTest, BlockSubspace) +{ + std::vector psi_run = psi; + std::vector eval(nband, 0.0); + + hsolver::DiagoPPCG solver( + /* diag_thr = */ 1e-12, + /* max_iter = */ 100, + /* sbsize = */ 3, + /* rr_step = */ 3); + + auto h_op = [this](T* in, T* out, int ld_in, int ncol) { + dense_h_multiply(H_mat.data(), n_dim, in, out, ld_in, ncol); + }; + + double avg_iter = solver.diag(h_op, nullptr, ld, nband, n_dim, psi_run.data(), eval.data(), ethr, prec.data()); + + // --- Eigenvalue check --- + for (int i = 0; i < nband; ++i) + { + EXPECT_NEAR(eval[i], exact[i], 1e-8) << "Eigenvec BLOCK: eigenvalue[" << i << "] mismatch"; + } + + // Independently check the returned eigenvectors using a fresh H application. + std::vector hpsi(n_dim); + std::vector res(n_dim); + for (int i = 0; i < nband; ++i) + { + dense_h_multiply(H_mat.data(), n_dim, psi_run.data() + i * ld, hpsi.data(), n_dim, 1); + for (int j = 0; j < n_dim; ++j) + { + res[j] = hpsi[j] - eval[i] * psi_run[j + i * ld]; + } + Real res_nrm = column_norm(res.data(), n_dim); + EXPECT_LT(res_nrm, 1e-4) << "Eigenvec BLOCK: residual[" << i << "] too large: " << res_nrm; + } + + // --- Orthogonality check: |ψ_i^H ψ_j - δ_ij| < 1e-8 --- + for (int i = 0; i < nband; ++i) + { + for (int j = 0; j < nband; ++j) + { + T dot = 0; + for (int k = 0; k < n_dim; ++k) + { + dot += std::conj(psi_run[k + i * ld]) * psi_run[k + j * ld]; + } + if (i == j) + { + EXPECT_NEAR(std::abs(dot), 1.0, 1e-8) + << "Eigenvec BLOCK: ψ[" << i << "] not normalized, |dot|=" << std::abs(dot); + } + else + { + EXPECT_LT(std::abs(dot), 1e-8) + << "Eigenvec BLOCK: ψ[" << i << "] not orthogonal to ψ[" << j << "], |dot|=" << std::abs(dot); + } + } + } + + EXPECT_LE(avg_iter, double(100)) << "Eigenvec BLOCK: too many iterations"; +} + +// ============================================================================= +// Test fixture: all eigenvalues of a small matrix (nband = n_dim) +// 3×3 tridiagonal Laplacian, compute all 3 eigenvalues. +// Exercises the degenerate case where every band is requested. +// ============================================================================= +class DiagoPPCGAllBandsTest : public ::testing::Test +{ + protected: + void SetUp() override + { + n_dim = 3; + nband = 3; + ld = n_dim; + + H_mat.assign(n_dim * n_dim, T(0)); + for (int i = 0; i < n_dim; ++i) + { + H_mat[i + i * n_dim] = T(2.0, 0); + if (i > 0) + { + H_mat[i + (i - 1) * n_dim] = T(-1.0, 0); + } + if (i < n_dim - 1) + { + H_mat[i + (i + 1) * n_dim] = T(-1.0, 0); + } + } + + prec.assign(n_dim, 2.0); + + exact.resize(nband); + for (int k = 0; k < nband; ++k) + { + exact[k] = 2.0 - 2.0 * std::cos(Real(k + 1) * M_PI / Real(n_dim + 1)); + } + + ethr.assign(nband, 1e-10); + + std::mt19937 rng(101); + std::uniform_real_distribution dist(-1.0, 1.0); + + psi.assign(ld * nband, T(0)); + for (int j = 0; j < nband; ++j) + { + for (int i = 0; i < n_dim; ++i) + { + psi[i + j * ld] = T(dist(rng), 0.0); + } + } + + for (int j = 0; j < nband; ++j) + { + for (int k = 0; k < j; ++k) + { + T dot = 0; + for (int i = 0; i < n_dim; ++i) + { + dot += std::conj(psi[i + k * ld]) * psi[i + j * ld]; + } + for (int i = 0; i < n_dim; ++i) + { + psi[i + j * ld] -= dot * psi[i + k * ld]; + } + } + Real nrm = 0; + for (int i = 0; i < n_dim; ++i) + { + nrm += std::norm(psi[i + j * ld]); + } + nrm = std::sqrt(nrm); + for (int i = 0; i < n_dim; ++i) + { + psi[i + j * ld] /= nrm; + } + } + } + + int n_dim; + int nband; + int ld; + std::vector H_mat; + std::vector prec; + std::vector exact; + std::vector ethr; + std::vector psi; +}; + +TEST_F(DiagoPPCGAllBandsTest, BlockSubspace) +{ + std::vector psi_run = psi; + std::vector eval(nband, 0.0); + + hsolver::DiagoPPCG solver( + /* diag_thr = */ 1e-12, + /* max_iter = */ 100, + /* sbsize = */ 3, + /* rr_step = */ 3); + + auto h_op = [this](T* in, T* out, int ld_in, int ncol) { + dense_h_multiply(H_mat.data(), n_dim, in, out, ld_in, ncol); + }; + + double avg_iter = solver.diag(h_op, nullptr, ld, nband, n_dim, psi_run.data(), eval.data(), ethr, prec.data()); + + for (int i = 0; i < nband; ++i) + { + EXPECT_NEAR(eval[i], exact[i], 1e-8) << "AllBands BLOCK: eigenvalue[" << i << "] mismatch"; + } + EXPECT_LE(avg_iter, double(100)) << "AllBands BLOCK: too many iterations"; +} + +// ============================================================================= +// Test fixture: medium-sized tridiagonal (15×15, nband=4) +// Bridges the gap between the 10×10 and 20×20 tests. +// ============================================================================= +class DiagoPPCGMediumTridiagTest : public ::testing::Test +{ + protected: + void SetUp() override + { + n_dim = 15; + nband = 4; + ld = n_dim; + + H_mat.assign(n_dim * n_dim, T(0)); + for (int i = 0; i < n_dim; ++i) + { + H_mat[i + i * n_dim] = T(2.0, 0); + if (i > 0) + { + H_mat[i + (i - 1) * n_dim] = T(-1.0, 0); + } + if (i < n_dim - 1) + { + H_mat[i + (i + 1) * n_dim] = T(-1.0, 0); + } + } + + prec.assign(n_dim, 2.0); + + exact.resize(nband); + for (int k = 0; k < nband; ++k) + { + exact[k] = 2.0 - 2.0 * std::cos(Real(k + 1) * M_PI / Real(n_dim + 1)); + } + + ethr.assign(nband, 1e-10); + + std::mt19937 rng(202); + std::uniform_real_distribution dist(-1.0, 1.0); + + psi.assign(ld * nband, T(0)); + for (int j = 0; j < nband; ++j) + { + for (int i = 0; i < n_dim; ++i) + { + psi[i + j * ld] = T(dist(rng), 0.0); + } + } + + for (int j = 0; j < nband; ++j) + { + for (int k = 0; k < j; ++k) + { + T dot = 0; + for (int i = 0; i < n_dim; ++i) + { + dot += std::conj(psi[i + k * ld]) * psi[i + j * ld]; + } + for (int i = 0; i < n_dim; ++i) + { + psi[i + j * ld] -= dot * psi[i + k * ld]; + } + } + Real nrm = 0; + for (int i = 0; i < n_dim; ++i) + { + nrm += std::norm(psi[i + j * ld]); + } + nrm = std::sqrt(nrm); + for (int i = 0; i < n_dim; ++i) + { + psi[i + j * ld] /= nrm; + } + } + } + + int n_dim; + int nband; + int ld; + std::vector H_mat; + std::vector prec; + std::vector exact; + std::vector ethr; + std::vector psi; +}; + +TEST_F(DiagoPPCGMediumTridiagTest, BlockSubspace) +{ + std::vector psi_run = psi; + std::vector eval(nband, 0.0); + + hsolver::DiagoPPCG solver( + /* diag_thr = */ 1e-12, + /* max_iter = */ 120, + /* sbsize = */ 4, + /* rr_step = */ 4); + + auto h_op = [this](T* in, T* out, int ld_in, int ncol) { + dense_h_multiply(H_mat.data(), n_dim, in, out, ld_in, ncol); + }; + + double avg_iter = solver.diag(h_op, nullptr, ld, nband, n_dim, psi_run.data(), eval.data(), ethr, prec.data()); + + for (int i = 0; i < nband; ++i) + { + EXPECT_NEAR(eval[i], exact[i], 1e-8) << "Medium Tridiag BLOCK: eigenvalue[" << i << "] mismatch"; + } + EXPECT_LE(avg_iter, double(120)) << "Medium Tridiag BLOCK: too many iterations"; +} + +// ============================================================================= +// Test fixture: pentadiagonal Toeplitz (discrete biharmonic operator) +// H[i,i]=6, H[i,i±1]=-4, H[i,i±2]=1. Eigenvalues: +// λ_k = 16·sin⁴(k·π / (2·(n+1))), k = 1,...,n +// Wider bandwidth (5 vs 3) tests the solver with more off-diagonal coupling. +// ============================================================================= +class DiagoPPCGPentaTest : public ::testing::Test +{ + protected: + void SetUp() override + { + n_dim = 8; + nband = 4; + ld = n_dim; + + // H = T² where T is the tridiagonal Laplacian (2 on diag, -1 on off-diag). + // The corners of T² have diag=5 (not 6) since (T²)[0,0] = 2² + (-1)² = 5. + // Interior: (T²)[i,i] = (-1)² + 2² + (-1)² = 6. + H_mat.assign(n_dim * n_dim, T(0)); + for (int i = 0; i < n_dim; ++i) + { + H_mat[i + i * n_dim] = T((i == 0 || i == n_dim - 1) ? 5.0 : 6.0, 0); + if (i >= 1) + { + H_mat[i + (i - 1) * n_dim] = T(-4.0, 0); + } + if (i < n_dim - 1) + { + H_mat[i + (i + 1) * n_dim] = T(-4.0, 0); + } + if (i >= 2) + { + H_mat[i + (i - 2) * n_dim] = T(1.0, 0); + } + if (i < n_dim - 2) + { + H_mat[i + (i + 2) * n_dim] = T(1.0, 0); + } + } + + prec.assign(n_dim, 6.0); + prec[0] = 5.0; + prec[n_dim - 1] = 5.0; + + exact.resize(nband); + for (int k = 0; k < nband; ++k) + { + Real theta = Real(k + 1) * M_PI / Real(2 * (n_dim + 1)); + Real s = std::sin(theta); + exact[k] = Real(16) * s * s * s * s; + } + + ethr.assign(nband, 1e-10); + + std::mt19937 rng(505); + std::uniform_real_distribution dist(-1.0, 1.0); + + psi.assign(ld * nband, T(0)); + for (int j = 0; j < nband; ++j) + { + for (int i = 0; i < n_dim; ++i) + { + psi[i + j * ld] = T(dist(rng), 0.0); + } + } + + for (int j = 0; j < nband; ++j) + { + for (int k = 0; k < j; ++k) + { + T dot = 0; + for (int i = 0; i < n_dim; ++i) + { + dot += std::conj(psi[i + k * ld]) * psi[i + j * ld]; + } + for (int i = 0; i < n_dim; ++i) + { + psi[i + j * ld] -= dot * psi[i + k * ld]; + } + } + Real nrm = 0; + for (int i = 0; i < n_dim; ++i) + { + nrm += std::norm(psi[i + j * ld]); + } + nrm = std::sqrt(nrm); + for (int i = 0; i < n_dim; ++i) + { + psi[i + j * ld] /= nrm; + } + } + } + + int n_dim; + int nband; + int ld; + std::vector H_mat; + std::vector prec; + std::vector exact; + std::vector ethr; + std::vector psi; +}; + +TEST_F(DiagoPPCGPentaTest, BlockSubspace) +{ + std::vector psi_run = psi; + std::vector eval(nband, 0.0); + + hsolver::DiagoPPCG solver( + /* diag_thr = */ 1e-12, + /* max_iter = */ 150, + /* sbsize = */ 4, + /* rr_step = */ 4); + + auto h_op = [this](T* in, T* out, int ld_in, int ncol) { + dense_h_multiply(H_mat.data(), n_dim, in, out, ld_in, ncol); + }; + + double avg_iter = solver.diag(h_op, nullptr, ld, nband, n_dim, psi_run.data(), eval.data(), ethr, prec.data()); + + for (int i = 0; i < nband; ++i) + { + EXPECT_NEAR(eval[i], exact[i], 1e-8) << "Penta BLOCK: eigenvalue[" << i << "] mismatch"; + } + EXPECT_LE(avg_iter, double(150)) << "Penta BLOCK: too many iterations"; +} + +// ============================================================================= +// Test fixture: gapped spectrum +// H = diag(0.1, 0.5, 5.0, 6.0, 10.0), nband=3. +// Large gaps between eigenvalue groups test the solver's band separation. +// ============================================================================= +class DiagoPCGGappedSpectrumTest : public ::testing::Test +{ + protected: + void SetUp() override + { + n_dim = 5; + nband = 3; + ld = n_dim; + + H_mat.assign(n_dim * n_dim, T(0)); + H_mat[0 + 0 * n_dim] = T(0.1, 0); + H_mat[1 + 1 * n_dim] = T(0.5, 0); + H_mat[2 + 2 * n_dim] = T(5.0, 0); + H_mat[3 + 3 * n_dim] = T(6.0, 0); + H_mat[4 + 4 * n_dim] = T(10.0, 0); + + prec.assign(n_dim, 1.0); + + exact = {0.1, 0.5, 5.0}; + + ethr.assign(nband, 1e-10); + + std::mt19937 rng(606); + std::uniform_real_distribution dist(-1.0, 1.0); + + psi.assign(ld * nband, T(0)); + for (int j = 0; j < nband; ++j) + { + for (int i = 0; i < n_dim; ++i) + { + psi[i + j * ld] = T(dist(rng), 0.0); + } + } + + for (int j = 0; j < nband; ++j) + { + for (int k = 0; k < j; ++k) + { + T dot = 0; + for (int i = 0; i < n_dim; ++i) + { + dot += std::conj(psi[i + k * ld]) * psi[i + j * ld]; + } + for (int i = 0; i < n_dim; ++i) + { + psi[i + j * ld] -= dot * psi[i + k * ld]; + } + } + Real nrm = 0; + for (int i = 0; i < n_dim; ++i) + { + nrm += std::norm(psi[i + j * ld]); + } + nrm = std::sqrt(nrm); + for (int i = 0; i < n_dim; ++i) + { + psi[i + j * ld] /= nrm; + } + } + } + + int n_dim; + int nband; + int ld; + std::vector H_mat; + std::vector prec; + std::vector exact; + std::vector ethr; + std::vector psi; +}; + +TEST_F(DiagoPCGGappedSpectrumTest, BlockSubspace) +{ + std::vector psi_run = psi; + std::vector eval(nband, 0.0); + + hsolver::DiagoPPCG solver( + /* diag_thr = */ 1e-12, + /* max_iter = */ 100, + /* sbsize = */ 3, + /* rr_step = */ 3); + + auto h_op = [this](T* in, T* out, int ld_in, int ncol) { + dense_h_multiply(H_mat.data(), n_dim, in, out, ld_in, ncol); + }; + + double avg_iter = solver.diag(h_op, nullptr, ld, nband, n_dim, psi_run.data(), eval.data(), ethr, prec.data()); + + for (int i = 0; i < nband; ++i) + { + EXPECT_NEAR(eval[i], exact[i], 1e-8) << "Gapped BLOCK: eigenvalue[" << i << "] mismatch"; + } + EXPECT_LE(avg_iter, double(100)) << "Gapped BLOCK: too many iterations"; +} + +// ============================================================================= +// Test fixture: preconditioner stress test +// Uses the 10×10 tridiagonal Laplacian but with a suboptimal preconditioner: +// prec[i] = 1.0 instead of 2.0 (the exact diagonal). The solver should still +// converge, just more slowly. Tests robustness against a bad preconditioner. +// ============================================================================= +class DiagoPPCGBadPrecTest : public ::testing::Test +{ + protected: + void SetUp() override + { + n_dim = 10; + nband = 3; + ld = n_dim; + + H_mat.assign(n_dim * n_dim, T(0)); + for (int i = 0; i < n_dim; ++i) + { + H_mat[i + i * n_dim] = T(2.0, 0); + if (i > 0) + { + H_mat[i + (i - 1) * n_dim] = T(-1.0, 0); + } + if (i < n_dim - 1) + { + H_mat[i + (i + 1) * n_dim] = T(-1.0, 0); + } + } + + // Bad preconditioner: use 1.0 instead of 2.0 + prec.assign(n_dim, 1.0); + + exact.resize(nband); + for (int k = 0; k < nband; ++k) + { + exact[k] = 2.0 - 2.0 * std::cos(Real(k + 1) * M_PI / Real(n_dim + 1)); + } + + ethr.assign(nband, 1e-10); + + std::mt19937 rng(707); + std::uniform_real_distribution dist(-1.0, 1.0); + + psi.assign(ld * nband, T(0)); + for (int j = 0; j < nband; ++j) + { + for (int i = 0; i < n_dim; ++i) + { + psi[i + j * ld] = T(dist(rng), 0.0); + } + } + + for (int j = 0; j < nband; ++j) + { + for (int k = 0; k < j; ++k) + { + T dot = 0; + for (int i = 0; i < n_dim; ++i) + { + dot += std::conj(psi[i + k * ld]) * psi[i + j * ld]; + } + for (int i = 0; i < n_dim; ++i) + { + psi[i + j * ld] -= dot * psi[i + k * ld]; + } + } + Real nrm = 0; + for (int i = 0; i < n_dim; ++i) + { + nrm += std::norm(psi[i + j * ld]); + } + nrm = std::sqrt(nrm); + for (int i = 0; i < n_dim; ++i) + { + psi[i + j * ld] /= nrm; + } + } + } + + int n_dim; + int nband; + int ld; + std::vector H_mat; + std::vector prec; + std::vector exact; + std::vector ethr; + std::vector psi; +}; + +TEST_F(DiagoPPCGBadPrecTest, BlockSubspace) +{ + std::vector psi_run = psi; + std::vector eval(nband, 0.0); + + hsolver::DiagoPPCG solver( + /* diag_thr = */ 1e-12, + /* max_iter = */ 200, // more iterations due to bad preconditioner + /* sbsize = */ 4, + /* rr_step = */ 4); + + auto h_op = [this](T* in, T* out, int ld_in, int ncol) { + dense_h_multiply(H_mat.data(), n_dim, in, out, ld_in, ncol); + }; + + double avg_iter = solver.diag(h_op, nullptr, ld, nband, n_dim, psi_run.data(), eval.data(), ethr, prec.data()); + + for (int i = 0; i < nband; ++i) + { + EXPECT_NEAR(eval[i], exact[i], 1e-8) << "BadPrec BLOCK: eigenvalue[" << i << "] mismatch"; + } + EXPECT_LE(avg_iter, double(200)) << "BadPrec BLOCK: too many iterations"; +} + +// ============================================================================= +// Test fixture: n_dim = 1, nband = 1 — absolute minimum +// H is a 1×1 matrix [5.0], eigenvalue = 5.0 +// ============================================================================= +class DiagoPPCG1x1Test : public ::testing::Test +{ + protected: + void SetUp() override + { + n_dim = 1; + nband = 1; + ld = n_dim; + H_mat = {T(5.0, 0)}; + prec = {5.0}; + exact = {5.0}; + ethr.assign(nband, 1e-10); + psi = {T(1.0, 0)}; // already normalized + } + int n_dim; + int nband; + int ld; + std::vector H_mat; + std::vector prec; + std::vector exact; + std::vector ethr; + std::vector psi; +}; + +TEST_F(DiagoPPCG1x1Test, BlockSubspace) +{ + std::vector psi_run = psi; + std::vector eval(nband, 0.0); + hsolver::DiagoPPCG solver(1e-12, 10, 1, 1 + ); + auto h_op = [this](T* in, T* out, int ldi, int nc) { dense_h_multiply(H_mat.data(), n_dim, in, out, ldi, nc); }; + double avg_iter = solver.diag(h_op, nullptr, ld, nband, n_dim, psi_run.data(), eval.data(), ethr, prec.data()); + EXPECT_NEAR(eval[0], exact[0], 1e-8) << "1x1 BLOCK: mismatch"; + EXPECT_LE(avg_iter, 10.0) << "1x1 BLOCK: too many iterations"; +} + +// ============================================================================= +// Test fixture: scaled tridiagonal (eigenvalues × 100) +// H = 100 × tridiag(2, -1, -1). Tests convergence with large eigenvalues. +// ============================================================================= +class DiagoPPCGScaledTest : public ::testing::Test +{ + protected: + void SetUp() override + { + n_dim = 8; + nband = 3; + ld = n_dim; + H_mat.assign(n_dim * n_dim, T(0)); + for (int i = 0; i < n_dim; ++i) + { + H_mat[i + i * n_dim] = T(200.0, 0); + if (i > 0) + { + H_mat[i + (i - 1) * n_dim] = T(-100.0, 0); + } + if (i < n_dim - 1) + { + H_mat[i + (i + 1) * n_dim] = T(-100.0, 0); + } + } + prec.assign(n_dim, 200.0); + exact.resize(nband); + for (int k = 0; k < nband; ++k) + { + exact[k] = 100.0 * (2.0 - 2.0 * std::cos(Real(k + 1) * M_PI / Real(n_dim + 1))); + } + ethr.assign(nband, 1e-8); + init_psi(808); + } + void init_psi(int seed) + { + std::mt19937 rng(seed); + std::uniform_real_distribution dist(-1.0, 1.0); + psi.assign(ld * nband, T(0)); + for (int j = 0; j < nband; ++j) + { + for (int i = 0; i < n_dim; ++i) + { + psi[i + j * ld] = T(dist(rng), 0.0); + } + } + for (int j = 0; j < nband; ++j) + { + for (int k = 0; k < j; ++k) + { + T d = 0; + for (int i = 0; i < n_dim; ++i) + { + d += std::conj(psi[i + k * ld]) * psi[i + j * ld]; + } + for (int i = 0; i < n_dim; ++i) + { + psi[i + j * ld] -= d * psi[i + k * ld]; + } + } + Real nr = 0; + for (int i = 0; i < n_dim; ++i) + { + nr += std::norm(psi[i + j * ld]); + } + nr = std::sqrt(nr); + for (int i = 0; i < n_dim; ++i) + { + psi[i + j * ld] /= nr; + } + } + } + int n_dim; + int nband; + int ld; + std::vector H_mat; + std::vector prec; + std::vector exact; + std::vector ethr; + std::vector psi; +}; + +TEST_F(DiagoPPCGScaledTest, BlockSubspace) +{ + std::vector psi_run = psi; + std::vector eval(nband, 0.0); + hsolver::DiagoPPCG solver(1e-10, 120, 4, 4 + ); + auto h_op = [this](T* in, T* out, int ldi, int nc) { dense_h_multiply(H_mat.data(), n_dim, in, out, ldi, nc); }; + double avg_iter = solver.diag(h_op, nullptr, ld, nband, n_dim, psi_run.data(), eval.data(), ethr, prec.data()); + for (int i = 0; i < nband; ++i) + { + EXPECT_NEAR(eval[i], exact[i], 1e-6) << "Scaled BLOCK: eigenvalue[" << i << "] mismatch"; + } + EXPECT_LE(avg_iter, 120.0) << "Scaled BLOCK: too many iterations"; +} + +// ============================================================================= +// Test fixture: many bands (n_dim=12, nband=4) +// Moderate band-to-dimension ratio (1:3). +// ============================================================================= +class DiagoPPCGManyBandsTest : public ::testing::Test +{ + protected: + void SetUp() override + { + n_dim = 12; + nband = 4; + ld = n_dim; + H_mat.assign(n_dim * n_dim, T(0)); + for (int i = 0; i < n_dim; ++i) + { + H_mat[i + i * n_dim] = T(2.0, 0); + if (i > 0) + { + H_mat[i + (i - 1) * n_dim] = T(-1.0, 0); + } + if (i < n_dim - 1) + { + H_mat[i + (i + 1) * n_dim] = T(-1.0, 0); + } + } + prec.assign(n_dim, 2.0); + exact.resize(nband); + for (int k = 0; k < nband; ++k) + { + exact[k] = 2.0 - 2.0 * std::cos(Real(k + 1) * M_PI / Real(n_dim + 1)); + } + ethr.assign(nband, 1e-10); + init_psi(909); + } + void init_psi(int seed) + { + std::mt19937 rng(seed); + std::uniform_real_distribution dist(-1.0, 1.0); + psi.assign(ld * nband, T(0)); + for (int j = 0; j < nband; ++j) + { + for (int i = 0; i < n_dim; ++i) + { + psi[i + j * ld] = T(dist(rng), 0.0); + } + } + for (int j = 0; j < nband; ++j) + { + for (int k = 0; k < j; ++k) + { + T d = 0; + for (int i = 0; i < n_dim; ++i) + { + d += std::conj(psi[i + k * ld]) * psi[i + j * ld]; + } + for (int i = 0; i < n_dim; ++i) + { + psi[i + j * ld] -= d * psi[i + k * ld]; + } + } + Real nr = 0; + for (int i = 0; i < n_dim; ++i) + { + nr += std::norm(psi[i + j * ld]); + } + nr = std::sqrt(nr); + for (int i = 0; i < n_dim; ++i) + { + psi[i + j * ld] /= nr; + } + } + } + int n_dim; + int nband; + int ld; + std::vector H_mat; + std::vector prec; + std::vector exact; + std::vector ethr; + std::vector psi; +}; + +TEST_F(DiagoPPCGManyBandsTest, BlockSubspace) +{ + std::vector psi_run = psi; + std::vector eval(nband, 0.0); + hsolver::DiagoPPCG solver(1e-12, 150, 4, 4 + ); + auto h_op = [this](T* in, T* out, int ldi, int nc) { dense_h_multiply(H_mat.data(), n_dim, in, out, ldi, nc); }; + double avg_iter = solver.diag(h_op, nullptr, ld, nband, n_dim, psi_run.data(), eval.data(), ethr, prec.data()); + for (int i = 0; i < nband; ++i) + { + EXPECT_NEAR(eval[i], exact[i], 1e-8) << "ManyBands BLOCK: eigenvalue[" << i << "] mismatch"; + } + EXPECT_LE(avg_iter, 150.0) << "ManyBands BLOCK: too many iterations"; +} + +// ============================================================================= +// Test fixture: rr_step = 1 (most frequent Rayleigh-Ritz refinement) +// 8×8 tridiagonal, nband=3. Tests aggressive subspace diagonalization. +// ============================================================================= +class DiagoPPCGRrStep1Test : public ::testing::Test +{ + protected: + void SetUp() override + { + n_dim = 8; + nband = 3; + ld = n_dim; + H_mat.assign(n_dim * n_dim, T(0)); + for (int i = 0; i < n_dim; ++i) + { + H_mat[i + i * n_dim] = T(2.0, 0); + if (i > 0) + { + H_mat[i + (i - 1) * n_dim] = T(-1.0, 0); + } + if (i < n_dim - 1) + { + H_mat[i + (i + 1) * n_dim] = T(-1.0, 0); + } + } + prec.assign(n_dim, 2.0); + exact.resize(nband); + for (int k = 0; k < nband; ++k) + { + exact[k] = 2.0 - 2.0 * std::cos(Real(k + 1) * M_PI / Real(n_dim + 1)); + } + ethr.assign(nband, 1e-10); + init_psi(111); + } + void init_psi(int seed) + { + std::mt19937 rng(seed); + std::uniform_real_distribution dist(-1.0, 1.0); + psi.assign(ld * nband, T(0)); + for (int j = 0; j < nband; ++j) + { + for (int i = 0; i < n_dim; ++i) + { + psi[i + j * ld] = T(dist(rng), 0.0); + } + } + for (int j = 0; j < nband; ++j) + { + for (int k = 0; k < j; ++k) + { + T d = 0; + for (int i = 0; i < n_dim; ++i) + { + d += std::conj(psi[i + k * ld]) * psi[i + j * ld]; + } + for (int i = 0; i < n_dim; ++i) + { + psi[i + j * ld] -= d * psi[i + k * ld]; + } + } + Real nr = 0; + for (int i = 0; i < n_dim; ++i) + { + nr += std::norm(psi[i + j * ld]); + } + nr = std::sqrt(nr); + for (int i = 0; i < n_dim; ++i) + { + psi[i + j * ld] /= nr; + } + } + } + int n_dim; + int nband; + int ld; + std::vector H_mat; + std::vector prec; + std::vector exact; + std::vector ethr; + std::vector psi; +}; + +TEST_F(DiagoPPCGRrStep1Test, BlockSubspace) +{ + std::vector psi_run = psi; + std::vector eval(nband, 0.0); + hsolver::DiagoPPCG solver(1e-12, 100, 3, 1 /*rr_step=1*/ + ); + auto h_op = [this](T* in, T* out, int ldi, int nc) { dense_h_multiply(H_mat.data(), n_dim, in, out, ldi, nc); }; + double avg_iter = solver.diag(h_op, nullptr, ld, nband, n_dim, psi_run.data(), eval.data(), ethr, prec.data()); + for (int i = 0; i < nband; ++i) + { + EXPECT_NEAR(eval[i], exact[i], 1e-8) << "RrStep1 BLOCK: eigenvalue[" << i << "] mismatch"; + } + EXPECT_LE(avg_iter, 100.0) << "RrStep1 BLOCK: too many iterations"; +} + +// ============================================================================= +// Test fixture: Neumann boundary tridiagonal Laplacian +// H[0,0]=H[n-1,n-1]=1, interior diag=2, off-diag=-1. +// Eigenvalues: λ_k = 2 - 2·cos(k·π/n), k=0,1,...,n-1 +// ============================================================================= +class DiagoPPCGNeumannTest : public ::testing::Test +{ + protected: + void SetUp() override + { + n_dim = 8; + nband = 4; + ld = n_dim; + H_mat.assign(n_dim * n_dim, T(0)); + for (int i = 0; i < n_dim; ++i) + { + H_mat[i + i * n_dim] = T((i == 0 || i == n_dim - 1) ? 1.0 : 2.0, 0); + if (i > 0) + { + H_mat[i + (i - 1) * n_dim] = T(-1.0, 0); + } + if (i < n_dim - 1) + { + H_mat[i + (i + 1) * n_dim] = T(-1.0, 0); + } + } + prec.assign(n_dim, 2.0); + prec[0] = 1.0; + prec[n_dim - 1] = 1.0; + exact.resize(nband); + for (int k = 0; k < nband; ++k) + { + exact[k] = 2.0 - 2.0 * std::cos(Real(k) * M_PI / Real(n_dim)); + } + ethr.assign(nband, 1e-10); + init_psi(222); + } + void init_psi(int seed) + { + std::mt19937 rng(seed); + std::uniform_real_distribution dist(-1.0, 1.0); + psi.assign(ld * nband, T(0)); + for (int j = 0; j < nband; ++j) + { + for (int i = 0; i < n_dim; ++i) + { + psi[i + j * ld] = T(dist(rng), 0.0); + } + } + for (int j = 0; j < nband; ++j) + { + for (int k = 0; k < j; ++k) + { + T d = 0; + for (int i = 0; i < n_dim; ++i) + { + d += std::conj(psi[i + k * ld]) * psi[i + j * ld]; + } + for (int i = 0; i < n_dim; ++i) + { + psi[i + j * ld] -= d * psi[i + k * ld]; + } + } + Real nr = 0; + for (int i = 0; i < n_dim; ++i) + { + nr += std::norm(psi[i + j * ld]); + } + nr = std::sqrt(nr); + for (int i = 0; i < n_dim; ++i) + { + psi[i + j * ld] /= nr; + } + } + } + int n_dim; + int nband; + int ld; + std::vector H_mat; + std::vector prec; + std::vector exact; + std::vector ethr; + std::vector psi; +}; + +TEST_F(DiagoPPCGNeumannTest, BlockSubspace) +{ + std::vector psi_run = psi; + std::vector eval(nband, 0.0); + hsolver::DiagoPPCG solver(1e-12, 100, 4, 4 + ); + auto h_op = [this](T* in, T* out, int ldi, int nc) { dense_h_multiply(H_mat.data(), n_dim, in, out, ldi, nc); }; + double avg_iter = solver.diag(h_op, nullptr, ld, nband, n_dim, psi_run.data(), eval.data(), ethr, prec.data()); + for (int i = 0; i < nband; ++i) + { + EXPECT_NEAR(eval[i], exact[i], 1e-8) << "Neumann BLOCK: eigenvalue[" << i << "] mismatch"; + } + EXPECT_LE(avg_iter, 100.0) << "Neumann BLOCK: too many iterations"; +} + +// ============================================================================= +// Test fixture: tight convergence threshold (ethr = 1e-14) +// 6×6 tridiagonal, nband=2. Tests deep convergence. +// ============================================================================= +class DiagoPPCGTightEthrTest : public ::testing::Test +{ + protected: + void SetUp() override + { + n_dim = 6; + nband = 2; + ld = n_dim; + H_mat.assign(n_dim * n_dim, T(0)); + for (int i = 0; i < n_dim; ++i) + { + H_mat[i + i * n_dim] = T(2.0, 0); + if (i > 0) + { + H_mat[i + (i - 1) * n_dim] = T(-1.0, 0); + } + if (i < n_dim - 1) + { + H_mat[i + (i + 1) * n_dim] = T(-1.0, 0); + } + } + prec.assign(n_dim, 2.0); + exact.resize(nband); + for (int k = 0; k < nband; ++k) + { + exact[k] = 2.0 - 2.0 * std::cos(Real(k + 1) * M_PI / Real(n_dim + 1)); + } + ethr.assign(nband, 1e-14); + init_psi(333); + } + void init_psi(int seed) + { + std::mt19937 rng(seed); + std::uniform_real_distribution dist(-1.0, 1.0); + psi.assign(ld * nband, T(0)); + for (int j = 0; j < nband; ++j) + { + for (int i = 0; i < n_dim; ++i) + { + psi[i + j * ld] = T(dist(rng), 0.0); + } + } + for (int j = 0; j < nband; ++j) + { + for (int k = 0; k < j; ++k) + { + T d = 0; + for (int i = 0; i < n_dim; ++i) + { + d += std::conj(psi[i + k * ld]) * psi[i + j * ld]; + } + for (int i = 0; i < n_dim; ++i) + { + psi[i + j * ld] -= d * psi[i + k * ld]; + } + } + Real nr = 0; + for (int i = 0; i < n_dim; ++i) + { + nr += std::norm(psi[i + j * ld]); + } + nr = std::sqrt(nr); + for (int i = 0; i < n_dim; ++i) + { + psi[i + j * ld] /= nr; + } + } + } + int n_dim; + int nband; + int ld; + std::vector H_mat; + std::vector prec; + std::vector exact; + std::vector ethr; + std::vector psi; +}; + +TEST_F(DiagoPPCGTightEthrTest, BlockSubspace) +{ + std::vector psi_run = psi; + std::vector eval(nband, 0.0); + hsolver::DiagoPPCG solver(1e-14, 200, 3, 3 + ); + auto h_op = [this](T* in, T* out, int ldi, int nc) { dense_h_multiply(H_mat.data(), n_dim, in, out, ldi, nc); }; + double avg_iter = solver.diag(h_op, nullptr, ld, nband, n_dim, psi_run.data(), eval.data(), ethr, prec.data()); + for (int i = 0; i < nband; ++i) + { + EXPECT_NEAR(eval[i], exact[i], 1e-8) << "TightEthr BLOCK: eigenvalue[" << i << "] mismatch"; + } + EXPECT_LE(avg_iter, 200.0) << "TightEthr BLOCK: too many iterations"; +} + +// ============================================================================= +// Test fixture: non-diagonal S matrix (tridiagonal overlap) +// H = tridiag(2,-1,-1), S = tridiag(1.0, 0.2, 0.2) — both tridiagonal. +// Verifies S-orthogonalization with a non-diagonal S. +// ============================================================================= +class DiagoPPCGTridiagSTest : public ::testing::Test +{ + protected: + void SetUp() override + { + n_dim = 6; + nband = 2; + ld = n_dim; + H_mat.assign(n_dim * n_dim, T(0)); + for (int i = 0; i < n_dim; ++i) + { + H_mat[i + i * n_dim] = T(2.0, 0); + if (i > 0) + { + H_mat[i + (i - 1) * n_dim] = T(-1.0, 0); + } + if (i < n_dim - 1) + { + H_mat[i + (i + 1) * n_dim] = T(-1.0, 0); + } + } + S_mat.assign(n_dim * n_dim, T(0)); + for (int i = 0; i < n_dim; ++i) + { + S_mat[i + i * n_dim] = T(1.0, 0); + if (i > 0) + { + S_mat[i + (i - 1) * n_dim] = T(0.2, 0); + S_mat[(i - 1) + i * n_dim] = T(0.2, 0); + } + } + prec.assign(n_dim, 2.0); + // H and S share sine eigenvectors. Their generalized eigenvalues are + // (2 - 2 cos(theta)) / (1 + 0.4 cos(theta)), theta = k*pi/(n_dim+1). + exact.resize(nband); + for (int k = 0; k < nband; ++k) + { + const Real theta = Real(k + 1) * M_PI / Real(n_dim + 1); + const Real cosine = std::cos(theta); + exact[k] = (2.0 - 2.0 * cosine) / (1.0 + 0.4 * cosine); + } + ethr.assign(nband, 1e-8); + init_psi(444); + } + void init_psi(int seed) + { + std::mt19937 rng(seed); + std::uniform_real_distribution dist(-1.0, 1.0); + psi.assign(ld * nband, T(0)); + for (int j = 0; j < nband; ++j) + { + for (int i = 0; i < n_dim; ++i) + { + psi[i + j * ld] = T(dist(rng), 0.0); + } + } + // S-orthonormalize: S = tridiag(1.0, 0.2, 0.2) + for (int j = 0; j < nband; ++j) + { + for (int k = 0; k < j; ++k) + { + T d = 0; + for (int i = 0; i < n_dim; ++i) + { + T si = 0; + si += T(1.0, 0) * psi[i + k * ld]; + if (i > 0) + { + si += T(0.2, 0) * psi[(i - 1) + k * ld]; + } + if (i < n_dim - 1) + { + si += T(0.2, 0) * psi[(i + 1) + k * ld]; + } + d += std::conj(si) * psi[i + j * ld]; + } + for (int i = 0; i < n_dim; ++i) + { + psi[i + j * ld] -= d * psi[i + k * ld]; + } + } + Real nr = 0; + for (int i = 0; i < n_dim; ++i) + { + T si = 0; + si += T(1.0, 0) * psi[i + j * ld]; + if (i > 0) + { + si += T(0.2, 0) * psi[(i - 1) + j * ld]; + } + if (i < n_dim - 1) + { + si += T(0.2, 0) * psi[(i + 1) + j * ld]; + } + nr += std::real(std::conj(psi[i + j * ld]) * si); + } + nr = std::sqrt(nr); + for (int i = 0; i < n_dim; ++i) + { + psi[i + j * ld] /= nr; + } + } + } + int n_dim; + int nband; + int ld; + std::vector H_mat; + std::vector S_mat; + std::vector prec; + std::vector exact; + std::vector ethr; + std::vector psi; +}; + +TEST_F(DiagoPPCGTridiagSTest, BlockSubspace) +{ + std::vector psi_run = psi; + std::vector eval(nband, 0.0); + auto spsi_func = [this](T* in, T* out, int ldi, int nc) { + for (int j = 0; j < nc; ++j) + { + for (int i = 0; i < n_dim; ++i) + { + out[i + j * ldi] = T(1.0, 0) * in[i + j * ldi]; + if (i > 0) + { + out[i + j * ldi] += T(0.2, 0) * in[(i - 1) + j * ldi]; + } + if (i < n_dim - 1) + { + out[i + j * ldi] += T(0.2, 0) * in[(i + 1) + j * ldi]; + } + } + } + }; + hsolver::DiagoPPCG solver(1e-10, 150, 3, 3); + auto h_op = [this](T* in, T* out, int ldi, int nc) { dense_h_multiply(H_mat.data(), n_dim, in, out, ldi, nc); }; + const double avg_iter = solver.diag(h_op, spsi_func, ld, nband, n_dim, psi_run.data(), eval.data(), ethr, prec.data()); + + EXPECT_TRUE(solver.converged()); + EXPECT_EQ(solver.active_band_count(), 0); + EXPECT_FALSE(solver.iteration_limit_reached()); + for (int i = 0; i < nband; ++i) + { + EXPECT_NEAR(eval[i], exact[i], ethr[i]) << "TridiagS BLOCK: wrong eigenvalue at band " << i; + } + + // Reapply the dense matrices independently of the solver's operator images. + const int image_size = ld * nband; + std::vector hpsi(image_size); + std::vector spsi(image_size); + std::vector res(n_dim); + dense_h_multiply(H_mat.data(), n_dim, psi_run.data(), hpsi.data(), ld, nband); + dense_h_multiply(S_mat.data(), n_dim, psi_run.data(), spsi.data(), ld, nband); + for (int i = 0; i < nband; ++i) + { + for (int j = 0; j < n_dim; ++j) + { + const int offset = j + i * ld; + res[j] = hpsi[offset] - eval[i] * spsi[offset]; + } + const Real residual_norm = column_norm(res.data(), n_dim); + EXPECT_LE(residual_norm, ethr[i]) << "TridiagS BLOCK: residual at band " << i; + + for (int k = 0; k < nband; ++k) + { + T overlap = 0; + for (int j = 0; j < n_dim; ++j) + { + overlap += std::conj(psi_run[j + k * ld]) * spsi[j + i * ld]; + } + const T expected_overlap = (i == k) ? T(1) : T(0); + const T overlap_difference = overlap - expected_overlap; + const Real overlap_error = std::abs(overlap_difference); + EXPECT_LE(overlap_error, 1e-10) << "TridiagS BLOCK: S-overlap at (" << k << ", " << i << ")"; + } + } + EXPECT_LE(avg_iter, 150.0) << "TridiagS BLOCK: too many iterations"; +} + +// ============================================================================= +// Performance benchmark: PPCG vs BLOCK vs Davidson comparison +// +// Runs PPCG on random sparse symmetric matrices of various sizes (matching +// the BLOCK/Davidson test sizes: npw=100,200,500) and reports avg_iter and +// wall-clock time. avg_iter is the primary metric — it counts the average +// number of H·psi applications per band. +// +// Typical results from the existing BLOCK/Davidson tests (for reference): +// BLOCK: avg_iter ~ 20-50 on random sparse, ~ 5-15 on tridiagonal +// Davidson: avg_iter ~ 15-40 on random sparse (but each iter is heavier) +// PPCG: avg_iter ~ 2-10 on tridiagonal/diagonal, varies with sparsity +// +// This test is DISABLED by default (too slow for CI). Run manually with +// --gtest_also_run_disabled_tests --gtest_filter='*Benchmark*' +// ============================================================================= +class DiagoPPCGBenchmarkTest : public ::testing::Test +{ + protected: + struct RunResult + { + double average_iterations; + double wall_time; + bool converged; + Real max_residual; + }; + + void SetUp() override + { + } + + // Generate a random sparse symmetric matrix of size n with given sparsity. + // sparsity=0 means dense, sparsity=80 means 80% zeros. + static void make_random_hamilt(int n, int sparsity_pct, std::vector& H, std::vector& prec) + { + H.assign(n * n, T(0)); + std::mt19937 rng(unsigned(n * 100 + sparsity_pct)); + std::uniform_real_distribution dist(-1.0, 1.0); + int nnz = 0; + for (int i = 0; i < n; ++i) + { + for (int j = i; j < n; ++j) + { + if (i != j && (rng() % 100) < sparsity_pct) + { + continue; + } + Real val = (i == j) ? std::abs(dist(rng)) * n + 1.0 : dist(rng) * 0.5; + H[i + j * n] = T(val, 0); + if (i != j) + { + H[j + i * n] = T(val, 0); + } + if (val != 0) + { + ++nnz; + } + } + } + // Simple diagonal preconditioner + prec.resize(n); + for (int i = 0; i < n; ++i) + { + prec[i] = std::max(std::real(H[i + i * n]), 1e-6); + } + } + + static RunResult run_ppcg(int n, int nband, const std::vector& H, const std::vector& prec) + { + int ld = n; + std::mt19937 rng(42); + std::uniform_real_distribution dist(-1.0, 1.0); + std::vector psi(ld * nband, T(0)); + for (int j = 0; j < nband; ++j) + { + for (int i = 0; i < n; ++i) + { + psi[i + j * ld] = T(dist(rng), 0.0); + } + } + // GS orthonormalize + for (int j = 0; j < nband; ++j) + { + for (int k = 0; k < j; ++k) + { + T d = 0; + for (int i = 0; i < n; ++i) + { + d += std::conj(psi[i + k * ld]) * psi[i + j * ld]; + } + for (int i = 0; i < n; ++i) + { + psi[i + j * ld] -= d * psi[i + k * ld]; + } + } + Real nr = 0; + for (int i = 0; i < n; ++i) + { + nr += std::norm(psi[i + j * ld]); + } + nr = std::sqrt(nr); + for (int i = 0; i < n; ++i) + { + psi[i + j * ld] /= nr; + } + } + + std::vector eval(nband, 0.0); + std::vector ethr(nband, 1e-4); + auto h_op = [&H, n](T* in, T* out, int ldi, int nc) { dense_h_multiply(H.data(), n, in, out, ldi, nc); }; + + const int rr_step = std::min(nband, 4); + hsolver::DiagoPPCG solver(1e-8, 500, nband, rr_step); + + auto t0 = std::chrono::high_resolution_clock::now(); + double avg_iter = solver.diag(h_op, nullptr, ld, nband, n, psi.data(), eval.data(), ethr, prec.data()); + auto t1 = std::chrono::high_resolution_clock::now(); + double wall = std::chrono::duration(t1 - t0).count(); + return {avg_iter, wall, solver.converged(), solver.max_residual()}; + } +}; + +TEST_F(DiagoPPCGBenchmarkTest, DISABLED_FullBenchmark) +{ + struct Case + { + int n; + int nband; + int sparsity; + }; + std::vector cases = { + {50, 10, 0}, {50, 10, 60}, {100, 10, 0}, {100, 10, 60}, + {100, 10, 80}, {200, 10, 60}, {200, 10, 80}, {500, 10, 80}, + }; + + std::cout << "\n========== PPCG Performance Benchmark ==========\n"; + std::cout << " n_dim nband sparsity avg_iter wall_time(s)\n"; + std::cout << "-------------------------------------------------\n"; + for (auto& c : cases) + { + std::vector H; + std::vector prec; + make_random_hamilt(c.n, c.sparsity, H, prec); + const RunResult result = run_ppcg(c.n, c.nband, H, prec); + const double avg_iter = result.average_iterations; + const double wall = result.wall_time; + printf(" %5d %3d %2d%% %6.1f %7.4f\n", c.n, c.nband, c.sparsity, avg_iter, wall); + } + std::cout << "=================================================\n"; + SUCCEED(); +} + +// Quick convergence smoke test: one representative case, fast enough for CI. +TEST_F(DiagoPPCGBenchmarkTest, QuickBenchmark) +{ + std::vector H; + std::vector prec; + make_random_hamilt(80, 60, H, prec); + const RunResult result = run_ppcg(80, 8, H, prec); + EXPECT_TRUE(result.converged); + EXPECT_TRUE(std::isfinite(result.max_residual)); + EXPECT_LE(result.max_residual, Real(1.0e-4)); + EXPECT_LE(result.average_iterations, 500.0); +} + +int main(int argc, char** argv) +{ +#ifdef __MPI + int nproc = 1; + int myrank = 0; + setupmpi(argc, argv, nproc, myrank); + int nproc_in_pool = 1; + int kpar = nproc; + int mypool = 0; + int rank_in_pool = 0; + divide_pools(nproc, myrank, nproc_in_pool, kpar, mypool, rank_in_pool); +#endif + testing::InitGoogleTest(&argc, argv); + const int result = RUN_ALL_TESTS(); +#ifdef __MPI + finishmpi(); +#endif + return result; +} diff --git a/source/source_hsolver/unittests/test_diago_ppcg_float.cpp b/source/source_hsolver/unittests/test_diago_ppcg_float.cpp new file mode 100644 index 00000000000..a2701cb9ec3 --- /dev/null +++ b/source/source_hsolver/unittests/test_diago_ppcg_float.cpp @@ -0,0 +1,306 @@ +/** + * diago_ppcg_float_test.cpp — single-precision unit test for DiagoPPCG. + * + * Exercises the std::complex instantiation of the BLOCK_SUBSPACE and + * BLOCK_SUBSPACE strategy on dense matrices with analytical reference + * eigenvalues. Tolerances are looser than the double-precision suite because + * single precision has roughly 7 significant digits. + */ + +#include "../diago_ppcg.h" +#ifdef __MPI +#include "source_basis/module_pw/test/test_tool.h" +#endif + +#include +#include +#include +#include +#include +#include + +#ifndef M_PI +#define M_PI 3.14159265358979323846 +#endif + +using T = std::complex; +using Real = float; + +// ----------------------------------------------------------------------------- +// Helper: dense H-matrix times a set of column vectors (column-major H). +// ----------------------------------------------------------------------------- +static void dense_h_multiply(const T* H_data, int n_dim, + const T* in, T* out, int ld, int ncol) +{ + for (int j = 0; j < ncol; ++j) + { + for (int i = 0; i < n_dim; ++i) + { + T sum = T(0.0f, 0.0f); + for (int k = 0; k < n_dim; ++k) + { + sum += H_data[i + k * n_dim] * in[k + j * ld]; + } + out[i + j * ld] = sum; + } + } +} + +// Orthonormalize columns of psi in-place (S = I). +static void gram_schmidt(std::vector& psi, int ld, int n_dim, int nband) +{ + for (int j = 0; j < nband; ++j) + { + for (int k = 0; k < j; ++k) + { + T dot = T(0.0f, 0.0f); + for (int i = 0; i < n_dim; ++i) + { + dot += std::conj(psi[i + k * ld]) * psi[i + j * ld]; + } + for (int i = 0; i < n_dim; ++i) + { + psi[i + j * ld] -= dot * psi[i + k * ld]; + } + } + Real nrm = 0.0f; + for (int i = 0; i < n_dim; ++i) + { + nrm += std::norm(psi[i + j * ld]); + } + nrm = std::sqrt(nrm); + for (int i = 0; i < n_dim; ++i) + { + psi[i + j * ld] /= nrm; + } + } +} + +// ----------------------------------------------------------------------------- +// Diagonal matrix: H = diag(1, 2, 3). All eigenvalues are computed (nband == +// n_dim), so there is no ambiguity about which end of the spectrum to converge +// to; single-precision Rayleigh-Ritz can otherwise drift toward the upper +// eigenvalues on some platforms. +// ----------------------------------------------------------------------------- +TEST(DiagoPPCGFloatTest, DiagonalBlockSubspace) +{ + const int n_dim = 3; + const int nband = 3; + const int ld = n_dim; + + std::vector H_mat(n_dim * n_dim, T(0.0f, 0.0f)); + for (int i = 0; i < n_dim; ++i) + { + H_mat[i + i * n_dim] = T(Real(i + 1), 0.0f); + } + + std::vector prec(n_dim); + for (int i = 0; i < n_dim; ++i) + { + prec[i] = Real(i + 1); + } + + const Real exact[3] = {1.0f, 2.0f, 3.0f}; + std::vector ethr(nband, 1e-4); + + std::mt19937 rng(42); + std::uniform_real_distribution dist(-1.0f, 1.0f); + std::vector psi(ld * nband, T(0.0f, 0.0f)); + for (int j = 0; j < nband; ++j) + { + for (int i = 0; i < n_dim; ++i) + { + psi[i + j * ld] = T(dist(rng), 0.0f); + } + } + gram_schmidt(psi, ld, n_dim, nband); + + std::vector psi_run = psi; + std::vector eval(nband, 0.0f); + + hsolver::DiagoPPCG solver( + /* diag_thr = */ 1e-5f, + /* max_iter = */ 100, + /* sbsize = */ 3, + /* rr_step = */ 3); + + auto h_op = [&](T* in, T* out, int ld_in, int ncol) { + dense_h_multiply(H_mat.data(), n_dim, in, out, ld_in, ncol); + }; + + double avg_iter = solver.diag(h_op, nullptr, ld, nband, n_dim, + psi_run.data(), eval.data(), ethr, prec.data()); + + for (int i = 0; i < nband; ++i) + { + EXPECT_NEAR(double(eval[i]), double(exact[i]), 1e-4) + << "Diagonal float BLOCK: eigenvalue[" << i << "] mismatch"; + } + EXPECT_LE(avg_iter, 100.0) << "Diagonal float BLOCK: too many iterations"; +} + +// ----------------------------------------------------------------------------- +// Tridiagonal Laplacian: H[i,i]=2, H[i,i±1]=-1, exact λ_k = 2 - 2cos(kπ/(n+1)) +// ----------------------------------------------------------------------------- +TEST(DiagoPPCGFloatTest, TridiagonalBlockSubspace) +{ + const int n_dim = 10; + const int nband = 3; + const int ld = n_dim; + + std::vector H_mat(n_dim * n_dim, T(0.0f, 0.0f)); + for (int i = 0; i < n_dim; ++i) + { + H_mat[i + i * n_dim] = T(2.0f, 0.0f); + if (i > 0) + { + H_mat[i + (i - 1) * n_dim] = T(-1.0f, 0.0f); + } + if (i < n_dim - 1) + { + H_mat[i + (i + 1) * n_dim] = T(-1.0f, 0.0f); + } + } + + std::vector prec(n_dim, 2.0f); + std::vector exact(nband); + for (int k = 0; k < nband; ++k) + { + exact[k] = 2.0f - 2.0f * std::cos(Real(k + 1) * M_PI + / Real(n_dim + 1)); + } + std::vector ethr(nband, 1e-4); + + std::mt19937 rng(42); + std::uniform_real_distribution dist(-1.0f, 1.0f); + std::vector psi(ld * nband, T(0.0f, 0.0f)); + for (int j = 0; j < nband; ++j) + { + for (int i = 0; i < n_dim; ++i) + { + psi[i + j * ld] = T(dist(rng), 0.0f); + } + } + gram_schmidt(psi, ld, n_dim, nband); + + std::vector psi_run = psi; + std::vector eval(nband, 0.0f); + + hsolver::DiagoPPCG solver( + /* diag_thr = */ 1e-5f, + /* max_iter = */ 100, + /* sbsize = */ 4, + /* rr_step = */ 4); + + auto h_op = [&](T* in, T* out, int ld_in, int ncol) { + dense_h_multiply(H_mat.data(), n_dim, in, out, ld_in, ncol); + }; + + double avg_iter = solver.diag(h_op, nullptr, ld, nband, n_dim, + psi_run.data(), eval.data(), ethr, prec.data()); + + for (int i = 0; i < nband; ++i) + { + EXPECT_NEAR(double(eval[i]), double(exact[i]), 1e-4) + << "Tridiagonal float BLOCK: eigenvalue[" << i << "] mismatch"; + } + EXPECT_LE(avg_iter, 100.0) << "Tridiagonal float BLOCK: too many iterations"; +} + +// ----------------------------------------------------------------------------- +// Non-finite input validation (throws). +// ----------------------------------------------------------------------------- +TEST(DiagoPPCGFloatTest, NonFiniteInputThrows) +{ + const int n_dim = 5; + const int nband = 3; + const int ld = n_dim; + + std::vector H_mat(n_dim * n_dim, T(0.0f, 0.0f)); + for (int i = 0; i < n_dim; ++i) + { + H_mat[i + i * n_dim] = T(Real(i + 1), 0.0f); + } + + std::vector prec(n_dim, 1.0f); + std::vector psi(ld * nband, T(1.0f, 0.0f)); + std::vector eval(nband, 0.0f); + std::vector ethr(nband, 1e-4); + + hsolver::DiagoPPCG solver( + /* diag_thr = */ 1e-5f, 100, 3, 3); + + auto h_op = [&](T* in, T* out, int ld_in, int ncol) { + dense_h_multiply(H_mat.data(), n_dim, in, out, ld_in, ncol); + }; + + std::vector bad_ethr = ethr; + bad_ethr[0] = std::numeric_limits::quiet_NaN(); + EXPECT_THROW(solver.diag(h_op, nullptr, ld, nband, n_dim, + psi.data(), eval.data(), bad_ethr, prec.data()), + std::invalid_argument); + + std::vector bad_prec = prec; + bad_prec[0] = std::numeric_limits::quiet_NaN(); + EXPECT_THROW(solver.diag(h_op, nullptr, ld, nband, n_dim, + psi.data(), eval.data(), ethr, bad_prec.data()), + std::invalid_argument); +} + +TEST(DiagoPPCGFloatTest, RankDeficientStartupSubspace) +{ + const int n_dim = 4; + const int nband = 2; + const int ld = n_dim; + const Real inv_sqrt_two = Real(1) / std::sqrt(Real(2)); + std::vector h_mat(n_dim * n_dim, T(0)); + std::vector prec(n_dim); + for (int i = 0; i < n_dim; ++i) + { + const Real value = Real(i + 1); + h_mat[i + i * n_dim] = T(value, 0); + prec[i] = value; + } + + std::vector psi(ld * nband, T(0)); + psi[0] = T(inv_sqrt_two, 0); + psi[1] = T(inv_sqrt_two, 0); + psi[3 + ld] = T(1, 0); + std::vector eval(nband, Real(0)); + const std::vector ethr(nband, 1e-4); + const auto h_op = [&h_mat, n_dim](T* in, T* out, int ld_in, int ncol) { + dense_h_multiply(h_mat.data(), n_dim, in, out, ld_in, ncol); + }; + hsolver::DiagoPPCG solver( + /* diag_thr = */ 1e-5f, + /* max_iter = */ 20, + /* sbsize = */ 1, + /* rr_step = */ 2); + + solver.diag(h_op, nullptr, ld, nband, n_dim, + psi.data(), eval.data(), ethr, prec.data()); + + EXPECT_NEAR(eval[0], Real(1), Real(2e-4)); + EXPECT_NEAR(eval[1], Real(2), Real(2e-4)); + EXPECT_TRUE(solver.converged()); +} + +int main(int argc, char** argv) +{ +#ifdef __MPI + int nproc = 1; + int myrank = 0; + setupmpi(argc, argv, nproc, myrank); + int nproc_in_pool = 1; + int kpar = nproc; + int mypool = 0; + int rank_in_pool = 0; + divide_pools(nproc, myrank, nproc_in_pool, kpar, mypool, rank_in_pool); +#endif + testing::InitGoogleTest(&argc, argv); + const int result = RUN_ALL_TESTS(); +#ifdef __MPI + finishmpi(); +#endif + return result; +} diff --git a/source/source_hsolver/unittests/test_diago_ppcg_parallel.cpp b/source/source_hsolver/unittests/test_diago_ppcg_parallel.cpp new file mode 100644 index 00000000000..cb59229e316 --- /dev/null +++ b/source/source_hsolver/unittests/test_diago_ppcg_parallel.cpp @@ -0,0 +1,346 @@ +/** + * diago_ppcg_parallel_test.cpp — MPI parallel test for DiagoPPCG. + * + * Distributes the rows of a diagonal matrix across MPI processes: each process + * owns a slice of the diagonal and the corresponding rows of psi, computes the + * partial Gram matrix / residual locally, and relies on the solver's pooled + * MPI reductions (reduce_pool) to sum the partial results. The eigenvalues of + * the global diagonal matrix must be recovered identically on every process. + * + * Run with: mpirun -np ./MODULE_HSOLVER_ppcg_parallel + */ + +#include "../diago_ppcg.h" + +#include "source_base/parallel_comm.h" +#include "source_base/parallel_reduce.h" +#include "source_basis/module_pw/test/test_tool.h" + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +// A local block can discover two low roots while another block already holds +// an exact higher root. The discarded low direction must survive global root +// selection, including when its rows and projection norms are distributed. +static bool check_competing_roots(const int rank_in_pool, const int nproc_in_pool) +{ + using Scalar = std::complex; + const int bands = 2; + const int local_rows = nproc_in_pool == 1 ? 4 : 2; + const int size = local_rows * bands; + std::vector diagonal(local_rows); + std::vector preconditioner(local_rows, 1.0); + std::vector thresholds(bands, 1.0e-10); + std::vector vectors(size, Scalar(0)); + const double scale = 1.0 / std::sqrt(3.0); + for (int row = 0; row < local_rows; ++row) + { + const int global_row = rank_in_pool * local_rows + row; + diagonal[row] = global_row < 3 ? global_row + 1.0 : 97.0 + global_row; + if (global_row == 0 || global_row == 1 || global_row == 3) + { + vectors[row] = Scalar(scale); + } + if (global_row == 2) + { + vectors[row + local_rows] = Scalar(1); + } + } + bool operator_capacity_ok = true; + const auto h_op = [&](Scalar* in, Scalar* out, int ld, int ncol) { + operator_capacity_ok = operator_capacity_ok && ncol <= bands; + for (int column = 0; column < ncol; ++column) + { + for (int row = 0; row < local_rows; ++row) + { + out[row + column * ld] = diagonal[row] * in[row + column * ld]; + } + } + }; + hsolver::DiagoPPCG solver(1.0e-10, 80, 1, 4); + std::vector eigenvalues(bands); + solver.diag(h_op, nullptr, local_rows, bands, local_rows, + vectors.data(), eigenvalues.data(), thresholds, preconditioner.data()); + + std::vector residuals(bands, 0.0); + std::vector overlap(bands * bands, Scalar(0)); + for (int column = 0; column < bands; ++column) + { + for (int row = 0; row < local_rows; ++row) + { + const Scalar value = vectors[row + column * local_rows]; + const Scalar residual = (diagonal[row] - eigenvalues[column]) * value; + residuals[column] += std::norm(residual); + for (int other = 0; other < bands; ++other) + { + overlap[other + column * bands] += std::conj(vectors[row + other * local_rows]) * value; + } + } + } + Parallel_Reduce::reduce_pool(residuals.data(), bands); + const int overlap_size = bands * bands; + Parallel_Reduce::reduce_pool(overlap.data(), overlap_size); + bool correct = solver.converged() && operator_capacity_ok && solver.root_recovery_count() > 0; + for (int column = 0; column < bands; ++column) + { + const double expected = column + 1.0; + correct = correct && std::isfinite(eigenvalues[column]); + correct = correct && std::abs(eigenvalues[column] - expected) < 1.0e-8; + correct = correct && std::isfinite(residuals[column]) && residuals[column] < 1.0e-16; + for (int other = 0; other < bands; ++other) + { + const Scalar expected_overlap = other == column ? Scalar(1) : Scalar(0); + const double error = std::abs(overlap[other + column * bands] - expected_overlap); + correct = correct && std::isfinite(error) && error < 1.0e-8; + } + } + if (!correct) + { + std::fprintf(stderr, "pool rank %d: competing-root regression returned %.12g, %.12g\n", + rank_in_pool, eigenvalues[0], eigenvalues[1]); + } + return correct; +} + +int main(int argc, char** argv) +{ + int nproc = 1; + int myrank = 0; + setupmpi(argc, argv, nproc, myrank); + + int nproc_in_pool = 0; + int kpar = argc > 1 ? std::atoi(argv[1]) : 1; + int mypool = 0; + int rank_in_pool = 0; + divide_pools(nproc, myrank, nproc_in_pool, kpar, mypool, rank_in_pool); + const bool world_root_is_correct = Parallel_Reduce::is_world_root() == (myrank == 0); + BP_WORLD = MPI_COMM_WORLD; + + using T = std::complex; + using Real = double; + + const int nband = 3; + const int n_local = 2 * nband; // each process owns this many rows + const int n_dim_total = nproc_in_pool * n_local; + + // With one pool, only rank zero requests a trace to exercise inconsistent + // process environments. With multiple pools, every rank requests the same + // path to verify that only the MPI world root writes it. + const char* trace_environment = "ABACUS_PPCG_RESIDUAL_TRACE"; + const std::string trace_path = "ppcg_residual_trace_parallel_test.csv"; + int trace_environment_ok = 1; + if (myrank == 0) + { + std::remove(trace_path.c_str()); + trace_environment_ok = ::setenv(trace_environment, trace_path.c_str(), 1) == 0 ? 1 : 0; + } + else if (kpar > 1) + { + trace_environment_ok = ::setenv(trace_environment, trace_path.c_str(), 1) == 0 ? 1 : 0; + } + else + { + trace_environment_ok = ::unsetenv(trace_environment) == 0 ? 1 : 0; + } + + // Diagonal H with entries 1, 2, ..., n_dim_total; process owns a slice. + std::vector diag_local(n_local); + std::vector prec(n_local); + for (int i = 0; i < n_local; ++i) + { + diag_local[i] = Real(rank_in_pool * n_local + i + 1); + prec[i] = diag_local[i]; + } + + std::vector ethr(nband, 1e-8); + + // Random initial guess (fixed seed for reproducibility). + std::mt19937 rng(42); + std::uniform_real_distribution dist(-1.0, 1.0); + std::vector psi(n_local * nband, T(0.0, 0.0)); + for (int j = 0; j < nband; ++j) + { + for (int i = 0; i < n_local; ++i) + { + psi[i + j * n_local] = T(dist(rng), 0.0); + } + } + + hsolver::DiagoPPCG solver( + /* diag_thr = */ 1e-12, + /* max_iter = */ 100, + /* sbsize = */ 2, + /* rr_step = */ nband); + + int h_calls = 0; + auto h_op = [&](T* in, T* out, int ld, int ncol) { + ++h_calls; + for (int j = 0; j < ncol; ++j) + { + for (int i = 0; i < n_local; ++i) + { + out[i + j * ld] = diag_local[i] * in[i + j * ld]; + } + } + }; + + std::vector eval(nband, 0.0); + // One invalid rank must stop every rank in its pool before H/S collectives. + // Otherwise the valid ranks would hang while the invalid rank unwinds. + std::vector invalid_prec = prec; + if (rank_in_pool == 0) + { + invalid_prec[0] = std::numeric_limits::quiet_NaN(); + } + bool collective_validation_failed = false; + try + { + solver.diag(h_op, nullptr, n_local, nband, n_local, + psi.data(), eval.data(), ethr, invalid_prec.data()); + } + catch (const std::invalid_argument&) + { + collective_validation_failed = true; + } + const bool validation_before_operator = collective_validation_failed && h_calls == 0; + + // Only one rank advertises an unrepresentable local [X,W,P] workspace. + // The wavefunction product still fits, and the tiny real buffers must not + // be read on that rank or passed to H on any of its peers. + int checked_ld = n_local; + if (rank_in_pool == 0) + { + checked_ld = std::numeric_limits::max() / nband; + } + bool collective_size_failed = false; + try + { + solver.diag(h_op, nullptr, checked_ld, nband, n_local, + psi.data(), eval.data(), ethr, prec.data()); + } + catch (const std::invalid_argument&) + { + collective_size_failed = true; + } + const bool size_validation_before_operator = collective_size_failed && h_calls == 0; + + solver.diag(h_op, nullptr, n_local, nband, n_local, + psi.data(), eval.data(), ethr, prec.data()); + + // The lowest nband eigenvalues of the global diagonal matrix are 1..nband. + int ok = 1; + if (!world_root_is_correct + || !validation_before_operator + || !size_validation_before_operator + || trace_environment_ok != 1) + { + ok = 0; + } + if (!world_root_is_correct) + { + std::fprintf(stderr, "rank %d: MPI world-root detection is inconsistent\n", myrank); + } + if (!validation_before_operator) + { + std::fprintf(stderr, "rank %d: invalid input was not rejected before operator calls\n", myrank); + } + if (!size_validation_before_operator) + { + std::fprintf(stderr, "rank %d: workspace overflow was not rejected before operator calls\n", myrank); + } + if (!solver.converged() + || solver.active_band_count() != 0 + || solver.iteration_limit_reached() + || !std::isfinite(solver.max_residual()) + || solver.max_residual() > 1.0e-8) + { + const int converged = solver.converged() ? 1 : 0; + const int active_band_count = solver.active_band_count(); + const Real max_residual = solver.max_residual(); + std::fprintf(stderr, + "rank %d: solver status is inconsistent (converged=%d, active=%d, max_residual=%.12e)\n", + myrank, + converged, + active_band_count, + max_residual); + ok = 0; + } + for (int i = 0; i < nband; ++i) + { + if (std::abs(eval[i] - Real(i + 1)) > 1e-6) + { + std::printf("rank %d: eval[%d] = %.12f != %d\n", myrank, i, eval[i], i + 1); + ok = 0; + } + } + + if (myrank == 0) + { + std::ifstream trace(trace_path.c_str()); + std::string line; + int header_count = 0; + int record_count = 0; + bool has_final_record = false; + while (std::getline(trace, line)) + { + if (line.find("iteration,stage,max_residual") == 0) + { + ++header_count; + } + else if (!line.empty()) + { + ++record_count; + if (line.find(",final,") != std::string::npos) + { + has_final_record = true; + } + } + } + if (!trace.eof() || header_count != 1 || record_count == 0 || !has_final_record) + { + const int final_record = has_final_record ? 1 : 0; + std::fprintf(stderr, + "rank 0: residual trace is incomplete (headers=%d, records=%d, final=%d)\n", + header_count, + record_count, + final_record); + ok = 0; + } + std::remove(trace_path.c_str()); + } + + // Complete the single-trace assertions before this additional solve writes + // its own trace. Each pool must independently recover both lowest roots. + const int trace_unset_status = ::unsetenv(trace_environment); + const bool competing_roots_ok = check_competing_roots(rank_in_pool, nproc_in_pool); + if (trace_unset_status != 0 || !competing_roots_ok) + { + ok = 0; + } + + int global_ok = ok; + Parallel_Reduce::reduce_min(global_ok); + + if (myrank == 0) + { + if (global_ok == 1) + { + std::printf("PPCG MPI parallel test PASSED\n"); + } + else + { + std::printf("PPCG MPI parallel test FAILED\n"); + } + } + const int exit_code = global_ok == 1 ? 0 : 1; + finishmpi(); + return exit_code; +} diff --git a/source/source_hsolver/unittests/test_hsolver.cpp b/source/source_hsolver/unittests/test_hsolver.cpp index 97629b3a002..59b21d73a84 100644 --- a/source/source_hsolver/unittests/test_hsolver.cpp +++ b/source/source_hsolver/unittests/test_hsolver.cpp @@ -20,6 +20,7 @@ * - non pw-ksdft returns 0 * - nscf, first scf iteration (by init_chg and calculation), later iterations * - single precision lower bound + * - optional SCF/relax upper bound, including every electronic-cycle reset * - hsolver::set_diagethr_sdft * - non pw-sdft returns 0 * - nscf, first iteration of the first/later ion step, later iterations @@ -45,10 +46,10 @@ TEST(HSolverCalHsolveErrorTest, OtherCasesAreZero) TEST(HSolverSetDiagethrKsTest, NonPwKsdftIsZero) { EXPECT_EQ( - hsolver::set_diagethr_ks("lcao", "ksdft", "scf", "atomic", "double", 0, 1, 1.0e-3, 1.0e-2, 1.0e-2, 8.0, 1.0e-6), + hsolver::set_diagethr_ks("lcao", "ksdft", "scf", "atomic", "double", 0, 1, 1.0e-3, 1.0e-2, 1.0e-2, 8.0, 1.0e-6, 0.0), 0.0); EXPECT_EQ( - hsolver::set_diagethr_ks("pw", "sdft", "scf", "atomic", "double", 0, 1, 1.0e-3, 1.0e-2, 1.0e-2, 8.0, 1.0e-6), + hsolver::set_diagethr_ks("pw", "sdft", "scf", "atomic", "double", 0, 1, 1.0e-3, 1.0e-2, 1.0e-2, 8.0, 1.0e-6, 0.0), 0.0); } @@ -56,11 +57,11 @@ TEST(HSolverSetDiagethrKsTest, Nscf) { // the default 1e-2 is tightened to 0.1 * min(1e-2, scf_thr / nelec) EXPECT_DOUBLE_EQ( - hsolver::set_diagethr_ks("pw", "ksdft", "nscf", "atomic", "double", 0, 1, 1.0e-3, 1.0e-2, 1.0e-2, 8.0, 1.0e-6), + hsolver::set_diagethr_ks("pw", "ksdft", "nscf", "atomic", "double", 0, 1, 1.0e-3, 1.0e-2, 1.0e-2, 8.0, 1.0e-6, 0.0), 1.25e-8); // a threshold the user already tightened is kept EXPECT_DOUBLE_EQ( - hsolver::set_diagethr_ks("pw", "ksdft", "nscf", "atomic", "double", 0, 1, 1.0e-3, 1.0e-2, 1.0e-6, 8.0, 1.0e-6), + hsolver::set_diagethr_ks("pw", "ksdft", "nscf", "atomic", "double", 0, 1, 1.0e-3, 1.0e-2, 1.0e-6, 8.0, 1.0e-6, 0.0), 1.0e-6); } @@ -68,18 +69,18 @@ TEST(HSolverSetDiagethrKsTest, FirstIteration) { // atomic starting charge keeps the loose 1e-2 EXPECT_DOUBLE_EQ( - hsolver::set_diagethr_ks("pw", "ksdft", "scf", "atomic", "double", 0, 1, 1.0e-3, 1.0e-2, 1.0e-2, 8.0, 1.0e-6), + hsolver::set_diagethr_ks("pw", "ksdft", "scf", "atomic", "double", 0, 1, 1.0e-3, 1.0e-2, 1.0e-2, 8.0, 1.0e-6, 0.0), 1.0e-2); // a charge read from file is trusted, so the first diagonalization is strict EXPECT_DOUBLE_EQ( - hsolver::set_diagethr_ks("pw", "ksdft", "scf", "file", "double", 0, 1, 1.0e-3, 1.0e-2, 1.0e-2, 8.0, 1.0e-6), + hsolver::set_diagethr_ks("pw", "ksdft", "scf", "file", "double", 0, 1, 1.0e-3, 1.0e-2, 1.0e-2, 8.0, 1.0e-6, 0.0), 1.0e-5); EXPECT_DOUBLE_EQ( - hsolver::set_diagethr_ks("pw", "ksdft", "scf", "wfc", "double", 0, 1, 1.0e-3, 1.0e-2, 1.0e-2, 8.0, 1.0e-6), + hsolver::set_diagethr_ks("pw", "ksdft", "scf", "wfc", "double", 0, 1, 1.0e-3, 1.0e-2, 1.0e-2, 8.0, 1.0e-6, 0.0), 1.0e-5); // relax/md never go below pw_diag_thr_init EXPECT_DOUBLE_EQ( - hsolver::set_diagethr_ks("pw", "ksdft", "relax", "file", "double", 0, 1, 1.0e-3, 1.0e-3, 1.0e-2, 8.0, 1.0e-6), + hsolver::set_diagethr_ks("pw", "ksdft", "relax", "file", "double", 0, 1, 1.0e-3, 1.0e-3, 1.0e-2, 8.0, 1.0e-6, 0.0), 1.0e-3); } @@ -87,24 +88,149 @@ TEST(HSolverSetDiagethrKsTest, LaterIterations) { // iteration 2 restarts from 1e-2 and is bounded by 0.1 * drho / nelec EXPECT_DOUBLE_EQ( - hsolver::set_diagethr_ks("pw", "ksdft", "scf", "atomic", "double", 0, 2, 1.0e-3, 1.0e-2, 1.0e-7, 8.0, 1.0e-6), + hsolver::set_diagethr_ks("pw", "ksdft", "scf", "atomic", "double", 0, 2, 1.0e-3, 1.0e-2, 1.0e-7, 8.0, 1.0e-6, 0.0), 1.25e-5); // later iterations only ever tighten the threshold EXPECT_DOUBLE_EQ( - hsolver::set_diagethr_ks("pw", "ksdft", "scf", "atomic", "double", 0, 3, 1.0e-3, 1.0e-2, 1.0e-6, 8.0, 1.0e-6), + hsolver::set_diagethr_ks("pw", "ksdft", "scf", "atomic", "double", 0, 3, 1.0e-3, 1.0e-2, 1.0e-6, 8.0, 1.0e-6, 0.0), 1.0e-6); EXPECT_DOUBLE_EQ( - hsolver::set_diagethr_ks("pw", "ksdft", "scf", "atomic", "double", 0, 3, 1.0e-3, 1.0e-2, 1.0e-4, 8.0, 1.0e-6), + hsolver::set_diagethr_ks("pw", "ksdft", "scf", "atomic", "double", 0, 3, 1.0e-3, 1.0e-2, 1.0e-4, 8.0, 1.0e-6, 0.0), 1.25e-5); } TEST(HSolverSetDiagethrKsTest, SinglePrecisionLowerBound) { EXPECT_DOUBLE_EQ( - hsolver::set_diagethr_ks("pw", "ksdft", "scf", "atomic", "single", 0, 3, 1.0e-3, 1.0e-2, 1.0e-6, 8.0, 1.0e-6), + hsolver::set_diagethr_ks("pw", "ksdft", "scf", "atomic", "single", 0, 3, 1.0e-3, 1.0e-2, 1.0e-6, 8.0, 1.0e-6, 0.0), 0.5e-4); } +TEST(HSolverSetDiagethrKsTest, ScfUpperBoundAppliesFromFirstIteration) +{ + for (const int iter : {1, 2, 3}) + { + const double threshold = hsolver::set_diagethr_ks( + "pw", "ksdft", "scf", "atomic", "double", 0, iter, 1.0, 1.0e-2, 1.0e-2, 8.0, 1.0e-6, 1.0e-12); + EXPECT_DOUBLE_EQ(threshold, 1.0e-12); + } +} + +TEST(HSolverSetDiagethrKsTest, ScfUpperBoundKeepsTighterThresholds) +{ + const double initial = hsolver::set_diagethr_ks( + "pw", "ksdft", "scf", "atomic", "double", 0, 1, 1.0, 1.0e-14, 1.0e-14, 8.0, 1.0e-6, 1.0e-12); + const double adaptive = hsolver::set_diagethr_ks( + "pw", "ksdft", "scf", "atomic", "double", 0, 2, 8.0e-14, 1.0e-2, 1.0e-12, 8.0, 1.0e-6, 1.0e-12); + const double previous = hsolver::set_diagethr_ks( + "pw", "ksdft", "scf", "atomic", "double", 0, 3, 1.0, 1.0e-2, 1.0e-14, 8.0, 1.0e-6, 1.0e-12); + EXPECT_DOUBLE_EQ(initial, 1.0e-14); + EXPECT_DOUBLE_EQ(adaptive, 1.0e-15); + EXPECT_DOUBLE_EQ(previous, 1.0e-14); +} + +TEST(HSolverSetDiagethrKsTest, ScfUpperBoundPreservesSinglePrecisionFloor) +{ + const double capped = hsolver::set_diagethr_ks( + "pw", "ksdft", "scf", "atomic", "single", 0, 2, 1.0, 1.0e-2, 1.0e-2, 8.0, 1.0e-6, 1.0e-4); + const double floor = hsolver::set_diagethr_ks( + "pw", "ksdft", "scf", "atomic", "single", 0, 3, 1.0e-12, 1.0e-2, 1.0e-8, 8.0, 1.0e-6, 5.0e-5); + EXPECT_DOUBLE_EQ(capped, 1.0e-4); + EXPECT_DOUBLE_EQ(floor, 5.0e-5); +} + +TEST(HSolverSetDiagethrKsTest, RelaxUpperBoundSurvivesEveryIonicCycle) +{ + // Carry the preceding electronic threshold into each new ionic cycle. + // Both the first-iteration initial floor and iteration-2 reset can loosen it. + for (const std::string& init_chg : {"atomic", "file"}) + { + double previous = 1.0e-2; + for (const int ionic_step : {0, 1, 2}) + { + for (const int iteration : {1, 2, 3}) + { + SCOPED_TRACE(::testing::Message() << init_chg << " ionic=" << ionic_step + << " electronic=" << iteration); + const double threshold = hsolver::set_diagethr_ks( + "pw", "ksdft", "relax", init_chg, "double", ionic_step, iteration, + 1.0, 1.0e-2, previous, 8.0, 1.0e-6, 1.0e-12); + EXPECT_DOUBLE_EQ(threshold, 1.0e-12); + previous = threshold; + } + } + } +} + +TEST(HSolverSetDiagethrKsTest, RelaxZeroUpperBoundPreservesLegacyIonicResets) +{ + const double residuals[] = {1.0, 8.0e-3, 8.0e-4}; + const double expected[] = {1.0e-3, 1.0e-4, 1.0e-5}; + double previous = 1.0e-2; + for (const int ionic_step : {0, 1, 2}) + { + for (int index = 0; index < 3; ++index) + { + const int iteration = index + 1; + SCOPED_TRACE(::testing::Message() << "ionic=" << ionic_step << " electronic=" << iteration); + const double threshold = hsolver::set_diagethr_ks( + "pw", "ksdft", "relax", "file", "double", ionic_step, iteration, + residuals[index], 1.0e-3, previous, 8.0, 1.0e-6, 0.0); + EXPECT_DOUBLE_EQ(threshold, expected[index]); + previous = threshold; + } + } +} + +TEST(HSolverSetDiagethrKsTest, RelaxUpperBoundKeepsTighterThresholds) +{ + for (const int ionic_step : {0, 1, 2}) + { + SCOPED_TRACE(ionic_step); + const double initial = hsolver::set_diagethr_ks( + "pw", "ksdft", "relax", "atomic", "double", ionic_step, 1, + 1.0, 1.0e-14, 1.0e-14, 8.0, 1.0e-6, 1.0e-12); + const double adaptive = hsolver::set_diagethr_ks( + "pw", "ksdft", "relax", "atomic", "double", ionic_step, 2, + 8.0e-14, 1.0e-2, 1.0e-12, 8.0, 1.0e-6, 1.0e-12); + const double previous = hsolver::set_diagethr_ks( + "pw", "ksdft", "relax", "atomic", "double", ionic_step, 3, + 1.0, 1.0e-2, 1.0e-14, 8.0, 1.0e-6, 1.0e-12); + EXPECT_DOUBLE_EQ(initial, 1.0e-14); + EXPECT_DOUBLE_EQ(adaptive, 1.0e-15); + EXPECT_DOUBLE_EQ(previous, 1.0e-14); + } +} + +TEST(HSolverSetDiagethrKsTest, RelaxUpperBoundPreservesSinglePrecisionFloor) +{ + const double capped = hsolver::set_diagethr_ks( + "pw", "ksdft", "relax", "atomic", "single", 1, 2, 1.0, 1.0e-2, 1.0e-2, 8.0, 1.0e-6, 1.0e-4); + const double floor = hsolver::set_diagethr_ks( + "pw", "ksdft", "relax", "atomic", "single", 1, 3, 1.0e-12, 1.0e-2, 1.0e-8, 8.0, 1.0e-6, 5.0e-5); + EXPECT_DOUBLE_EQ(capped, 1.0e-4); + EXPECT_DOUBLE_EQ(floor, 5.0e-5); +} + +TEST(HSolverSetDiagethrKsTest, UpperBoundDoesNotChangeOtherPolicies) +{ + const double nscf = hsolver::set_diagethr_ks( + "pw", "ksdft", "nscf", "atomic", "double", 0, 1, 1.0e-3, 1.0e-2, 1.0e-2, 8.0, 1.0e-6, 1.0e-12); + const double cell_relax = hsolver::set_diagethr_ks( + "pw", "ksdft", "cell-relax", "file", "double", 0, 1, 1.0e-3, 1.0e-3, 1.0e-2, 8.0, 1.0e-6, 1.0e-12); + const double md = hsolver::set_diagethr_ks( + "pw", "ksdft", "md", "file", "double", 0, 1, 1.0e-3, 1.0e-3, 1.0e-2, 8.0, 1.0e-6, 1.0e-12); + const double lcao = hsolver::set_diagethr_ks( + "lcao", "ksdft", "scf", "atomic", "double", 0, 1, 1.0e-3, 1.0e-2, 1.0e-2, 8.0, 1.0e-6, 1.0e-12); + const double sdft = hsolver::set_diagethr_ks( + "pw", "sdft", "scf", "atomic", "double", 0, 1, 1.0e-3, 1.0e-2, 1.0e-2, 8.0, 1.0e-6, 1.0e-12); + EXPECT_DOUBLE_EQ(nscf, 1.25e-8); + EXPECT_DOUBLE_EQ(cell_relax, 1.0e-3); + EXPECT_DOUBLE_EQ(md, 1.0e-3); + EXPECT_DOUBLE_EQ(lcao, 0.0); + EXPECT_DOUBLE_EQ(sdft, 0.0); +} + TEST(HSolverSetDiagethrSdftTest, NonPwSdftIsZero) { EXPECT_EQ( @@ -151,14 +277,24 @@ TEST(HSolverSetDiagethrSdftTest, LaterIterations) TEST(HSolverResetDiagEthrTest, PwKsdft) { - const std::string log_name = "test_hsolver_reset_diag_ethr.log"; + const std::string log_name = testing::TempDir() + "test_hsolver_reset_diag_ethr.log"; std::ofstream ofs(log_name.c_str()); - const double new_ethr = hsolver::reset_diag_ethr(ofs, "pw", "ksdft", "double", 1.0e-2, 1.0e-3, 1.0e-2, 8.0); - const double new_ethr_single = hsolver::reset_diag_ethr(ofs, "pw", "ksdft", "single", 1.0e-2, 1.0e-3, 1.0e-2, 8.0); + const double new_ethr = hsolver::reset_diag_ethr(ofs, "pw", "ksdft", "double", 1.0e-2, 1.0e-3, 1.0e-2, 8.0, 0.0); + const double new_ethr_single = hsolver::reset_diag_ethr(ofs, "pw", "ksdft", "single", 1.0e-2, 1.0e-3, 1.0e-2, 8.0, 0.0); + // A fractional electron count can otherwise make the reset exceed the bound. + const double capped = hsolver::reset_diag_ethr( + ofs, "pw", "ksdft", "double", 1.0e-8, 5.0e-9, 1.0e-8, 1.0e-2, 1.0e-8); + const double tighter = hsolver::reset_diag_ethr( + ofs, "pw", "ksdft", "double", 1.0e-2, 1.0e-12, 1.0e-8, 8.0, 1.0e-8); + const double single = hsolver::reset_diag_ethr( + ofs, "pw", "ksdft", "single", 1.0e-2, 1.0e-12, 1.0e-8, 8.0, 5.0e-5); ofs.close(); EXPECT_DOUBLE_EQ(new_ethr, 1.25e-5); EXPECT_DOUBLE_EQ(new_ethr_single, 0.5e-4); + EXPECT_DOUBLE_EQ(capped, 1.0e-8); + EXPECT_DOUBLE_EQ(tighter, 1.25e-14); + EXPECT_DOUBLE_EQ(single, 5.0e-5); std::ifstream ifs(log_name.c_str()); std::stringstream buffer; @@ -172,6 +308,6 @@ TEST(HSolverResetDiagEthrTest, PwKsdft) TEST(HSolverResetDiagEthrTest, OtherCasesAreZero) { std::ofstream ofs; - EXPECT_EQ(hsolver::reset_diag_ethr(ofs, "lcao", "ksdft", "double", 1.0e-2, 1.0e-3, 1.0e-2, 8.0), 0.0); - EXPECT_EQ(hsolver::reset_diag_ethr(ofs, "pw", "sdft", "double", 1.0e-2, 1.0e-3, 1.0e-2, 8.0), 0.0); + EXPECT_EQ(hsolver::reset_diag_ethr(ofs, "lcao", "ksdft", "double", 1.0e-2, 1.0e-3, 1.0e-2, 8.0, 0.0), 0.0); + EXPECT_EQ(hsolver::reset_diag_ethr(ofs, "pw", "sdft", "double", 1.0e-2, 1.0e-3, 1.0e-2, 8.0, 0.0), 0.0); } diff --git a/source/source_hsolver/unittests/test_hsolver_pw.cpp b/source/source_hsolver/unittests/test_hsolver_pw.cpp index 694349948d8..c41ac07aa29 100644 --- a/source/source_hsolver/unittests/test_hsolver_pw.cpp +++ b/source/source_hsolver/unittests/test_hsolver_pw.cpp @@ -6,12 +6,17 @@ #include "hsolver_pw_sup.h" #include "hsolver_supplementary_mock.h" #include "source_hsolver/diag_comm_info.h" +#include "source_hsolver/diago_ppcg.h" #include "source_hsolver/hs_operator.h" #include "source_hsolver/hsolver_lcaopw.h" #include "source_hsolver/hsolver_pw.h" +#include "source_basis/module_pw/test/test_tool.h" +#include "source_base/parallel_reduce.h" #include #include +#include +#include /// H = S = identity: the simplest operator the solvers can be handed template @@ -31,6 +36,27 @@ class IdentityHSOperator : public hsolver::HSOperator> +{ + public: + void hpsi(const std::complex* x, + std::complex* hx, + const int ld, + const int nvec) const override + { + block_sizes.push_back(nvec); + for (int ib = 0; ib < nvec; ++ib) + { + for (int ig = 0; ig < ld; ++ig) + { + hx[ib * ld + ig] = double(ig + 1) * x[ib * ld + ig]; + } + } + } + + mutable std::vector block_sizes; +}; + // Mock implementations for the template functions causing linking errors namespace ModulePW { // Mock implementation for recip_to_real @@ -172,14 +198,34 @@ class TestHSolverPW : public ::testing::Test { // class derived from it and friendship is not inherited, so the call into // the protected hamiltSolvePsiK() is routed through here. template - static void hamiltSolvePsiK(hsolver::HSolverPW& hs, + static bool hamiltSolvePsiK(hsolver::HSolverPW& hs, const hsolver::HSOperator& op, psi::Psi& ps, std::vector::type>& pre, typename GetTypeReal::type* eig, const int ntry) { - hs.hamiltSolvePsiK(op, ps, pre, eig, ntry); + return hs.hamiltSolvePsiK(op, ps, pre, eig, ntry).converged; + } + + template + static void set_thresholds(hsolver::HSolverPW& hs, const std::vector& thresholds) + { + hs.ethr_band = thresholds; + } + + static void report_failed_kpoints(hsolver::HSolverPW>& hs, std::ostream& log) + { + using Status = hsolver::HSolverPW>::DiagonalizationStatus; + std::vector statuses(3); + statuses[0].converged = false; + statuses[0].iteration_limit_reached = true; + statuses[0].active_band_count = 2; + statuses[0].max_residual = 0.125; + statuses[2].converged = false; + statuses[2].active_band_count = 1; + statuses[2].max_residual = 0.25; + hs.report_ppcg_convergence(statuses, log); } ModulePW::PW_Basis_K pwbk; @@ -198,6 +244,7 @@ class TestHSolverPW : public ::testing::Test { 0, false, 4, + 16, 0, 0); hsolver::HSolverPW, base_device::DEVICE_CPU> hs_d @@ -215,6 +262,7 @@ class TestHSolverPW : public ::testing::Test { 0, false, 4, + 16, 0, 0); @@ -473,3 +521,152 @@ TEST_F(TestHSolverPW, NpwxLessThanNbandsDeath) ".*" ); } + +TEST_F(TestHSolverPW, PpcgUsesIndependentConstructorConfiguration) +{ + using T = std::complex; + const int dim = 18; + const int nband = 3; + const int maxiter = 4; + const double residual_threshold = 1.0e-8; + const std::vector residual_thresholds = {1.0e-8, 4.0e-8, 9.0e-8}; + std::vector precondition(dim, 1.0); + std::vector initial(dim * nband); + for (size_t i = 0; i < initial.size(); ++i) + { + const double index = double(i + 1); + const double phase = index * 1.7; + const double real = std::sin(index); + const double imaginary = std::cos(phase); + initial[i] = T(real, imaginary); + } + + // Neither object calls setup_diago_params_pw. Dispatch must honor its own + // block size and RR interval, including after a different object was used. + hsolver::HSolverPW first(&pwbk, "scf", "pw", "ppcg", false, 1, 1, + maxiter, residual_threshold, false, nband, false, 1, 3, 0, 0); + hsolver::HSolverPW second(&pwbk, "scf", "pw", "ppcg", false, 1, 1, + maxiter, residual_threshold, false, nband, false, 2, 1, 0, 0); + set_thresholds(first, residual_thresholds); + set_thresholds(second, residual_thresholds); + hsolver::HSolverPW* solvers[] = {&first, &second, &first}; + const int block_sizes[] = {1, 2, 1}; + const int rr_steps[] = {3, 1, 3}; + for (int run = 0; run < 3; ++run) + { + psi::Psi wavefunctions(1, nband, dim, dim, true); + wavefunctions.set_all_psi(initial.data(), initial.size()); + std::vector reference = initial; + std::vector actual_eigenvalues(nband); + std::vector reference_eigenvalues(nband); + DiagonalPPCGOperator actual_operator; + DiagonalPPCGOperator reference_operator; + hsolver::DiagoPPCG direct( + residual_threshold, maxiter, block_sizes[run], rr_steps[run]); + const auto hpsi = [&](T* x, T* hx, const int ld, const int nvec) { + reference_operator.hpsi(x, hx, ld, nvec); + }; + const auto spsi = [&](T* x, T* sx, const int ld, const int nvec) { + reference_operator.spsi(x, sx, ld, nvec); + }; + direct.diag(hpsi, spsi, dim, nband, dim, reference.data(), + reference_eigenvalues.data(), residual_thresholds, precondition.data()); + const bool converged = hamiltSolvePsiK(*solvers[run], actual_operator, wavefunctions, + precondition, actual_eigenvalues.data(), 1); + EXPECT_EQ(converged, direct.converged()); + EXPECT_EQ(actual_operator.block_sizes, reference_operator.block_sizes); + for (int ib = 0; ib < nband; ++ib) + { + EXPECT_NEAR(actual_eigenvalues[ib], reference_eigenvalues[ib], 1.0e-12); + } + } +} + +TEST_F(TestHSolverPW, PpcgFailureUsesApplicationDiagnostic) +{ + using T = std::complex; + hsolver::HSolverPW solver(&pwbk, "scf", "pw", "ppcg", false, 1, 1, + 4, 1.0e-8, false, 1, false, 1, 3, 0, 0); + const std::vector thresholds(1, 1.0e-8); + set_thresholds(solver, thresholds); + psi::Psi wavefunctions(1, 1, 3, 3, true); + std::vector precondition(3, 1.0); + precondition[0] = std::numeric_limits::quiet_NaN(); + std::vector eigenvalues(1); + EXPECT_EXIT({ + std::cout.rdbuf(std::cerr.rdbuf()); + hamiltSolvePsiK(solver, hamilt_test_d, wavefunctions, + precondition, eigenvalues.data(), 1); + }, ::testing::ExitedWithCode(1), "PPCG diagonalization failed:.*preconditioner contains non-finite"); +} + +TEST_F(TestHSolverPW, PpcgPreservesRequestedResidualNorm) +{ + using T = std::complex; + const double threshold = 1.0e-8; + hsolver::HSolverPW solver(&pwbk, "scf", "pw", "ppcg", false, 1, 1, + 10, threshold, false, 1, false, 1, 3, 0, 0); + const std::vector thresholds(1, threshold); + set_thresholds(solver, thresholds); + psi::Psi wavefunctions(1, 1, 3, 3, true); + wavefunctions.zero_out(); + wavefunctions(0, 0) = T(1.0, 0.0); + wavefunctions(0, 1) = T(1.0e-5, 0.0); + std::vector precondition(3, 1.0); + std::vector eigenvalues(1); + DiagonalPPCGOperator op; + // This starting residual is below sqrt(threshold), but above threshold. + // Energy alone is already accurate to O(1e-10), so check the wavefunction. + EXPECT_TRUE(hamiltSolvePsiK(solver, op, wavefunctions, + precondition, eigenvalues.data(), 1)); + double squared_residual = 0.0; + for (int ig = 0; ig < 3; ++ig) + { + const T residual = (double(ig + 1) - eigenvalues[0]) * wavefunctions(0, ig); + squared_residual += std::norm(residual); + } + const double residual_norm = std::sqrt(squared_residual); + EXPECT_LE(residual_norm, threshold); +} + +TEST_F(TestHSolverPW, PpcgConvergenceWarningSummarizesKpoints) +{ + int contributing_pools = 1; + Parallel_Reduce::reduce_all(contributing_pools); + const int kpoints = 2 * contributing_pools; + const int bands = 3 * contributing_pools; + const std::string expected_kpoints = std::to_string(kpoints) + " k point(s)"; + const std::string expected_limits = std::to_string(contributing_pools) + " reached the iteration limit"; + const std::string expected_bands = "active bands across those k points = " + std::to_string(bands); + std::ostringstream log; + report_failed_kpoints(hs_d, log); + const std::string message = log.str(); + EXPECT_NE(message.find(expected_kpoints), std::string::npos); + EXPECT_NE(message.find(expected_limits), std::string::npos); + EXPECT_NE(message.find(expected_bands), std::string::npos); + EXPECT_NE(message.find("2.500000e-01"), std::string::npos); + EXPECT_EQ(std::count(message.begin(), message.end(), '\n'), 1); +} + +int main(int argc, char** argv) +{ +#ifdef __MPI + int nproc = 1; + int myrank = 0; + setupmpi(argc, argv, nproc, myrank); + int nproc_in_pool = 1; + int kpar = nproc; + int mypool = 0; + int rank_in_pool = 0; + divide_pools(nproc, myrank, nproc_in_pool, kpar, mypool, rank_in_pool); +#endif + testing::InitGoogleTest(&argc, argv); + // Re-execute death tests so they initialize their own MPI runtime instead + // of using MPI state inherited from a forked parent process. + testing::FLAGS_gtest_death_test_style = "threadsafe"; + const int result = RUN_ALL_TESTS(); +#ifdef __MPI + finishmpi(); +#endif + return result; +} diff --git a/source/source_io/module_parameter/input_parameter.h b/source/source_io/module_parameter/input_parameter.h index 672afba6478..6c008d2dc04 100644 --- a/source/source_io/module_parameter/input_parameter.h +++ b/source/source_io/module_parameter/input_parameter.h @@ -94,9 +94,11 @@ struct Input_para ///< use mesh, which is used in QE. int nspin = 1; ///< LDA ; LSDA ; non-linear spin int pw_diag_nmax = 50; - double pw_diag_thr = 0.01; ///< used in cg method + double pw_diag_thr = 0.01; ///< solver-dependent threshold for iterative PW solvers + double pw_diag_thr_max = 0.0; ///< optional upper bound on the adaptive PW SCF threshold; zero disables it bool diago_smooth_ethr = false; ///< smooth ethr for iter methods int pw_diag_ndim = 4; ///< dimension of workspace for Davidson diagonalization + int pw_diag_rr_step = 1; ///< full Rayleigh-Ritz interval for PPCG diagonalization int diago_cg_prec = 1; ///< mohan add 2012-03-31 int diag_subspace = 0; // 0: Lapack, 1: elpa, 2: scalapack bool use_k_continuity = false; ///< whether to use k-point continuity for initializing wave functions diff --git a/source/source_io/module_parameter/read_inp_estruc.cpp b/source/source_io/module_parameter/read_inp_estruc.cpp index 777dbdbfe78..a282e631e0c 100644 --- a/source/source_io/module_parameter/read_inp_estruc.cpp +++ b/source/source_io/module_parameter/read_inp_estruc.cpp @@ -3,6 +3,8 @@ #include "read_input.h" #include "read_input_tool.h" +#include + namespace ModuleIO { void ReadInput::item_elec_stru() @@ -37,9 +39,10 @@ void ReadInput::item_elec_stru() // Electronic Structure { Input_Item item("ks_solver"); - item.domain.set_choices({"cg", "dav", "bpcg", "dav_subspace", "genelpa", "elpa", "lapack", "scalapack_gvx", "cusolver", "cusolvermp", "pexsi"}); - item.domain.add_requirement("basis_type==pw", "ks_solver in [cg, dav, bpcg, dav_subspace]"); + item.domain.set_choices({"cg", "dav", "bpcg", "dav_subspace", "ppcg", "genelpa", "elpa", "lapack", "scalapack_gvx", "cusolver", "cusolvermp", "pexsi"}); + item.domain.add_requirement("basis_type==pw", "ks_solver in [cg, dav, bpcg, dav_subspace, ppcg]"); item.domain.add_requirement("basis_type==lcao", "ks_solver in [genelpa, elpa, lapack, scalapack_gvx, cusolver, cusolvermp, pexsi]"); + item.domain.add_requirement("ks_solver==ppcg", "esolver_type!=sdft"); item.domain.add_requirement("ks_solver==genelpa", "device==cpu"); item.annotation = "cg; dav; lapack; genelpa; elpa; scalapack_gvx; cusolver"; item.category = "Electronic structure"; @@ -50,7 +53,8 @@ For plane-wave basis, * cg: The conjugate-gradient (CG) method. * dav: The Davidson algorithm. -* dav_subspace: The Davidson algorithm without orthogonalization operation, this method is the most recommended for efficiency. `pw_diag_ndim` can be set to 2 for this method. +* dav_subspace: A block Davidson method. The CPU implementation orthogonalizes trial vectors in the overlap metric and filters linearly dependent expansion directions. `pw_diag_ndim` can be set to 2 for this method. +* ppcg: The projection preconditioned conjugate-gradient method. It is optimized and validated for CPU plane-wave calculations; non-CPU devices use a transitional host/device bridge. It is not supported for esolver_type=sdft, including mixed stochastic-deterministic calculations. * bpcg: The BPCG method, which is a block-parallel Conjugate Gradient (CG) method, typically exhibits higher acceleration in a GPU environment. The BPCG method is currently under testing and is not recommended for use. For numerical atomic orbitals basis, @@ -1000,10 +1004,10 @@ Use case: When experimental or high-level theoretical results suggest that the S } { Input_Item item("pw_diag_thr"); - item.annotation = "threshold for eigenvalues is cg electron iterations"; + item.annotation = "threshold for iterative PW diagonalization"; item.category = "Plane wave related variables"; item.type = "Real"; - item.description = "Only used when you use ks_solver = cg/dav/dav_subspace/bpcg. It indicates the threshold for the first electronic iteration, from the second iteration the pw_diag_thr will be updated automatically. For nscf calculations with planewave basis set, pw_diag_thr should be <= 1e-3."; + item.description = "Only used when you use ks_solver = cg/dav/dav_subspace/bpcg/ppcg. It indicates the threshold for the first electronic iteration, from the second iteration the pw_diag_thr will be updated automatically. The stopping criterion is solver-dependent: CG and Davidson methods compare eigenvalue changes against the threshold, whereas PPCG requires the norm of H psi - epsilon S psi to be at most max(ethr_band, ethr, 1e-14). The same numerical threshold therefore does not imply the same accuracy or cost across these methods. PPCG keeps the direct residual tolerance because an energy-only or square-root residual tolerance can leave forces and stresses insufficiently converged. For nscf calculations with planewave basis set, pw_diag_thr should be <= 1e-3."; item.default_value = "0.01"; item.unit = ""; read_sync_double(input.pw_diag_thr); @@ -1018,6 +1022,46 @@ Use case: When experimental or high-level theoretical results suggest that the S }; this->add_item(item); } + { + Input_Item item("pw_diag_thr_max"); + item.annotation = "optional upper bound on the adaptive PW SCF diagonalization threshold"; + item.category = "Plane wave related variables"; + item.type = "Real"; + item.description = R"(An optional upper bound on the diagonalization threshold selected during SCF iterations. Zero disables the bound and preserves the existing adaptive strategy. A positive value limits the threshold from the first electronic iteration of each self-consistent cycle onward, including every ionic step of relaxation, while allowing the adaptive strategy to tighten it further. It does not change the meaning or default value of pw_diag_thr, which sets the initial threshold. + +Positive values are supported only for basis_type=pw, esolver_type=ksdft, calculation=scf or relax, with diago_smooth_ethr=false. Smoothing is incompatible with this bound because it relaxes thresholds for weakly occupied bands. With precision=single, the bound must be at least 5e-5, the single-precision threshold floor. + +This bounds the SCF threshold supplied to the iterative solver, not a universal eigenvalue or residual error. Each solver retains its stopping criterion and numerical limits; in particular, PPCG has a minimum residual tolerance of 1e-14. Tight bounds can increase diagonalization cost or exhaust pw_diag_nmax. Validate convergence, energies, forces, and stresses for the intended calculation.)"; + item.default_value = "0"; + item.unit = ""; + item.set_availability("basis_type==pw and esolver_type==ksdft and calculation in [scf, relax]"); + read_sync_double(input.pw_diag_thr_max); + item.check_value = [](const Input_Item& item, const Parameter& para) { + const Input_para& input = para.inp; + if (!std::isfinite(input.pw_diag_thr_max) || input.pw_diag_thr_max < 0.0) + { + ModuleBase::WARNING_QUIT("ReadInput", "pw_diag_thr_max must be finite and nonnegative"); + } + if (input.pw_diag_thr_max == 0.0) + { + return; + } + const bool supports_bound = input.calculation == "scf" || input.calculation == "relax"; + if (input.basis_type != "pw" || input.esolver_type != "ksdft" || !supports_bound) + { + ModuleBase::WARNING_QUIT("ReadInput", "Positive pw_diag_thr_max requires PW KSDFT SCF or relaxation"); + } + if (input.diago_smooth_ethr) + { + ModuleBase::WARNING_QUIT("ReadInput", "Positive pw_diag_thr_max requires diago_smooth_ethr=false"); + } + if (input.precision == "single" && input.pw_diag_thr_max < 5.0e-5) + { + ModuleBase::WARNING_QUIT("ReadInput", "pw_diag_thr_max must be at least 5e-5 for single precision"); + } + }; + this->add_item(item); + } { Input_Item item("diago_smooth_ethr"); item.annotation = "smooth ethr for iter methods"; @@ -1044,27 +1088,46 @@ Use case: When experimental or high-level theoretical results suggest that the S } { Input_Item item("pw_diag_nmax"); - item.annotation = "max iteration number for cg"; + item.annotation = "maximum iteration count for iterative PW diagonalization"; item.category = "Plane wave related variables"; item.type = "Integer"; - item.description = "Only useful when you use ks_solver = cg/dav/dav_subspace/bpcg. It indicates the maximal iteration number for cg/david/dav_subspace/bpcg method."; + item.description = "Sets the iterative diagonalization limit for ks_solver = cg/dav/dav_subspace/ppcg in both SCF and NSCF calculations. NSCF uses this input, including its default of 50, rather than retaining a previous calculation's limit. CG and Davidson methods apply this limit within each internal diagonalization attempt and may retry; PPCG applies it to its subspace updates. It is not a common bound on total operator applications across solvers. Increasing the limit does not change convergence thresholds or guarantee convergence. BPCG currently does not use this parameter; its internal budget remains 24 updates in the first electronic iteration and 4 in later iterations."; item.default_value = "50"; item.unit = ""; - item.set_availability("basis_type==pw and ks_solver in [cg, dav, dav_subspace, bpcg]"); + item.set_availability("basis_type==pw and ks_solver in [cg, dav, dav_subspace, bpcg, ppcg]"); read_sync_int(input.pw_diag_nmax); this->add_item(item); } { Input_Item item("pw_diag_ndim"); - item.annotation = "dimension of workspace for Davidson diagonalization"; + item.annotation = "dimension of workspace for iterative PW diagonalization"; item.category = "Plane wave related variables"; item.type = "Integer"; - item.description = "Only useful when you use ks_solver = dav or ks_solver = dav_subspace. It indicates dimension of workspace(number of wavefunction packets, at least 2 needed) for the Davidson method. A larger value may yield a smaller number of iterations in the algorithm but uses more memory and more CPU time in subspace diagonalization."; + item.description = "Only useful when you use ks_solver = dav, dav_subspace, or ppcg. It indicates dimension of workspace(number of wavefunction packets, at least 2 needed) for the Davidson method, and the block size for the PPCG method. A larger value may yield a smaller number of iterations in the algorithm but uses more memory and more CPU time in subspace diagonalization."; item.default_value = "4"; item.unit = ""; + item.set_availability("basis_type==pw and ks_solver in [dav, dav_subspace, ppcg]"); read_sync_int(input.pw_diag_ndim); this->add_item(item); } + { + Input_Item item("pw_diag_rr_step"); + item.annotation = "full Rayleigh-Ritz interval for PPCG"; + item.category = "Plane wave related variables"; + item.type = "Integer"; + item.description = "Only useful when you use ks_solver = ppcg. It must be a positive integer and sets the interval, in subspace iterations, between full Rayleigh-Ritz diagonalizations. The conservative default of 1 preserves the established every-iteration Rayleigh-Ritz cadence. The solver also performs a full Rayleigh-Ritz step at initialization, after the startup global correction, at final validation or tentative convergence, or when overlap orthogonalization fails. Additional expanded global Rayleigh-Ritz steps may be triggered when discarded local-block directions compete with the requested lowest roots; this root recovery is conditional, not a separate fixed interval. Between these steps it maintains S-orthogonality without diagonalizing the full projected Hamiltonian. H and S images are refreshed independently every 16 iterations, with additional refreshes when search-direction projection loses precision or residuals approach a floating-point cancellation scale. Convergence and final residuals are checked using freshly applied operators. Larger Rayleigh-Ritz intervals reduce full projected eigensolves but can change the SCF trajectory or require more diagonalization iterations; validate energies, forces, stresses, and the iteration budget for the intended calculation."; + item.default_value = "1"; + item.unit = ""; + item.set_availability("basis_type==pw and ks_solver==ppcg"); + read_sync_int(input.pw_diag_rr_step); + item.check_value = [](const Input_Item& item, const Parameter& para) { + if (para.input.pw_diag_rr_step <= 0) + { + ModuleBase::WARNING_QUIT("ReadInput", "pw_diag_rr_step must be greater than 0"); + } + }; + this->add_item(item); + } { Input_Item item("diago_cg_prec"); item.annotation = "diago_cg_prec"; diff --git a/source/source_io/module_wf/write_wfc_pw.cpp b/source/source_io/module_wf/write_wfc_pw.cpp index 450bb596988..2192b73e6f0 100644 --- a/source/source_io/module_wf/write_wfc_pw.cpp +++ b/source/source_io/module_wf/write_wfc_pw.cpp @@ -213,14 +213,14 @@ void ModuleIO::write_wfc_pw( { ofs2 << "\n< Band " << ib + 1 << " >" << std::endl; } - ofs2 << std::scientific; + ofs2 << std::scientific << std::setprecision(16); for (int ig = 0; ig < ng; ig++) { if (ig % 4 == 0 && (ig != 0 || id != 0)) { ofs2 << "\n"; } - ofs2 << std::setw(15) << psi(ib, ig).real() << std::setw(15) << psi(ib, ig).imag(); + ofs2 << std::setw(25) << psi(ib, ig).real() << std::setw(25) << psi(ib, ig).imag(); } // end ig if (id == nproc_in_pool - 1 && npol == 1) { @@ -264,14 +264,14 @@ void ModuleIO::write_wfc_pw( { std::ofstream ofs2(wfilename[ik].c_str(), std::ios::app); - ofs2 << std::scientific; + ofs2 << std::scientific << std::setprecision(16); for (int ig = 0; ig < ng; ig++) { if (ig % 4 == 0 && (ig != 0 || id != 0)) { ofs2 << "\n"; } - ofs2 << std::setw(15) << psi(ib, ig + ng_max).real() << std::setw(15) + ofs2 << std::setw(25) << psi(ib, ig + ng_max).real() << std::setw(25) << psi(ib, ig + ng_max).imag(); } // end ig if (id == nproc_in_pool - 1) diff --git a/source/source_io/test_serial/read_input_item_test.cpp b/source/source_io/test_serial/read_input_item_test.cpp index b1221ff0b77..138e7955295 100644 --- a/source/source_io/test_serial/read_input_item_test.cpp +++ b/source/source_io/test_serial/read_input_item_test.cpp @@ -1,4 +1,5 @@ #include +#include #include #include #include @@ -126,6 +127,122 @@ TEST_F(InputTest, RelaxMethod) EXPECT_EQ(find_label("relax_new", readinput.input_lists), readinput.input_lists.end()); } +TEST_F(InputTest, PwDiagThresholdUpperBoundDefaultsAndParsing) +{ + ModuleIO::ReadInput readinput(0); + Parameter param; + Input_para& input = TestParameters::input(param); + auto it = find_label("pw_diag_thr_max", readinput.input_lists); + ASSERT_NE(it, readinput.input_lists.end()); + EXPECT_DOUBLE_EQ(input.pw_diag_thr_max, 0.0); + EXPECT_EQ(it->second.default_value, "0"); + EXPECT_EQ(it->second.type, "Real"); + EXPECT_EQ(it->second.get_availability(), "basis_type==pw and esolver_type==ksdft and calculation in [scf, relax]"); + it->second.check_value(it->second, param); + + it->second.str_values = {"1e-12"}; + it->second.read_value(it->second, param); + EXPECT_DOUBLE_EQ(input.pw_diag_thr_max, 1.0e-12); + EXPECT_DOUBLE_EQ(input.pw_diag_thr, 0.01); + it->second.check_value(it->second, param); + + it->second.str_values = {"0"}; + it->second.read_value(it->second, param); + input.basis_type = "lcao"; + input.esolver_type = "sdft"; + input.calculation = "nscf"; + input.diago_smooth_ethr = true; + input.precision = "single"; + it->second.check_value(it->second, param); + EXPECT_DOUBLE_EQ(input.pw_diag_thr_max, 0.0); +} + +TEST_F(InputTest, PwDiagThresholdUpperBoundAcceptsRelaxParsing) +{ + ModuleIO::ReadInput readinput(0); + Parameter param; + Input_para& input = TestParameters::input(param); + input.calculation = "relax"; + auto it = find_label("pw_diag_thr_max", readinput.input_lists); + ASSERT_NE(it, readinput.input_lists.end()); + it->second.str_values = {"1e-13"}; + it->second.read_value(it->second, param); + EXPECT_DOUBLE_EQ(input.pw_diag_thr_max, 1.0e-13); + EXPECT_DOUBLE_EQ(input.pw_diag_thr, 0.01); + // A rejected calculation terminates; isolate the validation so the old + // implementation reports a test failure without aborting the whole suite. + EXPECT_EXIT( + { + it->second.check_value(it->second, param); + std::exit(0); + }, + ::testing::ExitedWithCode(0), ""); +} + +TEST_F(InputTest, PwDiagThresholdUpperBoundRejectsInvalidNumbers) +{ + ModuleIO::ReadInput readinput(0); + Parameter param; + auto it = find_label("pw_diag_thr_max", readinput.input_lists); + ASSERT_NE(it, readinput.input_lists.end()); + for (const std::string& value : {"-1e-12", "nan", "inf", "-inf"}) + { + it->second.str_values = {value}; + it->second.read_value(it->second, param); + EXPECT_EXIT(it->second.check_value(it->second, param), ::testing::ExitedWithCode(1), ""); + } +} + +TEST_F(InputTest, PwDiagThresholdUpperBoundRejectsUnsupportedPolicies) +{ + ModuleIO::ReadInput readinput(0); + Parameter param; + Input_para& input = TestParameters::input(param); + input.pw_diag_thr_max = 1.0e-12; + const Input_para valid = input; + auto it = find_label("pw_diag_thr_max", readinput.input_lists); + ASSERT_NE(it, readinput.input_lists.end()); + it->second.check_value(it->second, param); + + for (const std::string& calculation : {"nscf", "cell-relax", "md"}) + { + input = valid; + input.calculation = calculation; + EXPECT_EXIT(it->second.check_value(it->second, param), ::testing::ExitedWithCode(1), ""); + } + for (const std::string& calculation : {"scf", "relax"}) + { + input = valid; + input.calculation = calculation; + input.basis_type = "lcao"; + EXPECT_EXIT(it->second.check_value(it->second, param), ::testing::ExitedWithCode(1), ""); + input = valid; + input.calculation = calculation; + input.esolver_type = "sdft"; + EXPECT_EXIT(it->second.check_value(it->second, param), ::testing::ExitedWithCode(1), ""); + input = valid; + input.calculation = calculation; + input.diago_smooth_ethr = true; + EXPECT_EXIT(it->second.check_value(it->second, param), ::testing::ExitedWithCode(1), ""); + } +} + +TEST_F(InputTest, PwDiagThresholdUpperBoundChecksSinglePrecisionFloor) +{ + ModuleIO::ReadInput readinput(0); + Parameter param; + Input_para& input = TestParameters::input(param); + input.precision = "single"; + auto it = find_label("pw_diag_thr_max", readinput.input_lists); + ASSERT_NE(it, readinput.input_lists.end()); + input.pw_diag_thr_max = 5.0e-5; + it->second.check_value(it->second, param); + input.pw_diag_thr_max = 1.0e-4; + it->second.check_value(it->second, param); + input.pw_diag_thr_max = 4.0e-5; + EXPECT_EXIT(it->second.check_value(it->second, param), ::testing::ExitedWithCode(1), ""); +} + TEST_F(InputTest, PccSelectionAndUnsupportedCombinations) { ModuleIO::ReadInput readinput(0); @@ -515,6 +632,20 @@ TEST_F(InputTest, Item_test) it->second.reset_value(it->second, param); EXPECT_EQ(TestParameters::input(param).pw_diag_thr, 1.0e-5); } + { // pw_diag_rr_step + auto it = find_label("pw_diag_rr_step", readinput.input_lists); + const Input_para defaults; + EXPECT_EQ(defaults.pw_diag_rr_step, 1); + EXPECT_EQ(it->second.default_value, "1"); + TestParameters::input(param).pw_diag_rr_step = 1; + EXPECT_NO_THROW(it->second.check_value(it->second, param)); + + TestParameters::input(param).pw_diag_rr_step = 16; + EXPECT_NO_THROW(it->second.check_value(it->second, param)); + + TestParameters::input(param).pw_diag_rr_step = 0; + EXPECT_EXIT(it->second.check_value(it->second, param), ::testing::ExitedWithCode(1), ""); + } { // nb2d auto it = find_label("nb2d", readinput.input_lists); TestParameters::input(param).nb2d = -1; @@ -988,6 +1119,18 @@ TEST_F(InputTest, Item_test) it->second.reset_value(it->second, param); EXPECT_EQ(TestParameters::input(param).ks_solver, "cg"); + TestParameters::input(param).ks_solver = "ppcg"; + TestParameters::input(param).basis_type = "pw"; + TestParameters::input(param).esolver_type = "ksdft"; + check_item(it->second, param); + EXPECT_EQ(TestParameters::input(param).ks_solver, "ppcg"); + + TestParameters::input(param).basis_type = "lcao"; + testing::internal::CaptureStdout(); + EXPECT_EXIT(check_item(it->second, param), ::testing::ExitedWithCode(1), ""); + output = testing::internal::GetCapturedStdout(); + EXPECT_THAT(output, testing::HasSubstr("Requires ks_solver in")); + TestParameters::input(param).ks_solver = "default"; TestParameters::input(param).basis_type = "lcao"; TestParameters::input(param).device = "gpu"; diff --git a/source/source_io/test_serial/test_read_input.cpp b/source/source_io/test_serial/test_read_input.cpp index 8dd695ec800..003536089b9 100644 --- a/source/source_io/test_serial/test_read_input.cpp +++ b/source/source_io/test_serial/test_read_input.cpp @@ -73,6 +73,50 @@ TEST_F(InputDomainReaderTest, KeepsExplicitValuesAndContextualDefaults) EXPECT_EQ(input.inp.gga_grad, 0); } +TEST_F(InputDomainReaderTest, ReadsAndWritesPpcgWithTheUnifiedDomainRules) +{ + write_input("ks_solver ppcg\n"); + ModuleIO::ReadInput reader(0); + reader.check_ntype_flag = false; + Parameter input; + reader.read_parameters(input, filename); + EXPECT_EQ(input.inp.ks_solver, "ppcg"); + + const std::string written = filename + ".written"; + reader.write_parameters(input, written); + ModuleIO::ReadInput rereader(0); + rereader.check_ntype_flag = false; + Parameter restored; + rereader.read_parameters(restored, written); + EXPECT_EQ(restored.inp.ks_solver, "ppcg"); + + const ModuleIO::ParameterMetadata metadata = ModuleIO::ParameterHelp::get_metadata("ks_solver"); + EXPECT_NE(metadata.allowed_values.find("ppcg"), std::string::npos); +} + +TEST_F(InputDomainReaderTest, RejectsStochasticPpcgWithOrWithoutDeterministicBands) +{ + for (const int bands : {0, 4}) + { + const std::string settings = "esolver_type sdft\nnbands " + std::to_string(bands) + + "\nnbands_sto 8\nsmearing_method fd\n"; + const std::string invalid = settings + "ks_solver ppcg\n"; + const std::string output = reject_input(invalid); + EXPECT_NE(output.find("ks_solver"), std::string::npos); + EXPECT_NE(output.find("Requires esolver_type!=sdft"), std::string::npos); + + const std::string supported = settings + "ks_solver cg\n"; + write_input(supported); + ModuleIO::ReadInput reader(0); + reader.check_ntype_flag = false; + Parameter input; + reader.read_parameters(input, filename); + EXPECT_EQ(input.inp.ks_solver, "cg"); + EXPECT_EQ(input.inp.esolver_type, "sdft"); + EXPECT_EQ(input.inp.nbands, bands); + } +} + TEST_F(InputDomainReaderTest, RejectsMalformedAndOutOfDomainExplicitValues) { for (const std::string value : {"3", "0", "-1", "2abc", "2.0", "2 4", "999999999999999999999"}) diff --git a/source/source_io/unittests/test_read_input_parallel.cpp b/source/source_io/unittests/test_read_input_parallel.cpp index 3c8ce4c54bd..c3a4804542c 100644 --- a/source/source_io/unittests/test_read_input_parallel.cpp +++ b/source/source_io/unittests/test_read_input_parallel.cpp @@ -155,6 +155,7 @@ TEST_F(InputParaTest, ParaRead) EXPECT_EQ(param.inp.pw_diag_nmax, 50); EXPECT_EQ(param.inp.diago_cg_prec, 1); EXPECT_EQ(param.inp.pw_diag_ndim, 4); + EXPECT_EQ(param.inp.pw_diag_rr_step, 1); EXPECT_DOUBLE_EQ(param.inp.pw_diag_thr, 1.0e-2); EXPECT_FALSE(param.inp.diago_smooth_ethr); EXPECT_EQ(param.inp.nb2d, 0); diff --git a/source/source_lcao/module_deltaspin/deltaspin_pw_mi.cpp b/source/source_lcao/module_deltaspin/deltaspin_pw_mi.cpp index d692ecda7eb..bcc93656dff 100644 --- a/source/source_lcao/module_deltaspin/deltaspin_pw_mi.cpp +++ b/source/source_lcao/module_deltaspin/deltaspin_pw_mi.cpp @@ -359,23 +359,25 @@ void update_psi_charge_pw_cpu(ScState& state, { // Full PW diagonalization: subspace rotation provides a good initial guess, // then HSolverPW iteratively refines psi in the full plane-wave space and calls psiToRho. + const Input_para& input = PARAM.inp; hsolver::HSolverPW, base_device::DEVICE_CPU> hsolver_pw_obj( pw_wfc, - PARAM.inp.calculation, - PARAM.inp.basis_type, - PARAM.inp.ks_solver, + input.calculation, + input.basis_type, + input.ks_solver, PARAM.globalv.use_uspp, - PARAM.inp.nspin, + input.nspin, hsolver::DiagoIterAssist>::SCF_ITER, hsolver::DiagoIterAssist>::PW_DIAG_NMAX, hsolver::DiagoIterAssist>::PW_DIAG_THR, hsolver::DiagoIterAssist>::need_subspace, - PARAM.inp.nbands, - PARAM.inp.diago_smooth_ethr, - PARAM.inp.pw_diag_ndim, - PARAM.inp.diag_subspace, - PARAM.inp.nb2d, - PARAM.inp.use_k_continuity); + input.nbands, + input.diago_smooth_ethr, + input.pw_diag_ndim, + input.pw_diag_rr_step, + input.diag_subspace, + input.nb2d, + input.use_k_continuity); hamilt::HamiltHSOperator, base_device::DEVICE_CPU> op(hamilt_t, pw_wfc); hsolver_pw_obj.solve(op, @@ -472,23 +474,25 @@ void update_psi_charge_pw_gpu(ScState& state, // STAGE 2: Full-space update (GPU) if (pw_solve) { + const Input_para& input = PARAM.inp; hsolver::HSolverPW, base_device::DEVICE_GPU> hsolver_pw_obj( pw_wfc, - PARAM.inp.calculation, - PARAM.inp.basis_type, - PARAM.inp.ks_solver, + input.calculation, + input.basis_type, + input.ks_solver, PARAM.globalv.use_uspp, - PARAM.inp.nspin, + input.nspin, hsolver::DiagoIterAssist, base_device::DEVICE_GPU>::SCF_ITER, hsolver::DiagoIterAssist, base_device::DEVICE_GPU>::PW_DIAG_NMAX, hsolver::DiagoIterAssist, base_device::DEVICE_GPU>::PW_DIAG_THR, hsolver::DiagoIterAssist, base_device::DEVICE_GPU>::need_subspace, - PARAM.inp.nbands, - PARAM.inp.diago_smooth_ethr, - PARAM.inp.pw_diag_ndim, - PARAM.inp.diag_subspace, - PARAM.inp.nb2d, - PARAM.inp.use_k_continuity); + input.nbands, + input.diago_smooth_ethr, + input.pw_diag_ndim, + input.pw_diag_rr_step, + input.diag_subspace, + input.nb2d, + input.use_k_continuity); hamilt::HamiltHSOperator, base_device::DEVICE_GPU> op(hamilt_t, pw_wfc); hsolver_pw_obj.solve(op, diff --git a/source/source_pw/module_stodft/sto_hsolver_pw.h b/source/source_pw/module_stodft/sto_hsolver_pw.h index 8c9b8cef9c2..abebc50fc57 100644 --- a/source/source_pw/module_stodft/sto_hsolver_pw.h +++ b/source/source_pw/module_stodft/sto_hsolver_pw.h @@ -27,6 +27,7 @@ class StoHSolverPW : public hsolver::HSolverPW const int nbands_in, const bool diago_smooth_ethr_in, const int pw_diag_ndim_in, + const int pw_diag_rr_step_in, const int diag_subspace_in, const int nb2d_in, const bool ks_run_in, @@ -45,6 +46,7 @@ class StoHSolverPW : public hsolver::HSolverPW nbands_in, diago_smooth_ethr_in, pw_diag_ndim_in, + pw_diag_rr_step_in, diag_subspace_in, nb2d_in), ks_run(ks_run_in), all_ks_run(all_ks_run_in), bndpar(bndpar_in) diff --git a/source/source_pw/module_stodft/test/CMakeLists.txt b/source/source_pw/module_stodft/test/CMakeLists.txt index 8abb1f159a0..2977f2ae70d 100644 --- a/source/source_pw/module_stodft/test/CMakeLists.txt +++ b/source/source_pw/module_stodft/test/CMakeLists.txt @@ -50,6 +50,7 @@ addtest( ../sto_hsolver_pw.cpp ../../../source_hsolver/hsolver_pw.cpp ../../../source_hsolver/diago_bpcg.cpp + ../../../source_hsolver/diago_ppcg.cpp ../../../source_hsolver/diago_dav_subspace.cpp ../../../source_hsolver/diag_const_nums.cpp ../../../source_hsolver/diago_iter_assist.cpp diff --git a/source/source_pw/module_stodft/test/test_sto_hsolver_pw.cpp b/source/source_pw/module_stodft/test/test_sto_hsolver_pw.cpp index 0a361000f7d..02e8a4801f4 100644 --- a/source/source_pw/module_stodft/test/test_sto_hsolver_pw.cpp +++ b/source/source_pw/module_stodft/test/test_sto_hsolver_pw.cpp @@ -299,6 +299,7 @@ class TestStoHSolverPW : public ::testing::Test 0, false, 4, + 16, 0, 0, false, diff --git a/tests/01_PW/006_PW_UPF201_Eu/INPUT b/tests/01_PW/006_PW_UPF201_Eu/INPUT index f2ba550f072..e848c120dee 100644 --- a/tests/01_PW/006_PW_UPF201_Eu/INPUT +++ b/tests/01_PW/006_PW_UPF201_Eu/INPUT @@ -5,10 +5,11 @@ ntype 1 nspin 1 pseudo_dir ../../PP_ORB symmetry 1 -nbands 10 +nbands 20 ecutwfc 20 -scf_thr 1.0e-7 // about iteration +scf_thr 1.0e-11 +pw_diag_nmax 200 basis_type pw smearing_method gauss diff --git a/tests/01_PW/006_PW_UPF201_Eu/README b/tests/01_PW/006_PW_UPF201_Eu/README index d502e127c02..006dc4bb2a1 100644 --- a/tests/01_PW/006_PW_UPF201_Eu/README +++ b/tests/01_PW/006_PW_UPF201_Eu/README @@ -1 +1,48 @@ -test upf201 pseudopotential which has no cutoff_radius_index and multiserial