DMRG using unexpectedly large memory

[cfengno1]

Dear Huanchen,

I am using DMRG to perform small bond dimension (200) calculations. But the program takes about 100GB memory just to get the random initial state. The state files are also very large for the first sites. Can you help to explain why?

from pyblock2.driver.core import DMRGDriver, SymmetryTypes, MPOAlgorithmTypes
import numpy as np
import os
Lx = 160
Ly = 2
N_db = 64
L = Lx * Ly
J = 0.4
V = 0.2
t = 1
N_ELEC = L
TWO_S = 0
Nthreads = 6
Mem = 204800
init_occ = 1

driver = DMRGDriver(scratch="./tmp", symm_type=SymmetryTypes.SAnySU2, n_threads=Nthreads, stack_mem=Mem << 20, restart_dir="./MPS", clean_scratch = False)
driver.set_symmetry_groups("U1Fermi", "U1", "SU2", "SU2") # electron number, doublon number, total spin, total spin
Q = driver.bw.SX

# [Part A] Set states and matrix representation of operators in local Hilbert space
site_basis, site_ops = [], []

for k in range(L):
    basis = [(Q(0, 0, 0, 0), 1), (Q(1, 0, 1, 1), 1), (Q(2, 1, 0, 0) ,1)] # [01]
    ops = {
        "": np.array([[1, 0, 0], [0, 1, 0], [0, 0, 1]]),        # identity
        "C": np.array([[0, 0, 0], [0, 0, 0], [0, -2**0.5, 0]]),       # create doublon
        "c": np.array([[0, 0, 0], [1, 0, 0], [0, 0, 0]]),  # annhilate holon
        "D": np.array([[0, 0, 0], [0, 0, 1], [0, 0, 0]]),       # annhilate doublon
        "d": np.array([[0, 2**0.5, 0], [0, 0, 0], [0, 0, 0]]),  # create holon
        "n": np.array([[0, 0, 0], [0, 1, 0], [0, 0, 2]]),  # electron number
    }
    site_basis.append(basis)
    site_ops.append(ops)

# [Part B] Set Hamiltonian terms
driver.initialize_system(n_sites=L, vacuum=Q(0, 0, 0, 0), target=Q(N_ELEC, N_db, TWO_S, TWO_S), hamil_init=False)
driver.ghamil = driver.get_custom_hamiltonian(site_basis, site_ops)
b = driver.expr_builder()

la = np.loadtxt("Lattice.dat",delimiter=",",dtype=int)
for i in range(0,la.shape[0]):
    s1 = la[i][0] - 1
    s2 = la[i][1] - 1
    b.add_term("(n+n)0", [s1, s2], V)
    b.add_term("(C+D)0", [s1, s2], -(2 ** 0.5) * t)
    b.add_term("(D+C)0", [s1, s2], -(2 ** 0.5) * t)
    b.add_term("(c+d)0", [s1, s2], -(2 ** 0.5) * t)
    b.add_term("(d+c)0", [s1, s2], -(2 ** 0.5) * t)
    b.add_term("((c+d)2+(c+d)2)0", [s1, s1, s2, s2], J * -(3 ** 0.5) / 2)
    b.add_term("((C+c)0+(d+D)0)0", [s1, s1, s2, s2], -J / 4)
    b.add_term("((d+D)0+(C+c)0)0", [s1, s1, s2, s2], -J / 4)
    b.add_term("(n+n)0", [s1, s2],  - J / 4)

# [Part C] Perform DMRG
mpo = driver.get_mpo(b.finalize(adjust_order=True), algo_type=MPOAlgorithmTypes.FastBipartite, iprint=1)
mps = driver.get_random_mps(tag="KET", bond_dim=200, nroots=1)

[hczhai]

1. When you have a large number of quantum numbers and a large number of sites, computing the symmetry blocks in the initial MPS is non-trivial. You have 160 sites and three quantum numbers, so each quantum number can take values between [0, 320], [0, 160], [0, 160], respectively. To prepare the initial MPS, block2 needs to consider all these possibilities for the left and right bonds for each site. For each considered quantum number in the SAny mode, 56 bytes of storage is required. In total this requires 2 * 56 * 160 * 320 * 160 * 160 B = 136.7 GB memory. You need to reduce the number of symmetry groups and/or number of sites to reduce the memory cost. MPS bond dimension is not relevant here.
2. You can use driver.get_random_mps(..., occs=[N_ELEC / L] * L) to reduce the size of the initial MPS tensors and save time in the first sweep (with the occs argument, block2 will discard some unnecessary quantum numbers). But the memory cost in "1" will not be reduced, because block2 still needs to compute the probability distribution of all quantum numbers to decide which one to discard.