Feature: Add diago_dav_subspace module to pyabacus

python/pyabacus/CMakeLists.txt

@@ -11,6 +11,9 @@ find_package(pybind11 CONFIG REQUIRED)
 set(ABACUS_SOURCE_DIR "${PROJECT_SOURCE_DIR}/../../source")
 set(BASE_PATH "${ABACUS_SOURCE_DIR}/module_base")
 set(NAO_PATH "${ABACUS_SOURCE_DIR}/module_basis/module_nao")
+set(HSOLVER_PATH "${ABACUS_SOURCE_DIR}/module_hsolver")
+set(HAMILT_PATH "${ABACUS_SOURCE_DIR}/module_hamilt_general")
+set(PSI_PATH "${ABACUS_SOURCE_DIR}/module_psi")
 set(ENABLE_LCAO ON)
 list(APPEND CMAKE_MODULE_PATH "${PROJECT_SOURCE_DIR}/../../cmake")
 
@@ -100,6 +103,29 @@ list(APPEND _naos
 add_library(naopack SHARED 
     ${_naos}
     )
+# add diago shared library
+list(APPEND _diago
+    ${HSOLVER_PATH}/diago_dav_subspace.cpp
+    ${HSOLVER_PATH}/diag_const_nums.cpp
+    ${HSOLVER_PATH}/diago_iter_assist.cpp
+
+    ${HSOLVER_PATH}/kernels/dngvd_op.cpp
+    ${HSOLVER_PATH}/kernels/math_kernel_op.cpp
+    # dependency
+    ${BASE_PATH}/module_device/device.cpp
+    ${BASE_PATH}/module_device/memory_op.cpp
+
+    ${HAMILT_PATH}/operator.cpp
+    ${PSI_PATH}/psi.cpp
+  )
+add_library(diagopack SHARED 
+    ${_diago}
+    )
+target_link_libraries(diagopack
+    PRIVATE
+    ${OpenBLAS_LIBRARIES} 
+    ${LAPACK_LIBRARIES}
+    )
 # link math_libs
 if(MKLROOT)
   target_link_libraries(naopack
@@ -125,9 +151,10 @@ list(APPEND _sources
     ${PROJECT_SOURCE_DIR}/src/py_abacus.cpp
     ${PROJECT_SOURCE_DIR}/src/py_base_math.cpp
     ${PROJECT_SOURCE_DIR}/src/py_m_nao.cpp
+    ${PROJECT_SOURCE_DIR}/src/py_diago_dav_subspace.cpp
     )
 pybind11_add_module(_core MODULE ${_sources})
-target_link_libraries(_core PRIVATE pybind11::headers naopack)
+target_link_libraries(_core PRIVATE pybind11::headers naopack diagopack)
 target_compile_definitions(_core PRIVATE VERSION_INFO=${PROJECT_VERSION})
 # set RPATH
 execute_process(
@@ -141,5 +168,7 @@ set(TARGET_PACK pyabacus)
 set(CMAKE_INSTALL_RPATH "${PYTHON_SITE_PACKAGES}/${TARGET_PACK}")
 set_target_properties(_core PROPERTIES INSTALL_RPATH "$ORIGIN")
 set_target_properties(naopack PROPERTIES INSTALL_RPATH "$ORIGIN")
+set_target_properties(diagopack PROPERTIES INSTALL_RPATH "$ORIGIN")
 install(TARGETS _core naopack DESTINATION ${TARGET_PACK})
+install(TARGETS _core diagopack DESTINATION ${TARGET_PACK})
 

python/pyabacus/examples/diago_matrix.py

@@ -0,0 +1,52 @@
+from pyabacus import hsolver
+import numpy as np
+
+h_mat = np.array(
+    [
+        4.0+0.0j, 2.0+0.0j, 2.0+0.0j,
+        2.0+0.0j, 4.0+0.0j, 2.0+0.0j,
+        2.0+0.0j, 2.0+0.0j, 4.0+0.0j
+    ], 
+dtype=np.complex128, order='C')
+# h_mat = np.array(
+#     [
+#         1.0+0.0j, 0.0+0.0j, 0.0+0.0j,
+#         0.0+0.0j, 1.0+0.0j, 0.0+0.0j,
+#         0.0+0.0j, 0.0+0.0j, 1.0+0.0j
+#     ],
+# dtype=np.complex128, order='C')
+pre_condition = np.ones(3, dtype=np.float64, order='C')
+nband = 1
+nbasis = 3
+dav_ndim = 2
+diag_thr = 1e-2
+diag_nmax = 1000
+need_subspace = False
+comm_info = hsolver.diag_comm_info(0, 1)
+
+psi = np.ones(nbasis * nband, dtype=np.complex128, order='C')
+eigenvalues = np.zeros(nband, dtype=np.float64, order='C')
+is_occupied = [True] * nband
+
+diago_dav_subspace = hsolver.Diago_DavSubspace(
+    pre_condition,
+    nband,
+    nbasis,
+    dav_ndim,
+    diag_thr,
+    diag_nmax,
+    need_subspace,
+    comm_info
+)
+
+res = diago_dav_subspace.diag(
+    h_mat,
+    psi,
+    nbasis,
+    eigenvalues,
+    is_occupied,
+    False
+)
+
+print(f'res: {res}')
+print(f'eigenvalues: {eigenvalues}')

python/pyabacus/src/py_abacus.cpp

@@ -5,10 +5,12 @@ namespace py = pybind11;
 
 void bind_base_math(py::module& m);
 void bind_m_nao(py::module& m);
+void bind_diago_dav_subspace(py::module& m);
 
 PYBIND11_MODULE(_core, m)
 {
     m.doc() = "Python extension for ABACUS built with pybind11 and scikit-build.";
     bind_base_math(m);
     bind_m_nao(m);
+    bind_diago_dav_subspace(m);
 }

python/pyabacus/src/py_diago_dav_subspace.cpp

@@ -0,0 +1,88 @@
+#include <complex>
+#include <functional>
+#include <pybind11/pybind11.h>
+#include <pybind11/stl.h>
+#include <pybind11/complex.h>
+#include <pybind11/numpy.h>
+
+#include "module_hsolver/diago_dav_subspace.h"
+#include "module_hsolver/kernels/math_kernel_op.h"
+#include "module_base/module_device/types.h"
+
+namespace py = pybind11;
+using namespace pybind11::literals;
+
+void bind_diago_dav_subspace(py::module& m)
+{
+    py::module hsolver = m.def_submodule("hsolver");
+
+    py::class_<hsolver::diag_comm_info>(hsolver, "diag_comm_info")
+        .def(py::init<const int, const int>(), "rank"_a, "nproc"_a)
+        .def_readonly("rank", &hsolver::diag_comm_info::rank)
+        .def_readonly("nproc", &hsolver::diag_comm_info::nproc);
+
+    py::class_<hsolver::Diago_DavSubspace<std::complex<double>, base_device::DEVICE_CPU>>(hsolver, "Diago_DavSubspace")
+        .def(py::init<const std::vector<double>, const int, const int,
+                    const int, const double, const int, const bool,
+                    const hsolver::diag_comm_info>(),
+        "precondition"_a, "nband"_a, "nbasis"_a, "david_ndim"_a, "diag_thr"_a, "diag_nmax"_a, "need_subspace"_a, "diag_comm"_a)
+        .def("diag",
+            [] (hsolver::Diago_DavSubspace<std::complex<double>, base_device::DEVICE_CPU>& self,
+                py::array_t<std::complex<double>> h_mat,
+                py::array_t<std::complex<double>> psi_in,
+                const int psi_in_dmax,
+                py::array_t<double> eigenvalue_in,
+                const std::vector<bool>& is_occupied,
+                const bool& scf_type)
+            {
+                py::buffer_info h_buf = h_mat.request();
+                if (h_buf.ndim != 1)
+                {
+                    throw std::runtime_error("h_mat must be 1D array representing a matrix");
+                }
+                py::buffer_info psi_buf = psi_in.request();
+                if (psi_buf.ndim != 1)
+                {
+                    throw std::runtime_error("psi_in must be 1D array representing a matrix");
+                }
+                py::buffer_info eigenvalue_buf = eigenvalue_in.request();
+                if (eigenvalue_buf.ndim != 1)
+                {
+                    throw std::runtime_error("eigenvalue_in must be 1D array storing eigenvalues");
+                }
+
+                std::complex<double>* h_mat_ptr = static_cast<std::complex<double>*>(h_buf.ptr);
+                std::complex<double>* psi_in_ptr = static_cast<std::complex<double>*>(psi_buf.ptr);
+                double* eigenvalue_in_ptr = static_cast<double*>(eigenvalue_buf.ptr);
+
+                // TODO: Wrap std::function<void(T*, T*, const int, const int, const int, const int)>
+                //       to a python callable type
+                auto hpsi_func = [h_mat_ptr] (std::complex<double> *hpsi_out,
+                            std::complex<double> *psi_in, const int nband_in,
+                            const int nbasis_in, const int band_index1,
+                            const int band_index2) 
+                {
+                    const std::complex<double> *one_ = nullptr, *zero_ = nullptr;
+
+                    one_ = new std::complex<double>(1.0, 0.0);
+                    zero_ = new std::complex<double>(0.0, 0.0);
+
+                    base_device::DEVICE_CPU *ctx = {};
+
+                    hsolver::gemm_op<std::complex<double>, base_device::DEVICE_CPU>()(
+                        ctx, 'N', 'N', 
+                        nbasis_in, band_index2 - band_index1 + 1, nbasis_in, 
+                        one_, h_mat_ptr, nbasis_in, 
+                        psi_in + band_index1 * nbasis_in, nbasis_in,
+                        zero_, hpsi_out + band_index1 * nbasis_in, nbasis_in);
+                };
+
+                return self.diag(hpsi_func, 
+                                 psi_in_ptr, 
+                                 psi_in_dmax, 
+                                 eigenvalue_in_ptr, 
+                                 is_occupied, 
+                                 scf_type);
+            }
+        );
+}

python/pyabacus/src/pyabacus/__init__.py

@@ -1,3 +1,3 @@
 from __future__ import annotations
-from ._core import ModuleBase, ModuleNAO
-__all__ = ["ModuleBase", "ModuleNAO"]
+from ._core import ModuleBase, ModuleNAO, hsolver
+__all__ = ["ModuleBase", "ModuleNAO", "hsolver"]