source_hsolver: decouple the eigensolvers from hamilt through HSOperator/HSMatrix interfaces

python/pyabacus/src/hsolver/diago_adapter.hpp

@@ -86,8 +86,7 @@ class PyDiagoDavidAdapter
         int max_iter,
         ::hsolver::diag_comm_info comm_info)
     {
-        auto hpsi_func = make_hpsi_func_fstyle<T>(mm_op);
-        auto spsi_func = make_spsi_func_identity<Traits>();
+        const PyHSOperator<T> op(mm_op);
 
         solver_ = std::make_unique<SolverType>(
             precond_vec.data(),
@@ -98,8 +97,7 @@ class PyDiagoDavidAdapter
         );
 
         return solver_->diag(
-            hpsi_func,
-            spsi_func,
+            op,
             nbasis_,
             storage_.psi_ptr(),
             storage_.eigenvalue_ptr(),
@@ -178,8 +176,7 @@ class PyDiagoDavSubspaceAdapter
         int diag_subspace,
         int nb2d)
     {
-        auto hpsi_func = make_hpsi_func_fstyle<T>(mm_op);
-        auto spsi_func = make_spsi_func_identity<Traits>();
+        const PyHSOperator<T> op(mm_op);
 
         solver_ = std::make_unique<SolverType>(
             precond_vec,
@@ -194,8 +191,7 @@ class PyDiagoDavSubspaceAdapter
         );
 
         return solver_->diag(
-            hpsi_func,
-            spsi_func,
+            op,
             storage_.psi_ptr(),
             nbasis_,
             storage_.eigenvalue_ptr(),

python/pyabacus/src/hsolver/diago_traits.hpp

@@ -22,6 +22,7 @@
 #include <pybind11/stl.h>
 
 #include "source_hsolver/diago_david.h"
+#include "source_hsolver/hs_operator.h"
 #include "source_hsolver/diago_dav_subspace.h"
 #include "source_hsolver/diago_cg.h"
 #include "source_base/module_device/memory_op.h"
@@ -384,52 +385,53 @@ struct DiagoCGTraits
 #endif // __ENABLE_ATEN
 
 // ============================================================================
-// Helper Functions for Creating HPsi/SPsi Lambdas
+// HSOperator view of a Python callable, plus tensor helpers for the ATen CG path
 // ============================================================================
 
 /**
- * @brief Create hpsi_func lambda for raw pointer interface (F-style)
+ * @brief The Python matrix-vector callable seen through hsolver::HSOperator.
  *
- * Wraps a Python callable to work with ABACUS raw pointer interface.
- * Handles array layout conversion between Python (row-major) and
- * ABACUS (column-major for Davidson methods).
+ * H is applied by the Python callable mm_op on a column-major (ld_psi x nvec)
+ * array; S is the identity. This is all the Davidson-type solvers need.
  */
 template <typename T>
-auto make_hpsi_func_fstyle(
-    std::function<py::array_t<T>(py::array_t<T>)> mm_op)
+class PyHSOperator : public ::hsolver::HSOperator<T, base_device::DEVICE_CPU>
 {
-    return [mm_op](T* psi_in, T* hpsi_out, const int ld_psi, const int nvec) {
+public:
+    explicit PyHSOperator(std::function<py::array_t<T>(py::array_t<T>)> mm_op)
+        : mm_op_(std::move(mm_op))
+    {
+    }
+
+    void update_k(const int ik) override
+    {
+    }
+
+    void hpsi(const T* x, T* hx, const int ld, const int nvec) const override
+    {
         // Create F-style numpy array (column-major)
-        py::array_t<T, py::array::f_style> psi({ld_psi, nvec});
+        py::array_t<T, py::array::f_style> psi({ld, nvec});
         py::buffer_info buf = psi.request();
         T* ptr = static_cast<T*>(buf.ptr);
-        std::copy(psi_in, psi_in + nvec * ld_psi, ptr);
+        std::copy(x, x + static_cast<size_t>(nvec) * ld, ptr);
 
         // Call Python function
-        py::array_t<T, py::array::f_style> hpsi = mm_op(psi);
+        py::array_t<T, py::array::f_style> hpsi_arr = mm_op_(psi);
 
         // Copy result back
-        py::buffer_info hpsi_buf = hpsi.request();
-        T* hpsi_ptr = static_cast<T*>(hpsi_buf.ptr);
-        std::copy(hpsi_ptr, hpsi_ptr + nvec * ld_psi, hpsi_out);
-    };
-}
+        py::buffer_info hpsi_buf = hpsi_arr.request();
+        const T* hpsi_ptr = static_cast<const T*>(hpsi_buf.ptr);
+        std::copy(hpsi_ptr, hpsi_ptr + static_cast<size_t>(nvec) * ld, hx);
+    }
 
-/**
- * @brief Create spsi_func lambda for raw pointer interface (identity)
- *
- * For non-orthogonal basis, S*psi = psi (identity operation).
- */
-template <typename Traits>
-auto make_spsi_func_identity()
-{
-    using T = typename Traits::T;
-    using syncmem_op = typename Traits::syncmem_op;
+    void spsi(const T* x, T* sx, const int ld, const int nvec) const override
+    {
+        std::copy(x, x + static_cast<size_t>(nvec) * ld, sx);
+    }
 
-    return [](const T* psi_in, T* spsi_out, const int nrow, const int nbands) {
-        syncmem_op()(spsi_out, psi_in, static_cast<size_t>(nbands * nrow));
-    };
-}
+private:
+    std::function<py::array_t<T>(py::array_t<T>)> mm_op_;
+};
 
 #ifdef __ENABLE_ATEN
 /**

source/source_esolver/esolver_ks_lcao.cpp

@@ -13,6 +13,7 @@
 #include "source_hamilt/module_gint/gint.h"
 #include "source_estate/elecstate_lcao.h"
 #include "source_lcao/hamilt_lcao.h"
+#include "source_hamilt/hamilt_hs_adapter.h"
 #include "source_hsolver/hsolver_lcao.h"
 #ifdef __EXX
 #include "../source_lcao/module_ri/exx_opt_orb.h"
@@ -480,7 +481,9 @@ void ESolver_KS_LCAO<TK, TR>::hamilt2rho_single(UnitCell& ucell, int istep, int
                                                   this->inp_->device == "gpu",
                                                   GlobalV::NPROC,
                                                   GlobalV::MY_RANK);
-        hsolver_lcao_obj.solve(static_cast<hamilt::Hamilt<TK>*>(this->p_hamilt), this->psi[0], this->pelec, *this->dmat.dm, 
+        // the eigensolvers only ever ask the Hamiltonian for H(k) and S(k)
+        hamilt::HamiltHSMatrix<TK> hs(static_cast<hamilt::Hamilt<TK>*>(this->p_hamilt));
+        hsolver_lcao_obj.solve(hs, this->psi[0], this->pelec, *this->dmat.dm, 
           this->chr, this->inp_->nspin, skip_charge);
     }
     else

source/source_esolver/esolver_ks_lcao_tddft.cpp

@@ -19,6 +19,7 @@
 #include "source_estate/module_pot/h_tddft_pw.h"
 #include "source_estate/module_pot/potential_new.h"
 #include "source_estate/module_pot/td_field_manager.h"
+#include "source_hamilt/hamilt_hs_adapter.h"
 #include "source_hsolver/hsolver_lcao.h"
 #include "source_lcao/module_rt/evolve_elec.h"
 #include "source_lcao/rho_tau_lcao.h"
@@ -365,7 +366,9 @@ void ESolver_KS_LCAO_TDDFT<TR, Device>::hamilt2rho_single(UnitCell& ucell, const
                                                                         this->inp_->device == "gpu",
                                                                         GlobalV::NPROC,
                                                                         GlobalV::MY_RANK);
-            hsolver_lcao_obj.solve(static_cast<hamilt::Hamilt<std::complex<double>>*>(this->p_hamilt),
+            hamilt::HamiltHSMatrix<std::complex<double>> hs(
+                static_cast<hamilt::Hamilt<std::complex<double>>*>(this->p_hamilt));
+            hsolver_lcao_obj.solve(hs,
                                    this->psi[0],
                                    this->pelec,
                                    *this->dmat.dm,

source/source_esolver/esolver_ks_lcaopw.cpp

@@ -155,16 +155,19 @@ namespace ModuleESolver
 #else
         const hsolver::diag_comm_info diag_comm(0, 1);
 #endif
-        hsolver_lip_obj.solve(static_cast<hamilt::Hamilt<T>*>(this->p_hamilt),
+        // the solver sees the Hamiltonian only through this operator; the EXX
+        // subspace hooks live in the adapter, not in the solver
+        hamilt::HamiltLIPHSOperator<T> op(dynamic_cast<hamilt::HamiltLIP<T>*>(this->p_hamilt),
+                                          this->pw_wfc,
+                                          this->general_exx_info_.cal_exx,
+                                          this->general_exx_info_.hybrid_alpha);
+        hsolver_lip_obj.solve(op,
                               *this->stp.template get_psi_t<T, base_device::DEVICE_CPU>(),
                               this->pelec,
                               *this->psi_local,
                               diag_comm,
                               GlobalV::ofs_running,
-                              skip_charge,
-                              ucell.tpiba,
-                              ucell.nat,
-                              this->general_exx_info_);
+                              skip_charge);
 
         // add exx
 #ifdef __EXX

source/source_esolver/esolver_ks_pw.cpp

@@ -4,6 +4,7 @@
 #include "source_estate/module_charge/symm_rho.h"
 #include "source_hsolver/diago_iter_assist.h"
 #include "source_hsolver/diago_params.h"
+#include "source_hamilt/hamilt_hs_adapter.h"
 #include "source_hsolver/hsolver_pw.h"
 #include "source_io/module_parameter/parameter.h"
 #include "source_pw/module_pwdft/force_pw.h"
@@ -264,16 +265,16 @@ void ESolver_KS_PW<T, Device>::hamilt2rho_single(UnitCell& ucell, const int iste
                                                      this->inp_->nb2d,
                                                      this->inp_->use_k_continuity);
 
-        hsolver_pw_obj.solve(static_cast<hamilt::Hamilt<T, Device>*>(this->p_hamilt),
+        // the iterative eigensolvers see the Hamiltonian only through this operator
+        hamilt::HamiltHSOperator<T, Device> op(static_cast<hamilt::Hamilt<T, Device>*>(this->p_hamilt), this->pw_wfc);
+        hsolver_pw_obj.solve(op,
                              *this->stp.template get_psi_t<T, Device>(),
                              this->pelec,
                              this->pelec->ekb.c,
                              GlobalV::RANK_IN_POOL,
                              GlobalV::NPROC_IN_POOL,
                              GlobalV::ofs_running,
-                             skip_charge,
-                             ucell.tpiba,
-                             ucell.nat);
+                             skip_charge);
     }
 
     // symmetrize the charge density

source/source_hamilt/hamilt_hs_adapter.h

@@ -0,0 +1,84 @@
+#ifndef HAMILT_HS_ADAPTER_H
+#define HAMILT_HS_ADAPTER_H
+
+#include "source_base/tool_quit.h"
+#include "source_basis/module_pw/pw_basis_k.h"
+#include "source_hamilt/hamilt.h"
+#include "source_hsolver/hs_matrix.h"
+#include "source_hsolver/hs_operator.h"
+#include "source_psi/psi.h"
+
+namespace hamilt
+{
+
+/**
+ * @brief Presents a Hamilt as the H/S block-vector operator the iterative
+ *        eigensolvers work on (hsolver::HSOperator).
+ *
+ * This is the only place that wraps raw pointers into Psi / hpsi_info for the
+ * operator chain; hsolver itself never sees Hamilt.
+ */
+template <typename T, typename Device = base_device::DEVICE_CPU>
+class HamiltHSOperator : public hsolver::HSOperator<T, Device>
+{
+  public:
+    HamiltHSOperator(Hamilt<T, Device>* hm, const ModulePW::PW_Basis_K* wfc_basis) : hm_(hm), wfc_basis_(wfc_basis)
+    {
+    }
+
+    void update_k(const int ik) override
+    {
+        hm_->updateHk(ik);
+        ik_ = ik;
+        npw_ = wfc_basis_->npwk[ik];
+    }
+
+    void hpsi(const T* x, T* hx, const int ld, const int nvec) const override
+    {
+        if (hm_->ops == nullptr)
+        {
+            ModuleBase::WARNING_QUIT("HamiltHSOperator::hpsi", "Operators in Hamilt are not allocated yet");
+        }
+        // non-owning view of x: one k point, nvec bands, leading dimension ld, npw valid rows
+        psi::Psi<T, Device> x_view(const_cast<T*>(x), 1, nvec, ld, npw_);
+        typename Operator<T, Device>::hpsi_info info(&x_view, psi::Range(true, 0, 0, nvec - 1), hx);
+        hm_->ops->hPsi(info);
+    }
+
+    void spsi(const T* x, T* sx, const int ld, const int nvec) const override
+    {
+        hm_->sPsi(x, sx, ld, npw_, nvec);
+    }
+
+  protected:
+    Hamilt<T, Device>* hm_ = nullptr;
+    const ModulePW::PW_Basis_K* wfc_basis_ = nullptr;
+    int ik_ = 0;  ///< k point set by the last update_k()
+    int npw_ = 0; ///< number of plane waves of that k point (without npol)
+};
+
+/**
+ * @brief Presents a Hamilt as the H(k)/S(k) matrix source the direct
+ *        eigensolvers work on (hsolver::HSMatrix).
+ */
+template <typename T>
+class HamiltHSMatrix : public hsolver::HSMatrix<T>
+{
+  public:
+    explicit HamiltHSMatrix(Hamilt<T>* hm) : hm_(hm)
+    {
+    }
+
+    void hs_at_k(const int ik, ModuleBase::MatrixBlock<T>& hk, ModuleBase::MatrixBlock<T>& sk) override
+    {
+        hm_->updateHk(ik);
+        hm_->matrix(hk, sk);
+    }
+
+  private:
+    Hamilt<T>* hm_ = nullptr;
+};
+
+} // namespace hamilt
+
+#endif // HAMILT_HS_ADAPTER_H