diff --git a/python/pyabacus/CMakeLists.txt b/python/pyabacus/CMakeLists.txt index 3fd07f724e..d2616366b0 100644 --- a/python/pyabacus/CMakeLists.txt +++ b/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}) diff --git a/python/pyabacus/README.md b/python/pyabacus/README.md index dc9a533613..e4b97d6ac3 100644 --- a/python/pyabacus/README.md +++ b/python/pyabacus/README.md @@ -1,11 +1,10 @@ Build Example: TwoCenterIntegral Section in ABACUS -============== +================================================== An example project built with [pybind11](https://github.com/pybind/pybind11) and scikit-build-core. Python 3.7+ (see older commits for older versions of Python). - Installation ------------ @@ -13,21 +12,22 @@ Installation - Clone ABACUS main repository and `cd abacus-develop/python/pyabacus`. - Build pyabacus by `pip install -v .` or install test dependencies & build pyabacus by `pip install .[test]`. (Use `pip install -v .[test] -i https://pypi.tuna.tsinghua.edu.cn/simple` to accelerate installation process.) - CI Examples ----------- There are examples for CI in `.github/workflows`. A simple way to produces binary "wheels" for all platforms is illustrated in the "wheels.yml" file, -using [`cibuildwheel`][]. +using . Use `pytest -v` to run all the unit tests for pyabacus in the local machine. + ```shell $ cd tests/ $ pytest -v ``` Run `python vis_nao.py` to visualize the numerical orbital. + ```shell $ cd examples/ $ python vis_nao.py @@ -41,6 +41,16 @@ $ python ex_s_rotate.py norm(S_e3 - S_numer) = 3.341208104032616e-15 ``` +Run `python diago_matrix.py` in `examples` to check the diagonalization of a matrix. + +```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] +``` + License ------- @@ -58,4 +68,4 @@ s.sphbesj(1, 0.0) 0.0 ``` -[`cibuildwheel`]: https://cibuildwheel.readthedocs.io \ No newline at end of file +[`cibuildwheel`]: https://cibuildwheel.readthedocs.io diff --git a/python/pyabacus/examples/Si2.mat b/python/pyabacus/examples/Si2.mat new file mode 100644 index 0000000000..cf83b44cf1 Binary files /dev/null and b/python/pyabacus/examples/Si2.mat differ diff --git a/python/pyabacus/examples/diago_matrix.py b/python/pyabacus/examples/diago_matrix.py new file mode 100644 index 0000000000..23ee1fbac6 --- /dev/null +++ b/python/pyabacus/examples/diago_matrix.py @@ -0,0 +1,38 @@ +from pyabacus import hsolver +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 diff --git a/python/pyabacus/pyproject.toml b/python/pyabacus/pyproject.toml index 85bd47833c..fd3fecc9e9 100644 --- a/python/pyabacus/pyproject.toml +++ b/python/pyabacus/pyproject.toml @@ -10,6 +10,7 @@ description="A minimal two_center_integral package (with pybind11)" readme = "README.md" authors = [ { name = "Jie Li", email = "lij@aisi.ac.cn" }, + { name = "Chenxu Bai", email = "chenxu.bai@stu.pku.edu.cn" }, ] requires-python = ">=3.7" classifiers = [ diff --git a/python/pyabacus/src/py_abacus.cpp b/python/pyabacus/src/py_abacus.cpp index 9313323593..db450cc005 100644 --- a/python/pyabacus/src/py_abacus.cpp +++ b/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); } \ No newline at end of file diff --git a/python/pyabacus/src/py_diago_dav_subspace.cpp b/python/pyabacus/src/py_diago_dav_subspace.cpp new file mode 100644 index 0000000000..c9a28b49ae --- /dev/null +++ b/python/pyabacus/src/py_diago_dav_subspace.cpp @@ -0,0 +1,93 @@ +#include +#include +#include +#include +#include +#include + +#include "module_hsolver/diago_dav_subspace.h" +#include "module_hsolver/kernels/math_kernel_op.h" +#include "module_base/module_device/types.h" + +#include "./py_diago_dav_subspace.hpp" + +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") + .def(py::init(), "rank"_a, "nproc"_a) + .def_readonly("rank", &hsolver::diag_comm_info::rank) + .def_readonly("nproc", &hsolver::diag_comm_info::nproc); + + py::class_(hsolver, "diago_dav_subspace") + .def(py::init(), R"pbdoc( + Constructor of diago_dav_subspace, a class for diagonalizing + a linear operator using the Davidson-Subspace 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::PyDiagoDavSubspace::diag, R"pbdoc( + Diagonalize the linear operator 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. + 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. + need_subspace : bool + Whether to use the subspace function. + is_occupied : list[bool] + A list of boolean values indicating whether the band is occupied, + meaning that the corresponding eigenvalue is to be calculated. + scf_type : bool + Whether to use the SCF type, which is used to determine the + convergence criterion. + If true, it indicates a self-consistent field (SCF) calculation, + where 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. + )pbdoc", + "mm_op"_a, + "precond_vec"_a, + "dav_ndim"_a, + "tol"_a, + "max_iter"_a, + "need_subspace"_a, + "is_occupied"_a, + "scf_type"_a, + "comm_info"_a) + .def("set_psi", &py_hsolver::PyDiagoDavSubspace::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::PyDiagoDavSubspace::get_psi, R"pbdoc( + Get the eigenvectors. + )pbdoc") + .def("init_eigenvalue", &py_hsolver::PyDiagoDavSubspace::init_eigenvalue, R"pbdoc( + Initialize the eigenvalues as zero. + )pbdoc") + .def("get_eigenvalue", &py_hsolver::PyDiagoDavSubspace::get_eigenvalue, R"pbdoc( + Get the eigenvalues. + )pbdoc"); +} diff --git a/python/pyabacus/src/py_diago_dav_subspace.hpp b/python/pyabacus/src/py_diago_dav_subspace.hpp new file mode 100644 index 0000000000..f6cf7e9a5b --- /dev/null +++ b/python/pyabacus/src/py_diago_dav_subspace.hpp @@ -0,0 +1,158 @@ +#ifndef PYTHON_PYABACUS_SRC_PY_DIAGO_DAV_SUBSPACE_HPP +#define PYTHON_PYABACUS_SRC_PY_DIAGO_DAV_SUBSPACE_HPP + +#include +#include + +#include +#include +#include +#include +#include + +#include "module_hsolver/diago_dav_subspace.h" + +namespace py = pybind11; + +namespace py_hsolver +{ + +class PyDiagoDavSubspace +{ +public: + PyDiagoDavSubspace(int nbasis, int nband) : nbasis(nbasis), nband(nband) + { + psi = new std::complex[nbasis * nband]; + eigenvalue = new double[nband]; + } + + PyDiagoDavSubspace(const PyDiagoDavSubspace&) = delete; + PyDiagoDavSubspace& operator=(const PyDiagoDavSubspace&) = delete; + PyDiagoDavSubspace(PyDiagoDavSubspace&& other) : nbasis(other.nbasis), nband(other.nband) + { + psi = other.psi; + eigenvalue = other.eigenvalue; + + other.psi = nullptr; + other.eigenvalue = nullptr; + } + + ~PyDiagoDavSubspace() + { + 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 need_subspace, + std::vector is_occupied, + bool scf_type, + 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); + }; + + obj = std::make_unique, base_device::DEVICE_CPU>>( + precond_vec, + nband, + nbasis, + dav_ndim, + tol, + max_iter, + need_subspace, + comm_info + ); + + return obj->diag(hpsi_func, psi, nbasis, eigenvalue, is_occupied, scf_type); + } + +private: + std::complex* psi = nullptr; + double* eigenvalue = nullptr; + + int nbasis; + int nband; + + std::unique_ptr, base_device::DEVICE_CPU>> obj; +}; + +} // namespace py_hsolver + +#endif \ No newline at end of file diff --git a/python/pyabacus/src/py_m_nao.cpp b/python/pyabacus/src/py_m_nao.cpp index 826f11369e..d7c7283614 100644 --- a/python/pyabacus/src/py_m_nao.cpp +++ b/python/pyabacus/src/py_m_nao.cpp @@ -48,7 +48,7 @@ void bind_m_nao(py::module& m) "update"_a = 0) .def("set_uniform_grid", &RadialCollection::set_uniform_grid, - "Sets a common uniform grid for all RadialSet objects", + "Sets a common uniform grid for all RadialSet objects.", "for_r_space"_a, "ngrid"_a, "cutoff"_a, @@ -68,7 +68,7 @@ void bind_m_nao(py::module& m) } self.set_grid(for_r_space, ngrid, static_cast(grid_info.ptr), mode); }, - "Sets a common grid for all RadialSet objects", + "Sets a common grid for all RadialSet objects.", "for_r_space"_a, "ngrid"_a, "grid"_a, diff --git a/python/pyabacus/src/pyabacus/ModuleBase/__init__.py b/python/pyabacus/src/pyabacus/ModuleBase/__init__.py new file mode 100644 index 0000000000..95bf3a0461 --- /dev/null +++ b/python/pyabacus/src/pyabacus/ModuleBase/__init__.py @@ -0,0 +1,4 @@ +from __future__ import annotations +from ._module_base import * + +__all__ = ["Sphbes", "Integral", "SphericalBesselTransformer"] \ No newline at end of file diff --git a/python/pyabacus/src/pyabacus/ModuleBase/_module_base.py b/python/pyabacus/src/pyabacus/ModuleBase/_module_base.py new file mode 100644 index 0000000000..fe96247e43 --- /dev/null +++ b/python/pyabacus/src/pyabacus/ModuleBase/_module_base.py @@ -0,0 +1,105 @@ +""" +pyabacus.ModuleBase +=================== +Basic math functions and integrals. +""" + +import numpy as np + +from typing import overload +from numpy.typing import NDArray + +class Sphbes: + def __init__(self) -> None: ... + @overload + @staticmethod + def sphbesj(l: int, x: float) -> float: ... + @overload + @staticmethod + def sphbesj( + n: int, + r: NDArray[np.float64], + q: int, + l: int, + jl: NDArray[np.float64] + ) -> None: ... + @overload + @staticmethod + def dsphbesj(l: int, x: float) -> float: ... + @overload + @staticmethod + def dsphbesj( + n: int, + r: NDArray[np.float64], + q: int, + l: int, + djl: NDArray[np.float64] + ) -> None: ... + @staticmethod + def sphbes_zeros(l: int, n: int, zeros: NDArray[np.float64]) -> None: ... + +class Integral: + def __init__(self) -> None: ... + @overload + @staticmethod + def Simpson_Integral( + mesh: int, + func: NDArray[np.float64], + rab: NDArray[np.float64], + asum: float + ) -> float: ... + @overload + @staticmethod + def Simpson_Integral( + mesh: int, + func: NDArray[np.float64], + dr: float, + asum: float + ) -> float: ... + @staticmethod + def Simpson_Integral_0toall( + mesh: int, + func: NDArray[np.float64], + rab: NDArray[np.float64], + asum: NDArray[np.float64] + ) -> None: ... + @staticmethod + def Simpson_Integral_alltoinf( + mesh: int, + func: NDArray[np.float64], + rab: NDArray[np.float64], + asum: NDArray[np.float64] + ) -> None: ... + @overload + @staticmethod + def simpson( + n: int, + f: NDArray[np.float64], + dx: float + ) -> float: ... + @overload + @staticmethod + def simpson( + n: int, + f: NDArray[np.float64], + h: NDArray[np.float64], + ) -> float: ... + @overload + @staticmethod + def Gauss_Legendre_grid_and_weight( + n: int, + x: NDArray[np.float64], + w: NDArray[np.float64], + ) -> None: ... + @overload + @staticmethod + def Gauss_Legendre_grid_and_weight( + xmin: float, + xmax: float, + n: int, + x: NDArray[np.float64], + w: NDArray[np.float64], + ) -> None: ... + +class SphericalBesselTransformer: + def __init__(self) -> None: ... \ No newline at end of file diff --git a/python/pyabacus/src/pyabacus/ModuleNAO/__init__.py b/python/pyabacus/src/pyabacus/ModuleNAO/__init__.py new file mode 100644 index 0000000000..6c43ea250f --- /dev/null +++ b/python/pyabacus/src/pyabacus/ModuleNAO/__init__.py @@ -0,0 +1 @@ +from __future__ import annotations \ No newline at end of file diff --git a/python/pyabacus/src/pyabacus/ModuleNAO/_module_nao.py b/python/pyabacus/src/pyabacus/ModuleNAO/_module_nao.py new file mode 100644 index 0000000000..c27f2ffe5e --- /dev/null +++ b/python/pyabacus/src/pyabacus/ModuleNAO/_module_nao.py @@ -0,0 +1,102 @@ +""" +pyabacus.ModuleNAO +================== +Module for Numerical Atomic Orbitals (NAO) in ABACUS +""" + +import numpy as np + +from numpy.typing import NDArray +from pyabacus.ModuleBase import SphericalBesselTransformer +from typing import overload, List + + +class RadialCollection: + def __init__(self) -> None: + """ + A class that holds all numerical radial functions of the same kind. + + An instance of this class could be the collection of all radial functions + of numerical atomic orbitals, or all Kleinman-Bylander beta functions from + all elements involved in a calculation. + """ + pass + + def build( + self, + nfile: int, + file_list: List[str], + ftype: str = '\0' + ) -> None: + """ + Builds the collection from (orbital) files. + """ + pass + + def set_transformer( + self, + sbt: SphericalBesselTransformer, + update: int = 0 + ) -> None: + """ + Sets a spherical Bessel transformers for all RadialSet objects. + """ + pass + + def set_uniform_grid( + self, + for_r_space: bool, + ngrid: int, + cutoff: float, + mode: str = 'i', + enable_fft: bool = False + ) -> None: + """ + Sets a common uniform grid for all RadialSet objects. + """ + pass + + def set_grid( + self, + for_r_space: bool, + ngrid: int, + grid: NDArray[np.float64], + mode: str = 'i' + ) -> None: + """ + Sets a common grid for all RadialSet objects + """ + pass + + def __call__(self, itype: int, l: int, izeta: int) -> 'NumericalRadial': ... + def symbol(self, itype: int) -> str: ... + @property + def ntype(self) -> int: ... + @overload + def lmax(self, itype: int) -> int: ... + @property + def lmax(self) -> int: ... + @overload + def rcut_max(self, itype: int) -> float: ... + @property + def rcut_max(self) -> float: ... + def nzeta(self, itype: int, l: int) -> int: ... + @overload + def nzeta_max(self, itype: int) -> int: ... + @overload + def nzeta_max(self) -> int: ... + @overload + def nchi(self, itype: int) -> int: ... + @overload + def nchi(self) -> int: ... + +class TwoCenterIntegrator: ... + +class NumericalRadial: ... + + + + + + + \ No newline at end of file diff --git a/python/pyabacus/src/pyabacus/__init__.py b/python/pyabacus/src/pyabacus/__init__.py index e3776b1616..d73557a7d4 100644 --- a/python/pyabacus/src/pyabacus/__init__.py +++ b/python/pyabacus/src/pyabacus/__init__.py @@ -1,3 +1,16 @@ from __future__ import annotations -from ._core import ModuleBase, ModuleNAO -__all__ = ["ModuleBase", "ModuleNAO"] \ No newline at end of file + +__submodules__ = ["ModuleBase", "ModuleNAO", "hsolver"] + +__all__ = list(__submodules__) + +def __getattr__(attr): + if attr == "ModuleBase": + from ._core import ModuleBase + return ModuleBase + elif attr == "ModuleNAO": + from ._core import ModuleNAO + return ModuleNAO + elif attr == "hsolver": + import pyabacus.hsolver as hsolver + return hsolver \ No newline at end of file diff --git a/python/pyabacus/src/pyabacus/hsolver/__init__.py b/python/pyabacus/src/pyabacus/hsolver/__init__.py new file mode 100644 index 0000000000..e6a2da7ea6 --- /dev/null +++ b/python/pyabacus/src/pyabacus/hsolver/__init__.py @@ -0,0 +1,4 @@ +from __future__ import annotations +from ._hsolver import * + +__all__ = ["dav_subspace"] \ No newline at end of file diff --git a/python/pyabacus/src/pyabacus/hsolver/_hsolver.py b/python/pyabacus/src/pyabacus/hsolver/_hsolver.py new file mode 100644 index 0000000000..dc9fa1441a --- /dev/null +++ b/python/pyabacus/src/pyabacus/hsolver/_hsolver.py @@ -0,0 +1,104 @@ +# pyabacus.hsolver + +import numpy as np +from numpy.typing import NDArray +from typing import Tuple, List, Union, Callable + +from .._core import hsolver + +class diag_comm_info: + def __init__(self, rank: int, nproc: int) -> None: ... + + @property + def rank(self) -> int: ... + + @property + def nproc(self) -> int: ... + +def dav_subspace( + 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, + need_subspace: 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. + need_subspace : bool, optional + Whether to use subspace function, by default False. + is_occupied : List[bool] | None, optional + The list of occupied bands, by default None. This indicates how many eigenvalues + need to be calculated, starting from the smallest eigenvalue. Only the energy levels + occupied by electrons (occupied) need to be calculated. + 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. + + 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 is_occupied is None: + is_occupied = [True] * num_eigs + + 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_dav_subspace(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) + assert dav_ndim > 1, "dav_ndim must be greater than 1." + assert dav_ndim * num_eigs < dim * comm_info.nproc, "dav_ndim * num_eigs must be less than dim * comm_info.nproc." + + _ = _diago_obj_dav_subspace.diag( + mm_op, + pre_condition, + dav_ndim, + tol, + max_iter, + need_subspace, + is_occupied, + scf_type, + 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_diag/Na5.mat b/python/pyabacus/tests/test_diag/Na5.mat new file mode 100644 index 0000000000..bf33a6ce1c Binary files /dev/null and b/python/pyabacus/tests/test_diag/Na5.mat differ diff --git a/python/pyabacus/tests/test_diag/Si2.mat b/python/pyabacus/tests/test_diag/Si2.mat new file mode 100644 index 0000000000..cf83b44cf1 Binary files /dev/null and b/python/pyabacus/tests/test_diag/Si2.mat differ diff --git a/python/pyabacus/tests/test_diago_dav_subspace.py b/python/pyabacus/tests/test_diago_dav_subspace.py new file mode 100644 index 0000000000..dcf70bdcde --- /dev/null +++ b/python/pyabacus/tests/test_diago_dav_subspace.py @@ -0,0 +1,56 @@ +from __future__ import annotations + +import pytest +from pyabacus import hsolver +import numpy as np +import scipy + +def diag_pyabacus(h_sparse, nband): + def mm_op(x): + return h_sparse.dot(x) + + nbasis = h_sparse.shape[0] + + v0 = np.random.rand(nbasis, nband) + + diag_elem = h_sparse.diagonal() + 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( + mm_op, + v0, + nbasis, + nband, + precond, + dav_ndim=8, + tol=1e-12, + max_iter=5000, + scf_type=True + ) + + return e + +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(): + 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_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) +]) +def test_diag(file_name, nband, atol): + h_sparse = scipy.io.loadmat(file_name)['Problem']['A'][0, 0] + e_pyabacus = diag_pyabacus(h_sparse, nband) + e_scipy = diag_eigsh(h_sparse, nband) + np.testing.assert_allclose(e_pyabacus, e_scipy, atol=atol)