Feature: Add CG algorithm to pyabacus.hsolver (#5398)

a1henu · web-flow · commit 5f9a1c64c175 · 2024-11-04T11:51:59.000+08:00
* fix some typos in `_hsolver.py` * fix some bugs caused by #5134 * Refactor hsolver module and remove unused code * refactor the structure of pythonization source code * fix some bug * Refactor __getattr__ function in __init__.py to handle attribute errors * fix some bugs * Add CONTRIBUTING.md to facilitate contributing to pyabacus project * fix typos * Update CONTRIBUTING.md * Update CONTRIBUTING.md * update README.md and CONTRIBUTING.md * update README.md * update CONTRIBUTING.md * update CONTRIBUTING.md * fix a bug caused by tuple in python3.8 * update * add basic framework for diagoCG * add cd to hsolver * change the signature of cg * add cg diagnolization method * add __all__ * remove unused code in example
diff --git a/python/pyabacus/examples/diago_matrix.py b/python/pyabacus/examples/diago_matrix.py
@@ -10,10 +10,13 @@ def load_mat(mat_file):
     return h_mat, nbasis, nband
 
 def calc_eig_pyabacus(mat_file, method):
-    algo = {
+    dav = {
         'dav_subspace': hsolver.dav_subspace,
         'davidson': hsolver.davidson
     }
+    cg = {
+        'cg': hsolver.cg
+    }
     
     h_mat, nbasis, nband = load_mat(mat_file)
     
@@ -22,25 +25,27 @@ def calc_eig_pyabacus(mat_file, method):
     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):
+    def mvv_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
-    )
+    if method in dav:
+        algo = dav[method]
+        # args: mvvop, init_v, dim, num_eigs, precondition, dav_ndim, tol, max_iter
+        args = (mvv_op, v0, nbasis, nband, precond, 8, 1e-12, 5000)
+    elif method in cg:
+        algo = cg[method]
+        # args: mvvop, init_v, dim, num_eigs, precondition, tol, max_iter
+        args = (mvv_op, v0, nbasis, nband, precond, 1e-12, 5000)
+    else:
+        raise ValueError(f"Method {method} not available")
+    
+    e, _ = algo(*args)
 
     print(f'eigenvalues calculated by pyabacus-{method} is: \n', e)
     
     return e
 
-def calc_eig_scipy(mat_file):
+def calc_eigsh(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)
@@ -50,13 +55,12 @@ def calc_eig_scipy(mat_file):
 
 if __name__ == '__main__':
     mat_file = './Si2.mat'
-    method = ['dav_subspace', 'davidson']
+    method = ['dav_subspace', 'davidson', 'cg']
     
     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)
+        e_scipy = calc_eigsh(mat_file)
         
         print('eigenvalues difference: \n', e_pyabacus - e_scipy)
-    
-    
+        
diff --git a/python/pyabacus/src/hsolver/CMakeLists.txt b/python/pyabacus/src/hsolver/CMakeLists.txt
@@ -2,6 +2,7 @@
 list(APPEND _diago
     ${HSOLVER_PATH}/diago_dav_subspace.cpp
     ${HSOLVER_PATH}/diago_david.cpp
+    ${HSOLVER_PATH}/diago_cg.cpp
     ${HSOLVER_PATH}/diag_const_nums.cpp
     ${HSOLVER_PATH}/diago_iter_assist.cpp
 
diff --git a/python/pyabacus/src/hsolver/py_diago_cg.hpp b/python/pyabacus/src/hsolver/py_diago_cg.hpp
@@ -0,0 +1,192 @@
+#ifndef PYTHON_PYABACUS_SRC_PY_DIAGO_CG_HPP
+#define PYTHON_PYABACUS_SRC_PY_DIAGO_CG_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 <ATen/core/tensor.h>
+#include <ATen/core/tensor_map.h>
+#include <ATen/core/tensor_types.h>
+
+#include "module_hsolver/diago_cg.h"
+#include "module_base/module_device/memory_op.h"
+
+namespace py = pybind11;
+
+namespace py_hsolver
+{
+
+class PyDiagoCG
+{
+public:
+    PyDiagoCG(int dim, int num_eigs) : dim{dim}, num_eigs{num_eigs} { }
+    PyDiagoCG(const PyDiagoCG&) = delete;
+    PyDiagoCG& operator=(const PyDiagoCG&) = delete;
+    PyDiagoCG(PyDiagoCG&& other)
+    {
+        psi = other.psi;
+        other.psi = nullptr;
+
+        eig = other.eig;
+        other.eig = nullptr;
+    }
+
+    ~PyDiagoCG() 
+    {
+        if (psi != nullptr) 
+        {
+            delete psi;
+            psi = nullptr;
+        }
+
+        if (eig != nullptr)
+        {
+            delete eig;
+            eig = nullptr;
+        }
+    }
+
+    void init_eig() 
+    {
+        eig = new ct::Tensor(ct::DataType::DT_DOUBLE, {num_eigs});
+        eig->zero();
+    }
+
+    py::array_t<double> get_eig() 
+    {
+        py::array_t<double> eig_out(eig->NumElements());
+        py::buffer_info eig_buf = eig_out.request();
+        double* eig_out_ptr = static_cast<double*>(eig_buf.ptr);
+
+        if (eig == nullptr) {
+            throw std::runtime_error("eig is not initialized");
+        }
+        double* eig_ptr = eig->data<double>();
+
+        std::copy(eig_ptr, eig_ptr + eig->NumElements(), eig_out_ptr);
+        return eig_out;
+    }
+
+    void set_psi(py::array_t<std::complex<double>> psi_in)
+    {
+        py::buffer_info psi_buf = psi_in.request();
+        std::complex<double>* psi_ptr = static_cast<std::complex<double>*>(psi_buf.ptr);
+
+        psi = new ct::TensorMap(
+            psi_ptr,
+            ct::DataType::DT_COMPLEX_DOUBLE,
+            ct::DeviceType::CpuDevice,
+            ct::TensorShape({num_eigs, dim})
+        );
+    }
+
+    py::array_t<std::complex<double>> get_psi()
+    {
+        py::array_t<std::complex<double>> psi_out({num_eigs, dim});
+        py::buffer_info psi_buf = psi_out.request();
+        std::complex<double>* psi_out_ptr = static_cast<std::complex<double>*>(psi_buf.ptr);
+
+        if (psi == nullptr) {
+            throw std::runtime_error("psi is not initialized");
+        }
+        std::complex<double>* psi_ptr = psi->data<std::complex<double>>();
+
+        std::copy(psi_ptr, psi_ptr + psi->NumElements(), psi_out_ptr);
+        return psi_out;
+    }
+
+    void set_prec(py::array_t<double> prec_in)
+    {
+        py::buffer_info prec_buf = prec_in.request();
+        double* prec_ptr = static_cast<double*>(prec_buf.ptr);
+
+        prec = new ct::TensorMap(
+            prec_ptr,
+            ct::DataType::DT_DOUBLE,
+            ct::DeviceType::CpuDevice,
+            ct::TensorShape({dim})
+        );
+    }
+
+    void diag(
+        std::function<py::array_t<std::complex<double>>(py::array_t<std::complex<double>>)> mm_op,
+        int diag_ndim,
+        double tol,
+        bool need_subspace,
+        bool scf_type,
+        int nproc_in_pool = 1
+    ) {
+        const std::string basis_type = "pw";
+        const std::string calculation = scf_type ? "scf" : "nscf";
+
+        auto hpsi_func = [mm_op] (const ct::Tensor& psi_in, ct::Tensor& hpsi_out) {
+            const auto ndim = psi_in.shape().ndim();
+            REQUIRES_OK(ndim <= 2, "dims of psi_in should be less than or equal to 2");
+            const int nvec   = ndim == 1 ? 1 : psi_in.shape().dim_size(0);
+            const int ld_psi = ndim == 1 ? psi_in.NumElements() : psi_in.shape().dim_size(1);
+
+            // Note: numpy's py::array_t is row-major, and
+            //       our tensor-array is row-major
+            py::array_t<std::complex<double>> psi({ld_psi, nvec});
+            py::buffer_info psi_buf = psi.request();
+            std::complex<double>* psi_ptr = static_cast<std::complex<double>*>(psi_buf.ptr);
+            std::copy(psi_in.data<std::complex<double>>(), psi_in.data<std::complex<double>>() + nvec * ld_psi, psi_ptr);
+
+            py::array_t<std::complex<double>> 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 + nvec * ld_psi, hpsi_out.data<std::complex<double>>());
+        };
+
+        auto subspace_func = [] (const ct::Tensor& psi_in, ct::Tensor& psi_out) { /*do nothing*/ };
+
+        auto spsi_func = [this] (const ct::Tensor& psi_in, ct::Tensor& spsi_out) {
+            const auto ndim = psi_in.shape().ndim();
+            REQUIRES_OK(ndim <= 2, "dims of psi_in should be less than or equal to 2");
+            const int nrow   = ndim == 1 ? psi_in.NumElements() : psi_in.shape().dim_size(1);
+            const int nbands = ndim == 1 ? 1 : psi_in.shape().dim_size(0);
+            syncmem_z2z_h2h_op()(
+                this->ctx,
+                this->ctx,
+                spsi_out.data<std::complex<double>>(), 
+                psi_in.data<std::complex<double>>(), 
+                static_cast<size_t>(nrow * nbands)
+            );
+        };
+
+        cg = std::make_unique<hsolver::DiagoCG<std::complex<double>, base_device::DEVICE_CPU>>(
+            basis_type,
+            calculation,
+            need_subspace,
+            subspace_func,
+            tol,
+            diag_ndim,
+            nproc_in_pool
+        );
+
+        cg->diag(hpsi_func, spsi_func, *psi, *eig, *prec);
+    }
+
+private:
+    base_device::DEVICE_CPU* ctx = {};
+
+    int dim;
+    int num_eigs;
+
+    ct::Tensor* psi = nullptr;
+    ct::Tensor* eig = nullptr;
+    ct::Tensor* prec = nullptr;
+
+    std::unique_ptr<hsolver::DiagoCG<std::complex<double>, base_device::DEVICE_CPU>> cg;
+};
+
+} // namespace py_hsolver
+
+#endif
diff --git a/python/pyabacus/src/hsolver/py_hsolver.cpp b/python/pyabacus/src/hsolver/py_hsolver.cpp
@@ -11,6 +11,7 @@
 
 #include "./py_diago_dav_subspace.hpp"
 #include "./py_diago_david.hpp"
+#include "./py_diago_cg.hpp"
 
 namespace py = pybind11;
 using namespace pybind11::literals;
@@ -144,6 +145,55 @@ void bind_hsolver(py::module& m)
         .def("get_eigenvalue", &py_hsolver::PyDiagoDavid::get_eigenvalue, R"pbdoc(
             Get the eigenvalues.        
         )pbdoc");
+
+    py::class_<py_hsolver::PyDiagoCG>(m, "diago_cg")
+        .def(py::init<int, int>(), R"pbdoc(
+            Constructor of diago_cg, a class for diagonalizing 
+            a linear operator using the Conjugate Gradient Method.
+
+            This class serves as a backend computation class. The interface 
+            for invoking this class is a function defined in _hsolver.py, 
+            which uses this class to perform the calculations.
+        )pbdoc")
+        .def("diag", &py_hsolver::PyDiagoCG::diag, R"pbdoc(
+            Diagonalize the linear operator using the Conjugate Gradient Method.
+
+            Parameters
+            ----------
+            mm_op : Callable[[NDArray[np.complex128]], NDArray[np.complex128]],
+                The operator to be diagonalized, which is a function that takes a matrix as input
+                and returns a matrix mv_op(X) = H * X as output.
+            max_iter : int
+                The maximum number of iterations.
+            tol : double
+                The tolerance for the convergence.
+            need_subspace : bool
+                Whether to use the subspace function.
+            scf_type : bool
+                Whether to use the SCF type, which is used to determine the
+                convergence criterion.
+        )pbdoc",
+        "mm_op"_a,
+        "max_iter"_a,
+        "tol"_a,
+        "need_subspace"_a,
+        "scf_type"_a,
+        "nproc_in_pool"_a)
+        .def("init_eig", &py_hsolver::PyDiagoCG::init_eig, R"pbdoc(
+            Initialize the eigenvalues.
+        )pbdoc")
+        .def("get_eig", &py_hsolver::PyDiagoCG::get_eig, R"pbdoc(
+            Get the eigenvalues.
+        )pbdoc")
+        .def("set_psi", &py_hsolver::PyDiagoCG::set_psi, R"pbdoc(
+            Set the eigenvectors.
+        )pbdoc", "psi_in"_a)
+        .def("get_psi", &py_hsolver::PyDiagoCG::get_psi, R"pbdoc(
+            Get the eigenvectors.
+        )pbdoc")
+        .def("set_prec", &py_hsolver::PyDiagoCG::set_prec, R"pbdoc(
+            Set the preconditioner.
+        )pbdoc", "prec_in"_a);
 }
 
 PYBIND11_MODULE(_hsolver_pack, m)
diff --git a/python/pyabacus/src/pyabacus/hsolver/__init__.py b/python/pyabacus/src/pyabacus/hsolver/__init__.py
@@ -1,4 +1,4 @@
 from __future__ import annotations
 from ._hsolver import *
 
-__all__ = ["diag_comm_info", "dav_subspace", "davidson"]
+__all__ = ["diag_comm_info", "dav_subspace", "davidson", "cg"]
diff --git a/python/pyabacus/src/pyabacus/hsolver/_hsolver.py b/python/pyabacus/src/pyabacus/hsolver/_hsolver.py
diff --git a/python/pyabacus/tests/test_hsolver.py b/python/pyabacus/tests/test_hsolver.py
