Obtaining spectral moments of a MPS state

[basilib]

Hi block2 developers,

I am interested in obtaining the spectral moments of an MPS. These are defined as:
$\langle \Psi | c_i H^n c^\dagger_j | \Psi \rangle$ and, $\langle \Psi | c^\dagger_j H^n c_i | \Psi \rangle $.

Since I require all of these quantities for up to some $n$, I think it would be most efficient to calculate the states $H^n c_i | \Psi \rangle$ by iteratively applying the Hamiltonian for $n=0,1,\ldots, N$, then taking the inner product with $\langle \Psi | c^\dagger_j$.

It might be possible to do this using the Green's function code or TD-DMRG functionality, however I am not familiar with the codebase and I would appreciate some advice on how to obtain these quantities.

[hczhai]

Please see if the following example script can solve your problem for computing $\langle\Psi|c_i^\dagger H^n c_j |\Psi\rangle$. Note methods 1 and 2 will be too expensive for $n\ge2$, so they are only for benchmarking purposes.

import numpy as np
from pyblock2._pyscf.ao2mo import integrals as itg
from pyblock2.driver.core import DMRGDriver, SymmetryTypes
from pyscf import gto, scf, lo

BOHR = 0.52917721092
R = 1.8 * BOHR
N = 6

mol = gto.M(atom=[['H', (i * R, 0, 0)] for i in range(N)], basis="sto6g", symmetry="c1", verbose=0)
mf = scf.RHF(mol).run(conv_tol=1E-14)

mf.mo_coeff = lo.orth.lowdin(mol.intor('cint1e_ovlp_sph'))
ncas, n_elec, spin, ecore, h1e, g2e, orb_sym = itg.get_rhf_integrals(mf, ncore=0, ncas=None, g2e_symm=1)

driver = DMRGDriver(scratch="./tmp", symm_type=SymmetryTypes.SZ, n_threads=4)
driver.initialize_system(n_sites=ncas, n_elec=n_elec, spin=spin, orb_sym=orb_sym)

bond_dims = [150] * 4 + [200] * 4
noises = [1e-4] * 4 + [1e-5] * 4 + [0]
thrds = [1e-10] * 8

mpo = driver.get_qc_mpo(h1e=h1e, g2e=g2e, ecore=ecore, integral_cutoff=1E-8, iprint=1)
ket = driver.get_random_mps(tag="KET", bond_dim=150, nroots=1)
energy = driver.dmrg(mpo, ket, n_sweeps=20, bond_dims=bond_dims, noises=noises,
    thrds=thrds, iprint=1)
print('Ground state energy = %20.15f' % energy)

# [Method 1] Get <i|H^n|j> by construction of MPO for "i H^n j"
def get_expt(ket, n, xi, xj, cutoff=1E-12):
    terms = [(1.0, "", [])]
    for _ in range(n):
        new_terms = []
        for f, op, ixs in terms:
            for i, j in np.ndindex(*h1e.shape):
                if abs(f * h1e[i, j] * 2 ** 0.5) > cutoff:
                    new_terms.append((f * h1e[i, j], op + "CD", ixs + [i, j]))
                    new_terms.append((f * h1e[i, j], op + "cd", ixs + [i, j]))
            for i, j, k, l in np.ndindex(*g2e.shape):
                if abs(f * g2e[i, j, k, l]) > cutoff:
                    new_terms.append((f * g2e[i, j, k, l] * 0.5, op + "CCDD", ixs + [i, k, l, j]))
                    new_terms.append((f * g2e[i, j, k, l] * 0.5, op + "ccdd", ixs + [i, k, l, j]))
                    new_terms.append((f * g2e[i, j, k, l] * 0.5, op + "CcdD", ixs + [i, k, l, j]))
                    new_terms.append((f * g2e[i, j, k, l] * 0.5, op + "cCDd", ixs + [i, k, l, j]))
            new_terms.append((f * ecore, op, ixs))
        terms = new_terms
    b = driver.expr_builder()
    for f, op, ixs in terms:
        b.add_term("c" + op + "d", [xi] + ixs + [xj], f)
        b.add_term("C" + op + "D", [xi] + ixs + [xj], f)
    xmpo = driver.get_mpo(b.finalize(fermionic_ops="cdCD"))
    return driver.expectation(ket, xmpo, ket)

h0mat = np.array([[get_expt(ket, 0, xi, xj, cutoff=1E-12) for xj in range(ncas)] for xi in range(ncas)])
h1mat = np.array([[get_expt(ket, 1, xi, xj, cutoff=1E-12) for xj in range(ncas)] for xi in range(ncas)])

# [Method 2] Get <i|H^n|j> by contraction of NPDM
ref_dm1 = np.sum(driver.get_1pdm(ket), axis=0)
print(np.linalg.norm(h0mat - ref_dm1))

dm2aa, dm2ab, dm2bb = driver.get_2pdm(ket)
ref_dm2 = np.sum([dm2aa, dm2ab, dm2ab.transpose(1, 0, 3, 2), dm2bb], axis=0)
dm3aaa, dm3aab, dm3abb, dm3bbb = driver.get_3pdm(ket)
ref_dm3 = np.sum([dm3aaa, dm3aab, dm3aab.transpose(0, 2, 1, 4, 3, 5), dm3abb, dm3aab.transpose(2, 0, 1, 4, 5, 3),
    dm3abb.transpose(1, 0, 2, 3, 5, 4), dm3abb.transpose(1, 2, 0, 5, 3, 4), dm3bbb], axis=0)
print(np.linalg.norm(h1mat - (ecore * ref_dm1 + np.einsum('ijkl,jk->il', ref_dm2, h1e, optimize=True)
    + 0.5 * np.einsum('ijklmn,jklm->in', ref_dm3, g2e.transpose(0, 2, 3, 1), optimize=True))))

# [Method 3] Get <i|H^n|j> by contraction of MPS for |i> and H^n|j>
def get_hn_ket(ket, n, xj):
    kets = []
    for xd, spin in zip("dD", "AB"):
        dmpo = driver.get_site_mpo(op=xd, site_index=xj, iprint=0)
        dket = driver.get_random_mps(tag="D%sKET-%s-H0" % (spin, xj), bond_dim=200,
            center=ket.center, target=dmpo.op.q_label + ket.info.target)
        driver.multiply(dket, dmpo, ket, n_sweeps=10, bond_dims=[200], thrds=[1E-10] * 10, iprint=0)
        for i in range(n):
            hdket = driver.copy_mps(dket, tag="D%sKET-%s-H%s" % (spin, xj, i + 1))
            driver.multiply(hdket, mpo, dket, n_sweeps=10, bond_dims=[200], thrds=[1E-10] * 10, iprint=0)
            dket = hdket
        kets.append(dket)
    return kets

impo = driver.get_identity_mpo()

h0kets = [get_hn_ket(ket, 0, xj) for xj in range(ncas)]
h0mat_mps = np.array([[np.sum([driver.expectation(h0kets[xi][s], impo, h0kets[xj][s])
    for s in [0, 1]]) for xj in range(ncas)] for xi in range(ncas)])
print(np.linalg.norm(h0mat - h0mat_mps))

h1kets = [get_hn_ket(ket, 1, xj) for xj in range(ncas)]
h1mat_mps = np.array([[np.sum([driver.expectation(h0kets[xi][s], impo, h1kets[xj][s])
    for s in [0, 1]]) for xj in range(ncas)] for xi in range(ncas)])
print(np.linalg.norm(h1mat - h1mat_mps))

h2kets = [get_hn_ket(ket, 2, xj) for xj in range(ncas)]