Add PPCG block subspace solver

source/source_hsolver/diago_ppcg.cpp

@@ -0,0 +1,18 @@
+#include "diago_ppcg.h"
+
+#include "ppcg/diago_ppcg_lapack.hpp"
+#include "ppcg/diago_ppcg_ops.hpp"
+#include "ppcg/diago_ppcg_subspace.hpp"
+#include "ppcg/diago_ppcg_orth.hpp"
+#include "ppcg/diago_ppcg_cg.hpp"
+#include "ppcg/diago_ppcg_diag.hpp"
+
+namespace hsolver {
+
+// =============================================================================
+// Explicit template instantiation (CPU only; extend for GPU as needed)
+// =============================================================================
+template class DiagoPPCG<std::complex<float>,  base_device::DEVICE_CPU>;
+template class DiagoPPCG<std::complex<double>, base_device::DEVICE_CPU>;
+
+} // namespace hsolver

source/source_hsolver/diago_ppcg.h

@@ -0,0 +1,228 @@
+#ifndef DIAGO_PPCG_H
+#define DIAGO_PPCG_H
+
+#include "source_base/module_device/types.h"
+
+#include <vector>
+#include <functional>
+#include <cmath>
+#include <algorithm>
+#include <limits>
+#include <stdexcept>
+#include <numeric>
+#include <complex>
+#include <type_traits>
+
+namespace hsolver {
+
+// -----------------------------------------------------------------------------
+// DiagoPPCG: Projection Preconditioned Conjugate Gradient solver
+// -----------------------------------------------------------------------------
+//
+// Supports two algorithmic strategies:
+//   CONJUGATE_GRADIENT — band-by-band Polak-Ribiere CG with line minimization
+//     (File 2 approach).
+//   BLOCK_SUBSPACE — block subspace diagonalization (File 1 approach).
+//
+// BLOCK_SUBSPACE is the production path used by ks_solver=ppcg.
+// CONJUGATE_GRADIENT is kept as an explicit fallback strategy.
+// -----------------------------------------------------------------------------
+
+enum class PpcgStrategy { BLOCK_SUBSPACE, CONJUGATE_GRADIENT };
+
+namespace base_device = ::base_device;
+
+template <typename T, typename Device>
+class DiagoPPCG
+{
+public:
+    // -------------------------------------------------------------------------
+    // Type aliases
+    // -------------------------------------------------------------------------
+    using Real = typename std::conditional<
+        std::is_same<T, std::complex<double>>::value, double,
+        float>::type;
+    using HPsiFunc = std::function<void(T*, T*, int, int)>;
+    using SPsiFunc = std::function<void(T*, T*, int, int)>;
+
+    // -------------------------------------------------------------------------
+    // Constructor
+    // -------------------------------------------------------------------------
+    DiagoPPCG(const Real& diag_thr,
+              const int& diag_iter_max,
+              const int& sbsize,
+              const int& rr_step,
+              const bool gamma_g0_real,
+              const PpcgStrategy strategy = PpcgStrategy::BLOCK_SUBSPACE);
+
+    // -------------------------------------------------------------------------
+    // Main entry point
+    //
+    // Returns average number of subspace iterations per band.
+    // -------------------------------------------------------------------------
+    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<double>& ethr_band,
+                const Real* prec);
+
+private:
+    // -------------------------------------------------------------------------
+    // Data members
+    // -------------------------------------------------------------------------
+    int maxiter_;
+    int sbsize_;
+    int rr_step_;
+    Real diag_thr_;
+    bool gamma_g0_real_;
+    PpcgStrategy strategy_;
+
+    // Problem dimensions (set in diag())
+    int ld_psi_ = 0;
+    int n_band_ = 0;
+    int n_dim_ = 0;
+
+    // Cached S-operator (null if identity).
+    SPsiFunc spsi_func_;
+
+    // Working storage (column-major: ld_psi_ rows, n_band_ columns).
+    std::vector<T> hpsi_;
+    std::vector<T> spsi_;
+    std::vector<T> w_;       // residual / preconditioned residual
+    std::vector<T> sw_;      // S * w
+    std::vector<T> hw_;      // H * w
+    std::vector<T> p_;       // previous search direction (for block subspace)
+    std::vector<T> sp_;      // S * p
+    std::vector<T> hp_;      // H * p
+
+    // Polak-Ribiere state (CONJUGATE_GRADIENT strategy)
+    std::vector<T> grad_old_;   // previous gradient
+    std::vector<T> z_old_;      // previous preconditioned residual
+    std::vector<Real> beta_denom_;
+
+    // -------------------------------------------------------------------------
+    // Internal helpers
+    // -------------------------------------------------------------------------
+    static inline int idx(int row, int col, int ld)
+    {
+        return row + col * ld;
+    }
+
+    void validate_input(const T* psi_in, const Real* eigenvalue_in,
+                        const std::vector<double>& ethr_band,
+                        const Real* prec) const;
+
+    void force_g0_real(T* x, int ncol) 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 <x|y> (real part only).
+    Real gamma_dot(const T* x, const T* y) const;
+    T complex_dot(const T* x, const T* y) const;
+
+    // Gram matrix: out[i, j] = <a_i | b_j>.
+    void gram(const T* a, const T* b,
+              int ncol_a, int ncol_b,
+              std::vector<T>& out, int ld_out) const;
+
+    // Gather / scatter columns.
+    void copy_cols(const T* src, const std::vector<int>& cols,
+                   std::vector<T>& dst) const;
+    void scatter_cols(T* dst, const std::vector<int>& cols,
+                      const std::vector<T>& src) const;
+
+    // Project x onto vectors orthogonal to the S-orthonormal basis.
+    void project_against(const T* basis, const T* sbasis,
+                         const std::vector<int>& basis_cols,
+                         std::vector<T>& x, std::vector<T>& sx,
+                         const std::vector<int>& x_cols) const;
+
+    // x[c] /= max(prec, eps)  for each active column c.
+    void divide_by_preconditioner(const std::vector<int>& active_cols,
+                                  const Real* prec,
+                                  std::vector<T>& x) const;
+
+    // -------------------------------------------------------------------------
+    // Block-subspace strategy helpers (File 1 style)
+    // -------------------------------------------------------------------------
+    struct SmallSubspace
+    {
+        std::vector<T> k;        // K matrix (projected H)
+        std::vector<T> m;        // M matrix (projected S)
+        std::vector<Real> eval;  // eigenvalues
+        std::vector<Real> w_scale;
+        std::vector<Real> p_scale;
+    };
+
+    void lock_epairs(const std::vector<T>& residual,
+                     const std::vector<double>& ethr_band,
+                     std::vector<int>& active_cols) const;
+
+    void build_small_subspace(const T* psi,
+                              const std::vector<int>& cols,
+                              bool use_p,
+                              SmallSubspace& subspace) const;
+
+    void solve_small_generalized(int dim, SmallSubspace& subspace) const;
+
+    void update_one_block(T* psi,
+                          const std::vector<int>& cols,
+                          int l,
+                          bool use_p,
+                          const SmallSubspace& subspace);
+
+    void right_solve_upper(const std::vector<T>& r, int n,
+                           std::vector<T>& x) const;
+
+    bool is_s_orthonormal(const T* psi, const T* spsi, int ncol) const;
+
+    void s_gram_schmidt(T* psi, T* hpsi, T* spsi, int ncol) const;
+
+    void chol_qr_active(T* psi, const std::vector<int>& active_cols);
+
+    void rayleigh_ritz(T* psi, Real* eigenvalue,
+                       std::vector<int>& active_cols,
+                       const std::vector<double>& ethr_band);
+
+    Real trace_of_active_projected(const T* psi,
+                                    const std::vector<int>& active_cols) const;
+
+    // -------------------------------------------------------------------------
+    // Conjugate-gradient strategy helpers (File 2 style)
+    // -------------------------------------------------------------------------
+    void calc_gradient(const Real* prec,
+                       const T* hpsi,
+                       const T* spsi,
+                       const T* psi,
+                       const Real* eigenvalue,
+                       std::vector<T>& grad) const;
+
+    void orth_gradient(const T* psi, const T* spsi,
+                       std::vector<T>& grad) const;
+
+    void update_polak_ribiere(const std::vector<T>& grad,
+                              std::vector<T>& p,
+                              std::vector<T>& grad_old,
+                              std::vector<T>& z_old,
+                              std::vector<Real>& beta_denom,
+                              const Real* prec) const;
+
+    void line_minimize(T* psi, T* hpsi, T* spsi,
+                       const T* p, const T* hp, const T* sp,
+                       int ncol) const;
+
+    void orth_cholesky(T* psi, T* hpsi, T* spsi, int ncol) const;
+};
+
+} // namespace hsolver
+
+#endif // DIAGO_PPCG_H

source/source_hsolver/hsolver_pw.cpp

@@ -11,18 +11,81 @@
 #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_io/module_parameter/parameter.h"
 #include "source_psi/psi.h"
 #include "source_estate/elecstate_tools.h"
 
 
 #include <algorithm>
+#include <type_traits>
 #include <vector>
 
 namespace hsolver
 {
 
+namespace
+{
+template <typename T, typename Device, typename Real, typename HPsiFunc, typename SPsiFunc>
+double 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<double>& ethr_band,
+                   const Real* pre_condition,
+                   const double diag_thr,
+                   const int diag_iter_max,
+                   const int pw_diag_ndim,
+                   const bool gamma_only,
+                   std::true_type)
+{
+    const int sbsize = std::max(1, std::min(nband, pw_diag_ndim));
+    const int rr_step = std::max(1, pw_diag_ndim);
+
+    DiagoPPCG<T, Device> ppcg(static_cast<Real>(diag_thr),
+                              diag_iter_max,
+                              sbsize,
+                              rr_step,
+                              gamma_only,
+                              PpcgStrategy::BLOCK_SUBSPACE);
+
+    return ppcg.diag(hpsi_func,
+                     spsi_func,
+                     ld_psi,
+                     nband,
+                     dim,
+                     psi,
+                     eigenvalue,
+                     ethr_band,
+                     pre_condition);
+}
+
+template <typename T, typename Device, typename Real, typename HPsiFunc, typename SPsiFunc>
+double run_ppcg_pw(const HPsiFunc&,
+                   const SPsiFunc&,
+                   const int,
+                   const int,
+                   const int,
+                   T*,
+                   Real*,
+                   const std::vector<double>&,
+                   const Real*,
+                   const double,
+                   const int,
+                   const int,
+                   const bool,
+                   std::false_type)
+{
+    ModuleBase::WARNING_QUIT("HSolverPW::hamiltSolvePsiK",
+                             "PPCG is currently implemented for CPU PW calculations only.");
+    return 0.0;
+}
+} // namespace
+
 template <typename T, typename Device>
 void HSolverPW<T, Device>::cal_smooth_ethr(const double& wk,
                                            const double* wg,
@@ -83,7 +146,7 @@ void HSolverPW<T, Device>::solve(hamilt::Hamilt<T, Device>* pHamilt,
     this->nproc_in_pool = nproc_in_pool_in;
 
     // report if the specified diagonalization method is not supported
-    const std::initializer_list<std::string> _methods = {"cg", "dav", "dav_subspace", "bpcg"};
+    const std::initializer_list<std::string> _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!");
@@ -379,6 +442,24 @@ void HSolverPW<T, Device>::hamiltSolvePsiK(hamilt::Hamilt<T, Device>* hm,
                         ntry_max,
                         notconv_max));
     }
+    else if (this->method == "ppcg")
+    {
+        DiagoIterAssist<T, Device>::avg_iter += run_ppcg_pw<T, Device, Real>(
+            hpsi_func,
+            spsi_func,
+            psi.get_nbasis(),
+            psi.get_nbands(),
+            psi.get_current_ngk(),
+            psi.get_pointer(),
+            eigenvalue,
+            this->ethr_band,
+            pre_condition.data(),
+            this->diag_thr,
+            this->diag_iter_max,
+            PARAM.inp.pw_diag_ndim,
+            PARAM.globalv.gamma_only_pw,
+            std::is_same<Device, base_device::DEVICE_CPU>());
+    }
     ModuleBase::timer::end("HSolverPW", "solve_psik");
     return;
 }