change the signature of cg

Author: a1henu

python/pyabacus/examples/diago_matrix.py

@@ -9,18 +9,20 @@ def load_mat(mat_file):
 
     return h_mat, nbasis, nband
 
-def calc_eig_pyabacus(mat_file, method):
+def gen_dense_mat(dim):
+    # generate a random symmetric and positive definite matrix
+    h_mat = np.random.rand(dim, dim)
+    h_mat = h_mat + h_mat.T
+    h_mat = h_mat + dim * np.eye(dim)
+    
+    return h_mat
+
+def calc_eig_dav(mat_file, method):
     algo = {
         'dav_subspace': hsolver.dav_subspace,
-        'davidson': hsolver.davidson,
-        'cg': hsolver.cg
+        'davidson': hsolver.davidson
     }
 
-    if method is not 'cg':
-        ndim = 8
-    else:
-        ndim = 30
-    
     h_mat, nbasis, nband = load_mat(mat_file)
 
     v0 = np.random.rand(nbasis, nband)
@@ -37,31 +39,70 @@ def mm_op(x):
         nbasis,
         nband,
         precond,
-        ndim,
-        1e-8     # tol
+        dav_ndim=8,
+        tol=1e-8
     )
 
     print(f'eigenvalues calculated by pyabacus-{method} is: \n', e)
 
     return e
 
-def calc_eig_scipy(mat_file):
+def calc_eig_cg(h_mat, num_eigs):
+    dim = h_mat.shape[0]
+    v0 = np.random.rand(dim, num_eigs)
+    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, _ = hsolver.cg(
+        mm_op,
+        v0,
+        dim,
+        num_eigs,
+        precond,
+        tol=1e-8
+    )
+    
+    print('eigenvalues calculated by pyabacus-cg is: \n', e)
+    
+    return e
+
+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)
     print('eigenvalues calculated by scipy is: \n', e)
 
     return e
 
+def calc_eigh(h_mat, num_eigs):
+    e, _ = scipy.linalg.eigh(h_mat)
+    e = np.sort(e)
+    print('eigenvalues calculated by scipy is: \n', e[:num_eigs])
+    
+    return e
+
 if __name__ == '__main__':
     mat_file = './Si2.mat'
-    method = ['dav_subspace', 'davidson', 'cg']
+    method = ['dav_subspace', 'davidson']
 
     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_pyabacus = calc_eig_dav(mat_file, m)
+        e_scipy = calc_eigsh(mat_file)
 
         print('eigenvalues difference: \n', e_pyabacus - e_scipy)
+        
+    print("\n====== davidson and dav_subspace method Done! ======")
+    print("\n====== CG method ======")
+    
+    h_mat = gen_dense_mat(100)
+    num_eigs = 8
+    e_cg = calc_eig_cg(h_mat, num_eigs)
+    e_scipy = calc_eigh(h_mat, num_eigs)
 
-    
+    print('eigenvalues difference: \n', e_cg - e_scipy[:num_eigs])
+    print("\n====== CG method Done! ======")

python/pyabacus/src/hsolver/py_diago_cg.hpp

@@ -131,6 +131,8 @@ class PyDiagoCG
             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);
@@ -169,7 +171,7 @@ class PyDiagoCG
             nproc_in_pool
         );
 
-        return cg->diag(hpsi_func, spsi_func, *psi, *eig, *prec);
+        cg->diag(hpsi_func, spsi_func, *psi, *eig, *prec);
     }
 
 private:

python/pyabacus/src/pyabacus/hsolver/_hsolver.py

@@ -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,
@@ -185,12 +185,50 @@ def cg(
     init_v: NDArray[np.complex128],
     dim: int,
     num_eigs: int,
-    pre_condition: NDArray[np.float64],
-    diag_ndim: int = 2,
+    precondition: NDArray[np.float64],
     tol: float = 1e-2,
+    max_iter: int = 1000,
     need_subspace: bool = False,
-    scf_type: 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.")
 
@@ -204,17 +242,18 @@ def cg(
     _diago_obj_cg.set_psi(init_v)
     _diago_obj_cg.init_eig()
 
-    _diago_obj_cg.set_prec(pre_condition)
+    _diago_obj_cg.set_prec(precondition)
 
-    _ = _diago_obj_cg.diag(
+    _diago_obj_cg.diag(
         mvv_op,
-        diag_ndim, 
+        max_iter, 
         tol,
         need_subspace,
         scf_type,
+        nproc_in_pool
     )
 
-    e = _diago_obj_cg.get_eigenvalue()
+    e = _diago_obj_cg.get_eig()
     v = _diago_obj_cg.get_psi()
 
     return e, v

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]