diff --git a/renormalizer/mps/backend.py b/renormalizer/mps/backend.py
index 02256565..427b365b 100644
--- a/renormalizer/mps/backend.py
+++ b/renormalizer/mps/backend.py
@@ -35,8 +35,10 @@
 
 if not USE_GPU:
     logger.info("use numpy as backend")
+    OE_BACKEND = "numpy"
 else:
     logger.info("use cupy as backend")
+    OE_BACKEND = "cupy"
 
 
 xp.random.seed(2019)
diff --git a/renormalizer/mps/gs.py b/renormalizer/mps/gs.py
index 9130a0a9..75a76d21 100644
--- a/renormalizer/mps/gs.py
+++ b/renormalizer/mps/gs.py
@@ -6,45 +6,25 @@
 
 """
 
-from typing import Tuple, List
+from typing import Tuple, List, Union
 import logging
 
-from renormalizer.lib import davidson
-from renormalizer.mps.backend import xp, USE_GPU
-from renormalizer.mps.matrix import (
-    multi_tensor_contract,
-    tensordot,
-    asnumpy, 
-    asxp)
-from renormalizer.mps import Mpo, Mps
-from renormalizer.mps.lib import Environ, cvec2cmat
-from renormalizer.utils import Quantity
-
 import numpy as np
 import scipy
 import opt_einsum as oe
 import primme
 
+from renormalizer.lib import davidson
+from renormalizer.mps.backend import xp, OE_BACKEND
+from renormalizer.mps.matrix import multi_tensor_contract, tensordot, asnumpy, asxp
+from renormalizer.mps import Mpo, Mps
+from renormalizer.mps.lib import Environ, cvec2cmat
+from renormalizer.utils import Quantity
+
 logger = logging.getLogger(__name__)
 
-def find_lowest_energy(h_mpo: Mpo, nexciton, Mmax):
-    logger.debug("begin finding lowest energy")
-    model = h_mpo.model
-    mps = Mps.random(model, nexciton, Mmax)
-    energies, _ = optimize_mps(mps, h_mpo)
-    return energies[-1]
-
-def find_highest_energy(h_mpo: Mpo, nexciton, Mmax):
-    logger.debug("begin finding highest energy")
-    model = h_mpo.model
-    mps = Mps.random(model, nexciton, Mmax)
-    mps.optimize_config.inverse = -1.0
-    energies, _ = optimize_mps(mps, h_mpo)
-    return -energies[-1]
-
-def construct_mps_mpo_2(
-    model, Mmax, nexciton, offset=Quantity(0)
-):
+
+def construct_mps_mpo(model, mmax, nexciton, offset=Quantity(0)):
     """
     MPO/MPS structure 2
     e1,ph11,ph12,..e2,ph21,ph22,...en,phn1,phn2...
@@ -58,20 +38,20 @@ def construct_mps_mpo_2(
     """
     initialize MPS according to quantum number
     """
-    mps = Mps.random(model, nexciton, Mmax, percent=1)
+    mps = Mps.random(model, nexciton, mmax, percent=1)
     # print("initialize left-canonical:", mps.check_left_canonical())
 
     return mps, mpo
 
 
-def optimize_mps(mps: Mps, mpo: Mpo, omega :float =None) -> Tuple[List, Mps]:
-    r""" DMRG ground state algorithm and state-averaged excited states algorithm
-    
+def optimize_mps(mps: Mps, mpo: Mpo, omega: float = None) -> Tuple[List, Mps]:
+    r"""DMRG ground state algorithm and state-averaged excited states algorithm
+
     Parameters
     ----------
     mps : renormalizer.mps.Mps
-        initial guess of mps
-    mpo : renormalizer.mps.Mpo 
+        initial guess of mps. The MPS is overwritten during the optimization.
+    mpo : renormalizer.mps.Mpo
         mpo of Hamiltonian
     omega: float, optional
         target the eigenpair near omega with special variational function
@@ -84,30 +64,20 @@ def optimize_mps(mps: Mps, mpo: Mpo, omega :float =None) -> Tuple[List, Mps]:
         :math:`[e_0, e_0, \cdots, e_0]` if ``nroots=1``.
         :math:`[[e_0, \cdots, e_n], \dots, [e_0, \cdots, e_n]]` if ``nroots=n``.
     mps : renormalizer.mps.Mps
-        optimized ground state mps. The input mps is overwritten and could not
-        be used anymore.
-    
+        optimized ground state MPS.
+            Note it's not the same with the overwritten input MPS.
+
     See Also
     --------
     renormalizer.utils.configs.OptimizeConfig : The optimization configuration.
-    
+
     """
-    algo = mps.optimize_config.algo
-    method = mps.optimize_config.method
-    procedure = mps.optimize_config.procedure
-    inverse = mps.optimize_config.inverse
-    nroots = mps.optimize_config.nroots
 
-    assert method in ["2site", "1site"]
-    logger.info(f"optimization method: {method}")
+    assert mps.optimize_config.method in ["2site", "1site"]
+    logger.info(f"optimization method: {mps.optimize_config.method}")
     logger.info(f"e_rtol: {mps.optimize_config.e_rtol}")
     logger.info(f"e_atol: {mps.optimize_config.e_atol}")
-    
-    if USE_GPU:
-        oe_backend = "cupy"
-    else:
-        oe_backend = "numpy"
-    
+
     # ensure that mps is left or right-canonical
     # TODO: start from a mix-canonical MPS
     if mps.is_left_canon:
@@ -116,15 +86,7 @@ def optimize_mps(mps: Mps, mpo: Mpo, omega :float =None) -> Tuple[List, Mps]:
     else:
         mps.ensure_left_canon()
         env = "L"
-    
-    # in state-averged calculation, contains C of each state for better initial
-    # guess
-    averaged_ms = None   
-    
-    # the index of active site of the returned mps
-    res_mps_idx = None
-    
-    # target eigenstate close to omega with (H-omega)^2
+
     # construct the environment matrix
     if omega is not None:
         identity = Mpo.identity(mpo.model)
@@ -134,365 +96,424 @@ def optimize_mps(mps: Mps, mpo: Mpo, omega :float =None) -> Tuple[List, Mps]:
         environ = Environ(mps, mpo, env)
 
     macro_iteration_result = []
-    converged = False
-    for isweep, (mmax, percent) in enumerate(procedure):
+    # Idx of the active site with lowest energy for each sweep
+    # determines the index of active site of the returned mps
+    opt_e_idx: int = None
+    res_mps: Union[Mps, List[Mps]] = None
+    for isweep, (mmax, percent) in enumerate(mps.optimize_config.procedure):
         logger.debug(f"isweep: {isweep}")
         logger.debug(f"mmax, percent: {mmax}, {percent}")
         logger.debug(f"{mps}")
-        
-        micro_iteration_result = []
-        for imps in mps.iter_idx_list(full=True):
-            if method == "2site" and \
-                ((mps.to_right and imps == mps.site_num-1)
-                or ((not mps.to_right) and imps == 0)):
+
+        micro_iteration_result, res_mps = single_sweep(mps, mpo, environ, omega, mmax, percent, opt_e_idx)
+
+        opt_e = min(micro_iteration_result)
+        macro_iteration_result.append(opt_e[0])
+        opt_e_idx = opt_e[1]
+
+        logger.debug(
+            f"{isweep+1} sweeps are finished, lowest energy = {min(macro_iteration_result)}"
+        )
+        # check if convergence
+        if isweep > 0 and percent == 0:
+            v1, v2 = sorted(macro_iteration_result)[:2]
+            if np.allclose(
+                v1, v2, rtol=mps.optimize_config.e_rtol, atol=mps.optimize_config.e_atol
+            ):
+                logger.info("DMRG has converged!")
                 break
-            
-            if mps.to_right:
-                lmethod, rmethod = "System", "Enviro"
-            else:
-                lmethod, rmethod = "Enviro", "System"
+    else:
+        logger.warning("DMRG did not converge! Please increase the procedure!")
+        logger.info(f"The lowest two energies: {sorted(macro_iteration_result)[:2]}.")
 
-            if method == "1site":
+    assert res_mps is not None
+    # remove the redundant basis near the edge
+    if mps.optimize_config.nroots == 1:
+        res_mps = res_mps.normalize().ensure_left_canon().canonicalise()
+        logger.info(f"{res_mps}")
+    else:
+        res_mps = [mp.normalize().ensure_left_canon().canonicalise() for mp in res_mps]
+        logger.info(f"{res_mps[0]}")
+    return macro_iteration_result, res_mps
+
+
+def single_sweep(
+    mps: Mps,
+    mpo: Mpo,
+    environ: Environ,
+    omega: float,
+    mmax: int,
+    percent: float,
+    last_opt_e_idx: int
+):
+
+    method = mps.optimize_config.method
+    nroots = mps.optimize_config.nroots
+
+    # in state-averaged calculation, contains C of each state for better initial guess
+    averaged_ms = []
+    # optmized mps
+    res_mps: Union[Mps, List[Mps]] = None
+    # energies after optimizing each site
+    micro_iteration_result = []
+    for imps in mps.iter_idx_list(full=True):
+        if method == "2site" and (
+            (mps.to_right and imps == mps.site_num - 1)
+            or ((not mps.to_right) and imps == 0)
+        ):
+            break
+
+        if mps.to_right:
+            lmethod, rmethod = "System", "Enviro"
+        else:
+            lmethod, rmethod = "Enviro", "System"
+
+        if method == "1site":
+            lidx = imps - 1
+            cidx = [imps]
+            ridx = imps + 1
+        elif method == "2site":
+            if mps.to_right:
                 lidx = imps - 1
-                cidx= [imps]
-                ridx = imps + 1
-            elif method == "2site":
-                if mps.to_right:
-                    lidx = imps - 1
-                    cidx = [imps, imps+1] 
-                    ridx = imps + 2
-                else:
-                    lidx = imps - 2
-                    cidx = [imps-1, imps]  # center site
-                    ridx = imps + 1
+                cidx = [imps, imps + 1]
+                ridx = imps + 2
             else:
-                assert False
-            logger.debug(f"optimize site: {cidx}")
-            
-            if omega is None:
-                operator = mpo
-            else:
-                operator = [mpo, mpo]
+                lidx = imps - 2
+                cidx = [imps - 1, imps]  # center site
+                ridx = imps + 1
+        else:
+            assert False
+        logger.debug(f"optimize site: {cidx}")
 
-            ltensor = environ.GetLR(
-                "L", lidx, mps, operator, itensor=None, method=lmethod
-            )
-            rtensor = environ.GetLR(
-                "R", ridx, mps, operator, itensor=None, method=rmethod
-            )
+        if omega is None:
+            operator = mpo
+        else:
+            operator = [mpo, mpo]
 
-            # get the quantum number pattern
-            qnbigl, qnbigr, qnmat = mps._get_big_qn(cidx)
-            cshape = qnmat.shape
-            nonzeros = np.sum(qnmat == mps.qntot)
-            logger.debug(f"Hmat dim: {nonzeros}")
-            
-            # center mo
-            cmo = [asxp(mpo[idx]) for idx in cidx] 
-            
-            if qnmat.size > 1000 and algo != "direct":
-                # iterative algorithm
-                
-                # diagonal elements of H
-                if omega is None:
-                    tmp_ltensor = xp.einsum("aba -> ba", ltensor)
-                    tmp_cmo0 = xp.einsum("abbc -> abc", cmo[0])
-                    tmp_rtensor = xp.einsum("aba -> ba", rtensor)
-                    if method == "1site":
-                        #   S-a c f-S
-                        #   O-b-O-g-O
-                        #   S-a c f-S
-                        path = [([0, 1], "ba, bcg -> acg"), ([1, 0], "acg, gf -> acf")]
-                        hdiag = multi_tensor_contract(
-                            path, tmp_ltensor, tmp_cmo0, tmp_rtensor
-                        )[(qnmat == mps.qntot)]
-                    else:
-                        #   S-a c   d f-S
-                        #   O-b-O-e-O-g-O
-                        #   S-a c   d f-S
-                        tmp_cmo1 = xp.einsum("abbc -> abc", cmo[1])
-                        path = [
-                            ([0, 1], "ba, bce -> ace"),
-                            ([0, 1], "edg, gf -> edf"),
-                            ([0, 1], "ace, edf -> acdf"),
-                        ]
-                        hdiag = multi_tensor_contract(
-                            path, tmp_ltensor, tmp_cmo0, tmp_cmo1, tmp_rtensor
-                        )[(qnmat == mps.qntot)]
-                else:
-                    if method == "1site":
-                        #   S-a d h-S
-                        #   O-b-O-f-O
-                        #   |   e   |
-                        #   O-c-O-g-O
-                        #   S-a d h-S
-                        hdiag = oe.contract("abca, bdef, cedg, hfgh -> adh",
-                                ltensor, cmo[0], cmo[0], rtensor, backend=oe_backend)[(qnmat == mps.qntot)]
-                    else:
-                        #   S-a d   h l-S
-                        #   O-b-O-f-O-j-O
-                        #   |   e   i   |
-                        #   O-c-O-g-O-k-O
-                        #   S-a d   h l-S
-                        hdiag = oe.contract("abca, bdef, cedg, fhij, gihk, ljkl -> adhl",
-                                ltensor, cmo[0], cmo[0], cmo[1], cmo[1], rtensor, backend=oe_backend)[(qnmat == mps.qntot)]
-
-                hdiag = asnumpy(hdiag * inverse)
-                
-                # initial guess
+        ltensor = environ.GetLR("L", lidx, mps, operator, itensor=None, method=lmethod)
+        rtensor = environ.GetLR("R", ridx, mps, operator, itensor=None, method=rmethod)
+
+        # get the quantum number pattern
+        qnbigl, qnbigr, qnmat = mps._get_big_qn(cidx)
+        cshape = qnmat.shape
+
+        # center mo
+        cmo = [asxp(mpo[idx]) for idx in cidx]
+
+        use_direct_eigh = qnmat.size < 1000 or mps.optimize_config.algo == "direct"
+        if use_direct_eigh:
+            e, c = eigh_direct(mps, qnmat, ltensor, rtensor, cmo, omega)
+        else:
+            # the iterative approach
+            # generate initial guess
+            if nroots == 1:
                 if method == "1site":
                     # initial guess   b-S-c
                     #                   a
-                    if nroots == 1:
-                        cguess = [asnumpy(mps[cidx[0]])[qnmat == mps.qntot]]
-                    else:
-                        cguess = []
-                        if averaged_ms is not None:
-                            for ms in averaged_ms:
-                                cguess.append(asnumpy(ms)[qnmat == mps.qntot])
+                    raw_cguess = mps[cidx[0]]
                 else:
                     # initial guess b-S-c-S-e
                     #                 a   d
-                    if nroots == 1:
-                        cguess = [asnumpy(tensordot(mps[cidx[0]], mps[cidx[1]],
-                            axes=1)[qnmat == mps.qntot])]
-                    else:
-                        cguess = []
-                        if averaged_ms is not None:
-                            for ms in averaged_ms:
-                                if mps.to_right:
-                                    cguess.append(asnumpy(tensordot(ms, mps[cidx[1]],
-                                        axes=1)[qnmat == mps.qntot]))
-                                else:
-                                    cguess.append(asnumpy(tensordot(mps[cidx[0]],
-                                        ms,
-                                        axes=1)[qnmat == mps.qntot]))
-                if omega is not None:
+                    raw_cguess = tensordot(mps[cidx[0]], mps[cidx[1]], axes=1)
+                cguess = [asnumpy(raw_cguess)[qnmat == mps.qntot]]
+            else:
+                cguess = []
+                for ms in averaged_ms:
                     if method == "1site":
-                        #   S-a e j-S
-                        #   O-b-O-g-O
-                        #   |   f   |
-                        #   O-c-O-i-O
-                        #   S-d h k-S
-                        expr = oe.contract_expression(
-                                "abcd, befg, cfhi, jgik, aej -> dhk",
-                                ltensor, cmo[0], cmo[0], rtensor, cshape,
-                                constants=[0,1,2,3])
-                    else:
-                        #   S-a e   j o-S
-                        #   O-b-O-g-O-l-O
-                        #   |   f   k   |
-                        #   O-c-O-i-O-n-O
-                        #   S-d h   m p-S
-                        expr = oe.contract_expression(
-                                "abcd, befg, cfhi, gjkl, ikmn, olnp, aejo -> dhmp",
-                                ltensor, cmo[0], cmo[0], cmo[1], cmo[1],
-                                rtensor, cshape,
-                                constants=[0,1,2,3,4,5])
-
-                count = 0
-                def hop(x):
-                    nonlocal count
-                    count += 1
-                    clist = []
-                    if x.ndim == 1:
-                        clist.append(x)
+                        raw_cguess = asnumpy(ms)
                     else:
-                        for icol in range(x.shape[1]):
-                            clist.append(x[:,icol])
-                    res = []
-                    for c in clist:
-                        # convert c to initial structure according to qn pattern
-                        cstruct = asxp(cvec2cmat(cshape, c, qnmat, mps.qntot))
-                        
-                        if omega is None:
-                            if method == "1site":
-                                # S-a   l-S
-                                #     d
-                                # O-b-O-f-O
-                                #     e
-                                # S-c   k-S
-
-                                path = [
-                                    ([0, 1], "abc, adl -> bcdl"),
-                                    ([2, 0], "bcdl, bdef -> clef"),
-                                    ([1, 0], "clef, lfk -> cek"),
-                                ]
-                                cout = multi_tensor_contract(
-                                    path, ltensor, cstruct, cmo[0], rtensor
-                                )
-                            else:
-                                # S-a       l-S
-                                #     d   g
-                                # O-b-O-f-O-j-O
-                                #     e   h
-                                # S-c       k-S
-                                path = [
-                                    ([0, 1], "abc, adgl -> bcdgl"),
-                                    ([3, 0], "bcdgl, bdef -> cglef"),
-                                    ([2, 0], "cglef, fghj -> clehj"),
-                                    ([1, 0], "clehj, ljk -> cehk"),
-                                ]
-                                cout = multi_tensor_contract(
-                                    path,
-                                    ltensor,
-                                    cstruct,
-                                    cmo[0],
-                                    cmo[1],
-                                    rtensor,
-                                )
+                        if mps.to_right:
+                            raw_cguess = tensordot(ms, mps[cidx[1]], axes=1)
                         else:
-                            cout = expr(cstruct, backend=oe_backend)   
-                        
-                    # convert structure c to 1d according to qn
-                        res.append(asnumpy(cout)[qnmat == mps.qntot])
-                    
-                    if len(res) == 1:
-                        return inverse * res[0]
-                    else:
-                        return inverse * np.stack(res,axis=1)
-
-                if len(cguess) < nroots:    
-                    cguess.extend([np.random.random([nonzeros]) - 0.5 for i in
-                        range(len(cguess), nroots)])
-                
-                if algo == "davidson": 
-                    precond = lambda x, e, *args: x / (hdiag - e + 1e-4)
-                    
-                    e, c = davidson(
-                        hop, cguess, precond, max_cycle=100,
-                        nroots=nroots, max_memory=64000
-                    )
-                    # if one root, davidson return e as np.float
-
-                #elif algo == "arpack":
-                #    # scipy arpack solver : much slower than pyscf/davidson
-                #    A = scipy.sparse.linalg.LinearOperator((nonzeros,nonzeros), matvec=hop)
-                #    e, c = scipy.sparse.linalg.eigsh(A, k=nroots, which="SA", v0=cguess)
-                #    # scipy return numpy.array
-                #    if nroots == 1:
-                #        e = e[0]
-                #elif algo == "lobpcg":
-                #    precond = lambda x: scipy.sparse.diags(1/(hdiag+1e-4)) @ x
-                #    A = scipy.sparse.linalg.LinearOperator((nonzeros,nonzeros),
-                #            matvec=hop, matmat=hop)
-                #    M = scipy.sparse.linalg.LinearOperator((nonzeros,nonzeros),
-                #            matvec=precond, matmat=hop)
-                #    e, c = scipy.sparse.linalg.lobpcg(A, np.array(cguess).T,
-                #            M=M, largest=False)
-                elif algo == "primme":
-                    precond = lambda x: scipy.sparse.diags(1/(hdiag+1e-4)) @ x
-                    A = scipy.sparse.linalg.LinearOperator((nonzeros,nonzeros),
-                            matvec=hop, matmat=hop)
-                    M = scipy.sparse.linalg.LinearOperator((nonzeros,nonzeros),
-                            matvec=precond, matmat=hop)
-                    e, c = primme.eigsh(A, k=min(nroots,nonzeros), which="SA", 
-                            v0=np.array(cguess).T, OPinv=M,
-                            method="PRIMME_DYNAMIC", 
-                            tol=1e-6)
-                else:
-                    assert False
-                logger.debug(f"use {algo}, HC hops: {count}")
+                            raw_cguess = tensordot(mps[cidx[0]], ms, axes=1)
+                    cguess.append(asnumpy(raw_cguess)[qnmat == mps.qntot])
+
+            guess_dim = np.sum(qnmat == mps.qntot)
+            cguess.extend(
+                [np.random.rand(guess_dim) - 0.5 for i in range(len(cguess), nroots)]
+            )
+            e, c = eigh_iterative(mps, qnmat, ltensor, rtensor, cmo, omega, cguess)
+
+        # if multi roots, both davidson and primme return np.ndarray
+        if nroots > 1:
+            e = e.tolist()
+        logger.debug(f"energy: {e}")
+        micro_iteration_result.append((e, cidx))
+
+        cstruct = cvec2cmat(cshape, c, qnmat, mps.qntot, nroots=nroots)
+
+        # store the "optimal" mps (usually in the middle of each sweep)
+        if cidx == last_opt_e_idx:
+            if nroots == 1:
+                res_mps = mps.copy()
+                res_mps._update_mps(cstruct, cidx, qnbigl, qnbigr, mmax, percent)
             else:
-                logger.debug(f"use direct eigensolver")
-                
-                # direct algorithm
-                if omega is None:
-                    if method == "1site":
-                        # S-a   l-S
-                        #     d
-                        # O-b-O-f-O
-                        #     e
-                        # S-c   k-S
-                        ham = oe.contract("abc,bdef,lfk->adlcek",
-                                ltensor, cmo[0], rtensor, backend=oe_backend)
-                        ham = ham[:,:,:,qnmat==mps.qntot][qnmat==mps.qntot,:] * inverse
-                    else:
-                        # S-a       l-S
-                        #     d   g
-                        # O-b-O-f-O-j-O
-                        #     e   h
-                        # S-c       k-S
-                        ham = oe.contract("abc,bdef,fghj,ljk->adglcehk",
-                                ltensor, cmo[0], cmo[1], rtensor)
-                        ham = ham[:,:,:,:,qnmat==mps.qntot][qnmat==mps.qntot,:] * inverse
-                else:
-                    if method == "1site":
-                        #   S-a e j-S
-                        #   O-b-O-g-O
-                        #   |   f   |
-                        #   O-c-O-i-O
-                        #   S-d h k-S
-                        ham = oe.contract(
-                                "abcd, befg, cfhi, jgik -> aejdhk",
-                                ltensor, cmo[0], cmo[0], rtensor)
-                        ham = ham[:,:,:,qnmat==mps.qntot][qnmat==mps.qntot,:] * inverse
-                    else:
-                        #   S-a e   j o-S
-                        #   O-b-O-g-O-l-O
-                        #   |   f   k   |
-                        #   O-c-O-i-O-n-O
-                        #   S-d h   m p-S
-                        ham = oe.contract(
-                                "abcd, befg, cfhi, gjkl, ikmn, olnp -> aejodhmp",
-                                ltensor, cmo[0], cmo[0], cmo[1], cmo[1],
-                                rtensor)
-                        ham = ham[:,:,:,:,qnmat==mps.qntot][qnmat==mps.qntot,:] * inverse
-                
-                w, v = scipy.linalg.eigh(asnumpy(ham))
-                if nroots == 1:
-                    e = w[0]
-                    c = v[:,0]
-                else:
-                    e = w[:nroots]
-                    c = [v[:,iroot] for iroot in range(min(nroots,v.shape[1]))]
-            # if multi roots, both davidson and primme return np.ndarray
-            if nroots > 1:
-                e = e.tolist()
-            logger.debug(f"energy: {e}")
-            micro_iteration_result.append(e)
-                
-            cstruct = cvec2cmat(cshape, c, qnmat, mps.qntot, nroots=nroots)
-            # store the "optimal" mps (usually in the middle of each sweep)
-            if res_mps_idx is not None and res_mps_idx == imps:
-                if nroots == 1:
-                    res_mps = mps.copy()
-                    res_mps._update_mps(cstruct, cidx, qnbigl, qnbigr, mmax, percent)
-                else:
-                    res_mps = [mps.copy() for i in range(len(cstruct))]
-                    for iroot in range(len(cstruct)):
-                        res_mps[iroot]._update_mps(cstruct[iroot], cidx, qnbigl,
-                                qnbigr, mmax, percent)
-            
-            averaged_ms = mps._update_mps(cstruct, cidx, qnbigl, qnbigr, mmax, percent)
-            
-            
-        mps._switch_direction()
-        
-        res_mps_idx = micro_iteration_result.index(min(micro_iteration_result)) 
-        macro_iteration_result.append(micro_iteration_result[res_mps_idx])
-        # check if convergence
-        if isweep > 0 and percent == 0:
-            v1, v2 = sorted(macro_iteration_result)[:2]
-            if np.allclose(v1, v2, rtol=mps.optimize_config.e_rtol,
-                    atol=mps.optimize_config.e_atol):
-                converged = True
-                break
-        
+                res_mps = [mps.copy() for i in range(len(cstruct))]
+                for iroot in range(len(cstruct)):
+                    res_mps[iroot]._update_mps(
+                        cstruct[iroot], cidx, qnbigl, qnbigr, mmax, percent
+                    )
+
+        averaged_ms = mps._update_mps(cstruct, cidx, qnbigl, qnbigr, mmax, percent)
+
+    mps._switch_direction()
+    return micro_iteration_result, res_mps
+
+
+def eigh_direct(
+    mps: Mps,
+    qnmat: np.ndarray,
+    ltensor: xp.ndarray,
+    rtensor: xp.ndarray,
+    cmo: List[xp.ndarray],
+    omega: float,
+):
+    logger.debug(f"use direct eigensolver")
 
-    logger.debug(f"{isweep+1} sweeps are finished, lowest energy = {sorted(macro_iteration_result)[0]}")
-    if converged:
-        logger.info("DMRG is converged!")
+    # direct algorithm
+    if omega is None:
+        if mps.optimize_config.method == "1site":
+            # S-a   l-S
+            #     d
+            # O-b-O-f-O
+            #     e
+            # S-c   k-S
+            ham = oe.contract(
+                "abc,bdef,lfk->adlcek",
+                ltensor, cmo[0], rtensor,
+                backend=OE_BACKEND
+            )
+            ham = ham[:, :, :, qnmat == mps.qntot][qnmat == mps.qntot, :]
+        else:
+            # S-a       l-S
+            #     d   g
+            # O-b-O-f-O-j-O
+            #     e   h
+            # S-c       k-S
+            ham = oe.contract(
+                "abc,bdef,fghj,ljk->adglcehk",
+                ltensor, cmo[0], cmo[1], rtensor,
+                backend=OE_BACKEND,
+            )
+            ham = ham[:, :, :, :, qnmat == mps.qntot][qnmat == mps.qntot, :]
     else:
-        logger.warning("DMRG is not converged! Please increase the procedure!")
-        logger.info(f"The lowest two energies: {sorted(macro_iteration_result)[:2]}.")
-        
-    # remove the redundant basis near the edge
+        if mps.optimize_config.method == "1site":
+            #   S-a e j-S
+            #   O-b-O-g-O
+            #   |   f   |
+            #   O-c-O-i-O
+            #   S-d h k-S
+            ham = oe.contract(
+                "abcd, befg, cfhi, jgik -> aejdhk",
+                ltensor, cmo[0], cmo[0], rtensor,
+                backend=OE_BACKEND,
+            )
+            ham = ham[:, :, :, qnmat == mps.qntot][qnmat == mps.qntot, :]
+        else:
+            #   S-a e   j o-S
+            #   O-b-O-g-O-l-O
+            #   |   f   k   |
+            #   O-c-O-i-O-n-O
+            #   S-d h   m p-S
+            ham = oe.contract(
+                "abcd, befg, cfhi, gjkl, ikmn, olnp -> aejodhmp",
+                ltensor, cmo[0], cmo[0], cmo[1], cmo[1], rtensor,
+                backend=OE_BACKEND,
+            )
+            ham = ham[:, :, :, :, qnmat == mps.qntot][qnmat == mps.qntot, :]
+
+    inverse = mps.optimize_config.inverse
+    w, v = scipy.linalg.eigh(asnumpy(ham) * inverse)
+
+    nroots = mps.optimize_config.nroots
     if nroots == 1:
-        res_mps = res_mps.normalize().ensure_left_canon().canonicalise()
-        logger.info(f"{res_mps}")
+        e = w[0]
+        c = v[:, 0]
     else:
-        res_mps = [mp.normalize().ensure_left_canon().canonicalise() for mp in res_mps]
-        logger.info(f"{res_mps[0]}")
-    return macro_iteration_result, res_mps
+        e = w[:nroots]
+        c = [v[:, iroot] for iroot in range(min(nroots, v.shape[1]))]
+    return e, c
+
+
+def eigh_iterative(
+    mps: Mps,
+    qnmat: np.ndarray,
+    ltensor: xp.ndarray,
+    rtensor: xp.ndarray,
+    cmo: List[xp.ndarray],
+    omega: float,
+    cguess: List[np.ndarray],
+):
+    # iterative algorithm
+    method = mps.optimize_config.method
+    inverse = mps.optimize_config.inverse
+
+    # diagonal elements of H for preconditioning
+    if omega is None:
+        tmp_ltensor = xp.einsum("aba -> ba", ltensor)
+        tmp_cmo0 = xp.einsum("abbc -> abc", cmo[0])
+        tmp_rtensor = xp.einsum("aba -> ba", rtensor)
+        if method == "1site":
+            #   S-a c f-S
+            #   O-b-O-g-O
+            #   S-a c f-S
+            path = [([0, 1], "ba, bcg -> acg"), ([1, 0], "acg, gf -> acf")]
+            hdiag = multi_tensor_contract(path, tmp_ltensor, tmp_cmo0, tmp_rtensor)
+        else:
+            #   S-a c   d f-S
+            #   O-b-O-e-O-g-O
+            #   S-a c   d f-S
+            tmp_cmo1 = xp.einsum("abbc -> abc", cmo[1])
+            path = [
+                ([0, 1], "ba, bce -> ace"),
+                ([0, 1], "edg, gf -> edf"),
+                ([0, 1], "ace, edf -> acdf"),
+            ]
+            hdiag = multi_tensor_contract(
+                path, tmp_ltensor, tmp_cmo0, tmp_cmo1, tmp_rtensor
+            )
+    else:
+        if method == "1site":
+            #   S-a d h-S
+            #   O-b-O-f-O
+            #   |   e   |
+            #   O-c-O-g-O
+            #   S-a d h-S
+            hdiag = oe.contract(
+                "abca, bdef, cedg, hfgh -> adh",
+                ltensor, cmo[0], cmo[0], rtensor,
+                backend=OE_BACKEND,
+            )
+        else:
+            #   S-a d   h l-S
+            #   O-b-O-f-O-j-O
+            #   |   e   i   |
+            #   O-c-O-g-O-k-O
+            #   S-a d   h l-S
+            hdiag = oe.contract(
+                "abca, bdef, cedg, fhij, gihk, ljkl -> adhl",
+                ltensor, cmo[0], cmo[0], cmo[1], cmo[1], rtensor,
+                backend=OE_BACKEND,
+            )
 
+    hdiag = asnumpy(hdiag[(qnmat == mps.qntot)] * inverse)
+
+    # Define the H operator
+
+    # contraction expression
+    cshape = qnmat.shape
+    if omega is None:
+        if method == "1site":
+            # S-a   l-S
+            #     d
+            # O-b-O-f-O
+            #     e
+            # S-c   k-S
+            expr = oe.contract_expression(
+                "abc, bdef, lfk, adl -> cek",
+                ltensor, cmo[0], rtensor, cshape,
+                constants=[0, 1, 2],
+            )
+        else:
+            # S-a       l-S
+            #     d   g
+            # O-b-O-f-O-j-O
+            #     e   h
+            # S-c       k-S
+            expr = oe.contract_expression(
+                "abc, bdef, fghj, ljk, adgl -> cehk",
+                ltensor, cmo[0], cmo[1], rtensor, cshape,
+                constants=[0, 1, 2, 3],
+            )
+    else:
+        if method == "1site":
+            #   S-a e j-S
+            #   O-b-O-g-O
+            #   |   f   |
+            #   O-c-O-i-O
+            #   S-d h k-S
+            expr = oe.contract_expression(
+                "abcd, befg, cfhi, jgik, aej -> dhk",
+                ltensor, cmo[0], cmo[0], rtensor, cshape,
+                constants=[0, 1, 2, 3],
+            )
+        else:
+            #   S-a e   j o-S
+            #   O-b-O-g-O-l-O
+            #   |   f   k   |
+            #   O-c-O-i-O-n-O
+            #   S-d h   m p-S
+            expr = oe.contract_expression(
+                "abcd, befg, cfhi, gjkl, ikmn, olnp, aejo -> dhmp",
+                ltensor, cmo[0], cmo[0], cmo[1], cmo[1], rtensor, cshape,
+                constants=[0, 1, 2, 3, 4, 5],
+            )
+
+    count = 0
+
+    def hop(x):
+        nonlocal count
+        count += 1
+        clist = []
+        if x.ndim == 1:
+            clist.append(x)
+        else:
+            for icol in range(x.shape[1]):
+                clist.append(x[:, icol])
+        res = []
+        for c in clist:
+            # convert c to initial structure according to qn pattern
+            cstruct = asxp(cvec2cmat(cshape, c, qnmat, mps.qntot))
+            cout = expr(cstruct) * inverse
+            # convert structure c to 1d according to qn
+            res.append(asnumpy(cout)[qnmat == mps.qntot])
+
+        if len(res) == 1:
+            return res[0]
+        else:
+            return np.stack(res, axis=1)
+
+    # Find the eigenvectors
+    algo = mps.optimize_config.algo
+    nroots = mps.optimize_config.nroots
+    if algo == "davidson":
+        precond = lambda x, e, *args: x / (hdiag - e + 1e-4)
+
+        e, c = davidson(
+            hop, cguess, precond, max_cycle=100, nroots=nroots, max_memory=64000
+        )
+        # if one root, davidson return e as np.float
+
+    # elif algo == "arpack":
+    #    # scipy arpack solver : much slower than pyscf/davidson
+    #    A = scipy.sparse.linalg.LinearOperator((nonzeros,nonzeros), matvec=hop)
+    #    e, c = scipy.sparse.linalg.eigsh(A, k=nroots, which="SA", v0=cguess)
+    #    # scipy return numpy.array
+    #    if nroots == 1:
+    #        e = e[0]
+    # elif algo == "lobpcg":
+    #    precond = lambda x: scipy.sparse.diags(1/(hdiag+1e-4)) @ x
+    #    A = scipy.sparse.linalg.LinearOperator((nonzeros,nonzeros),
+    #            matvec=hop, matmat=hop)
+    #    M = scipy.sparse.linalg.LinearOperator((nonzeros,nonzeros),
+    #            matvec=precond, matmat=hop)
+    #    e, c = scipy.sparse.linalg.lobpcg(A, np.array(cguess).T,
+    #            M=M, largest=False)
+    elif algo == "primme":
+        h_dim = np.sum(qnmat == mps.qntot)
+        precond = lambda x: scipy.sparse.diags(1 / (hdiag + 1e-4)) @ x
+        A = scipy.sparse.linalg.LinearOperator((h_dim, h_dim), matvec=hop, matmat=hop)
+        M = scipy.sparse.linalg.LinearOperator((h_dim, h_dim), matvec=precond, matmat=hop)
+        e, c = primme.eigsh(
+            A,
+            k=min(nroots, h_dim),
+            which="SA",
+            v0=np.array(cguess).T,
+            OPinv=M,
+            method="PRIMME_DYNAMIC",
+            tol=1e-6,
+        )
+    else:
+        assert False
+    logger.debug(f"use {algo}, HC hops: {count}")
+    return e, c
diff --git a/renormalizer/mps/tests/test_gs.py b/renormalizer/mps/tests/test_gs.py
index 0fe2c49c..5b51fa11 100644
--- a/renormalizer/mps/tests/test_gs.py
+++ b/renormalizer/mps/tests/test_gs.py
@@ -4,7 +4,7 @@
 import numpy as np
 import pytest
 
-from renormalizer.mps.gs import construct_mps_mpo_2, optimize_mps
+from renormalizer.mps.gs import construct_mps_mpo, optimize_mps
 from renormalizer.tests.parameter import holstein_model, custom_model
 from renormalizer.utils import constant
 
@@ -21,7 +21,7 @@
         "2site",
 ))
 def test_optimization(scheme, method):
-    mps, mpo = construct_mps_mpo_2(holstein_model.switch_scheme(scheme), procedure[0][0], nexciton)
+    mps, mpo = construct_mps_mpo(holstein_model.switch_scheme(scheme), procedure[0][0], nexciton)
     mps.optimize_config.procedure = procedure
     mps.optimize_config.method = method
     energies, mps_opt = optimize_mps(mps.copy(), mpo)
@@ -37,7 +37,7 @@ def test_optimization(scheme, method):
         "primme",
 ))
 def test_multistate(method, algo):
-    mps, mpo = construct_mps_mpo_2(holstein_model, procedure[0][0], nexciton)
+    mps, mpo = construct_mps_mpo(holstein_model, procedure[0][0], nexciton)
     mps.optimize_config.procedure = procedure
     mps.optimize_config.nroots = 4
     mps.optimize_config.method = method
@@ -59,7 +59,7 @@ def test_multistate(method, algo):
         4,
 ))
 def test_ex(method, nroots):
-    mps, mpo = construct_mps_mpo_2(holstein_model, procedure[0][0], nexciton)
+    mps, mpo = construct_mps_mpo(holstein_model, procedure[0][0], nexciton)
     mps.optimize_config.procedure = procedure
     mps.optimize_config.nroots = nroots
     mps.optimize_config.method = method