MPO/MPS Export (quimb)

Open in Colab

[1]:

!pip install block2==0.5.3 -qq --progress-bar off
!pip install pyscf==2.3.0 -qq --progress-bar off
!pip install quimb==1.10.0 -qq --progress-bar off
!pip install git+https://github.com/jcmgray/symmray.git@939288ebf52ff8903fca988a0b72cb540574c6b1 -qq --progress-bar off

  Installing build dependencies ... done
  Getting requirements to build wheel ... done
  Preparing metadata (pyproject.toml) ... done

Introduction

In this tutorial we explain how to transform the MPO and MPS objects from block2 to quimb. quimb (https://github.com/jcmgray/quimb) is an easy but fast Python library for quantum information many-body calculations, focusing primarily on tensor networks.

Depending on the symmetry of the problem, we consider two cases.

In the first case the MPO and MPS tensors are stored as dense arrays (which is more common for quantum information simulations). We use the “no-symmetry” mode in block2 to perform DMRG with these dense tensors. The local Hilbert basis is the qubit basis.

In the second case the MPO and MPS tensors are stored as block-sparse arrays (which is more common for quantum chemistry applications) with fermion statistics. This corresponds to the SymmetryTypes.SZ mode in block2. In quimb, we use the symmray library to handle block-sparse data structure and fermion signs.

Qubit basis with no symmetry

The following is an examaple script that performs quantum chemistry DMRG in the qubit basis with no symmetry in block2. The MPS and MPO can then be expressed as dense tensors and contracted in quimb, producing the same energy expectation.

[2]:

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

bond_dims = [50] * 4 + [100] * 6 + [200] * 6
noises = [1e-4] * 4 + [1e-5] * 10 + [0]
thrds = [1e-10] * 16

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.SAny | SymmetryTypes.SGB, n_threads=4)

n_sites = ncas * 2

driver.set_symmetry_groups()
Q = driver.bw.SX

site_basis = [[(Q(), 2)]] * n_sites
ops = {
    "": np.array([[1, 0], [0, 1]]),
    "P": np.array([[0, 1], [0, 0]]),
    "M": np.array([[0, 0], [1, 0]]),
    "Z": np.array([[0.5, 0], [0, -0.5]]),
}
site_ops = [ops] * n_sites

driver.initialize_system(n_sites=n_sites, vacuum=Q(), target=Q(), hamil_init=False)
driver.ghamil = driver.get_custom_hamiltonian(site_basis, site_ops)

mpo = driver.get_qc_mpo(h1e=h1e, g2e=driver.unpack_g2e(g2e, n_sites=ncas), ecore=ecore, add_ident=False, simple_const=True, iprint=0)
ket = driver.get_random_mps(tag="KET", bond_dim=50, nroots=1)
energy = driver.dmrg(mpo, ket, n_sweeps=len(bond_dims), bond_dims=bond_dims, noises=noises, thrds=thrds, cutoff=-1, iprint=1)
print('DMRG energy = %20.15f' % energy)

driver.align_mps_center(ket, ref=0)
ket = driver.adjust_mps(ket, dot=1)[0]
pympo = MPOTools.from_block2(mpo.prim_mpo)
pyket = MPSTools.from_block2(ket)

import quimb.tensor as qtn

def mps_fill_fn(_, counter=[0]):
    counter[0] += 1
    return pyket.tensors[counter[0] - 1].blocks[0].reduced

def mpo_fill_fn(_, counter=[0]):
    counter[0] += 1
    return pympo.tensors[counter[0] - 1].blocks[0].reduced

qumpo = qtn.MatrixProductOperator.from_fill_fn(mpo_fill_fn, len(pyket.tensors), bond_dim=None, phys_dim=2, shape="ludr")
qumps = qtn.MatrixProductState.from_fill_fn(mps_fill_fn, len(pympo.tensors), bond_dim=None, phys_dim=2, shape="lpr")

qumps.show()
qumpo.show()

qumpsH = qumps.H
qumps.align_(qumpo, qumpsH)
print('norm =', qumps.H @ qumps)
print('expt =', (qumpsH & qumpo & qumps) ^ ...)


Sweep =    0 | Direction =  forward | Bond dimension =   50 | Noise =  1.00e-04 | Dav threshold =  1.00e-10
Time elapsed =     19.629 | E =    -107.6446337899 | DW = 7.91470e-05

Sweep =    1 | Direction = backward | Bond dimension =   50 | Noise =  1.00e-04 | Dav threshold =  1.00e-10
Time elapsed =     21.973 | E =    -107.6516392903 | DE = -7.01e-03 | DW = 3.25122e-05

Sweep =    2 | Direction =  forward | Bond dimension =   50 | Noise =  1.00e-04 | Dav threshold =  1.00e-10
Time elapsed =     23.938 | E =    -107.6522750481 | DE = -6.36e-04 | DW = 5.68993e-05

Sweep =    3 | Direction = backward | Bond dimension =   50 | Noise =  1.00e-04 | Dav threshold =  1.00e-10
Time elapsed =     25.968 | E =    -107.6525192835 | DE = -2.44e-04 | DW = 5.70002e-05

Sweep =    4 | Direction =  forward | Bond dimension =  100 | Noise =  1.00e-05 | Dav threshold =  1.00e-10
Time elapsed =     34.137 | E =    -107.6534665811 | DE = -9.47e-04 | DW = 1.88522e-17

Sweep =    5 | Direction = backward | Bond dimension =  100 | Noise =  1.00e-05 | Dav threshold =  1.00e-10
Time elapsed =     48.520 | E =    -107.6540805479 | DE = -6.14e-04 | DW = 3.85541e-07

Sweep =    6 | Direction =  forward | Bond dimension =  100 | Noise =  1.00e-05 | Dav threshold =  1.00e-10
Time elapsed =     55.746 | E =    -107.6541134665 | DE = -3.29e-05 | DW = 5.96280e-07

Sweep =    7 | Direction = backward | Bond dimension =  100 | Noise =  1.00e-05 | Dav threshold =  1.00e-10
Time elapsed =     64.993 | E =    -107.6541176322 | DE = -4.17e-06 | DW = 6.03228e-07

Sweep =    8 | Direction =  forward | Bond dimension =  100 | Noise =  1.00e-05 | Dav threshold =  1.00e-10
Time elapsed =     72.634 | E =    -107.6541196617 | DE = -2.03e-06 | DW = 5.98280e-07

Sweep =    9 | Direction = backward | Bond dimension =  100 | Noise =  1.00e-05 | Dav threshold =  1.00e-10
Time elapsed =     79.336 | E =    -107.6541195922 | DE = 6.95e-08 | DW = 5.61042e-07

Sweep =   10 | Direction =  forward | Bond dimension =  200 | Noise =  1.00e-05 | Dav threshold =  1.00e-10
Time elapsed =    105.850 | E =    -107.6541215992 | DE = -2.01e-06 | DW = 1.28080e-18

Sweep =   11 | Direction = backward | Bond dimension =  200 | Noise =  1.00e-05 | Dav threshold =  1.00e-10
Time elapsed =    152.107 | E =    -107.6541224475 | DE = -8.48e-07 | DW = 4.69341e-14

Sweep =   12 | Direction =  forward | Bond dimension =  200 | Noise =  1.00e-05 | Dav threshold =  1.00e-10
Time elapsed =    175.686 | E =    -107.6541224475 | DE = -2.31e-11 | DW = 6.24418e-16

Sweep =   13 | Direction = backward | Bond dimension =  200 | Noise =  1.00e-05 | Dav threshold =  1.00e-10
Time elapsed =    193.797 | E =    -107.6541224475 | DE = -5.68e-14 | DW = 7.25652e-17

Sweep =   14 | Direction =  forward | Bond dimension =  200 | Noise =  0.00e+00 | Dav threshold =  1.00e-10
Time elapsed =    213.780 | E =    -107.6541224475 | DE = 0.00e+00 | DW = 9.54936e-19

DMRG energy = -107.654122447516784

/usr/local/lib/python3.11/dist-packages/cotengra/hyperoptimizers/hyper.py:57: UserWarning: Couldn't find `optuna`, `cmaes`, `baytune (btb)`, `chocolate`, or `nevergrad` so will use completely random sampling in place of hyper-optimization.
  warnings.warn(
/usr/local/lib/python3.11/dist-packages/cotengra/hyperoptimizers/hyper.py:39: UserWarning: Couldn't import `kahypar` - skipping from default hyper optimizer and using basic `labels` method instead.
  warnings.warn(
/usr/local/lib/python3.11/dist-packages/cotengra/hyperoptimizers/hyper.py:76: UserWarning: Couldn't find `optuna`, `cmaes`, `baytune (btb)`, `chocolate`, or `nevergrad` so will use completely random sampling in place of hyper-optimization.
  warnings.warn(

 2 4 8 16 32 64 128 200 200 200 200 200 128 64 32 16 8 4 2
●─<─<─<──<──<──<━━━<━━━<━━━<━━━<━━━<━━━<━━━<──<──<──<─<─<─<
│ │ │ │  │  │  │   │   │   │   │   │   │   │  │  │  │ │ │ │
│7│20│45│62│81│104│125│126│151│148│177│178│147│120│97│70│43│20│7│
●─●──●──●──●──●━━━●━━━●━━━●━━━●━━━●━━━●━━━●━━━●━━━●──●──●──●──●─●
│ │  │  │  │  │   │   │   │   │   │   │   │   │   │  │  │  │  │ │
norm = 0.9999999999999901
expt = -107.65412244751523

Fermionic MPO and MPS

The following is an examaple script that performs quantum chemistry DMRG in the standard \(U(1)\otimes U(1)\otimes \text{Point-Group}\) symmetry (SymmetryTypes.SZ) in block2. The MPS and MPO are exported as block-sparse fermionic tensors and contracted in quimb using symmray as the backend.

[3]:

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

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

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, simple_const=True, iprint=0)
ket = driver.get_random_mps(tag="KET", bond_dim=100, nroots=1)
energy = driver.dmrg(mpo, ket, n_sweeps=len(bond_dims), bond_dims=bond_dims, noises=noises, thrds=thrds, cutoff=-1, iprint=1)
print('DMRG energy = %20.15f' % energy)

driver.align_mps_center(ket, ref=0)
ket = driver.adjust_mps(ket, dot=1)[0]
pympo = MPOTools.from_block2(mpo.prim_mpo)
pyket = MPSTools.from_block2(ket)

import symmray as sr
import quimb.tensor as qtn

bidx = lambda z, d: sr.BlockIndex(chargemap={(x.n, x.twos, x.pg): v for x, v in z.items()}, dual=d)
basis, basis_d = [[[bidx(bz, d)] for bz in driver.basis] for d in [False, True]]
l_bonds = [[[]] + [[bidx(ts.get_state_info(0), True)] for ts in tn.tensors[1:]] for tn in [pyket, pympo]]
r_bonds = [[[bidx(ts.get_state_info(-1), False)] for ts in tn.tensors[:-1]] + [[]] for tn in [pyket, pympo]]

class SZ(sr.Symmetry):
    __slots__ = ()
    def valid(self, *charges):
        return all(isinstance(charge, tuple) and all(isinstance(ch, int) for ch in charge) for charge in charges)
    def combine(self, *charges):
        return (sum(charge[0] for charge in charges), sum(charge[1] for charge in charges),
            int(np.bitwise_xor.reduce([charge[2] for charge in charges])))
    def sign(self, charge, dual=True):
        return (charge[0] * [1, -1][dual], charge[1] * [1, -1][dual], charge[2])
    def parity(self, charge):
        return charge[0] % 2

def mps_fill_fn(_, counter=[0]):
    counter[0] += 1
    ix, ts = counter[0] - 1, pyket.tensors[counter[0] - 1]
    indices = tuple(l_bonds[0][ix] + basis_d[ix] + r_bonds[0][ix])
    f = lambda qs: -1 if ix != 0 and abs(qs[0].n) % 2 == 1 and abs(qs[1].n) % 2 == 1 else 1
    blocks = {tuple((q.n, q.twos, q.pg) for q in blk.q_labels): blk.reduced * f(blk.q_labels) for blk in ts.blocks}
    return sr.FermionicArray(indices=indices, charge=(0, 0, 0), symmetry=SZ(), blocks=blocks)

def mpo_fill_fn(_, counter=[0]):
    counter[0] += 1
    ix, ts = counter[0] - 1, pympo.tensors[counter[0] - 1]
    indices = tuple(l_bonds[1][ix] + basis[ix] + basis_d[ix] + r_bonds[1][ix])
    blocks = {tuple((q.n, q.twos, q.pg) for q in blk.q_labels): blk.reduced for blk in ts.blocks}
    return sr.FermionicArray(indices, charge=(0, 0, 0, 0), symmetry=SZ(), blocks=blocks)

qumpo = qtn.MatrixProductOperator.from_fill_fn(mpo_fill_fn, len(pyket.tensors), bond_dim=None, phys_dim=4, shape="ludr")
qumps = qtn.MatrixProductState.from_fill_fn(mps_fill_fn, len(pyket.tensors), bond_dim=None, phys_dim=4, shape="lpr")

qumpsH = qumps.H
qumps.align_(qumpo, qumpsH)
print('norm =', qumps.H @ qumps)
print('expt =', (qumpsH & qumpo & qumps) ^ ...)


Sweep =    0 | Direction =  forward | Bond dimension =  100 | Noise =  1.00e-04 | Dav threshold =  1.00e-10
Time elapsed =      0.676 | E =    -107.6540799442 | DW = 7.37486e-07

Sweep =    1 | Direction = backward | Bond dimension =  100 | Noise =  1.00e-04 | Dav threshold =  1.00e-10
Time elapsed =      0.960 | E =    -107.6541182041 | DE = -3.83e-05 | DW = 7.19646e-07

Sweep =    2 | Direction =  forward | Bond dimension =  100 | Noise =  1.00e-04 | Dav threshold =  1.00e-10
Time elapsed =      1.244 | E =    -107.6541211363 | DE = -2.93e-06 | DW = 6.92037e-07

Sweep =    3 | Direction = backward | Bond dimension =  100 | Noise =  1.00e-04 | Dav threshold =  1.00e-10
Time elapsed =      1.515 | E =    -107.6541211363 | DE = -1.44e-12 | DW = 6.92000e-07

Sweep =    4 | Direction =  forward | Bond dimension =  150 | Noise =  1.00e-05 | Dav threshold =  1.00e-10
Time elapsed =      1.825 | E =    -107.6541216080 | DE = -4.72e-07 | DW = 2.29437e-14

Sweep =    5 | Direction = backward | Bond dimension =  150 | Noise =  1.00e-05 | Dav threshold =  1.00e-10
Time elapsed =      2.162 | E =    -107.6541224011 | DE = -7.93e-07 | DW = 1.20768e-14

Sweep =    6 | Direction =  forward | Bond dimension =  150 | Noise =  1.00e-05 | Dav threshold =  1.00e-10
Time elapsed =      2.468 | E =    -107.6541224011 | DE = 1.42e-14 | DW = 1.79174e-19

Sweep =    7 | Direction = backward | Bond dimension =  150 | Noise =  1.00e-05 | Dav threshold =  1.00e-10
Time elapsed =      2.773 | E =    -107.6541224011 | DE = -8.53e-14 | DW = 7.72014e-19

Sweep =    8 | Direction =  forward | Bond dimension =  150 | Noise =  1.00e-05 | Dav threshold =  1.00e-10
Time elapsed =      3.090 | E =    -107.6541224011 | DE = 4.26e-14 | DW = 1.01747e-19

Sweep =    9 | Direction = backward | Bond dimension =  150 | Noise =  1.00e-05 | Dav threshold =  1.00e-10
Time elapsed =      3.568 | E =    -107.6541224011 | DE = 0.00e+00 | DW = 2.49515e-18

Sweep =   10 | Direction =  forward | Bond dimension =  200 | Noise =  1.00e-05 | Dav threshold =  1.00e-10
Time elapsed =      4.159 | E =    -107.6541224011 | DE = -1.42e-14 | DW = 2.29806e-23

Sweep =   11 | Direction = backward | Bond dimension =  200 | Noise =  1.00e-05 | Dav threshold =  1.00e-10
Time elapsed =      4.759 | E =    -107.6541224012 | DE = -2.23e-11 | DW = 3.27264e-37

Sweep =   12 | Direction =  forward | Bond dimension =  200 | Noise =  1.00e-05 | Dav threshold =  1.00e-10
Time elapsed =      5.331 | E =    -107.6541224012 | DE = 0.00e+00 | DW = 3.23965e-23

Sweep =   13 | Direction = backward | Bond dimension =  200 | Noise =  1.00e-05 | Dav threshold =  1.00e-10
Time elapsed =      5.931 | E =    -107.6541224012 | DE = 7.11e-14 | DW = 2.35826e-35

Sweep =   14 | Direction =  forward | Bond dimension =  200 | Noise =  0.00e+00 | Dav threshold =  1.00e-10
Time elapsed =      6.534 | E =    -107.6541224012 | DE = -4.26e-14 | DW = 1.41390e-23

DMRG energy = -107.654122401168493
norm = 0.9999999999999998
expt = -107.65412240116848