Change SU2 -> SZ symmetry for excited states

[piariane]

Dear Block2 developers,

First of all, thanks for an awesome program and excellent documentation!

I would like to calculate orbital entropies for a set of states but I run into problems. My approach was to first split the MPS into the different states kets and subsequently loop over the states and perform the transformation, see below.

ncas, n_elec, spin, ecore, h1e, g2e, orb_sym = itg.get_rhf_integrals(mf,
    ncore=ncore, ncas=ncas, g2e_symm=8)

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

# Define inital state
mpo = driver.get_qc_mpo(h1e=h1e, g2e=g2e, ecore=ecore, iprint=0)
ket = driver.get_random_mps(tag="KET", bond_dim=250, nroots=2)

# Define bond dimension and stuff
bond_dims = [250] * 4 + [250] * 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('State-averaged MPS energies = [%s]' % " ".join("%20.15f" % x for x in energies))
print('1st exc. state [eV]:', (energies[1]-energies[0])*27.2114)

# Split states
kets = [driver.split_mps(ket, ir, tag="KET-%d" % ir) for ir in range(ket.nroots)]

for e,iket in enumerate(kets):
    # Change from SU(2) to SZ symmetry in order to calculate orbital entropies
    zket = driver.mps_change_to_sz(iket, "ZKET-"+str(e))
    driver.symm_type = SymmetryTypes.SZ
    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)
    impo = driver.get_identity_mpo()
    print("SZ energy :", driver.expectation(zket, mpo, zket) / driver.expectation(zket, impo, zket))
    
    # Calculate orbital entanglement
    ordm1 = driver.get_orbital_entropies(zket, orb_type=1)
    print("Orbital entropies:")
    print(ordm1)

This approach works well for the first state but for the second state, I get an assertion error:

pyblock2/driver/core.py", line 4170, in mps_change_to_sz
    assert SymmetryTypes.SU2 in bw.symm_type
AssertionError

Do you know of a way how to do this?
Many thanks,
Pi

[hczhai]

Thanks for using the block2 package.

As the error message indicates, mps_change_to_sz is supposed to be used in the SU2 mode. But since you put driver.symm_type = SymmetryTypes.SZ in the first iteration in the for loop, driver will be in the SZ mode during the second iteration. To solve the problem, you may put driver.symm_type = SymmetryTypes.SZ outside the loop, and make sure all mps_change_to_sz is done before driver.symm_type = SymmetryTypes.SZ. The following is an example script:

# Split states
kets = [driver.split_mps(ket, ir, tag="KET-%d" % ir) for ir in range(ket.nroots)]

# Change from SU(2) to SZ symmetry in order to calculate orbital entropies
zkets = [driver.mps_change_to_sz(ket, "ZKET-%d" % ir) for ir, ket in enumerate(kets)]

driver.symm_type = SymmetryTypes.SZ
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)
impo = driver.get_identity_mpo()

for zket in zkets:
    print("SZ energy :", driver.expectation(zket, mpo, zket) / driver.expectation(zket, impo, zket))
    # Calculate orbital entanglement
    ordm1 = driver.get_orbital_entropies(zket, orb_type=1)
    print("Orbital entropies:")
    print(ordm1)

Of course, you can always use print(driver.symm_type) to check the current mode, and use driver.symm_type = SymmetryTypes.SU2 to change back to the SU2 mode. After every mode change, driver.initialize_system, mpo, and impo, etc. need to be repeated to make sure that these objects are defined in the new mode.

[piariane]

Thanks a lot for your thorough explanation and for the suggested solution.

It works like a charm!
Pi