diff --git a/python/pyabacus/src/py_diago_dav_subspace.hpp b/python/pyabacus/src/py_diago_dav_subspace.hpp index f6cf7e9a5ba..bd8bbb3e414 100644 --- a/python/pyabacus/src/py_diago_dav_subspace.hpp +++ b/python/pyabacus/src/py_diago_dav_subspace.hpp @@ -110,11 +110,14 @@ class PyDiagoDavSubspace 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) - { + auto hpsi_func = [mm_op] ( + std::complex *psi_in, + std::complex *hpsi_out, + 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}); diff --git a/python/pyabacus/src/py_diago_david.hpp b/python/pyabacus/src/py_diago_david.hpp index a53147f6ce3..2db5284f4bc 100644 --- a/python/pyabacus/src/py_diago_david.hpp +++ b/python/pyabacus/src/py_diago_david.hpp @@ -109,8 +109,8 @@ class PyDiagoDavid hsolver::diag_comm_info comm_info ) { auto hpsi_func = [mm_op] ( - std::complex *hpsi_out, - std::complex *psi_in, + std::complex *psi_in, + std::complex *hpsi_out, const int nband_in, const int nbasis_in, const int band_index1, diff --git a/python/pyabacus/src/pyabacus/hsolver/_hsolver.py b/python/pyabacus/src/pyabacus/hsolver/_hsolver.py index 5f5875a0554..3dea18ef368 100644 --- a/python/pyabacus/src/pyabacus/hsolver/_hsolver.py +++ b/python/pyabacus/src/pyabacus/hsolver/_hsolver.py @@ -16,7 +16,7 @@ def rank(self) -> int: ... def nproc(self) -> int: ... def dav_subspace( - mm_op: Callable[[NDArray[np.complex128]], NDArray[np.complex128]], + mvv_op: Callable[[NDArray[np.complex128]], NDArray[np.complex128]], init_v: NDArray[np.complex128], dim: int, num_eigs: int, @@ -32,9 +32,10 @@ def dav_subspace( 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. + 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 @@ -68,8 +69,8 @@ def dav_subspace( 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 not callable(mvv_op): + raise TypeError("mvv_op must be a callable object.") if is_occupied is None: is_occupied = [True] * num_eigs @@ -86,7 +87,7 @@ def dav_subspace( 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, + mvv_op, pre_condition, dav_ndim, tol, @@ -103,7 +104,7 @@ def dav_subspace( return e, v def davidson( - mm_op: Callable[[NDArray[np.complex128]], NDArray[np.complex128]], + mvv_op: Callable[[NDArray[np.complex128]], NDArray[np.complex128]], init_v: NDArray[np.complex128], dim: int, num_eigs: int, @@ -119,9 +120,10 @@ def davidson( 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. + 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 @@ -146,8 +148,8 @@ def davidson( 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 not callable(mvv_op): + raise TypeError("mvv_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') @@ -159,7 +161,7 @@ def davidson( comm_info = hsolver.diag_comm_info(0, 1) _ = _diago_obj_dav_subspace.diag( - mm_op, + mvv_op, pre_condition, dav_ndim, tol,