Feature: Add davidson function to pyabacus (#5112)

* add PyDiagoDavid class

* add pybind11 wrapper for `hsolver::DiagoDavid`

* add `PyDiagoDavid` class as the pybind wrapper for `hsolver::DiagoDavid` class

* add python function signature and python tests

* fix some bugs in pytest:test_hsolver.py

* Add example for diagonalizing sparse matrices using Davidson method and update README.md with relevant documentation

* update example for hsolver and README.md

* Revert unintended changes to file

Author: a1henu

python/pyabacus/CMakeLists.txt

@@ -106,6 +106,7 @@ add_library(naopack SHARED
 # add diago shared library
 list(APPEND _diago
     ${HSOLVER_PATH}/diago_dav_subspace.cpp
+    ${HSOLVER_PATH}/diago_david.cpp
     ${HSOLVER_PATH}/diag_const_nums.cpp
     ${HSOLVER_PATH}/diago_iter_assist.cpp
 
@@ -151,7 +152,7 @@ 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
+    ${PROJECT_SOURCE_DIR}/src/py_hsolver.cpp
     )
 pybind11_add_module(_core MODULE ${_sources})
 target_link_libraries(_core PRIVATE pybind11::headers naopack diagopack)

python/pyabacus/README.md

@@ -46,9 +46,29 @@ Run `python diago_matrix.py` in `examples` to check the diagonalization of a mat
 ```shell
 $ cd examples/
 $ python diago_matrix.py
-eigenvalues calculated by pyabacus: [-0.38440611 0.24221155 0.31593272 0.53144616 0.85155108 1.06950155 1.11142051 1.12462152]
-eigenvalues calculated by scipy: [-0.38440611 0.24221155 0.31593272 0.53144616 0.85155108 1.06950154 1.11142051 1.12462151]
-error: [9.26164700e-12 2.42959514e-10 2.96529468e-11 7.77933273e-12 7.53686002e-12 2.95628810e-09 1.04678111e-09 7.79106313e-09]
+
+====== Calculating eigenvalues using davidson method... ======
+eigenvalues calculated by pyabacus-davidson is: 
+ [-0.38440611  0.24221155  0.31593272  0.53144616  0.85155108  1.06950154
+  1.11142053  1.12462153]
+eigenvalues calculated by scipy is: 
+ [-0.38440611  0.24221155  0.31593272  0.53144616  0.85155108  1.06950154
+  1.11142051  1.12462151]
+eigenvalues difference: 
+ [4.47258897e-12 5.67104697e-12 8.48299209e-12 1.08900666e-11
+ 1.87927451e-12 3.15688586e-10 2.11438165e-08 2.68884972e-08]
+
+====== Calculating eigenvalues using dav_subspace method... ======
+enter diag... is_subspace = 0, ntry = 0
+eigenvalues calculated by pyabacus-dav_subspace is: 
+ [-0.38440611  0.24221155  0.31593272  0.53144616  0.85155108  1.06950154
+  1.11142051  1.12462153]
+eigenvalues calculated by scipy is: 
+ [-0.38440611  0.24221155  0.31593272  0.53144616  0.85155108  1.06950154
+  1.11142051  1.12462151]
+eigenvalues difference: 
+ [ 4.64694949e-12  2.14706031e-12  1.09236509e-11  4.66293670e-13
+ -8.94295749e-12  4.71351846e-11  5.39378986e-10  1.97244101e-08]
 ```
 
 License

python/pyabacus/examples/diago_matrix.py

@@ -2,37 +2,61 @@
 import numpy as np
 import scipy
 
-h_mat = scipy.io.loadmat('./Si2.mat')['Problem']['A'][0, 0]
-
-nbasis = h_mat.shape[0]
-nband = 8
-
-v0 = np.random.rand(nbasis, nband)
-
-diag_elem = h_mat.diagonal()
-diag_elem = np.where(np.abs(diag_elem) < 1e-8, 1e-8, diag_elem)
-precond = 1.0 / np.abs(diag_elem)
-
-
-def mm_op(x):
-    return h_mat.dot(x)
-
-e, v = hsolver.dav_subspace(
-    mm_op,
-    v0,
-    nbasis,
-    nband,
-    precond,
-    dav_ndim=8,
-    tol=1e-8,
-    max_iter=1000,
-    scf_type=False
-)
-
-print('eigenvalues calculated by pyabacus: ', e)
-
-e_scipy, v_scipy = scipy.sparse.linalg.eigsh(h_mat, k=nband, which='SA', maxiter=1000)
-e_scipy = np.sort(e_scipy)
-print('eigenvalues calculated by scipy: ', e_scipy)
-
-print('error:', e - e_scipy)
+def load_mat(mat_file):
+    h_mat = scipy.io.loadmat(mat_file)['Problem']['A'][0, 0]
+    nbasis = h_mat.shape[0]
+    nband = 8
+    
+    return h_mat, nbasis, nband
+
+def calc_eig_pyabacus(mat_file, method):
+    algo = {
+        'dav_subspace': hsolver.dav_subspace,
+        'davidson': hsolver.davidson
+    }
+    
+    h_mat, nbasis, nband = load_mat(mat_file)
+    
+    v0 = np.random.rand(nbasis, nband)
+    diag_elem = h_mat.diagonal()
+    diag_elem = np.where(np.abs(diag_elem) < 1e-8, 1e-8, diag_elem)
+    precond = 1.0 / np.abs(diag_elem)
+
+    def mm_op(x):
+        return h_mat.dot(x)
+
+    e, _ = algo[method](
+        mm_op,
+        v0,
+        nbasis,
+        nband,
+        precond,
+        dav_ndim=8,
+        tol=1e-8,
+        max_iter=1000
+    )
+
+    print(f'eigenvalues calculated by pyabacus-{method} is: \n', e)
+    
+    return e
+
+def calc_eig_scipy(mat_file):
+    h_mat, _, nband = load_mat(mat_file)
+    e, _ = scipy.sparse.linalg.eigsh(h_mat, k=nband, which='SA', maxiter=1000)
+    e = np.sort(e)
+    print('eigenvalues calculated by scipy is: \n', e)
+    
+    return e
+
+if __name__ == '__main__':
+    mat_file = './Si2.mat'
+    method = ['davidson', 'dav_subspace']
+    
+    for m in method:
+        print(f'\n====== Calculating eigenvalues using {m} method... ======')
+        e_pyabacus = calc_eig_pyabacus(mat_file, m)
+        e_scipy = calc_eig_scipy(mat_file)
+        
+        print('eigenvalues difference: \n', e_pyabacus - e_scipy)
+    
+    

python/pyabacus/src/py_abacus.cpp

@@ -5,12 +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);
+void bind_hsolver(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);
+    bind_hsolver(m);
 }

python/pyabacus/src/py_diago_david.hpp

@@ -0,0 +1,170 @@
+#ifndef PYTHON_PYABACUS_SRC_PY_DIAGO_DAVID_HPP
+#define PYTHON_PYABACUS_SRC_PY_DIAGO_DAVID_HPP
+
+#include <complex>
+#include <functional>
+
+#include <pybind11/pybind11.h>
+#include <pybind11/complex.h>
+#include <pybind11/functional.h>
+#include <pybind11/numpy.h>
+#include <pybind11/stl.h>
+
+#include "module_hsolver/diago_david.h"
+
+namespace py = pybind11;
+
+namespace py_hsolver
+{
+
+class PyDiagoDavid
+{
+public:
+    PyDiagoDavid(int nbasis, int nband) : nbasis(nbasis), nband(nband)
+    {
+        psi = new std::complex<double>[nbasis * nband];
+        eigenvalue = new double[nband];
+    }
+
+    PyDiagoDavid(const PyDiagoDavid&) = delete;
+    PyDiagoDavid& operator=(const PyDiagoDavid&) = delete;
+    PyDiagoDavid(PyDiagoDavid&& other) : nbasis(other.nbasis), nband(other.nband)
+    {
+        psi = other.psi;
+        eigenvalue = other.eigenvalue;
+
+        other.psi = nullptr;
+        other.eigenvalue = nullptr;
+    }
+
+    ~PyDiagoDavid()
+    {
+        if (psi != nullptr) 
+        { 
+            delete[] psi; 
+            psi = nullptr; 
+        }
+        if (eigenvalue != nullptr) 
+        { 
+            delete[] eigenvalue; 
+            eigenvalue = nullptr; 
+        }
+    }
+
+    void set_psi(py::array_t<std::complex<double>> psi_in)
+    {
+        assert(psi_in.size() == nbasis * nband);
+
+        for (size_t i = 0; i < nbasis * nband; ++i)
+        {
+            psi[i] = psi_in.at(i);
+        }
+    }
+
+    py::array_t<std::complex<double>> get_psi()
+    {
+        py::array_t<std::complex<double>> psi_out(nband * nbasis);
+        py::buffer_info psi_out_buf = psi_out.request();
+
+        std::complex<double>* psi_out_ptr = static_cast<std::complex<double>*>(psi_out_buf.ptr);
+
+        for (size_t i = 0; i < nband * nbasis; ++i)
+        {
+            psi_out_ptr[i] = psi[i];
+        }
+
+        return psi_out;   
+    }
+
+    void init_eigenvalue()
+    {
+        for (size_t i = 0; i < nband; ++i)
+        {
+            eigenvalue[i] = 0.0;
+        }
+    }
+
+    py::array_t<double> get_eigenvalue()
+    {
+        py::array_t<double> eigenvalue_out(nband);
+        py::buffer_info eigenvalue_out_buf = eigenvalue_out.request();
+
+        double* eigenvalue_out_ptr = static_cast<double*>(eigenvalue_out_buf.ptr);
+
+        for (size_t i = 0; i < nband; ++i)
+        {
+            eigenvalue_out_ptr[i] = eigenvalue[i];
+        }
+
+        return eigenvalue_out;
+    }
+
+    int diag(
+        std::function<py::array_t<std::complex<double>>(py::array_t<std::complex<double>>)> mm_op,
+        std::vector<double> precond_vec,
+        int dav_ndim,
+        double tol,
+        int max_iter,
+        bool use_paw,
+        hsolver::diag_comm_info comm_info
+    ) {
+        auto hpsi_func = [mm_op] (
+            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
+        ) {
+            // Note: numpy's py::array_t is row-major, but
+            //       our raw pointer-array is column-major
+            py::array_t<std::complex<double>, py::array::f_style> psi({nbasis_in, band_index2 - band_index1 + 1});
+            py::buffer_info psi_buf = psi.request();
+            std::complex<double>* psi_ptr = static_cast<std::complex<double>*>(psi_buf.ptr);
+            std::copy(psi_in + band_index1 * nbasis_in, psi_in + (band_index2 + 1) * nbasis_in, psi_ptr);
+
+            py::array_t<std::complex<double>, py::array::f_style> hpsi = mm_op(psi);
+
+            py::buffer_info hpsi_buf = hpsi.request();
+            std::complex<double>* hpsi_ptr = static_cast<std::complex<double>*>(hpsi_buf.ptr);
+            std::copy(hpsi_ptr, hpsi_ptr + (band_index2 - band_index1 + 1) * nbasis_in, hpsi_out);
+        };
+
+        auto spsi_func = [this] (
+            const std::complex<double> *psi_in, 
+            std::complex<double> *spsi_out, 
+            const int nrow, 
+            const int npw, 
+            const int nbands
+        ) {
+            syncmem_op()(this->ctx, this->ctx, spsi_out, psi_in, static_cast<size_t>(nbands * nrow));
+        };
+
+        obj = std::make_unique<hsolver::DiagoDavid<std::complex<double>, base_device::DEVICE_CPU>>(
+            precond_vec.data(), 
+            nband, 
+            nbasis, 
+            dav_ndim, 
+            use_paw,
+            comm_info
+        );
+
+        return obj->diag(hpsi_func, spsi_func, nbasis, psi, eigenvalue, tol, max_iter);
+    }
+
+private:
+    std::complex<double>* psi = nullptr;
+    double* eigenvalue = nullptr;
+
+    int nbasis;
+    int nband;
+
+    std::unique_ptr<hsolver::DiagoDavid<std::complex<double>, base_device::DEVICE_CPU>> obj;
+
+    base_device::DEVICE_CPU* ctx = {};
+    using syncmem_op = base_device::memory::synchronize_memory_op<std::complex<double>, base_device::DEVICE_CPU, base_device::DEVICE_CPU>;
+};
+
+} // namespace py_hsolver
+
+#endif