diff --git a/docs/advanced/scf/hsolver.md b/docs/advanced/scf/hsolver.md index cc85b4d87be..184524c7d18 100644 --- a/docs/advanced/scf/hsolver.md +++ b/docs/advanced/scf/hsolver.md @@ -4,7 +4,7 @@ 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 targets the many-eigenpair regime where the bounded memory and block operations pay off; it is optimized and validated for CPU plane-wave calculations, and non-CPU devices use a transitional host/device bridge. For PPCG, `pw_diag_ndim` controls the block size and `pw_diag_rr_step` controls how often the Hamiltonian and overlap images are refreshed after Rayleigh--Ritz rotations. 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 43a4072889c..91dfb0d987c 100644 --- a/source/Makefile.Objects +++ b/source/Makefile.Objects @@ -474,6 +474,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_pw.cpp b/source/source_esolver/esolver_ks_pw.cpp index 7355d60a8c8..03ce5809ff2 100644 --- a/source/source_esolver/esolver_ks_pw.cpp +++ b/source/source_esolver/esolver_ks_pw.cpp @@ -287,6 +287,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 457c9483c89..854ca22f1b6 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/test/elecstate_print_test.cpp b/source/source_estate/test/elecstate_print_test.cpp index c2d5bec1743..6eb6d5bfcab 100644 --- a/source/source_estate/test/elecstate_print_test.cpp +++ b/source/source_estate/test/elecstate_print_test.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_hsolver/CMakeLists.txt b/source/source_hsolver/CMakeLists.txt index e7adc3fb5ad..f8edba71f45 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_pw_tddft.cpp diff --git a/source/source_hsolver/diago_params.cpp b/source/source_hsolver/diago_params.cpp index 28e1040a974..c2dd860d9f1 100644 --- a/source/source_hsolver/diago_params.cpp +++ b/source/source_hsolver/diago_params.cpp @@ -15,7 +15,6 @@ 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; diff --git a/source/source_hsolver/diago_ppcg.cpp b/source/source_hsolver/diago_ppcg.cpp new file mode 100644 index 00000000000..001d085c21c --- /dev/null +++ b/source/source_hsolver/diago_ppcg.cpp @@ -0,0 +1,2023 @@ +#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 + +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 + +// ============================================================================= +// 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 (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 +void 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; + 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 < nstates; ++state) + { + subspace.eval[state] = reduced_eigenvalues[state]; + } + ModuleBase::gemm_op()('N', + 'N', + dim, + nstates, + retained_dim, + &one, + retained_basis.data(), + dim, + reduced_hamiltonian.data(), + retained_dim, + &zero, + subspace.k.data(), + dim); + local_rank_truncation_succeeded = true; + } + } + catch (const std::runtime_error&) + { + // The shifted generalized solve below remains the fallback. + } + if (all_pool_operations_succeeded(local_rank_truncation_succeeded)) + { + return; + } + + // 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; + } + } + // 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)); + } +} + +// --------------------------------------------------------------------------- +// 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]; + } + } + 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()); + 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; + + 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_) + { + // 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 = std::max(1, (nact + sbsize_ - 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); + solve_small_generalized(local_dim, l, subspace); + 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); + 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..5d7cbd4821a --- /dev/null +++ b/source/source_hsolver/diago_ppcg.h @@ -0,0 +1,286 @@ +#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; + +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_; + } + +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; + + // 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; + }; + + 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; + + void solve_small_generalized(const int dim, + const int nstates, + SmallSubspace& subspace) const; + + 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_pw.cpp b/source/source_hsolver/hsolver_pw.cpp index 1b316b567e0..8b2c1d1699c 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,194 @@ #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 #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); + // The host solver uses int indices. Check the product before allocation or + // copying, and reject it on every pool rank before any operator callback. + const int max_elements = std::numeric_limits::max(); + int invalid_dimensions = ld_psi <= 0 || nband <= 0 || dim <= 0 || dim > ld_psi + || nband > max_elements / ld_psi ? 1 : 0; + Parallel_Reduce::reduce_pool(&invalid_dimensions, 1); + if (invalid_dimensions != 0) + { + ModuleBase::WARNING_QUIT("PPCG GPU bridge", "Invalid or oversized wavefunction 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 +259,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 +270,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 +319,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 +357,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 +372,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 +417,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 +542,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 d5cbb772a8c..44c24c70049 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/test/CMakeLists.txt b/source/source_hsolver/test/CMakeLists.txt index f9e8749a12a..451ef22d531 100644 --- a/source/source_hsolver/test/CMakeLists.txt +++ b/source/source_hsolver/test/CMakeLists.txt @@ -40,7 +40,7 @@ if (ENABLE_MPI) AddTest( TARGET MODULE_HSOLVER_bpcg LIBS parameter base psi device container - SOURCES diago_bpcg_test.cpp ../diago_bpcg.cpp ../para_lin_tf.cpp ../diago_iter_assist.cpp + SOURCES diago_bpcg_test.cpp ../diago_bpcg.cpp ../para_lin_tf.cpp ../diago_iter_assist.cpp ../../source_basis/module_pw/test/test_tool.cpp ) AddTest( @@ -91,9 +91,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) @@ -130,6 +132,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() + install(FILES H-KPoints-Si2.dat DESTINATION ${CMAKE_CURRENT_BINARY_DIR}) install(FILES H-GammaOnly-Si2.dat DESTINATION ${CMAKE_CURRENT_BINARY_DIR}) install(FILES S-KPoints-Si2.dat DESTINATION ${CMAKE_CURRENT_BINARY_DIR}) @@ -147,6 +182,7 @@ install(FILES KPoints-Si64-Solution.dat DESTINATION ${CMAKE_CURRENT_BINARY_DIR}) install(FILES diago_cg_parallel_test.sh DESTINATION ${CMAKE_CURRENT_BINARY_DIR}) install(FILES diago_david_parallel_test.sh DESTINATION ${CMAKE_CURRENT_BINARY_DIR}) install(FILES diago_lcao_parallel_test.sh DESTINATION ${CMAKE_CURRENT_BINARY_DIR}) +install(FILES diago_ppcg_parallel_test.sh DESTINATION ${CMAKE_CURRENT_BINARY_DIR}) install(FILES PEXSI-H-GammaOnly-Si2.dat DESTINATION ${CMAKE_CURRENT_BINARY_DIR}) install(FILES PEXSI-S-GammaOnly-Si2.dat DESTINATION ${CMAKE_CURRENT_BINARY_DIR}) @@ -206,6 +242,10 @@ if (ENABLE_MPI) COMMAND ${BASH} diago_david_parallel_test.sh 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/test/diago_ppcg_parallel_test.sh b/source/source_hsolver/test/diago_ppcg_parallel_test.sh new file mode 100644 index 00000000000..de1ec4f1afe --- /dev/null +++ b/source/source_hsolver/test/diago_ppcg_parallel_test.sh @@ -0,0 +1,43 @@ +#!/bin/bash + +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 ./MODULE_HSOLVER_ppcg_parallel + 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 ./MODULE_HSOLVER_ppcg_parallel 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 ./MODULE_HSOLVER_ppcg_parallel + 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/test/test_diago_compare.cpp b/source/source_hsolver/test/test_diago_compare.cpp new file mode 100644 index 00000000000..26db45bb6a4 --- /dev/null +++ b/source/source_hsolver/test/test_diago_compare.cpp @@ -0,0 +1,523 @@ +/** + * 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 per-band convergence threshold, and wall time is measured to the SAME + * reference accuracy (max_eval_err < err_target) so the timings are directly + * comparable despite the solvers' differing internal stopping rules. + * + * 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 +#include + +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; + +// --------------------------------------------------------------------------- +// Unified stopping criterion. Wall time is measured as "time to reach +// max_eval_err < err_target", the same LAPACK-reference accuracy for every +// solver. Each solver's own stopping rule (eigenvalue-change for PPCG/CG, +// residual for BPCG/Davidson) is looser than the reference error, so every +// solver is re-driven for up to max_outer_passes calls to diag(). max_err +// always reports the accuracy actually reached, so a solver that fails to hit +// err_target within the budget stays visible. +// +// NOTE: err_target must be strictly coarser than every solver's own stopping +// point. The CG solver stops on |eigenvalue change| < ethr, which in practice +// leaves max_eval_err ~ (2-4)x ethr; a target tighter than that (e.g. 1e-6 +// with ethr=1e-6) can never be reached, so CG burns all max_outer_passes doing +// an expensive subspace restart + per-band CG each round. A relaxed target +// keeps the cross-solver comparison honest instead of penalizing CG for being +// re-driven into repeated subspace restarts. +// --------------------------------------------------------------------------- +const double err_target = 1e-5; // relaxed: reachable by every solver +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; +} + +// Total heap memory currently allocated (bytes). Used to compare the peak +// working memory of the solvers: PPCG keeps a bounded subspace, while +// Davidson grows its basis with the number of iterations. +static long heap_bytes() +{ + struct mallinfo2 mi = mallinfo2(); + return static_cast(mi.uordblks) + static_cast(mi.hblkhd); +} + +// 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 + long mem_bytes = 0; // peak heap memory allocated by the solver + 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); + long mem0 = 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.mem_bytes = heap_bytes() - mem0; + 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); + long mem0 = 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.mem_bytes = heap_bytes() - mem0; + 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); + long mem0 = 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); }); + // BPCG::diag() is a single block-CG sweep; iterate until convergence. + 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.mem_bytes = heap_bytes() - mem0; + 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); + long mem0 = 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's diag() already iterates its growing subspace to convergence, + // so it must be called exactly once; a re-drive loop would reuse the stale + // Ritz basis and trigger a rank-deficient Schmidt orthogonalization. + 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.mem_bytes = heap_bytes() - mem0; + 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("%-5s %-5s %-6s %-10s %-14s %-10s %-12s\n", "n", "nband", "bw", "solver", "wall_time(s)", + "max_err", "mem(MB)"); + 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; + + std::printf("%-5d %-5d %-6d %-10s %-14.5f %-10.2e %-12.2f\n", c.n, c.nband, c.bw, "PPCG", r_ppcg.wall_s, + r_ppcg.max_err, r_ppcg.mem_bytes / 1048576.0); + std::printf("%-5s %-5s %-6s %-10s %-14.5f %-10.2e %-12.2f\n", "", "", "", "CG", r_cg.wall_s, + r_cg.max_err, r_cg.mem_bytes / 1048576.0); + std::printf("%-5s %-5s %-6s %-10s %-14.5f %-10.2e %-12.2f\n", "", "", "", "BPCG", r_bpcg.wall_s, + r_bpcg.max_err, r_bpcg.mem_bytes / 1048576.0); + std::printf("%-5s %-5s %-6s %-10s %-14.5f %-10.2e %-12.2f\n", "", "", "", "Davidson", r_dav.wall_s, + r_dav.max_err, r_dav.mem_bytes / 1048576.0); + std::printf("-----------------------------------------------------------------\n"); + } + + const int exit_code = all_ok ? 0 : 1; + finishmpi(); + return exit_code; +} diff --git a/source/source_hsolver/test/test_diago_params.cpp b/source/source_hsolver/test/test_diago_params.cpp new file mode 100644 index 00000000000..1d038f97f41 --- /dev/null +++ b/source/source_hsolver/test/test_diago_params.cpp @@ -0,0 +1,51 @@ +#include "../diago_iter_assist.h" +#include "../diago_params.h" + +#include + +#include + +TEST(DiagoParamsPWTest, AppliesPwDiagNmaxForScf) +{ + Input_para input; + input.calculation = "scf"; + input.ks_solver = "ppcg"; + input.pw_diag_nmax = 73; + + typedef hsolver::DiagoIterAssist, base_device::DEVICE_CPU> Assist; + const int saved_nmax = Assist::PW_DIAG_NMAX; + Assist::PW_DIAG_NMAX = 19; + + hsolver::setup_diago_params_pw, base_device::DEVICE_CPU>( + 0, + 1, + 1.0e-8, + input); + + EXPECT_EQ(Assist::PW_DIAG_NMAX, 73); + Assist::PW_DIAG_NMAX = saved_nmax; +} + +TEST(DiagoParamsPWTest, PreservesSharedNmaxForNscf) +{ + Input_para input; + input.calculation = "nscf"; + input.pw_diag_nmax = 73; + + typedef hsolver::DiagoIterAssist, base_device::DEVICE_CPU> Assist; + const int saved_nmax = Assist::PW_DIAG_NMAX; + Assist::PW_DIAG_NMAX = 19; + + const char* solvers[] = {"cg", "dav", "dav_subspace", "bpcg", "ppcg"}; + for (const char* solver : solvers) + { + input.ks_solver = solver; + hsolver::setup_diago_params_pw, base_device::DEVICE_CPU>( + 0, + 1, + 1.0e-8, + input); + EXPECT_EQ(Assist::PW_DIAG_NMAX, 19) << solver; + } + Assist::PW_DIAG_NMAX = saved_nmax; +} diff --git a/source/source_hsolver/test/test_diago_ppcg.cpp b/source/source_hsolver/test/test_diago_ppcg.cpp new file mode 100644 index 00000000000..3156bf82eea --- /dev/null +++ b/source/source_hsolver/test/test_diago_ppcg.cpp @@ -0,0 +1,3847 @@ +/** + * 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_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, ReusesImagesForWellSeparatedDirection) +{ + check_projection(false, false); +} + +TEST_F(DiagoPPCGProjectionTest, ClearsStaleImagesForZeroDirection) +{ + check_projection(false, true); +} + +} // 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, nband, 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. + 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 expected_ritz = 1 + scheduled + final_extra + startup_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.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, nband, 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(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(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, nband, 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, nband, 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, nband, 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, nband, 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, nband, 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), spsi(n_dim), 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, nband, 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, nband, 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"; + } + + // --- Residual check: ||Hψ_i - ε_i ψ_i|| < sqrt(ethr) --- + // The eigenvalue-change convergence criterion targets eigenvalue error ~ethr, + // so the eigenvector residual is naturally ~sqrt(ethr) = 1e-4 for ethr=1e-8. + std::vector hpsi(n_dim), 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, nband, 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, nband, 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, nband, 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, nband, 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, nband, 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, nband, 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, nband, 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, nband, 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, nband, 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, nband, 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, nband, 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); + // Exact eigenvalues unknown analytically for generalized problem + // with non-diagonal S. Just check convergence via residual. + exact = {0.0, 0.0}; + 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, nband, 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); }; + double avg_iter = solver.diag(h_op, spsi_func, ld, nband, n_dim, psi_run.data(), eval.data(), ethr, prec.data()); + // Check eigenvalues are positive and reasonable + for (int i = 0; i < nband; ++i) + { + EXPECT_GT(eval[i], 0.0); + EXPECT_LT(eval[i], 5.0); + } + // Residual check + std::vector hpsi(n_dim), spsi(n_dim), 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 rn = column_norm(res.data(), n_dim); + EXPECT_LT(rn, 1e-4) << "TridiagS BLOCK: residual[" << i << "] too large: " << rn; + } + 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/test/test_diago_ppcg_float.cpp b/source/source_hsolver/test/test_diago_ppcg_float.cpp new file mode 100644 index 00000000000..a2701cb9ec3 --- /dev/null +++ b/source/source_hsolver/test/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/test/test_diago_ppcg_parallel.cpp b/source/source_hsolver/test/test_diago_ppcg_parallel.cpp new file mode 100644 index 00000000000..b93b37f37b5 --- /dev/null +++ b/source/source_hsolver/test/test_diago_ppcg_parallel.cpp @@ -0,0 +1,229 @@ +/** + * 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 + +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; + + 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 + || 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 (!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()); + } + + 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/test/test_hsolver_pw.cpp b/source/source_hsolver/test/test_hsolver_pw.cpp index 694349948d8..c41ac07aa29 100644 --- a/source/source_hsolver/test/test_hsolver_pw.cpp +++ b/source/source_hsolver/test/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 8f07b087a3f..39eb1b194c9 100644 --- a/source/source_io/module_parameter/input_parameter.h +++ b/source/source_io/module_parameter/input_parameter.h @@ -94,9 +94,10 @@ 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 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 5da535085ca..63c3c2b5163 100644 --- a/source/source_io/module_parameter/read_inp_estruc.cpp +++ b/source/source_io/module_parameter/read_inp_estruc.cpp @@ -54,6 +54,7 @@ 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. +* 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. * 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, @@ -129,7 +130,7 @@ Then the user has to correct the input file and restart the calculation.)"; }; item.check_value = [](const Input_Item& item, const Parameter& para) { const std::string& ks_solver = para.input.ks_solver; - const std::vector pw_solvers = {"cg", "dav", "bpcg", "dav_subspace"}; + const std::vector pw_solvers = {"cg", "dav", "bpcg", "dav_subspace", "ppcg"}; const std::vector lcao_solvers = { "genelpa", "elpa", @@ -1048,10 +1049,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); @@ -1108,27 +1109,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 = "Only useful when you use ks_solver = cg/dav/dav_subspace/bpcg/ppcg. It indicates the maximal iteration number for cg/david/dav_subspace/bpcg/ppcg method. The shared PW NSCF setup retains the existing iteration limit (30 in a fresh run) for all solvers instead of updating it from this input."; 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. 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/test/read_input_ptest.cpp b/source/source_io/test/read_input_ptest.cpp index 30398cd7796..b042d66c03e 100644 --- a/source/source_io/test/read_input_ptest.cpp +++ b/source/source_io/test/read_input_ptest.cpp @@ -154,6 +154,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_io/test_serial/read_input_item_test.cpp b/source/source_io/test_serial/read_input_item_test.cpp index 6107585e733..451687242f4 100644 --- a/source/source_io/test_serial/read_input_item_test.cpp +++ b/source/source_io/test_serial/read_input_item_test.cpp @@ -412,6 +412,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; @@ -885,6 +899,11 @@ 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"; + it->second.check_value(it->second, param); + EXPECT_EQ(TestParameters::input(param).ks_solver, "ppcg"); + TestParameters::input(param).ks_solver = "default"; TestParameters::input(param).basis_type = "lcao"; TestParameters::input(param).device = "gpu"; 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 3e933f7d1cd..76bd3c1359a 100644 --- a/source/source_pw/module_stodft/test/CMakeLists.txt +++ b/source/source_pw/module_stodft/test/CMakeLists.txt @@ -22,6 +22,7 @@ AddTest( LIBS parameter psi device base container MPI::MPI_CXX SOURCES test_sto_hsolver_pw.cpp ../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 ../../../source_hsolver/para_lin_tf.cpp ../../../source_estate/elecstate_tools.cpp ../../../source_estate/occupy.cpp @@ -29,4 +30,4 @@ AddTest( # This test calls MPI_Init in main() and its mocks take MPI_Comm arguments # unconditionally, so it must keep __MPI even though this directory disables it. KEEP_FEATURE_DEFINITIONS __MPI -) \ No newline at end of file +) 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 22292183184..a6e3c136f76 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/817_PW_PPCG/INPUT b/tests/01_PW/817_PW_PPCG/INPUT new file mode 100644 index 00000000000..49cf788d32d --- /dev/null +++ b/tests/01_PW/817_PW_PPCG/INPUT @@ -0,0 +1,32 @@ +INPUT_PARAMETERS +#Parameters (General) +suffix autotest +pseudo_dir ../../PP_ORB +pw_seed 1 + +gamma_only 0 +calculation scf +symmetry 1 +out_level ie +smearing_method gaussian +smearing_sigma 0.02 + +#Parameters (3.PW) +ecutwfc 40 +scf_thr 1e-6 +scf_nmax 50 + +#Parameters (LCAO) +basis_type pw +ks_solver ppcg +device cpu +chg_extrap second-order +pw_diag_thr 0.00001 +pw_diag_ndim 4 + +cal_force 1 +cal_stress 1 + +mixing_type broyden +mixing_beta 0.4 +mixing_gg0 1.5 diff --git a/tests/01_PW/817_PW_PPCG/KPT b/tests/01_PW/817_PW_PPCG/KPT new file mode 100644 index 00000000000..b5b3bdb1ae2 --- /dev/null +++ b/tests/01_PW/817_PW_PPCG/KPT @@ -0,0 +1,4 @@ +K_POINTS +0 +Gamma +1 1 2 0 0 0 diff --git a/tests/01_PW/817_PW_PPCG/README b/tests/01_PW/817_PW_PPCG/README new file mode 100644 index 00000000000..c3a6155993b --- /dev/null +++ b/tests/01_PW/817_PW_PPCG/README @@ -0,0 +1 @@ +pw basis for GaAs with ks_solver ppcg (projection preconditioned conjugate-gradient), multi k diff --git a/tests/01_PW/817_PW_PPCG/STRU b/tests/01_PW/817_PW_PPCG/STRU new file mode 100644 index 00000000000..b03baadd25e --- /dev/null +++ b/tests/01_PW/817_PW_PPCG/STRU @@ -0,0 +1,23 @@ +ATOMIC_SPECIES +As 1 As_dojo.upf upf201 +Ga 1 Ga_dojo.upf upf201 + +LATTICE_CONSTANT +1 // add lattice constant, 10.58 ang + +LATTICE_VECTORS +5.33 5.33 0.0 +0.0 5.33 5.33 +5.33 0.0 5.33 +ATOMIC_POSITIONS +Direct //Cartesian or Direct coordinate. + +As +0 +1 +0.300000 0.3300000 0.27000000 0 0 0 + +Ga //Element Label +0 +1 //number of atom +0.00000 0.00000 0.000000 0 0 0 diff --git a/tests/01_PW/817_PW_PPCG/result.ref b/tests/01_PW/817_PW_PPCG/result.ref new file mode 100644 index 00000000000..2bb5340777f --- /dev/null +++ b/tests/01_PW/817_PW_PPCG/result.ref @@ -0,0 +1,8 @@ +etotref -4862.3309704772291298 +etotperatomref -2431.1654852386 +totalforceref 9.105146 +totalstressref 37223.049159 +pointgroupref C_1 +spacegroupref C_1 +nksibzref 2 +totaltimeref 4.64 diff --git a/tests/01_PW/818_PW_PPCG_Si/INPUT b/tests/01_PW/818_PW_PPCG_Si/INPUT new file mode 100644 index 00000000000..38436bc1e41 --- /dev/null +++ b/tests/01_PW/818_PW_PPCG_Si/INPUT @@ -0,0 +1,33 @@ +INPUT_PARAMETERS +#Parameters (General) +suffix autotest +pseudo_dir ../../PP_ORB +pw_seed 1 + +gamma_only 0 +calculation scf +symmetry 1 +out_level ie +smearing_method gaussian +smearing_sigma 0.02 + +#Parameters (3.PW) +ecutwfc 40 +scf_thr 1e-6 +scf_nmax 50 + +#Parameters (LCAO) +basis_type pw +ks_solver ppcg +device cpu +nbands 6 +chg_extrap second-order +pw_diag_thr 0.00001 +pw_diag_ndim 4 + +cal_force 1 +cal_stress 1 + +mixing_type broyden +mixing_beta 0.4 +mixing_gg0 1.5 diff --git a/tests/01_PW/818_PW_PPCG_Si/KPT b/tests/01_PW/818_PW_PPCG_Si/KPT new file mode 100644 index 00000000000..28006d5e2df --- /dev/null +++ b/tests/01_PW/818_PW_PPCG_Si/KPT @@ -0,0 +1,4 @@ +K_POINTS +0 +Gamma +2 2 2 0 0 0 diff --git a/tests/01_PW/818_PW_PPCG_Si/README b/tests/01_PW/818_PW_PPCG_Si/README new file mode 100644 index 00000000000..5f6c7b60ade --- /dev/null +++ b/tests/01_PW/818_PW_PPCG_Si/README @@ -0,0 +1 @@ +pw basis for Si with ks_solver ppcg (projection preconditioned conjugate-gradient), multi k diff --git a/tests/01_PW/818_PW_PPCG_Si/STRU b/tests/01_PW/818_PW_PPCG_Si/STRU new file mode 100644 index 00000000000..f3bade0c3b3 --- /dev/null +++ b/tests/01_PW/818_PW_PPCG_Si/STRU @@ -0,0 +1,19 @@ +ATOMIC_SPECIES +Si 1.000 Si_dojo_nsoc.upf + +LATTICE_CONSTANT +10.2 + +LATTICE_VECTORS +0.0 0.5 0.5 +0.5 0.0 0.5 +0.5 0.5 0.0 + +ATOMIC_POSITIONS +Direct + +Si +0.0 +2 +0.00 0.00 0.00 0 0 0 +0.25 0.25 0.25 0 0 0 diff --git a/tests/01_PW/818_PW_PPCG_Si/result.ref b/tests/01_PW/818_PW_PPCG_Si/result.ref new file mode 100644 index 00000000000..bd11216ce9f --- /dev/null +++ b/tests/01_PW/818_PW_PPCG_Si/result.ref @@ -0,0 +1,8 @@ +etotref -227.6796325586651903 +etotperatomref -113.8398162793 +totalforceref 0.000000 +totalstressref 342.892641 +pointgroupref T_d +spacegroupref O_h +nksibzref 3 +totaltimeref 1.45 diff --git a/tests/01_PW/819_PW_PPCG_Al/INPUT b/tests/01_PW/819_PW_PPCG_Al/INPUT new file mode 100644 index 00000000000..af15e1c0cb0 --- /dev/null +++ b/tests/01_PW/819_PW_PPCG_Al/INPUT @@ -0,0 +1,33 @@ +INPUT_PARAMETERS +#Parameters (General) +suffix autotest +pseudo_dir ../../PP_ORB +pw_seed 1 + +gamma_only 0 +calculation scf +symmetry 1 +out_level ie +smearing_method gaussian +smearing_sigma 0.02 + +#Parameters (3.PW) +ecutwfc 40 +scf_thr 1e-6 +scf_nmax 50 + +#Parameters (LCAO) +basis_type pw +ks_solver ppcg +device cpu +nbands 8 +chg_extrap second-order +pw_diag_thr 0.00001 +pw_diag_ndim 4 + +cal_force 1 +cal_stress 1 + +mixing_type broyden +mixing_beta 0.4 +mixing_gg0 1.5 diff --git a/tests/01_PW/819_PW_PPCG_Al/KPT b/tests/01_PW/819_PW_PPCG_Al/KPT new file mode 100644 index 00000000000..2df8740b6fa --- /dev/null +++ b/tests/01_PW/819_PW_PPCG_Al/KPT @@ -0,0 +1,4 @@ +K_POINTS +0 +Gamma +4 4 4 0 0 0 diff --git a/tests/01_PW/819_PW_PPCG_Al/README b/tests/01_PW/819_PW_PPCG_Al/README new file mode 100644 index 00000000000..a1ff49b6998 --- /dev/null +++ b/tests/01_PW/819_PW_PPCG_Al/README @@ -0,0 +1 @@ +pw basis for Al metal with ks_solver ppcg (projection preconditioned conjugate-gradient), multi k diff --git a/tests/01_PW/819_PW_PPCG_Al/STRU b/tests/01_PW/819_PW_PPCG_Al/STRU new file mode 100644 index 00000000000..1db1a774bdf --- /dev/null +++ b/tests/01_PW/819_PW_PPCG_Al/STRU @@ -0,0 +1,18 @@ +ATOMIC_SPECIES +Al 26.98 Al_ONCV_PBE-1.0.upf + +LATTICE_CONSTANT +7.63 + +LATTICE_VECTORS +0.0 0.5 0.5 +0.5 0.0 0.5 +0.5 0.5 0.0 + +ATOMIC_POSITIONS +Direct + +Al +0.0 +1 +0.00 0.00 0.00 1 1 1 diff --git a/tests/01_PW/819_PW_PPCG_Al/result.ref b/tests/01_PW/819_PW_PPCG_Al/result.ref new file mode 100644 index 00000000000..7d103ca9c8c --- /dev/null +++ b/tests/01_PW/819_PW_PPCG_Al/result.ref @@ -0,0 +1,8 @@ +etotref -1882.1779470693431904 +etotperatomref -1882.1779470693 +totalforceref 0.000000 +totalstressref 1108.971264 +pointgroupref O_h +spacegroupref O_h +nksibzref 8 +totaltimeref 1.84 diff --git a/tests/01_PW/820_PW_PPCG_Cutoff/INPUT b/tests/01_PW/820_PW_PPCG_Cutoff/INPUT new file mode 100644 index 00000000000..be2db4fddb4 --- /dev/null +++ b/tests/01_PW/820_PW_PPCG_Cutoff/INPUT @@ -0,0 +1,31 @@ +INPUT_PARAMETERS +# One outer iteration isolates the PPCG cutoff-degeneracy regression. +suffix autotest +calculation scf + +nbands 6 +symmetry 0 +pseudo_dir ../../PP_ORB + +ecutwfc 20 +init_chg atomic +# Stop after the deterministic first Hamiltonian; eigensolver tolerance stays strict. +scf_thr 1 +scf_nmax 1 +kpar 1 + +out_band 1 +out_app_flag 0 +dft_functional scan +pw_diag_thr 1e-12 +ks_solver ppcg +pw_diag_ndim 4 + +basis_type pw + +smearing_method gauss +smearing_sigma 0.002 + +mixing_type plain +mixing_beta 0.7 +mixing_gg0 0.0 diff --git a/tests/01_PW/820_PW_PPCG_Cutoff/KPT b/tests/01_PW/820_PW_PPCG_Cutoff/KPT new file mode 100644 index 00000000000..c289c0158aa --- /dev/null +++ b/tests/01_PW/820_PW_PPCG_Cutoff/KPT @@ -0,0 +1,4 @@ +K_POINTS +0 +Gamma +1 1 1 0 0 0 diff --git a/tests/01_PW/820_PW_PPCG_Cutoff/README b/tests/01_PW/820_PW_PPCG_Cutoff/README new file mode 100644 index 00000000000..9de604b2073 --- /dev/null +++ b/tests/01_PW/820_PW_PPCG_Cutoff/README @@ -0,0 +1 @@ +PPCG balanced-block regression for degenerate roots at the requested cutoff. diff --git a/tests/01_PW/820_PW_PPCG_Cutoff/STRU b/tests/01_PW/820_PW_PPCG_Cutoff/STRU new file mode 100644 index 00000000000..9058cabd7d1 --- /dev/null +++ b/tests/01_PW/820_PW_PPCG_Cutoff/STRU @@ -0,0 +1,19 @@ +ATOMIC_SPECIES +Si 14 Si.pz-vbc.UPF + +LATTICE_CONSTANT +10.2 + +LATTICE_VECTORS +0.0 0.5 0.5 +0.5 0.0 0.5 +0.5 0.5 0.0 + +ATOMIC_POSITIONS +Direct + +Si +0.0 +2 +0.00 0.00 0.00 1 1 1 +0.25 0.25 0.25 1 1 1 diff --git a/tests/01_PW/820_PW_PPCG_Cutoff/band.txt.ref b/tests/01_PW/820_PW_PPCG_Cutoff/band.txt.ref new file mode 100644 index 00000000000..1f19ed780b3 --- /dev/null +++ b/tests/01_PW/820_PW_PPCG_Cutoff/band.txt.ref @@ -0,0 +1 @@ + 1 0.00000000 -6.74890088 6.03406247 6.03406247 6.03406247 8.62917470 8.62917470 diff --git a/tests/01_PW/820_PW_PPCG_Cutoff/result.ref b/tests/01_PW/820_PW_PPCG_Cutoff/result.ref new file mode 100644 index 00000000000..3b54dd6439b --- /dev/null +++ b/tests/01_PW/820_PW_PPCG_Cutoff/result.ref @@ -0,0 +1,4 @@ +etotref -198.4293070843724 +etotperatomref -99.2146535422 +CompareBand_pass 0 +totaltimeref 0.54 diff --git a/tests/01_PW/CASES_CPU.txt b/tests/01_PW/CASES_CPU.txt index b0d6cade3e8..752071b15e3 100644 --- a/tests/01_PW/CASES_CPU.txt +++ b/tests/01_PW/CASES_CPU.txt @@ -121,6 +121,10 @@ scf_brav_bco_smear_mv 814_PW_LT_triclinic relax_dftu_spin2 relax_dftu_spin4 +817_PW_PPCG +818_PW_PPCG_Si +819_PW_PPCG_Al +820_PW_PPCG_Cutoff scf_deltaspin2 scf_deltaspin4 221_PW_PCC0D_H3O diff --git a/tests/11_PW_GPU/CASES_GPU.txt b/tests/11_PW_GPU/CASES_GPU.txt index a7a12920b4d..68969d84535 100644 --- a/tests/11_PW_GPU/CASES_GPU.txt +++ b/tests/11_PW_GPU/CASES_GPU.txt @@ -3,6 +3,7 @@ scf_cg scf_cg_single scf_dav scf_dav_sub +scf_ppcg scf_out_wf scf_out_wf_norm scf_out_wf_spinor diff --git a/tests/11_PW_GPU/scf_ppcg/INPUT b/tests/11_PW_GPU/scf_ppcg/INPUT new file mode 100644 index 00000000000..9af711ac742 --- /dev/null +++ b/tests/11_PW_GPU/scf_ppcg/INPUT @@ -0,0 +1,34 @@ +INPUT_PARAMETERS +#Parameters (General) +suffix autotest +pseudo_dir ../../PP_ORB + +gamma_only 0 +calculation scf +symmetry 1 +relax_nmax 1 +out_level ie +smearing_method gaussian +smearing_sigma 0.02 + +#Parameters (3.PW) +ecutwfc 40 +scf_thr 1e-7 +scf_nmax 100 + +#Parameters (LCAO) +basis_type pw +ks_solver ppcg +device gpu +chg_extrap second-order +pw_diag_thr 0.00001 +pw_diag_ndim 4 + +cal_force 1 +cal_stress 1 + +mixing_type broyden +mixing_beta 0.4 +mixing_gg0 1.5 + +pw_seed 1 diff --git a/tests/11_PW_GPU/scf_ppcg/KPT b/tests/11_PW_GPU/scf_ppcg/KPT new file mode 100644 index 00000000000..28006d5e2df --- /dev/null +++ b/tests/11_PW_GPU/scf_ppcg/KPT @@ -0,0 +1,4 @@ +K_POINTS +0 +Gamma +2 2 2 0 0 0 diff --git a/tests/11_PW_GPU/scf_ppcg/README b/tests/11_PW_GPU/scf_ppcg/README new file mode 100644 index 00000000000..314b6e1ec20 --- /dev/null +++ b/tests/11_PW_GPU/scf_ppcg/README @@ -0,0 +1 @@ +This test for: PPCG method for GaAs on GPU (transitional host/device bridge) diff --git a/tests/11_PW_GPU/scf_ppcg/STRU b/tests/11_PW_GPU/scf_ppcg/STRU new file mode 100644 index 00000000000..b03baadd25e --- /dev/null +++ b/tests/11_PW_GPU/scf_ppcg/STRU @@ -0,0 +1,23 @@ +ATOMIC_SPECIES +As 1 As_dojo.upf upf201 +Ga 1 Ga_dojo.upf upf201 + +LATTICE_CONSTANT +1 // add lattice constant, 10.58 ang + +LATTICE_VECTORS +5.33 5.33 0.0 +0.0 5.33 5.33 +5.33 0.0 5.33 +ATOMIC_POSITIONS +Direct //Cartesian or Direct coordinate. + +As +0 +1 +0.300000 0.3300000 0.27000000 0 0 0 + +Ga //Element Label +0 +1 //number of atom +0.00000 0.00000 0.000000 0 0 0 diff --git a/tests/11_PW_GPU/scf_ppcg/result.ref b/tests/11_PW_GPU/scf_ppcg/result.ref new file mode 100644 index 00000000000..d851d7baa46 --- /dev/null +++ b/tests/11_PW_GPU/scf_ppcg/result.ref @@ -0,0 +1,8 @@ +etotref -4869.7470517268475305 +etotperatomref -2434.8735258634 +totalforceref 5.206812 +totalstressref 37241.543010 +pointgroupref C_1 +spacegroupref C_1 +nksibzref 8 +totaltimeref 4.30 diff --git a/tests/11_PW_GPU/scf_ppcg/threshold b/tests/11_PW_GPU/scf_ppcg/threshold new file mode 100644 index 00000000000..ef64fd978a6 --- /dev/null +++ b/tests/11_PW_GPU/scf_ppcg/threshold @@ -0,0 +1,4 @@ +threshold 1 +force_threshold 1 +stress_threshold 1 +fatal_threshold 1