diff --git a/python/pyabacus/examples/diago_matrix.py b/python/pyabacus/examples/diago_matrix.py index 6a3d1eb75ae..657a3a679e8 100644 --- 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) - - \ No newline at end of file + \ No newline at end of file diff --git a/python/pyabacus/src/hsolver/CMakeLists.txt b/python/pyabacus/src/hsolver/CMakeLists.txt index 6bf690ecec1..853973e6536 100644 --- 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 new file mode 100644 index 00000000000..2499b96e113 --- /dev/null +++ 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 +#include + +#include +#include +#include +#include +#include + +#include +#include +#include + +#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 get_eig() + { + py::array_t eig_out(eig->NumElements()); + py::buffer_info eig_buf = eig_out.request(); + double* eig_out_ptr = static_cast(eig_buf.ptr); + + if (eig == nullptr) { + throw std::runtime_error("eig is not initialized"); + } + double* eig_ptr = eig->data(); + + std::copy(eig_ptr, eig_ptr + eig->NumElements(), eig_out_ptr); + return eig_out; + } + + void set_psi(py::array_t> psi_in) + { + py::buffer_info psi_buf = psi_in.request(); + std::complex* psi_ptr = static_cast*>(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> get_psi() + { + py::array_t> psi_out({num_eigs, dim}); + py::buffer_info psi_buf = psi_out.request(); + std::complex* psi_out_ptr = static_cast*>(psi_buf.ptr); + + if (psi == nullptr) { + throw std::runtime_error("psi is not initialized"); + } + std::complex* psi_ptr = psi->data>(); + + std::copy(psi_ptr, psi_ptr + psi->NumElements(), psi_out_ptr); + return psi_out; + } + + void set_prec(py::array_t prec_in) + { + py::buffer_info prec_buf = prec_in.request(); + double* prec_ptr = static_cast(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>)> 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> psi({ld_psi, nvec}); + py::buffer_info psi_buf = psi.request(); + std::complex* psi_ptr = static_cast*>(psi_buf.ptr); + std::copy(psi_in.data>(), psi_in.data>() + nvec * ld_psi, psi_ptr); + + py::array_t> hpsi = mm_op(psi); + + py::buffer_info hpsi_buf = hpsi.request(); + std::complex* hpsi_ptr = static_cast*>(hpsi_buf.ptr); + std::copy(hpsi_ptr, hpsi_ptr + nvec * ld_psi, hpsi_out.data>()); + }; + + 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>(), + psi_in.data>(), + static_cast(nrow * nbands) + ); + }; + + cg = std::make_unique, 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, base_device::DEVICE_CPU>> cg; +}; + +} // namespace py_hsolver + +#endif \ No newline at end of file diff --git a/python/pyabacus/src/hsolver/py_hsolver.cpp b/python/pyabacus/src/hsolver/py_hsolver.cpp index 2401832f725..7e76d1f874b 100644 --- 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_(m, "diago_cg") + .def(py::init(), 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 index d9d954d9ee0..207fcd785f7 100644 --- 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"] \ No newline at end of file +__all__ = ["diag_comm_info", "dav_subspace", "davidson", "cg"] \ No newline at end of file diff --git a/python/pyabacus/src/pyabacus/hsolver/_hsolver.py b/python/pyabacus/src/pyabacus/hsolver/_hsolver.py index 9dece524a74..d1cebad5ddc 100644 --- a/python/pyabacus/src/pyabacus/hsolver/_hsolver.py +++ b/python/pyabacus/src/pyabacus/hsolver/_hsolver.py @@ -9,7 +9,7 @@ from typing import Tuple, List, Union, Callable from ._hsolver_pack import diag_comm_info as _diag_comm_info -from ._hsolver_pack import diago_dav_subspace, diago_david +from ._hsolver_pack import diago_dav_subspace, diago_david, diago_cg class diag_comm_info(_diag_comm_info): def __init__(self, rank: int, nproc: int): @@ -28,7 +28,7 @@ def dav_subspace( init_v: NDArray[np.complex128], dim: int, num_eigs: int, - pre_condition: NDArray[np.float64], + precondition: NDArray[np.float64], dav_ndim: int = 2, tol: float = 1e-2, max_iter: int = 1000, @@ -50,7 +50,7 @@ def dav_subspace( The number of basis, i.e. the number of rows/columns in the matrix. num_eigs : int The number of bands to calculate, i.e. the number of eigenvalues to calculate. - pre_condition : NDArray[np.float64] + precondition : NDArray[np.float64] The preconditioner. dav_ndim : int, optional The number of vectors in the subspace, by default 2. @@ -94,7 +94,7 @@ def dav_subspace( _ = _diago_obj_dav_subspace.diag( mvv_op, - pre_condition, + precondition, dav_ndim, tol, max_iter, @@ -114,7 +114,7 @@ def davidson( init_v: NDArray[np.complex128], dim: int, num_eigs: int, - pre_condition: NDArray[np.float64], + precondition: NDArray[np.float64], dav_ndim: int = 2, tol: float = 1e-2, max_iter: int = 1000, @@ -135,7 +135,7 @@ def davidson( The number of basis, i.e. the number of rows/columns in the matrix. num_eigs : int The number of bands to calculate, i.e. the number of eigenvalues to calculate. - pre_condition : NDArray[np.float64] + precondition : NDArray[np.float64] The preconditioner. dav_ndim : int, optional The number of vectors in the subspace, by default 2. @@ -167,7 +167,7 @@ def davidson( _ = _diago_obj_david.diag( mvv_op, - pre_condition, + precondition, dav_ndim, tol, max_iter, @@ -179,4 +179,81 @@ def davidson( v = _diago_obj_david.get_psi() return e, v - \ No newline at end of file + +def cg( + mvv_op: Callable[[NDArray[np.complex128]], NDArray[np.complex128]], + init_v: NDArray[np.complex128], + dim: int, + num_eigs: int, + precondition: NDArray[np.float64], + tol: float = 1e-2, + max_iter: int = 1000, + need_subspace: bool = False, + scf_type: bool = False, + nproc_in_pool: int = 1 +) -> Tuple[NDArray[np.float64], NDArray[np.complex128]]: + """ A function to diagonalize a matrix using the Conjugate Gradient method. + + Parameters + ---------- + mvv_op : Callable[[NDArray[np.complex128]], NDArray[np.complex128]], + The operator to be diagonalized, which is a function that takes a set of + vectors X = [x1, ..., xN] as input and returns a matrix(vector block) + mvv_op(X) = H * X ([Hx1, ..., HxN]) as output. + init_v : NDArray[np.complex128] + The initial guess for the eigenvectors. + dim : int + The number of basis, i.e. the number of rows/columns in the matrix. + num_eigs : int + The number of bands to calculate, i.e. the number of eigenvalues to calculate. + precondition : NDArray[np.float64] + The preconditioner. + max_iter : int, optional + The maximum number of iterations, by default 1000. + tol : float, optional + The tolerance for the convergence, by default 1e-2. + need_subspace : bool, optional + Whether to use subspace function, by default False. + scf_type : bool, optional + Indicates whether the calculation is a self-consistent field (SCF) calculation. + If True, the initial precision of eigenvalue calculation can be coarse. + If False, it indicates a non-self-consistent field (non-SCF) calculation, + where high precision in eigenvalue calculation is required from the start. + nproc_in_pool : int, optional + The number of processes in the pool, by default 1. + + Returns + ------- + e : NDArray[np.float64] + The eigenvalues. + v : NDArray[np.complex128] + The eigenvectors corresponding to the eigenvalues. + """ + if not callable(mvv_op): + raise TypeError("mvv_op must be a callable object.") + + if init_v.ndim != 1 or init_v.dtype != np.complex128: + # the shape of init_v is (num_eigs, dim) = (dim, num_eigs).T + if init_v.ndim == 2: + init_v = init_v.T + init_v = init_v.flatten().astype(np.complex128, order='C') + + _diago_obj_cg = diago_cg(dim, num_eigs) + _diago_obj_cg.set_psi(init_v) + _diago_obj_cg.init_eig() + + _diago_obj_cg.set_prec(precondition) + + _diago_obj_cg.diag( + mvv_op, + max_iter, + tol, + need_subspace, + scf_type, + nproc_in_pool + ) + + e = _diago_obj_cg.get_eig() + v = _diago_obj_cg.get_psi() + + return e, v \ No newline at end of file diff --git a/python/pyabacus/tests/test_hsolver.py b/python/pyabacus/tests/test_hsolver.py index 9106bdbf7e2..82b54c5a3ae 100644 --- a/python/pyabacus/tests/test_hsolver.py +++ b/python/pyabacus/tests/test_hsolver.py @@ -6,10 +6,13 @@ import scipy def diag_pyabacus(h_sparse, nband, method): - algo = { + dav = { 'dav_subspace': hsolver.dav_subspace, 'davidson': hsolver.davidson } + cg = { + 'cg': hsolver.cg + } def mm_op(x): return h_sparse.dot(x) @@ -21,16 +24,16 @@ def mm_op(x): diag_elem = np.where(np.abs(diag_elem) < 1e-8, 1e-8, diag_elem) precond = 1.0 / np.abs(diag_elem) - e, _ = algo[method]( - mm_op, - v0, - nbasis, - nband, - precond, - dav_ndim=8, - tol=1e-12, - max_iter=5000 - ) + if method in dav: + algo = dav[method] + args = (mm_op, v0, nbasis, nband, precond, 8, 1e-12, 5000) + elif method in cg: + algo = cg[method] + args = (mm_op, v0, nbasis, nband, precond, 1e-12, 5000) + else: + raise ValueError(f"Method {method} not available") + + e, _ = algo(*args) return e @@ -40,7 +43,8 @@ def diag_eigsh(h_sparse, nband): @pytest.mark.parametrize("method", [ ('dav_subspace'), - ('davidson') + ('davidson'), + ('cg') ]) def test_random_matrix_diag(method): np.random.seed(12) @@ -55,8 +59,10 @@ def test_random_matrix_diag(method): @pytest.mark.parametrize("file_name, nband, atol, method", [ ('./test_diag/Si2.mat', 16, 1e-8, 'dav_subspace'), ('./test_diag/Si2.mat', 16, 1e-8, 'davidson'), + ('./test_diag/Si2.mat', 16, 1e-8, 'cg'), ('./test_diag/Na5.mat', 16, 1e-8, 'dav_subspace'), ('./test_diag/Na5.mat', 16, 1e-8, 'davidson'), + ('./test_diag/Na5.mat', 16, 1e-8, 'cg'), ]) def test_diag(file_name, nband, atol, method): h_sparse = scipy.io.loadmat(file_name)['Problem']['A'][0, 0]