Is there a way to get partial trace of density matrix?

[fqian5]

I am interested in getting the partial trace of the MPS density matrix for certain sites. Seems like getting NPDM is a way to do it but it's too expensive to calculate all possible combination. Is there a way that I can input a list of sites I want to trace out and get the reduced density matrix? Thanks!

[hczhai]

Thanks for your interest in the block2 package.

The description of the problem that you provided is too general, and it would be more helpful if you can provide specifically what kind of NPDM you are computing and how you will trace it. For example, "partial trace of an NPDM tensor" normally means something like $$T_{ik} = \sum_{j} \langle a^\dagger_i a^\dagger_j a_j a_k \rangle $$

Is this what you need (note that for this case there is no need to "input a list of sites")? If not, please provide an alternative specific operator expression. If you also need the restriction on indices, you may need to provide information including what indices (the traced or untraced ones) should be restricted, and how many sites you have in total and in the restricted set. Without these we cannot estimate the "expensiveness" of any computational schemes.

[fqian5]

Thank you for your reply! Sorry, I think I didn't use the term NPDM correctly. What I meant is that after I finished the DMRG and get the ground state MPS. I want to do following:
The list contain the site number. I wonder if there is an easy way to do it in block2. Thanks!

[hczhai]

Thanks for providing the detailed information. What your figure suggests is the contraction of MPS tensors, which can be done using the following script:

import numpy as np
from pyblock2._pyscf.ao2mo import integrals as itg
from pyblock2.driver.core import DMRGDriver, SymmetryTypes
from pyblock2.algebra.io import MPSTools
from pyblock2.algebra.core import Tensor, MPS

from pyscf import gto, scf

mol = gto.M(atom="N 0 0 0; N 0 0 1.1", basis="sto3g", symmetry="d2h", verbose=0)
mf = scf.RHF(mol).run(conv_tol=1E-14)
ncas, n_elec, spin, ecore, h1e, g2e, orb_sym = itg.get_rhf_integrals(mf,
    ncore=0, ncas=None, g2e_symm=8)

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)

mpo = driver.get_qc_mpo(h1e=h1e, g2e=g2e, ecore=ecore, add_ident=False, iprint=1)

ket = driver.get_random_mps(tag="GS", bond_dim=250, nroots=1)
bond_dims = [250] * 4 + [500] * 4
noises = [1e-4] * 4 + [1e-5] * 4 + [0]
thrds = [1e-10] * 8
energies = driver.dmrg(mpo, ket, n_sweeps=5, bond_dims=bond_dims,
    noises=noises, thrds=thrds, iprint=1)
print('DMRG energy = %20.15f' % energies)

ket = driver.adjust_mps(ket, dot=1)[0]
print(ket.n_sites, ket.canonical_form)

pyket = MPSTools.from_block2(ket)

def ket_to_rho(ket, traced_idx):
    tensors = [None] * ket.n_sites
    for i in range(ket.n_sites):
        r = ket.tensors[i].rank
        if i in traced_idx:
            phys_idx = [1] if i != 0 else [0]
            tensors[i] = Tensor.contract(ket.tensors[i], ket.tensors[i], phys_idx, phys_idx, None, None,
                tuple(j + k for j in range(r - 1) for k in [0, r - 1]))
        else:
            tensors[i] = Tensor.contract(ket.tensors[i], ket.tensors[i], [], [], None, None,
                tuple(j + k for j in range(r) for k in [0, r]))
    ket = MPS(tensors=tensors)
    ket.merge_virtual_dims()
    tensors = []
    for i in range(ket.n_sites):
        if i in traced_idx and i != 0:
            tensors[-1] = Tensor.contract(tensors[-1], ket.tensors[i], [-1], [0])
        else:
            tensors.append(ket.tensors[i])
    if 0 in traced_idx and len(tensors) > 1:
        tensors[:2] = [Tensor.contract(tensors[0], tensors[1], [-1], [0])]
    ket.tensors = tensors
    return ket

rho = ket_to_rho(pyket, traced_idx=range(0, pyket.n_sites, 2))
print('Sites 0, 2, ... traced = ', [x.rank for x in rho.tensors])
print('All traced (norm)      = ', ket_to_rho(pyket, traced_idx=range(0, pyket.n_sites))[0])

Note that the resulting object (if not fully traced) is a collection of tensors (not a standard MPS).

[fqian5]

Thanks a lot!

[fqian5]

Actually, I have one more question. I was trying to get the matrix representation of the reduced density matrix by using <bra|rho|ket>but it doesn't work. When I try to get the matrix representation of the sub tensor using MPOTools, it returns the error 'MPO' object has no attribute 'schemer'. Could you give me some insights about how to do it? Thanks!

[hczhai]

I was trying to get the matrix representation of the reduced density matrix by using <bra|rho|ket>but it doesn't work.

How this is related to your previous question about "partial trace"? You do not have a trace notation here in "the matrix representation of the reduced density matrix by using <bra|rho|ket>". When you say "matrix representation", what is your definition of rows and columns of the "matrix"? If you can provide a clear definition of the quantity that you want to compute, it will be helpful for people to understand your request.

When I try to get the matrix representation of the sub tensor using MPOTools, it returns the error 'MPO' object has no attribute 'schemer'.

I did not use MPOTools in the script I provided above. If you modified the script or wrote a new script which generates the error, please at least copy your script here which can reproduce the error, and explain the motivation of the modification.

[fqian5]

Thank you for reply! The purpose of my code is to find the kernel space of the reduced density matrix. But I don't need all n particle rdm, I just need those consecutive sites such as 123, 234, 345. That's why I asked how to find the partial trace. After I got the reduced density matrix, I want to find the matrix representation of the reduced density matrix in CSF |I>,|j>such as <i|rho|j>. I noticed that there is expectation function can calculate <i|rho|j>. But I don't know how to find the MPS form of CSF. Hopefully, I explain my question more clear this time. Sorry for any confusion before.

[hczhai]

Thanks for your explanation. Now I see you eventually need the partially traced density operator in the basis of determinants/CSFs. For this purpose the most efficient way is to perform the partial trace in the determinant basis. We do support the construction of MPS from CSF (see https://block2.readthedocs.io/en/latest/tutorial/qc-hamiltonians.html#Construct-MPS-from-CSFs-or-Determinants), but if you do this it will be too slow since there can be a large number of CSFs. So this is not recommended.

The following script demonstrates the efficient way to generate the quantity <i|rho|j> you described:

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

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

mol = gto.M(atom="N 0 0 0; N 0 0 1.1", basis="sto3g", symmetry="d2h", verbose=0)
mf = scf.UHF(mol).run(conv_tol=1E-14)
ncas, n_elec, spin, ecore, h1e, g2e, orb_sym = itg.get_uhf_integrals(mf, ncore=0, ncas=None, g2e_symm=8)

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)

mpo = driver.get_qc_mpo(h1e=h1e, g2e=g2e, ecore=ecore, iprint=0)

ket = driver.get_random_mps(tag="GS", bond_dim=250, nroots=1)
energy = driver.dmrg(mpo, ket, n_sweeps=20, bond_dims=bond_dims, noises=noises,
    thrds=thrds, iprint=1)
print('DMRG energy = %20.15f' % energy)

dets, coeffs = driver.get_csf_coefficients(ket, cutoff=0.005, iprint=1)

def partial_trace(n_sites, dets, coeffs, traced_idxs):
    untraced_idxs = [x for x in range(n_sites) if x not in traced_idxs]
    traced_dets, rows = np.unique(dets[:, np.array(traced_idxs, int)], return_inverse=True, axis=0)
    untraced_dets, cols  = np.unique(dets[:, np.array(untraced_idxs, int)], return_inverse=True, axis=0)
    wfn = np.zeros((len(traced_dets), len(untraced_dets)), dtype=coeffs.dtype)
    wfn[rows, cols] = coeffs
    return wfn.T.conj() @ wfn, ["".join(["0ab2"[x] for x in det]) for det in untraced_dets]

rho, idxs = partial_trace(ncas, dets, coeffs, traced_idxs=[2, 3, 4, 5])

print("\n%16s%16s%20s" % ("row", "col", "val"))
for i in range(len(rho)):
    for j in range(len(rho)):
        print("%4d%12s%4d%12s%20.10f" % (i, idxs[i], j, idxs[j], rho[i, j]))

[fqian5]

Cool! Thank you so much for your help!