Ejemplo n.º 1
0
     def hop(x):
         # H*X
         nonlocal count
         count += 1
         
         assert len(x) == xsize
         tda_coeff = reshape_x(x)
 
         res = [np.zeros_like(coeff) if coeff is not None else None for coeff in tda_coeff]
         
         # fix ket and sweep bra and accumulate into res
         for ims in range(site_num):
             if tda_coeff[ims] is None:
                 assert tangent_u[ims] is None
                 continue
             
             # mix-canonical mps
             mps_tangent = merge(mps_l_cano, mps_r_cano, ims+1)
             mps_tangent[ims] = tensordot(tangent_u[ims], tda_coeff[ims], (-1, 0))
             
             mps_tangent_conj = mps_r_cano.copy()
             environ = Environ(mps_tangent, mpo, "R", mps_conj=mps_tangent_conj)
             
             for ims_conj in range(site_num):
                 ltensor = environ.GetLR(
                     "L", ims_conj-1, mps_tangent, mpo, itensor=None,
                     mps_conj=mps_tangent_conj,
                     method="System"
                 )
                 rtensor = environ.GetLR(
                     "R", ims_conj+1, mps_tangent, mpo, itensor=None,
                     mps_conj=mps_tangent_conj,
                     method="Enviro"
                 )
                 if tda_coeff[ims_conj] is not None:
                     # S-a   l-S
                     #     d
                     # O-b-O-f-O
                     #     e
                     # S-c   k-S
 
                     path = [
                         ([0, 1], "abc, cek -> abek"),
                         ([2, 0], "abek, bdef -> akdf"),
                         ([1, 0], "akdf, lfk -> adl"),
                     ]
                     out = multi_tensor_contract(
                         path, ltensor, asxp(mps_tangent[ims_conj]),
                         asxp(mpo[ims_conj]), rtensor
                     )
                     res[ims_conj] += asnumpy(tensordot(tangent_u[ims_conj], out,
                         ([0,1], [0,1])))
                 
                 # mps_conj combine 
                 mps_tangent_conj[ims_conj] = mps_l_cano[ims_conj]    
         
         res = [mat for mat in res if mat is not None]
 
         return np.concatenate(res, axis=None)
Ejemplo n.º 2
0
def optimize_mps_dmrg(mps, mpo):
    """
    1 or 2 site optimization procedure
    """

    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"]
    # print("optimization method", method)

    nexciton = mps.nexciton

    # construct the environment matrix
    environ = Environ()
    environ.construct(mps, mps, mpo, "L")

    nMPS = len(mps)
    # construct each sweep cycle scheme
    if method == "1site":
        loop = [["R", i]
                for i in range(nMPS - 1, -1, -1)] + [["L", i]
                                                     for i in range(0, nMPS)]
    else:
        loop = [["R", i]
                for i in range(nMPS - 1, 0, -1)] + [["L", i]
                                                    for i in range(1, nMPS)]

    # initial matrix
    ltensor = ones((1, 1, 1))
    rtensor = ones((1, 1, 1))

    energies = []
    for isweep, (mmax, percent) in enumerate(procedure):
        logger.debug(f"mmax, percent: {mmax}, {percent}")
        logger.debug(f"energy: {mps.expectation(mpo)}")
        logger.debug(f"{mps}")

        for system, imps in loop:
            if system == "R":
                lmethod, rmethod = "Enviro", "System"
            else:
                lmethod, rmethod = "System", "Enviro"

            if method == "1site":
                lsite = imps - 1
                addlist = [imps]
            else:
                lsite = imps - 2
                addlist = [imps - 1, imps]

            ltensor = environ.GetLR("L",
                                    lsite,
                                    mps,
                                    mps,
                                    mpo,
                                    itensor=ltensor,
                                    method=lmethod)
            rtensor = environ.GetLR("R",
                                    imps + 1,
                                    mps,
                                    mps,
                                    mpo,
                                    itensor=rtensor,
                                    method=rmethod)

            # get the quantum number pattern
            qnmat, qnbigl, qnbigr = construct_qnmat(mps, mpo.ephtable,
                                                    mpo.pbond_list, addlist,
                                                    method, system)
            cshape = qnmat.shape

            # hdiag
            tmp_ltensor = einsum("aba -> ba", ltensor)
            tmp_MPOimps = einsum("abbc -> abc", mpo[imps])
            tmp_rtensor = 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_MPOimps,
                                              tmp_rtensor)[(qnmat == nexciton)]
                # initial guess   b-S-c
                #                   a
                cguess = mps[imps][qnmat == nexciton]
            else:
                #   S-a c   d f-S
                #   O-b-O-e-O-g-O
                #   S-a c   d f-S
                tmp_MPOimpsm1 = einsum("abbc -> abc", mpo[imps - 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_MPOimpsm1,
                                              tmp_MPOimps,
                                              tmp_rtensor)[(qnmat == nexciton)]
                # initial guess b-S-c-S-e
                #                 a   d
                cguess = tensordot(mps[imps - 1], mps[imps],
                                   axes=1)[qnmat == nexciton]
            cguess = cguess.asnumpy()
            hdiag *= inverse
            nonzeros = np.sum(qnmat == nexciton)

            # print("Hmat dim", nonzeros)

            def hop(c):
                # convert c to initial structure according to qn pattern
                cstruct = cvec2cmat(cshape, c, qnmat, nexciton)

                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,
                                                 Matrix(cstruct), mpo[imps],
                                                 rtensor)
                    # for small matrices, check hermite:
                    # a=tensordot(ltensor, mpo[imps], ((1), (0)))
                    # b=tensordot(a, rtensor, ((4), (1)))
                    # c=b.transpose((0, 2, 4, 1, 3, 5))
                    # d=c.reshape(16, 16)
                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,
                        Matrix(cstruct),
                        mpo[imps - 1],
                        mpo[imps],
                        rtensor,
                    )
                # convert structure c to 1d according to qn
                return inverse * cout.asnumpy()[qnmat == nexciton]

            if nroots != 1:
                cguess = [cguess]
                for iroot in range(nroots - 1):
                    cguess.append(np.random.random([nonzeros]) - 0.5)

            precond = lambda x, e, *args: x / (hdiag.asnumpy() - e + 1e-4)

            e, c = davidson(hop,
                            cguess,
                            precond,
                            max_cycle=100,
                            nroots=nroots,
                            max_memory=64000)
            # scipy arpack solver : much slower than davidson
            # A = spslinalg.LinearOperator((nonzeros,nonzeros), matvec=hop)
            # e, c = spslinalg.eigsh(A,k=1, which="SA",v0=cguess)
            # print("HC loops:", count[0])
            # logger.debug(f"isweep: {isweep}, e: {e}")

            energies.append(e)

            cstruct = cvec2cmat(cshape, c, qnmat, nexciton, nroots=nroots)

            if nroots == 1:
                # direct svd the coefficient matrix
                mt, mpsdim, mpsqn, compmps = renormalization_svd(
                    cstruct,
                    qnbigl,
                    qnbigr,
                    system,
                    nexciton,
                    Mmax=mmax,
                    percent=percent,
                )
            else:
                # diagonalize the reduced density matrix
                mt, mpsdim, mpsqn, compmps = renormalization_ddm(
                    cstruct,
                    qnbigl,
                    qnbigr,
                    system,
                    nexciton,
                    Mmax=mmax,
                    percent=percent,
                )

            if method == "1site":
                mps[imps] = mt
                if system == "L":
                    if imps != len(mps) - 1:
                        mps[imps + 1] = tensordot(compmps,
                                                  mps[imps + 1],
                                                  axes=1)
                        mps.qn[imps + 1] = mpsqn
                    else:
                        mps[imps] = tensordot(mps[imps], compmps, axes=1)
                        mps.qn[imps + 1] = [0]

                else:
                    if imps != 0:
                        mps[imps - 1] = tensordot(mps[imps - 1],
                                                  compmps,
                                                  axes=1)
                        mps.qn[imps] = mpsqn
                    else:
                        mps[imps] = tensordot(compmps, mps[imps], axes=1)
                        mps.qn[imps] = [0]
            else:
                if system == "L":
                    mps[imps - 1] = mt
                    mps[imps] = compmps
                else:
                    mps[imps] = mt
                    mps[imps - 1] = compmps

                # mps.dim_list[imps] = mpsdim
                mps.qn[imps] = mpsqn

    energies = np.array(energies)
    if nroots == 1:
        logger.debug("Optimization complete, lowest energy = %g",
                     energies.min())

    return energies
Ejemplo n.º 3
0
    def _evolve_dmrg_tdvp_ps(self, mpo, evolve_dt) -> "Mps":
        # PhysRevB.94.165116
        # TDVP projector splitting
        imag_time = np.iscomplex(evolve_dt)
        if imag_time:
            mps = self.copy()
            mps_conj = mps
        else:
            mps = self.to_complex()
            mps_conj = mps.conj()  # another copy, so 3x memory is used.

        # construct the environment matrix
        environ = Environ()
        # almost half is not used. Not a big deal.
        environ.construct(mps, mps_conj, mpo, "L")
        environ.construct(mps, mps_conj, mpo, "R")

        # a workaround for https://github.com/scipy/scipy/issues/10164
        if imag_time:
            evolve_dt = -evolve_dt.imag
            # used in calculating derivatives
            coef = -1
        else:
            coef = 1j

        # statistics for debug output
        cmf_rk_steps = []
        USE_RK = self.evolve_config.tdvp_ps_rk4
        # sweep for 2 rounds
        for i in range(2):
            for imps in mps.iter_idx_list(full=True):
                system = "L" if mps.left else "R"
                ltensor = environ.read("L", imps - 1)
                rtensor = environ.read("R", imps + 1)

                shape = list(mps[imps].shape)
                l_array = ltensor.array
                r_array = rtensor.array

                hop = hop_factory(l_array, r_array, mpo[imps].array, len(shape))

                def hop_svt(ms):
                    # S-a   l-S
                    #
                    # O-b - b-O
                    #
                    # S-c   k-S

                    path = [([0, 1], "abc, ck -> abk"), ([1, 0], "abk, lbk -> al")]
                    HC = multi_tensor_contract(path, l_array, ms, r_array)
                    return HC

                if USE_RK:
                    def func(t, y):
                        return hop(y.reshape(shape)).ravel() / coef
                    sol = solve_ivp(
                        func, (0, evolve_dt / 2.0), mps[imps].ravel().array, method="RK45"
                    )
                    cmf_rk_steps.append(len(sol.t))
                    mps_t = sol.y[:, -1]
                else:
                    # Can't use the same func because here H should be Hermitian
                    def func(y):
                        return hop(y.reshape(shape)).ravel()
                    mps_t = expm_krylov(func, (evolve_dt / 2) / coef, mps[imps].ravel().array)
                mps_t = mps_t.reshape(shape)

                qnbigl, qnbigr = mps._get_big_qn(imps)
                u, qnlset, v, qnrset = svd_qn.Csvd(
                    asnumpy(mps_t),
                    qnbigl,
                    qnbigr,
                    mps.qntot,
                    QR=True,
                    system=system,
                    full_matrices=False,
                )
                vt = v.T

                if mps.is_left_canon and imps != 0:
                    mps[imps] = vt.reshape([-1] + shape[1:])
                    mps_conj[imps] = mps[imps].conj()
                    mps.qn[imps] = qnrset

                    rtensor = environ.GetLR(
                        "R", imps, mps, mps_conj, mpo, itensor=rtensor, method="System"
                    )
                    r_array = rtensor.array

                    # reverse update u site
                    shape_u = u.shape

                    if USE_RK:
                        def func_u(t, y):
                            return hop_svt(y.reshape(shape_u)).ravel() / coef
                        sol_u = solve_ivp(
                            func_u, (0, -evolve_dt / 2), u.ravel(), method="RK45"
                        )
                        cmf_rk_steps.append(len(sol_u.t))
                        mps_t = sol_u.y[:, -1]
                    else:
                        def func_u(y):
                            return hop_svt(y.reshape(shape_u)).ravel()
                        mps_t = expm_krylov(func_u, (-evolve_dt / 2) / coef, u.ravel())
                    mps_t = mps_t.reshape(shape_u)
                    mps[imps - 1] = tensordot(
                        mps[imps - 1].array,
                        mps_t,
                        axes=(-1, 0),
                    )
                    mps_conj[imps - 1] = mps[imps - 1].conj()

                elif mps.is_right_canon and imps != len(mps) - 1:
                    mps[imps] = u.reshape(shape[:-1] + [-1])
                    mps_conj[imps] = mps[imps].conj()
                    mps.qn[imps + 1] = qnlset

                    ltensor = environ.GetLR(
                        "L", imps, mps, mps_conj, mpo, itensor=ltensor, method="System"
                    )
                    l_array = ltensor.array

                    # reverse update svt site
                    shape_svt = vt.shape

                    if USE_RK:
                        def func_svt(t, y):
                            return hop_svt(y.reshape(shape_svt)).ravel() / coef
                        sol_svt = solve_ivp(
                            func_svt, (0, -evolve_dt / 2), vt.ravel(), method="RK45"
                        )
                        cmf_rk_steps.append(len(sol_svt.t))
                        mps_t = sol_svt.y[:, -1]
                    else:
                        def func_svt(y):
                            return hop_svt(y.reshape(shape_svt)).ravel()
                        mps_t = expm_krylov(func_svt, (-evolve_dt / 2) / coef, vt.ravel())
                    mps_t = mps_t.reshape(shape_svt)
                    mps[imps + 1] = tensordot(
                        mps_t,
                        mps[imps + 1].array,
                        axes=(1, 0),
                    )
                    mps_conj[imps + 1] = mps[imps + 1].conj()

                else:
                    mps[imps] = mps_t
                    mps_conj[imps] = mps[imps].conj()
            mps._switch_direction()

        if USE_RK:
            steps_stat = stats.describe(cmf_rk_steps)
            logger.debug(f"TDVP-PS CMF steps: {steps_stat}")
            mps.evolve_config.stat = steps_stat

        return mps
Ejemplo n.º 4
0
    def _evolve_dmrg_tdvp_mctdhnew(self, mpo, evolve_dt) -> "Mps":
        # new regularization scheme
        # JCP 148, 124105 (2018)
        # JCP 149, 044119 (2018)

        # a workaround for https://github.com/scipy/scipy/issues/10164
        imag_time = np.iscomplex(evolve_dt)
        if imag_time:
            evolve_dt = -evolve_dt.imag
            # used in calculating derivatives
            coef = -1
        else:
            coef = 1j

        if self.is_left_canon:
            assert self.check_left_canonical()
            self.canonicalise()

        mps = self.to_complex(inplace=True)

        # construct the environment matrix
        environ = Environ()
        environ.construct(mps, mps.conj(), mpo, "R")

        # initial matrix
        ltensor = ones((1, 1, 1))
        rtensor = ones((1, 1, 1))

        new_mps = mps.metacopy()

        # statistics for debug output
        cmf_rk_steps = []

        for imps in range(len(mps)):
            shape = list(mps[imps].shape)

            system = "L" if mps.left else "R"
            qnbigl, qnbigr = mps._get_big_qn(imps)
            u, s, qnlset, v, s, qnrset = svd_qn.Csvd(
                mps[imps].asnumpy(),
                qnbigl,
                qnbigr,
                mps.qntot,
                system=system,
                full_matrices=False,
            )
            vt = v.T

            mps[imps] = u.reshape(shape[:-1] + [-1])

            ltensor = environ.GetLR(
                "L", imps - 1, mps, mps.conj(), mpo, itensor=ltensor, method="System"
            )
            rtensor = environ.GetLR(
                "R", imps + 1, mps, mps.conj(), mpo, itensor=rtensor, method="Enviro"
            )

            epsilon = 1e-10
            epsilon = np.sqrt(epsilon)
            s = s + epsilon * np.exp(-s / epsilon)

            svt = Matrix(np.diag(s).dot(vt))

            rtensor = tensordot(rtensor, svt, axes=(2, 1))
            rtensor = tensordot(Matrix(vt).conj(), rtensor, axes=(1, 0))

            if imps != len(mps) - 1:
                mps[imps + 1] = tensordot(svt, mps[imps + 1], axes=(-1, 0))
                mps.qn[imps + 1] = qnlset
                new_mps.qn[imps + 1] = qnlset.copy()

            S_inv = xp.diag(1.0 / s)

            hop = hop_factory(ltensor, rtensor, mpo[imps], len(shape))

            func = integrand_func_factory(shape, hop, imps == len(mps) - 1, S_inv, coef)

            sol = solve_ivp(
                func, (0, evolve_dt), mps[imps].ravel().array, method="RK45"
            )
            cmf_rk_steps.append(len(sol.t))
            ms = sol.y[:, -1].reshape(shape)

            if imps == len(mps) - 1:
                new_mps[imps] = ms * s[0]
            else:
                new_mps[imps] = ms
        mps._switch_direction()
        new_mps._switch_direction()
        new_mps.canonicalise()

        steps_stat = stats.describe(cmf_rk_steps)
        logger.debug(f"TDVP-MCTDH CMF steps: {steps_stat}")
        # new_mps.evolve_config.stat = steps_stat

        return new_mps
Ejemplo n.º 5
0
    def _evolve_dmrg_tdvp_mctdh(self, mpo, evolve_dt) -> "Mps":
        # TDVP for original MCTDH
        if self.is_right_canon:
            assert self.check_right_canonical()
            self.canonicalise()

        # a workaround for https://github.com/scipy/scipy/issues/10164
        imag_time = np.iscomplex(evolve_dt)
        if imag_time:
            evolve_dt = -evolve_dt.imag
            # used in calculating derivatives
            coef = -1
        else:
            coef = 1j

        # qn for this method has not been implemented
        self.use_dummy_qn = True
        self.clear_qn()
        mps = self.to_complex(inplace=True)
        mps_conj = mps.conj()
        environ = Environ()
        environ.construct(mps, mps_conj, mpo, "R")

        # initial matrix
        ltensor = np.ones((1, 1, 1))
        rtensor = np.ones((1, 1, 1))

        new_mps = self.metacopy()

        cmf_rk_steps = []

        for imps in range(len(mps)):
            ltensor = environ.GetLR(
                "L", imps - 1, mps, mps_conj, mpo, itensor=ltensor, method="System"
            )
            rtensor = environ.GetLR(
                "R", imps + 1, mps, mps_conj, mpo, itensor=rtensor, method="Enviro"
            )
            # density matrix
            S = transferMat(mps, mps_conj, "R", imps + 1).asnumpy()

            epsilon = 1e-8
            w, u = scipy.linalg.eigh(S)
            try:
                w = w + epsilon * np.exp(-w / epsilon)
            except FloatingPointError:
                logger.warning(f"eigenvalue of density matrix contains negative value")
                w -= 2 * w.min()
                w = w + epsilon * np.exp(-w / epsilon)
            # print
            # "sum w=", np.sum(w)
            # S  = u.dot(np.diag(w)).dot(np.conj(u.T))
            S_inv = xp.asarray(u.dot(np.diag(1.0 / w)).dot(np.conj(u.T)))

            # pseudo inverse
            # S_inv = scipy.linalg.pinvh(S,rcond=1e-2)

            shape = mps[imps].shape

            hop = hop_factory(ltensor, rtensor, mpo[imps], len(shape))

            func = integrand_func_factory(shape, hop, imps == len(mps) - 1, S_inv, coef)

            sol = solve_ivp(
                func, (0, evolve_dt), mps[imps].ravel().array, method="RK45"
            )
            # print
            # "CMF steps:", len(sol.t)
            cmf_rk_steps.append(len(sol.t))
            new_mps[imps] = sol.y[:, -1].reshape(shape)
            new_mps[imps].check_lortho()
            # print
            # "orthogonal1", np.allclose(np.tensordot(MPSnew[imps],
            #                                        np.conj(MPSnew[imps]), axes=([0, 1], [0, 1])),
            #                           np.diag(np.ones(MPSnew[imps].shape[2])))
        steps_stat = stats.describe(cmf_rk_steps)
        logger.debug(f"TDVP-MCTDH CMF steps: {steps_stat}")

        return new_mps
Ejemplo n.º 6
0
    def kernel(self, restart=False, include_psi0=False):
        r"""calculate the roots

        Parameters
        ----------
        restart: bool, optional
            if restart from the former converged root. Default is ``False``.
            If ``restart = True``, ``include_psi0`` must be the same as the
            former calculation.
        include_psi0: bool, optional
            if the basis of Hamiltonian includes the ground state
                :math:`\Psi_0`. Default is ``False``.

        Returns
        -------
        e: np.ndarray
            the energy of the states, if ``include_psi0 = True``, the first
            element is the ground state energy, otherwise, it is the energy of
            the first excited state.

        """
        # right canonical mps
        mpo = self.hmpo
        nroots = self.nroots
        algo = self.algo
        site_num = mpo.site_num

        if not restart:
            # make sure that M is not redundant near the edge
            mps = self.mps.ensure_right_canon().canonicalise().normalize().canonicalise()
            logger.debug(f"reference mps shape, {mps}")
            mps_r_cano = mps.copy()
            assert mps.to_right 
            
            tangent_u = []
    
            for ims, ms in enumerate(mps):
                
                shape = list(ms.shape)
                u, s, vt = scipy.linalg.svd(ms.l_combine(), full_matrices=True)
                rank = len(s)
                if include_psi0 and ims == site_num-1: 
                    tangent_u.append(u.reshape(shape[:-1]+[-1]))
                else:
                    if rank < u.shape[1]:
                        tangent_u.append(u[:,rank:].reshape(shape[:-1]+[-1]))
                    else:
                        tangent_u.append(None)  # the tangent space is None

                mps[ims] = u[:,:rank].reshape(shape[:-1]+[-1])
                
                vt = xp.einsum("i, ij -> ij", asxp(s), asxp(vt))
                if ims == site_num-1:
                    assert vt.size == 1 and xp.allclose(vt, 1)
                else:
                    mps[ims+1] = asnumpy(tensordot(vt, mps[ims+1], ([-1],[0])))
                
            mps_l_cano = mps.copy() 
            mps_l_cano.to_right = False
            mps_l_cano.qnidx = site_num-1

        else:
            mps_l_cano, mps_r_cano, tangent_u, tda_coeff_list = self.wfn
            cguess = []
            for iroot in range(len(tda_coeff_list)):
                tda_coeff = tda_coeff_list[iroot]
                x = [c.flatten() for c in tda_coeff if c is not None]
                x = np.concatenate(x,axis=None)
                cguess.append(x)
            cguess = np.stack(cguess, axis=1)

        xshape = [] 
        xsize = 0
        for ims in range(site_num):
            if tangent_u[ims] is None:
                xshape.append((0,0))
            else:
                if ims == site_num-1:
                    xshape.append((tangent_u[ims].shape[-1], 1))
                else:    
                    xshape.append((tangent_u[ims].shape[-1], mps_r_cano[ims+1].shape[0]))
                xsize += np.prod(xshape[-1])
        
        logger.debug(f"DMRG-TDA H dimension: {xsize}")
        
        if USE_GPU:
            oe_backend = "cupy"
        else:
            oe_backend = "numpy"
        
        mps_tangent = mps_r_cano.copy()
        environ = Environ(mps_tangent, mpo, "R")
        hdiag = []
        for ims in range(site_num):
            ltensor = environ.GetLR(
                "L", ims-1, mps_tangent, mpo, itensor=None,
                method="System"
            )
            rtensor = environ.GetLR(
                "R", ims+1, mps_tangent, mpo, itensor=None,
                method="Enviro"
            )
            if tangent_u[ims] is not None:
                u = asxp(tangent_u[ims])
                tmp = oe.contract("abc, ded, bghe, agl, chl -> ld", ltensor, rtensor,
                        asxp(mpo[ims]), u, u, backend=oe_backend)   
                hdiag.append(asnumpy(tmp))
            mps_tangent[ims] = mps_l_cano[ims]
        hdiag = np.concatenate(hdiag, axis=None)
    
        count = 0
        
        # recover the vector-like x back to the ndarray tda_coeff
        def reshape_x(x):
            tda_coeff = []
            offset = 0
            for shape in xshape:
                if shape == (0,0):
                    tda_coeff.append(None)
                else:
                    size = np.prod(shape)
                    tda_coeff.append(x[offset:size+offset].reshape(shape))
                    offset += size
            
            assert offset == xsize
            return tda_coeff
            
        def hop(x):
            # H*X
            nonlocal count
            count += 1
            
            assert len(x) == xsize
            tda_coeff = reshape_x(x)
    
            res = [np.zeros_like(coeff) if coeff is not None else None for coeff in tda_coeff]
            
            # fix ket and sweep bra and accumulate into res
            for ims in range(site_num):
                if tda_coeff[ims] is None:
                    assert tangent_u[ims] is None
                    continue
                
                # mix-canonical mps
                mps_tangent = merge(mps_l_cano, mps_r_cano, ims+1)
                mps_tangent[ims] = tensordot(tangent_u[ims], tda_coeff[ims], (-1, 0))
                
                mps_tangent_conj = mps_r_cano.copy()
                environ = Environ(mps_tangent, mpo, "R", mps_conj=mps_tangent_conj)
                
                for ims_conj in range(site_num):
                    ltensor = environ.GetLR(
                        "L", ims_conj-1, mps_tangent, mpo, itensor=None,
                        mps_conj=mps_tangent_conj,
                        method="System"
                    )
                    rtensor = environ.GetLR(
                        "R", ims_conj+1, mps_tangent, mpo, itensor=None,
                        mps_conj=mps_tangent_conj,
                        method="Enviro"
                    )
                    if tda_coeff[ims_conj] is not None:
                        # S-a   l-S
                        #     d
                        # O-b-O-f-O
                        #     e
                        # S-c   k-S
    
                        path = [
                            ([0, 1], "abc, cek -> abek"),
                            ([2, 0], "abek, bdef -> akdf"),
                            ([1, 0], "akdf, lfk -> adl"),
                        ]
                        out = multi_tensor_contract(
                            path, ltensor, asxp(mps_tangent[ims_conj]),
                            asxp(mpo[ims_conj]), rtensor
                        )
                        res[ims_conj] += asnumpy(tensordot(tangent_u[ims_conj], out,
                            ([0,1], [0,1])))
                    
                    # mps_conj combine 
                    mps_tangent_conj[ims_conj] = mps_l_cano[ims_conj]    
            
            res = [mat for mat in res if mat is not None]
    
            return np.concatenate(res, axis=None)
        
        if algo == "davidson":
            if restart:
                cguess = [cguess[:,i] for i in range(cguess.shape[1])]
            else:
                cguess = [np.random.random(xsize) - 0.5]
            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 nroots == 1:
                c = [c]
            c = np.stack(c, axis=1)

        elif algo == "primme":
            if not restart:
                cguess = None

            def multi_hop(x):
                if x.ndim == 1:
                    return hop(x)
                elif x.ndim == 2:
                    return np.stack([hop(x[:,i]) for i in range(x.shape[1])],axis=1)
                else:
                    assert False
    
            def precond(x): 
                if x.ndim == 1:
                    return np.einsum("i, i -> i", 1/(hdiag+1e-4), x)
                elif x.ndim ==2:
                    return np.einsum("i, ij -> ij", 1/(hdiag+1e-4), x)
                else:
                    assert False
            A = scipy.sparse.linalg.LinearOperator((xsize,xsize),
                    matvec=multi_hop, matmat=multi_hop)
            M = scipy.sparse.linalg.LinearOperator((xsize,xsize),
                    matvec=precond, matmat=precond)
            e, c = primme.eigsh(A, k=min(nroots,xsize), which="SA", 
                    v0=cguess,
                    OPinv=M,
                    method="PRIMME_DYNAMIC", 
                    tol=1e-6)
        else:
            assert False

        logger.debug(f"H*C times: {count}")
        
        tda_coeff_list = []
        for iroot in range(nroots):
            tda_coeff_list.append(reshape_x(c[:,iroot])) 
        
        self.e = np.array(e)
        self.wfn = [mps_l_cano, mps_r_cano, tangent_u, tda_coeff_list]
        
        return self.e
Ejemplo n.º 7
0
    def variational_compress(self, mpo=None, guess=None):
        r"""Variational compress an mps/mpdm/mpo
        
        Parameters
        ----------
        mpo : renormalizer.mps.Mpo, optional 
            Default is ``None``. if mpo is not ``None``, the returned mps is
            an approximation of ``mpo @ self``
        guess : renormalizer.mps.MatrixProduct, optional
            Initial guess of compressed mps/mpdm/mpo. Default is ``None``. 
        
        Note
        ----
        the variational compress related configurations is defined in
        ``self`` if ``guess=None``, otherwise is defined in ``guess``

        Returns
        -------
        mp : renormalizer.mps.MatrixProduct
            a new compressed mps/mpdm/mpo, ``self`` is not overwritten.
            ``guess`` is overwritten.
        
        """

        if mpo is None:
            logger.info(
                "Recommend to use svd to compress a single mps/mpo/mpdm.")
            raise NotImplementedError

        if guess is None:
            # a minimal representation of self and mpo
            compressed_mpo = mpo.copy().canonicalise().compress(
                temp_m_trunc=self.compress_config.vguess_m[0])
            compressed_mps = self.copy().canonicalise().compress(
                temp_m_trunc=self.compress_config.vguess_m[1])
            # the attributes of guess would be the same as self
            guess = compressed_mpo.apply(compressed_mps)
        mps = guess
        mps.ensure_left_canon()
        logger.info(f"initial guess bond dims: {mps.bond_dims}")

        procedure = mps.compress_config.vprocedure
        method = mps.compress_config.vmethod

        environ = Environ(self, mpo, "L", mps_conj=mps.conj())

        converged = False
        for isweep, (mmax, percent) in enumerate(procedure):
            logger.debug(f"isweep: {isweep}")
            logger.debug(f"mmax, percent: {mmax}, {percent}")
            logger.debug(f"mps bond dims: {mps.bond_dims}")

            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, imps + 1]
                        ridx = imps + 2
                    else:
                        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,
                                        self,
                                        mpo,
                                        itensor=None,
                                        method=lmethod,
                                        mps_conj=mps.conj())
                rtensor = environ.GetLR("R",
                                        ridx,
                                        self,
                                        mpo,
                                        itensor=None,
                                        method=rmethod,
                                        mps_conj=mps.conj())

                # get the quantum number pattern
                qnbigl, qnbigr, qnmat = mps._get_big_qn(cidx)

                # center mo
                cmo = [asxp(mpo[idx]) for idx in cidx]
                cms = [asxp(self[idx]) for idx in cidx]
                if method == "1site":
                    if cms[0].ndim == 3:
                        # S-a   l-S
                        #     d
                        # O-b-O-f-O
                        #     e
                        # S-c   k-S

                        path = [
                            ([0, 1], "abc, cek -> abek"),
                            ([2, 0], "abek, bdef -> akdf"),
                            ([1, 0], "akdf, lfk -> adl"),
                        ]
                    elif cms[0].ndim == 4:
                        # S-a   l-S
                        #     d
                        # O-b-O-f-O
                        #     e
                        # S-c   k-S
                        #     g
                        path = [
                            ([0, 2], "abc, bdef -> acdef"),
                            ([2, 0], "acdef, cegk -> adfgk"),
                            ([1, 0], "adfgk, lfk -> adgl"),
                        ]
                    cout = multi_tensor_contract(path, ltensor, cms[0], cmo[0],
                                                 rtensor)
                else:
                    if USE_GPU:
                        oe_backend = "cupy"
                    else:
                        oe_backend = "numpy"
                    if cms[0].ndim == 3:
                        # S-a       l-S
                        #     d   g
                        # O-b-O-f-O-j-O
                        #     e   h
                        # S-c   m   k-S

                        cout = oe.contract(
                            "abc, bdef, fghj, ljk, cem, mhk -> adgl",
                            ltensor,
                            cmo[0],
                            cmo[1],
                            rtensor,
                            cms[0],
                            cms[1],
                            backend=oe_backend)
                    elif cms[0].ndim == 4:
                        # S-a       l-S
                        #     d   g
                        # O-b-O-f-O-j-O
                        #     e   h
                        # S-c   m   k-S
                        #     n   p
                        cout = oe.contract(
                            "abc, bdef, fghj, ljk, cenm, mhpk -> adngpl",
                            ltensor,
                            cmo[0],
                            cmo[1],
                            rtensor,
                            cms[0],
                            cms[1],
                            backend=oe_backend)
                # clean up the elements which do not meet the qn requirements
                cout[qnmat != mps.qntot] = 0
                mps._update_mps(cout, cidx, qnbigl, qnbigr, mmax, percent)

            mps._switch_direction()

            # check if convergence
            if isweep > 0 and percent == 0 and \
                    mps.distance(mps_old) / np.sqrt(mps.dot(mps.conj()).real) < mps.compress_config.vrtol:
                converged = True
                break

            mps_old = mps.copy()

        if converged:
            logger.info("Variational compress is converged!")
        else:
            logger.warning(
                "Variational compress is not converged! Please increase the procedure!"
            )

        # remove the redundant bond dimension near the boundary of the MPS
        mps.canonicalise()
        logger.info(f"{mps}")

        return mps
Ejemplo n.º 8
0
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 
        mpo of Hamiltonian
    omega: float, optional
        target the eigenpair near omega with special variational function
        :math:(\hat{H}-\omega)^2. Default is `None`.

    Returns
    -------
    energy : list
        list of energy of each marco sweep.
        :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.
    
    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}")
    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:
        mps.ensure_right_canon()
        env = "R"
    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)
        mpo = mpo.add(identity.scale(-omega))
        environ = Environ(mps, [mpo, mpo], env)
    else:
        environ = Environ(mps, mpo, env)

    macro_iteration_result = []
    converged = False
    for isweep, (mmax, percent) in enumerate(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)):
                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, imps + 1]
                    ridx = imps + 2
                else:
                    lidx = imps - 2
                    cidx = [imps - 1, imps]  # center site
                    ridx = imps + 1
            else:
                assert False
            logger.debug(f"optimize site: {cidx}")

            if omega is None:
                operator = mpo
            else:
                operator = [mpo, mpo]

            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
            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
                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])
                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:
                    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))

                        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,
                                )
                        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}")
            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

    logger.debug(
        f"{isweep+1} sweeps are finished, lowest energy = {sorted(macro_iteration_result)[0]}"
    )
    if converged:
        logger.info("DMRG is converged!")
    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 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