Applying a general (unitary) orbital rotation to MPS

[yannra]

Hey there,
First of all: Thanks a lot for publishing this super performant code, we have already been able to make good use of it.

We do currently have a pretty specific application in mind and I was wondering if you could point me in the direction for the specific details how best to approach this problem with block2.
Overall, we want to evaluate expressions of the form $\langle \Psi_1 | \hat{U} | \Psi_2 \rangle$ where $\Psi_1$ and $\Psi_2$ are MPS which we obtained from previous DMRG runs, and $\hat{U}$ is a unitary corresponding to a specific orbital rotation. I have been following the orbital rotation tutorial (https://block2.readthedocs.io/en/latest/developer/orbital-rotation.html) which is really great and contains almost everything I need. However, the tutorial explicitly focuses on rotations with a +1 determinant, which is not always the case for us (i.e., we also want to consider non-orientation preserving orbital rotations with determinant -1).

The bottom line of my question is: What is the easiest/best approach to apply such orbital rotations, for which the logarithm of the single-body rotation contains complex entries, to an MPS?

What I can obviously do is to split up the unitary so that we need to find an MPS representation for $\hat{U}_1 \hat{U}_2 | \Psi_2 \rangle$ (where $\hat{U}_1 \hat{U}_2 = \hat{U}$) and I can choose $\hat{U}_2$ such that this single-body rotation has a +1 determinant, and $\hat{U}_1$ is trivial to apply. If the determinant of $U$ is equal to $-1$, one might, e.g., choose $\hat{U}_2$ such that it corresponds to $\hat{U}$ but with a multiplication of the first orbital with a (-1) sign which then has $det(U_2) = +1$. The action of $\hat{U}_1$ would then merely be another sign flip of the first orbital. In principle, it is then easy to apply $\hat{U}_2$ with the time evolution approach outlined in the tutorial and afterwards the action of the trivial transformation $\hat{U}_1$ can be evaluated (in this case simply corresponding to flipping the sign of the matrices associated with single occupancies of the first orbital). However, I don't really know how to do that in block2. Is there any way to directly manipulate the MPS tensors to get the transformed MPS incorporating a simple orbital sign flip (in a way which allows us to reuse the MPS for subsequent overlap evaluations in block2)?

If there is an easier approach than the one outlined above for our use case, I would obviously also be keen to learn about that. One other potential approach I see is that I could simply use complex parameters for the application of the orbital rotation, allowing us to rotate the MPS into the basis we want even if the full orbital rotation "Hamiltonian" contains complex parameters. However, this seems unnecessarily complicated as, I believe, this would require switching to complex-valued MPS which shouldn't be necessary.

Either way, probably you can easily point me towards the most straightforward implementation for our problem. I'd really appreciate your input!

Many thanks for your support and best wishes,
Yannic

[hczhai]

Hi Yannic,

Thanks for your interest in using block2 for your research. Here is an example python script for doing orbital rotation with -1 determinant rotation matrices:

from pyscf import gto, scf, lo
import scipy.linalg
import numpy as np

mol = gto.M(atom='N 0 0 0; N 0 0 1.1', basis='sto3g', symmetry='c1')
mf = scf.RHF(mol).run()

# localize orbitals
def scdm(coeff, overlap):
    aux = lo.orth.lowdin(overlap)
    no = coeff.shape[1]
    ova = coeff.T @ overlap @ aux
    piv = scipy.linalg.qr(ova, pivoting=True)[2]
    bc = ova[:, piv[:no]]
    ova = np.dot(bc.T, bc)
    s12inv = lo.orth.lowdin(ova)
    return coeff @ bc @ s12inv

# HF orbitals (old basis)
mo_coeff = mf.mo_coeff

# Localized orbitals (new basis)
lmo_coeff = scdm(mo_coeff, mol.intor('cint1e_ovlp_sph'))

# orbital transform rot[old, new]
orb_rot = np.linalg.pinv(mo_coeff) @ lmo_coeff
assert np.linalg.norm(orb_rot.T - np.linalg.inv(orb_rot)) < 1E-12
assert np.linalg.norm(lmo_coeff - mo_coeff @ orb_rot) < 1E-12

# whether det(orb_rot) is +1 or -1
neg_det = True

if neg_det:
    if np.linalg.det(orb_rot) > 0:
        orb_rot[:, 0] = -orb_rot[:, 0]
        lmo_coeff = mo_coeff @ orb_rot

    print(np.linalg.det(orb_rot))

    # apply orbital sign change at orbital 1
    orb_rot[:, 1] = -orb_rot[:, 1]
else:
    if np.linalg.det(orb_rot) < 0:
        orb_rot[:, 0] = -orb_rot[:, 0]
        lmo_coeff = mo_coeff @ orb_rot

    print(np.linalg.det(orb_rot))

kappa = scipy.linalg.logm(orb_rot)

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

if kappa.dtype == np.float64:
    print('use real')
    symm_type = SymmetryTypes.SU2
else:
    print('use complex')
    symm_type = SymmetryTypes.SU2 | SymmetryTypes.CPX

driver = DMRGDriver(scratch="./tmp", symm_type=symm_type, stack_mem=4 << 30, n_threads=8)

# DMRG to find ground state MPS in the new basis
mf.mo_coeff = lmo_coeff
ncas, n_elec, spin, ecore, h1e, g2e, orb_sym = itg.get_rhf_integrals(mf)
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_lo = driver.get_random_mps(tag="KET-LO", bond_dim=500, nroots=1)
energy = driver.dmrg(mpo, ket_lo, n_sweeps=10, bond_dims=[500] * 10, noises=[1e-5] * 4 + [0],
    thrds=[1e-9] * 10, tol=1E-12, cutoff=1E-24, iprint=0)
print('DMRG energy (LO) = %20.15f' % energy)

# DMRG to find ground state MPS in the old basis
mf.mo_coeff = mo_coeff
ncas, n_elec, spin, ecore, h1e, g2e, orb_sym = itg.get_rhf_integrals(mf)
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_hf = driver.get_random_mps(tag="KET-HF", bond_dim=500, nroots=1)
energy = driver.dmrg(mpo, ket_hf, n_sweeps=10, bond_dims=[500] * 10, noises=[1e-5] * 4 + [0],
    thrds=[1e-9] * 10, tol=1E-12, cutoff=1E-24, iprint=0)
print('DMRG energy (HF) = %20.15f' % energy)

# perform orbital rotation (with positive det) from old to new basis
mpo_kappa = driver.get_qc_mpo(h1e=kappa, g2e=None, ecore=0.0, iprint=0)
ket_rot = ket_hf.deep_copy("KET-ROT")
driver.td_dmrg(mpo_kappa, ket_rot, delta_t=0.05, target_t=1.0, final_mps_tag="KET-ROT", iprint=0)

if neg_det:
    # apply orbital sign change at orbital 1
    orb_idx = 1
    if SymmetryTypes.SU2 in driver.symm_type:
        b = driver.expr_builder()
        b.add_term("((C+(C+D)0)1+D)0", [orb_idx] * 4, 4)
        b.add_term("(C+D)0", [orb_idx] * 2, -2 * 2 ** 0.5)
        b.add_term("", [], 1)
        mpo_flip_sign = driver.get_mpo(b.finalize())
    else:
        b = driver.expr_builder()
        b.add_term("cdCD", [orb_idx] * 4, 4)
        b.add_term("cd", [orb_idx] * 2, -2)
        b.add_term("CD", [orb_idx] * 2, -2)
        b.add_term("", [], 1)
        mpo_flip_sign = driver.get_mpo(b.finalize())
    ket_s = ket_rot.deep_copy("KET-S")
    driver.multiply(ket_s, mpo_flip_sign, ket_rot, iprint=0)
    ket_rot = ket_s.deep_copy("KET-ROT")

# test if ket_rot is close to the ground state in the new basis
# the expt can be +1 or -1 because we did two separate DMRG in different basis and there can potentially be a phase difference
impo = driver.get_identity_mpo()
expt = driver.expectation(ket_rot, impo, ket_lo)
print(expt, abs(expt))

A few notes:

1. You can apply a +1 det rotation and then create an MPO called mpo_flip_sign which will flip the sign of a specific orbital in the MPS, when you use driver.multipy to apply this MPO to the MPS. The grammar for spin-adapted MPO construction is explained in https://block2.readthedocs.io/en/latest/tutorial/hubbard.html#id2 and if you feel the SU(2) notation difficult to understand, you can change symm_type = SymmetryTypes.SU2 to symm_type = SymmetryTypes.SZ to use the non-spin-adapted version and the grammar for non-spin-adapted MPO construction is explained in https://block2.readthedocs.io/en/latest/tutorial/hubbard.html#Build-Hamiltonian.
2. When the rotation matrix has a +1 det, it is not always true that kappa will be real. You can run the script multiple times and there is a small chance that the kappa is complex, then the code will switch to the complex mode, and everything can still work. Therefore, if you use the complex mode with symm_type = SymmetryTypes.SU2 | SymmetryTypes.CPX, you can directly use the -1 det rotation matrix and there is no need to apply mpo_flip_sign.
3. If you do not want to use the complex mode, then you have to make sure that from a +1 det matrix you always get a real kappa. For this you can change kappa = scipy.linalg.logm(orb_rot) to kappa = my_logm(orb_rot) where my_logm can be found in https://github.com/block-hczhai/block2-preview/blob/p0.5.2rc8/pyblock2/driver/block2main#L3022-L3040. my_logm will always produce a real kappa as long as the input has +1 det.

Please let me know if you have any further questions.

[yannra]

Thank you so much for this very swift and detailed response - this is exactly what I was looking for.

Thanks and best wishes!