diff --git a/python/pyabacus/CMakeLists.txt b/python/pyabacus/CMakeLists.txt index d2616366b06..e21f5e14463 100644 --- a/python/pyabacus/CMakeLists.txt +++ b/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) diff --git a/python/pyabacus/README.md b/python/pyabacus/README.md index e4b97d6ac3d..f751e52c1bb 100644 --- a/python/pyabacus/README.md +++ b/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 diff --git a/python/pyabacus/examples/diago_matrix.py b/python/pyabacus/examples/diago_matrix.py index 23ee1fbac6a..b7dba3865c5 100644 --- a/python/pyabacus/examples/diago_matrix.py +++ b/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) \ No newline at end of file +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) + + \ No newline at end of file diff --git a/python/pyabacus/src/py_abacus.cpp b/python/pyabacus/src/py_abacus.cpp index db450cc0052..5874667761a 100644 --- a/python/pyabacus/src/py_abacus.cpp +++ b/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); } \ No newline at end of file diff --git a/python/pyabacus/src/py_diago_david.hpp b/python/pyabacus/src/py_diago_david.hpp new file mode 100644 index 00000000000..a53147f6ce3 --- /dev/null +++ b/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 +#include + +#include +#include +#include +#include +#include + +#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[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> 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> get_psi() + { + py::array_t> psi_out(nband * nbasis); + py::buffer_info psi_out_buf = psi_out.request(); + + std::complex* psi_out_ptr = static_cast*>(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 get_eigenvalue() + { + py::array_t eigenvalue_out(nband); + py::buffer_info eigenvalue_out_buf = eigenvalue_out.request(); + + double* eigenvalue_out_ptr = static_cast(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>)> mm_op, + std::vector 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 *hpsi_out, + std::complex *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, py::array::f_style> psi({nbasis_in, band_index2 - band_index1 + 1}); + py::buffer_info psi_buf = psi.request(); + std::complex* psi_ptr = static_cast*>(psi_buf.ptr); + std::copy(psi_in + band_index1 * nbasis_in, psi_in + (band_index2 + 1) * nbasis_in, psi_ptr); + + py::array_t, py::array::f_style> 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 + (band_index2 - band_index1 + 1) * nbasis_in, hpsi_out); + }; + + auto spsi_func = [this] ( + const std::complex *psi_in, + std::complex *spsi_out, + const int nrow, + const int npw, + const int nbands + ) { + syncmem_op()(this->ctx, this->ctx, spsi_out, psi_in, static_cast(nbands * nrow)); + }; + + obj = std::make_unique, 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* psi = nullptr; + double* eigenvalue = nullptr; + + int nbasis; + int nband; + + std::unique_ptr, base_device::DEVICE_CPU>> obj; + + base_device::DEVICE_CPU* ctx = {}; + using syncmem_op = base_device::memory::synchronize_memory_op, base_device::DEVICE_CPU, base_device::DEVICE_CPU>; +}; + +} // namespace py_hsolver + +#endif \ No newline at end of file diff --git a/python/pyabacus/src/py_diago_dav_subspace.cpp b/python/pyabacus/src/py_hsolver.cpp similarity index 61% rename from python/pyabacus/src/py_diago_dav_subspace.cpp rename to python/pyabacus/src/py_hsolver.cpp index c9a28b49ae4..df8f6782562 100644 --- a/python/pyabacus/src/py_diago_dav_subspace.cpp +++ b/python/pyabacus/src/py_hsolver.cpp @@ -10,11 +10,12 @@ #include "module_base/module_device/types.h" #include "./py_diago_dav_subspace.hpp" +#include "./py_diago_david.hpp" namespace py = pybind11; using namespace pybind11::literals; -void bind_diago_dav_subspace(py::module& m) +void bind_hsolver(py::module& m) { py::module hsolver = m.def_submodule("hsolver"); @@ -90,4 +91,60 @@ void bind_diago_dav_subspace(py::module& m) .def("get_eigenvalue", &py_hsolver::PyDiagoDavSubspace::get_eigenvalue, R"pbdoc( Get the eigenvalues. )pbdoc"); + + py::class_(hsolver, "diago_david") + .def(py::init(), R"pbdoc( + Constructor of diago_david, a class for diagonalizing + a linear operator using the Davidson 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. + + Parameters + ---------- + nbasis : int + The number of basis functions. + nband : int + The number of bands to be calculated. + )pbdoc", "nbasis"_a, "nband"_a) + .def("diag", &py_hsolver::PyDiagoDavid::diag, R"pbdoc( + Diagonalize the linear operator using the Davidson 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. + precond_vec : np.ndarray + The preconditioner vector. + dav_ndim : int + The number of vectors, which is a multiple of the number of + eigenvectors to be calculated. + tol : double + The tolerance for the convergence. + max_iter : int + The maximum number of iterations. + use_paw : bool + Whether to use the projector augmented wave method. + )pbdoc", + "mm_op"_a, + "precond_vec"_a, + "dav_ndim"_a, + "tol"_a, + "max_iter"_a, + "use_paw"_a, + "comm_info"_a) + .def("set_psi", &py_hsolver::PyDiagoDavid::set_psi, R"pbdoc( + Set the initial guess of the eigenvectors, i.e. the wave functions. + )pbdoc", "psi_in"_a) + .def("get_psi", &py_hsolver::PyDiagoDavid::get_psi, R"pbdoc( + Get the eigenvectors. + )pbdoc") + .def("init_eigenvalue", &py_hsolver::PyDiagoDavid::init_eigenvalue, R"pbdoc( + Initialize the eigenvalues as zero. + )pbdoc") + .def("get_eigenvalue", &py_hsolver::PyDiagoDavid::get_eigenvalue, R"pbdoc( + Get the eigenvalues. + )pbdoc"); } diff --git a/python/pyabacus/src/pyabacus/hsolver/_hsolver.py b/python/pyabacus/src/pyabacus/hsolver/_hsolver.py index dc9fa1441ad..5f5875a0554 100644 --- a/python/pyabacus/src/pyabacus/hsolver/_hsolver.py +++ b/python/pyabacus/src/pyabacus/hsolver/_hsolver.py @@ -101,4 +101,75 @@ def dav_subspace( v = _diago_obj_dav_subspace.get_psi() return e, v + +def davidson( + mm_op: Callable[[NDArray[np.complex128]], NDArray[np.complex128]], + init_v: NDArray[np.complex128], + dim: int, + num_eigs: int, + pre_condition: NDArray[np.float64], + dav_ndim: int = 2, + tol: float = 1e-2, + max_iter: int = 1000, + use_paw: bool = False, + # is_occupied: Union[List[bool], None] = None, + # scf_type: bool = False +) -> Tuple[NDArray[np.float64], NDArray[np.complex128]]: + """ A function to diagonalize a matrix using the Davidson-Subspace 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. + 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. + pre_condition : NDArray[np.float64] + The preconditioner. + dav_ndim : int, optional + The number of vectors in the subspace, by default 2. + tol : float, optional + The tolerance for the convergence, by default 1e-2. + max_iter : int, optional + The maximum number of iterations, by default 1000. + use_paw : bool, optional + Whether to use projector augmented wave (PAW) method, by default False. + + Returns + ------- + e : NDArray[np.float64] + The eigenvalues. + v : NDArray[np.complex128] + The eigenvectors corresponding to the eigenvalues. + """ + if not callable(mm_op): + raise TypeError("mm_op must be a callable object.") + + if init_v.ndim != 1 or init_v.dtype != np.complex128: + init_v = init_v.flatten().astype(np.complex128, order='C') + + _diago_obj_dav_subspace = hsolver.diago_david(dim, num_eigs) + _diago_obj_dav_subspace.set_psi(init_v) + _diago_obj_dav_subspace.init_eigenvalue() + + comm_info = hsolver.diag_comm_info(0, 1) + + _ = _diago_obj_dav_subspace.diag( + mm_op, + pre_condition, + dav_ndim, + tol, + max_iter, + use_paw, + comm_info + ) + + e = _diago_obj_dav_subspace.get_eigenvalue() + v = _diago_obj_dav_subspace.get_psi() + + return e, v \ No newline at end of file diff --git a/python/pyabacus/tests/test_diago_dav_subspace.py b/python/pyabacus/tests/test_hsolver.py similarity index 59% rename from python/pyabacus/tests/test_diago_dav_subspace.py rename to python/pyabacus/tests/test_hsolver.py index dcf70bdcdee..9106bdbf7e2 100644 --- a/python/pyabacus/tests/test_diago_dav_subspace.py +++ b/python/pyabacus/tests/test_hsolver.py @@ -5,7 +5,11 @@ import numpy as np import scipy -def diag_pyabacus(h_sparse, nband): +def diag_pyabacus(h_sparse, nband, method): + algo = { + 'dav_subspace': hsolver.dav_subspace, + 'davidson': hsolver.davidson + } def mm_op(x): return h_sparse.dot(x) @@ -17,7 +21,7 @@ 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, _ = hsolver.dav_subspace( + e, _ = algo[method]( mm_op, v0, nbasis, @@ -25,8 +29,7 @@ def mm_op(x): precond, dav_ndim=8, tol=1e-12, - max_iter=5000, - scf_type=True + max_iter=5000 ) return e @@ -35,22 +38,28 @@ def diag_eigsh(h_sparse, nband): e, _ = scipy.sparse.linalg.eigsh(h_sparse, k=nband, which='SA', maxiter=5000, tol=1e-12) return e -def test_random_matrix_diag(): +@pytest.mark.parametrize("method", [ + ('dav_subspace'), + ('davidson') +]) +def test_random_matrix_diag(method): np.random.seed(12) n = 500 h_sparse = np.random.rand(n,n) h_sparse = h_sparse + h_sparse.conj().T + np.diag(np.random.random(n))*10 - e_pyabacus = diag_pyabacus(h_sparse, 8) + e_pyabacus = diag_pyabacus(h_sparse, 8, method) e_scipy = diag_eigsh(h_sparse, 8) np.testing.assert_allclose(e_pyabacus, e_scipy, atol=1e-8) -@pytest.mark.parametrize("file_name, nband, atol", [ - ('./test_diag/Si2.mat', 16, 1e-8), - ('./test_diag/Na5.mat', 16, 1e-8) +@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/Na5.mat', 16, 1e-8, 'dav_subspace'), + ('./test_diag/Na5.mat', 16, 1e-8, 'davidson'), ]) -def test_diag(file_name, nband, atol): +def test_diag(file_name, nband, atol, method): h_sparse = scipy.io.loadmat(file_name)['Problem']['A'][0, 0] - e_pyabacus = diag_pyabacus(h_sparse, nband) + e_pyabacus = diag_pyabacus(h_sparse, nband, method) e_scipy = diag_eigsh(h_sparse, nband) np.testing.assert_allclose(e_pyabacus, e_scipy, atol=atol)