Explicit form of tensors from MPS/MPO in numpy arrays

[Kowalikov]

Hi everyone, I'm using pyblock2 for the entanglement assessment between the molecule parts. I calculate the ground state MPS with DMRG and get entanglement entropy and some expectation values with implemented functions, but is there a way to get MPS or MPO tensors in explicit form e.g. numpy array? I'm starting my journey with quantum chemistry and tensor networks and I'd like to play with basics step by step to better understand what's going on.

[hczhai]

Hi @Kowalikov, thanks for your interest in using block2. Here is an example python script for translating the MPS and MPO into explicit numpy arrays:

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, MPOTools

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)
print('MPS tensor at site 1 =\n', pyket.tensors[1])

pympo = MPOTools.from_block2(mpo.prim_mpo)
print('MPO tensor at site 1 =\n', pympo.tensors[1])

print('MPS overlap using numpy arrays =', pyket @ pyket)
print('Energy expectation using numpy arrays =', pyket @ (pympo @ pyket))

A few notes:

1. To run this script you need to set an extra argument add_ident=False in the get_qc_mpo function. You may patch pyblock2/driver/core.py with this commit 70f9d7d to make it work.
2. This script only works with symm_type=SymmetryTypes.SZ. The spin-adapted quantum chemistry with symm_type=SymmetryTypes.SU2 is more complicated, for which we may add some pedagogical support in the future, but if you are new to quantum chemistry DMRG it is recommended to start from the non-spin-adapted (SymmetryTypes.SZ) version.
3. The printed MPS and MPO tensors are block-sparse tensors. So each tensor has many blocks and every block is associated with a set of quantum labels (one quantum label for one leg).
4. The last two lines were implemented in pure python, so you can have a look at pyblock2/algebra/core.py to see how the contraction is performed. Note that this python implementation is mostly for educational and debugging purposes, and in block2 things are implemented in a much more complicated way for efficiency considerations.

I am happy to help if you have any further questions.

[meandmytram]

That is super useful, thanks a lot!

[meandmytram]

@Kowalikov @hczhai In fact, I'm still having trouble trying to convert pympo tensors to dense numpy arrays. I can access subtensors in the dense form but it's not obvious for me how to convert pairs (quantum number, subtensor) to a numpy array. Could you please point me to how that could be done? Thank you!

[hczhai]

In pyblock3 which is a separate package, we provide a more complete tool set for manipulating MPS and MPO as numpy arrays. Instead of

from pyblock2.algebra.io import MPSTools, MPOTools

You can install pyblock3, and do

from pyblock3.block2.io import MPSTools, MPOTools

To get MPS and MPO in the pyblock3 package. Then you can find the to_dense() method (https://github.com/block-hczhai/pyblock3-preview/blob/master/pyblock3/algebra/core.py#L2265) on the tensor object.

Alternatively, you can follow the documentation in pyblock3 so that you can directly create your MPO from the Hamiltonian expression. The quantum chemistry DMRG implementation in pyblock3 can be fully understood in python. Then there is no need to use block2.

[meandmytram]

Thank you so much! @hczhai