g_vumps.py

on github (download).

"""Toy code implementing the variational uniform matrix product states (VUMPS) algorithm.

This implementation closely follows Laurens Vanderstraeten, Jutho Haegeman and Frank Verstraete, 
Tangent-space methods for uniform matrix product states, SciPost Physics Lecture Notes 007, 2019, 
https://arxiv.org/abs/1810.07006.
"""

import numpy as np
from scipy.sparse.linalg import LinearOperator, gmres, eigsh
from scipy.linalg import polar

from .f_umps import UniformMPS, TransferMatrix


def vumps_algorithm(h, guess_psi0, tol, maxruns=1_000):
    """Find the uMPS ground state of the Hamiltonian h with an initial guess up to tolerance tol."""
    vumps_engine = VUMPSEngine(guess_psi0, h, tol/10)
    for i in range(maxruns):
        vumps_engine.run()
        if vumps_engine.err <= tol:
            psi0 = vumps_engine.psi
            e0 = psi0.get_bond_expectation_value(h)
            var0 = vumps_engine.get_energy_variance()
            print(f"uMPS ground state converged with VUMPS up to tol={tol} in gradient norm. " \
                  + f"Final error after {i+1} iterations: {vumps_engine.err}.\n" \
                  + f"Ground state energy density: {e0}. \n"
                  + f"Ground state variance density: {var0}.")
            return e0, psi0, var0
    raise RuntimeError(f"Ground state did not converge up to tol={tol}. \
                        Final error after {maxruns} runs: {vumps_engine.err}.")


class VUMPSEngine:
    """Simple class for the VUMPS engine to perform the variational ground state optimization for h.

    For P_A the tangent-space projector and H = sum_n h_{n,n+1}, the tangent-space gradient reads
    P_A * H|psi(A)> = |psi(G,A)> with G = Heff1(AC) - AL Heff0(C) = Heff1(AC) - Heff0(C) AR.
    The variational ground state optimum (corresponding to G=0) in the uMPS manifold satisfies
        1) AC ground state of Heff1,
        2) C ground state of Heff0,
        3) AC = AL C = C AR.
    
    Parameters
    ----------
    psi, h, tol: Same as attributes.

    Attributes
    ----------
    psi: UniformMPS
         The current state to be iteratively optimized towards the ground state.
    h: np.array[ndim=4]
       The two-site Hamiltonian of which the ground state is searched.
    tol: float
         Tolerance up to which the geometric sum environments Lh and Rh are computed with gmres.
    Lh: np.array[ndim=2]
        Left environment computed from geometric sum of transfer matrix TL.
    Rh: np.array[ndim=2]
        Right environment computed from geometric sum of transfer matrix TR.
    err: float
         The error measure for psi, equal to the gradient norm ||Heff_AC(AC) - AL Heff_C(C)||.
    """
    def __init__(self, psi, h, tol):
        self.psi = psi.copy()
        self.h = subtract_energy_offset(self.psi, h, canonical_form=False)
        self.tol = tol
        self.Lh = get_Lh(self.psi, self.h, canonical_form=False, guess=None, tol=self.tol)
        self.Rh = get_Rh(self.psi, self.h, canonical_form=False, guess=None, tol=self.tol)
        self.err = self.get_gradient_norm()
        
    def run(self):
        """Perform one update of self.psi in the variational ground state optimization for self.h.
        
        1) AC -> ground state of Heff1,
        2) C -> ground state of Heff0,
        3) AL/AR from left/right polar decompositions of AC and C.
        """
        H_eff_1 = Heff1(self.h, self.Lh, self.psi.AL, self.psi.AR, self.Rh)
        H_eff_0 = Heff0(self.h, self.Lh, self.psi.AL, self.psi.AR, self.Rh)
        AC_new = self.get_theta_gs(Heff=H_eff_1, guess=self.psi.AC)  
        C_new = self.get_theta_gs(Heff=H_eff_0, guess=self.psi.C)  
        AL_new, AR_new = get_AL_AR(AC_new, C_new)  
        self.psi = UniformMPS(AL_new, AR_new, AC_new, C_new)
        self.h = subtract_energy_offset(self.psi, self.h, canonical_form=False)
        self.Lh = get_Lh(self.psi, self.h, canonical_form=False, guess=self.Lh, tol=self.tol)
        self.Rh = get_Rh(self.psi, self.h, canonical_form=False, guess=self.Rh, tol=self.tol)
        self.err = self.get_gradient_norm()

    @staticmethod
    def get_theta_gs(Heff, guess):
        """Find the ground state of Heff with an initial guess."""
        guess = np.reshape(guess, Heff.shape[1])
        _, theta_gs = eigsh(Heff, k=1, which="SA", v0=guess)
        theta_gs = np.reshape(theta_gs[:, 0], Heff.shape_theta)
        return theta_gs

    def get_gradient_norm(self):
        """Compute the gradient norm ||Heff1(AC) - AL Heff0(C)|| for self.psi."""
        H_eff_1 = Heff1(self.h, self.Lh, self.psi.AL, self.psi.AR, self.Rh)
        H_eff_0 = Heff0(self.h, self.Lh, self.psi.AL, self.psi.AR, self.Rh)
        AC = H_eff_1._matvec(np.reshape(self.psi.AC, H_eff_1.shape[1]))  # vL.p.vR
        AC = np.reshape(AC, H_eff_1.shape_theta)  # vL p vR
        C = H_eff_0._matvec(np.reshape(self.psi.C, H_eff_0.shape[1]))  # vL.vR
        C = np.reshape(C, H_eff_0.shape_theta)  # vL vR
        ALC = np.tensordot(self.psi.AL, C, axes=(2, 0))  # vL p [vR], [vL] vR
        gradient_norm = np.linalg.norm(AC - ALC)
        return gradient_norm
    
    def get_energy_variance(self):
        """For the Hamiltonian H = sum_n h_{n,n+1}, compute the energy variance density of self.psi.

        var(H) = <psi|H^2|psi> - <psi|H|psi>^2 = <psi|(H-<psi|H|psi>)^2|psi>.

        Diagrammatic representation for the variance density (after h -> h - <psi|h|psi>):

        .----(AC)--(AR)--.     .--(AL)--(AC)--(AR)--.     .--(AC)--(AR)--.     .--(AC)--(AR)--(AR)--.     
        |     |     |    |     |   |     |     |    |     |   |     |    |     |   |     |     |    |     
        |   (----h----)  |     |   |   (----h----)  |     | (----h----)  |     | (----h----)   |    |     
        |     |     |    |  +  |   |     |     |    |  +  |   |     |    |  +  |   |     |     |    |   
       (Lh)   |     |    |     | (----h----)   |    |     | (----h----)  |     |   |   (----h----)  |     
        |     |     |    |     |   |     |     |    |     |   |     |    |     |   |     |     |    |      
        .---(AC*)--(AR*)-.     .-(AL*)-(AC*)--(AR*)-.     .-(AC*)--(AR*)-.     .-(AC*)--(AR*)-(AR*)-.     

            .--(AC)--(AR)----.
            |   |     |      |
            | (----h----)    |
        +   |   |     |      |
            |   |     |    (Rh)
            |   |     |      |
            .-(AC*)--(AR*)---.
        """
        h = self.h
        Lh = self.Lh
        AL = self.psi.AL
        AC = self.psi.AC
        AR = self.psi.AR
        Rh = self.Rh
        ACAR = np.tensordot(AC, AR, axes=(2, 0))  # vL p1 [vR], [vL] p2 vR  -> vL p1 p2 vR
        ACAR_h = np.tensordot(ACAR, h, axes=((1, 2), (2, 3)))  # vL [p1] [p2] vR, p1 p2 [p1*] [p2*]
        ACAR_h = np.transpose(ACAR_h, (0, 2, 3, 1))  # vL p1 p2 vR
        var1 = np.tensordot(ACAR_h, np.conj(ACAR), axes=((1, 2, 3), (1, 2, 3)))  
        # vL [p1] [p2] [vR], vL* [p1*] [p2*] [vR*]
        var1 = np.tensordot(Lh, var1, axes=((0, 1), (0, 1)))  # [vR] [vR*], [vL] [vL*]
        var2 = np.tensordot(AL, ACAR_h, axes=(2, 0))  # vL p1 [vR], [vL] p2 p3 vR
        var2 = np.tensordot(h, var2, axes=((2, 3), (1, 2)))  # p1 p2 [p1*] [p2*], vL [p1] [p2] p3 vR
        var2 = np.tensordot(var2, np.conj(AL), axes=((0, 2), (1, 0)))  
        # [p1]  p2 [vL] p3 vR, [vL*] [p1*] vR*
        var2 = np.tensordot(var2, np.conj(ACAR), axes=((3, 0, 1, 2), (0, 1, 2, 3)))  
        # [p2] [p3] [vR] [vR*], [vL*] [p2*] [p3*] [vR*]
        var3 = np.tensordot(h, ACAR_h, axes=((2, 3), (1, 2)))  # p1 p2 [p1*] [p2*], vL [p1] [p2] vR
        var3 = np.tensordot(var3, np.conj(ACAR), axes=((2, 0, 1, 3), (0, 1, 2, 3)))  
        # [p1] [p2] [vL] [vR], [vL*] [p1*] [p2*] [vR*]
        var4 = np.tensordot(ACAR_h, AR, axes=(3, 0))  # vL p1 p2 [vR], [vL] p3 vR
        var4 = np.tensordot(h, var4, axes=((2, 3), (2, 3)))  # p2 p3 [p2*] [p3*], vL p1 [p2] [p3] vR
        var4 = np.tensordot(var4, np.conj(ACAR), axes=((2, 3, 0), (0, 1, 2)))  
        # [p2] p3 [vL] [p1] vR, [vL*] [p1*] [p2*] vR*
        var4 = np.tensordot(var4, np.conj(AR), axes=((2, 0, 1), (0, 1, 2)))  
        # [p3] [vR] [vR*], [vL*] [p3*] [vR*]
        var5 = np.tensordot(ACAR_h, np.conj(ACAR), axes=((0, 1, 2), (0, 1, 2)))  
        # [vL] [p1] [p2] vR, [vL*] [p1*] [p2*] vR*
        var5 = np.tensordot(var5, Rh, axes=((0, 1), (0, 1)))  # [vR] [vR*], [vL] [vL*]
        var = var1 + var2 + var3 + var4 + var5
        return var


# Classes and functions used both vor VUMPS and uTDVP

class Heff1(LinearOperator):
    """Class for the effective Hamiltonian acting on the center site tensor AC.
    
                            .---(AC)--.     .--(AL)--(AC)--.     .--(AC)--(AR)--.     .--(AC)---.
                            |    |    |     |   |     |    |     |   |     |    |     |   |     |
    matvec:  --(AC)--  ->  (Lh)  |    |  +  | (----h----)  |  +  | (----h----)  |  +  |   |   (Rh)
                |           |    |    |     |   |     |    |     |   |     |    |     |   |     |
                            .---    --.     .-(AL*)--    --.     .--   --(AR*)--.     .--    ---.
    """
    def __init__(self, h, Lh, AL, AR, Rh):
        self.h = h  # p1 p2 p1* p2*
        self.Lh = Lh  # vR vR*
        self.AL = AL  # vL p vR
        self.AR = AR  # vL p vR
        self.Rh = Rh  # vL vL*
        D = np.shape(self.AL)[0]
        d = np.shape(self.AL)[1]
        self.shape = (D * d * D, D * d * D)
        self.shape_theta = (D, d, D)
        self.dtype = self.AL.dtype
        
    def _matvec(self, AC):
        """Perform the matvec multiplication diagrammatically shown above."""
        AC = np.reshape(AC, self.shape_theta)  # vL p vR
        AC1 = np.tensordot(self.Lh, AC, axes=(0, 0))  # [vR] vR*, [vL] p vR -> vL p vR
        AC2 = np.tensordot(self.AL, AC, axes=(2, 0))  # vL p1 [vR], [vL] p2 vR
        AC2 = np.tensordot(self.h, AC2, axes=((2, 3), (1, 2)))  # p1 p2 [p1*] [p2*], vL [p1] [p2] vR
        AC2 = np.tensordot(np.conj(self.AL), AC2, axes=((0, 1), (2, 0)))  
        # [vL*] [p1*] vR*, [p1] p2 [vL] vR -> vL p vR  
        AC3 = np.tensordot(AC, self.AR, axes=(2, 0))  # vL p1 [vR], [vL] p2 vR
        AC3 = np.tensordot(self.h, AC3, axes=((2, 3), (1, 2)))  # p1 p2 [p1*] [p2*], vL [p1] [p2] vR
        AC3 = np.tensordot(np.conj(self.AR), AC3, axes=((1, 2), (1, 3)))
        # vL* [p2*] [vR*], p1 [p2] vL [vR] -> vR p vL
        AC3 = np.transpose(AC3, (2, 1, 0))  # vL p vR                        
        AC4 = np.tensordot(AC, self.Rh, axes=(2, 0))  # vL p [vR], [vL] vL* -> vL p vR
        AC_new = AC1 + AC2 + AC3 + AC4
        AC_new = np.reshape(AC_new, self.shape[1])  # vL.p.vR
        return AC_new
    
    def _adjoint(self):
        return self
    
    def trace(self):
        """Compute the trace of Heff1 (only needed for TDVP).

         .---.             .--(AL)-. .---.             .---. .--(AR)-.              .---.
         |   |             |   |   | |   |             |   | |   |   |              |   |
        (Lh) | * d * D  +  | (----h----) | * D  +  D * | (----h----) |  +  D * d *  | (Rh)
         |   |             |   |   | |   |             |   | |   |   |              |   |
         .---.             .-(AL*)-. .---.             .---. .-(AR*)-.              .---.
        """
        trace1 = np.trace(self.Lh) * self.shape_theta[1] * self.shape_theta[2]
        trace2 = np.tensordot(self.AL, np.conj(self.AL), axes=((0, 2), (0, 2)))
        # [vL] p [vR], [vL*] p* [vR*]
        trace2 = np.tensordot(trace2, np.trace(self.h, axis1=1, axis2=3), axes=((0, 1), (1, 0)))
        # [p] [p*], [p1] [p2] [p1*] [p2*]
        trace2 *= self.shape_theta[2]
        trace3 = np.tensordot(self.AR, np.conj(self.AR), axes=((0, 2), (0, 2)))
        # [vL] p [vR], [vL*] p* [vR*]
        trace3 = np.tensordot(trace3, np.trace(self.h, axis1=0, axis2=2), axes=((0, 1), (1, 0)))
        # [p] [p*], [p1] [p2] [p1*] [p2*]
        trace3 *= self.shape_theta[0]
        trace4 =  self.shape_theta[0] * self.shape_theta[1] * np.trace(self.Rh)
        trace = trace1 + trace2 + trace3 + trace4
        return trace


class Heff0(LinearOperator):
    """Class for the effective Hamiltonian acting on the center matrix C.

                           .---(C)--.     .--(AL)--(C)--(AR)--.     .--(C)---.
                           |        |     |   |          |    |     |        |
    matvec:  --(C)--  ->  (Lh)      |  +  | (------h-------)  |  +  |      (Rh)
                           |        |     |   |          |    |     |        |
                           .---   --.     .-(AL*)--   --(AR*)-.     .--   ---.
    """                    
    def __init__(self, h, Lh, AL, AR, Rh):
        self.h = h  # p1 p2 p1* p2*
        self.Lh = Lh  # vR vR*
        self.AL = AL  # vL p vR
        self.AR = AR  # vL p vR
        self.Rh = Rh  # vL vL*
        D = np.shape(self.AL)[0]
        self.shape = (D * D, D * D)
        self.shape_theta = (D, D)
        self.dtype = self.AL.dtype
        
    def _matvec(self, C):
        """Perform the matvec multiplication diagrammatically shown above."""
        C = np.reshape(C, self.shape_theta)  # vL vR
        C1 = np.tensordot(self.Lh, C, axes=(0, 0))  # [vR] vR*, [vL] vR -> vL vR
        C2 = np.tensordot(self.AL, C, axes=(2, 0))  # vL p1 [vR], [vL] vR
        C2 = np.tensordot(C2, self.AR, axes=(2, 0))  # vL p1 [vR], [vL] p2 vR
        C2 = np.tensordot(self.h, C2, axes=((2, 3), (1, 2)))  # p1 p2 [p1*] [p2*], vL [p1] [p2] vR
        C2 = np.tensordot(np.conj(self.AL), C2, axes=((0, 1), (2, 0)))  
        # [vL*] [p1*] vR*, [p1] p2 [vL] vR
        C2 = np.tensordot(C2, np.conj(self.AR), axes=((1, 2), (1, 2)))  
        # vR* [p2] [vR], vL* [p2*] [vR*] -> vL vR
        C3 = np.tensordot(C, self.Rh, axes=(1, 0))  # vL [vR], [vL] vL* -> vL vR
        C_new = C1 + C2 + C3
        C_new = np.reshape(C_new, self.shape[1])  # vL.vR
        return C_new
    
    def _adjoint(self):
        return self
    
    def trace(self):
        """Compute the trace of Heff0 (only needed for TDVP).

         .---.         .--(AL)-. .-(AR)--.         .---.
         |   |         |   |   | |   |   |         |   |
        (Lh) | * D  +  | (------h------) |  +  D * | (Rh)
         |   |         |   |   | |   |   |         |   |
         .---.         .-(AL*)-. .-(AR*)-.         .---.
        """
        trace1 = np.trace(self.Lh) * self.shape_theta[1]
        trace2 = np.tensordot(self.AL, np.conj(self.AL), axes=((0, 2), (0, 2)))
        # [vL] p1 [vR], [vL*] p1* [vR*]
        trace2 = np.tensordot(trace2, self.h, axes=((0, 1), (2, 0)))  
        # [p1] [p1*], [p1] p2 [p1*] p2*
        trace2 = np.tensordot(trace2, self.AR, axes=(1, 1))  # p2 [p2*], vL [p2] vR
        trace2 = np.tensordot(trace2, np.conj(self.AR), axes=((1, 0, 2), (0, 1, 2)))
        # [p2] [vL] [vR], [vL*] [p2*] [vR*]
        trace3 = self.shape_theta[0] * np.trace(self.Rh)
        trace = trace1 + trace2 + trace3
        return trace
    
    
class InverseGeometricSum(LinearOperator):
    """Class for the inverse of the geometric sum of a transfer matrix T_AB with leading right/left
    eigenvector |R>/<L|.
    
    sum_{n=0}^{infty} (alpha * T_AB)^n = [1 - (alpha * T_AB) + (alpha * |R><L|)]^{-1}  (pseudo=True)
                                       = [1 - (alpha * T_AB)]^{-1}  (pseudo=False)

    The case pseudo=True applies for transfer matrices T_AB with leading eigenvalue 1. We assume 
    that the term proportional to |N||R><L| cancels out after multiplication to a vector on the
    right/left. The case pseudo=False applies for transfer matrices T_AB with spectral radius 
    strictly smaller than 1. 
                                                                            
                                 ---.       ---.            ---(A)---.             ---.   .----.    
                                    |          |                |    |                |   |    |      
    matvec for transpose=False:    (X)  ->    (X)  -  alpha     |   (X)  [+  alpha   (R) (L)  (X)]    
                                    |          |                |    |                |   |    | 
                                 ---.       ---.            ---(B*)--.             ---.   .----.  
                                                                                        
                                                                         [iff pseudo=True]

                                 .---        .---           .---(A)---             .----.   .---
                                 |           |              |    |                 |    |   |
    matvec for transpose=True:  (X)    ->   (X)   -  alpha (X)   |      [+  alpha (X)  (R) (L)  ]
                                 |           |              |    |                 |    |   |
                                 .---        .---           .--(B*)---             .----.   .---
    """
    def __init__(self, A, B, R, L, transpose=False, alpha=1.,  pseudo=True):
        D = np.shape(A)[0]
        self.T_AB = TransferMatrix([A], [B], transpose=transpose)
        if pseudo:
            self.R = np.reshape(R, (D * D))
            self.L = np.reshape(L, (D * D))
        self.alpha = alpha
        self.transpose = transpose
        self.pseudo = pseudo
        self.shape = (D * D, D * D)
        self.shape_X = (D, D)
        self.dtype = A.dtype

    def _matvec(self, X):
        if not self.transpose:
            X1 = X
            X2 = -self.alpha * self.T_AB._matvec(X)
            if self.pseudo:
                X2 += self.alpha * self.R * np.inner(self.L, X)
            return X1 + X2
        else:
            X1 = X
            X2 = -self.alpha * self.T_AB._matvec(X)
            if self.pseudo:
                X2 += np.inner(X, self.R) * self.L
            return X1 + X2
            
    def multiply_geometric_sum(self, b, guess, tol):
        """Solve the linear equation self|X> = |b>."""
        b = np.reshape(b, self.shape[1])
        if guess is not None:
            guess = np.reshape(guess, self.shape[1])
        X = gmres(self, b, x0=guess, rtol=tol, atol=0)[0]
        X = np.reshape(X, self.shape_X)  # vL vL* or vR vR*
        return X

    
def get_Lh(psi, h, canonical_form, guess, tol):
    """Compute left environment Lh from geometric sum of transfer matrix TL (psi not necessarily in 
    canonical form).

     .---     .--(AL)--(AL)---                   ---(AL)---^n
     |        |   |     |                            |
    (Lh)   =  | (----h----)    sum_{n=0}^{infty}     | 
     |        |   |     |                            |
     .---     .-(AL*)-(AL*)---                   --(AL*)---
    """
    AL = psi.AL  # vL p vR
    D = np.shape(AL)[0]
    C = psi.C
    R = np.matmul(C, np.conj(C).T)  # vL vL*
    if not canonical_form:
        TL = TransferMatrix([AL], [AL])
        _, R = TL.get_leading_eigenpairs(k=1, guess=R)
        R /= np.trace(R)  # vL vL*   
    IGS = InverseGeometricSum(AL, AL, R=R, L=np.eye(D), transpose=True)
    theta2 = np.tensordot(AL, AL, axes=(2, 0))  # vL p1 p2 vR
    lh = np.tensordot(h, theta2, axes=((2, 3), (1, 2)))
    # p1 p2 [p1*] [p2*], vL [p1] [p2] vR
    lh = np.tensordot(lh, np.conj(theta2), axes=((2, 0, 1), (0, 1, 2)))
    # [p1] [p2] [vL] vR, [vL*] [p1*] [p2*] vR*
    Lh = IGS.multiply_geometric_sum(lh, guess, tol)  # vR vR*
    return Lh

def get_Rh(psi, h, canonical_form, guess, tol):
    """Compute right environment Rh from geometric sum of transfer matrix TR (psi not necessarily in 
    canonical form).

    ---.                        ---(AR)---^n  ---(AR)--(AR)--.
       |                            |             |     |    |
     (Rh)  =  sum_{n=0}^{infty}     |           (----h----)  |
       |                            |             |     |    |
    ---.                        --(AR*)---    --(AR*)-(AR*)--.
    """
    AR = psi.AR  # vL p vR
    D = np.shape(AR)[0]
    C = psi.C
    L = np.matmul(C.T, np.conj(C))  # vR vR*
    if not canonical_form:
        TR = TransferMatrix([AR], [AR], transpose=True)
        _, L = TR.get_leading_eigenpairs(k=1, guess=L)
        L /= np.trace(L)  # vR vR*
    IGS = InverseGeometricSum(AR, AR, R=np.eye(D), L=L)
    theta2 = np.tensordot(AR, AR, axes=(2, 0))  # vL p1 p2 vR
    rh = np.tensordot(h, theta2, axes=((2, 3), (1, 2)))
    # p1 p2 [p1*] [p2*], vL [p1] [p2] vR
    rh = np.tensordot(rh, np.conj(theta2), axes=((0, 1, 3), (1, 2, 3)))
    # [p1] [p2] vL [vR], vL* [p1*] [p2*] [vR*]
    Rh = IGS.multiply_geometric_sum(rh, guess, tol)  # vL vL*
    return Rh


def subtract_energy_offset(psi, h, canonical_form):
    """Subtract energy of psi from two site Hamiltonian h (psi not necessarily in canonical form).
    
                                           .--(AL)--(AL)--.
                                           |   |     |    |
    h -> h - e * Id with e = <psi|h|psi> = | (----h----) (R)    
                                           |   |     |    |
                                           .-(AL*)--(AL*)-.  
    """
    AL = psi.AL  # vL p vR
    R = np.matmul(psi.C, np.conj(psi.C).T)  # vL vL*
    if not canonical_form:
        T = TransferMatrix([AL], [AL])
        _, R = T.get_leading_eigenpairs(k=1, guess=R)
        R /= np.trace(R)  # vL vL*
    theta2 = np.tensordot(AL, AL, axes=(2, 0))  # vL p1 p2 vR
    lh = np.tensordot(h, theta2, axes=((2, 3), (1, 2)))
    # p1 p2 [p1*] [p2*], vL [p1] [p2] vR
    lh = np.tensordot(lh, np.conj(theta2), axes=((2, 0, 1), (0, 1, 2)))
    # [p1] [p2] [vL] vR, [vL*] [p1*] [p2*] vR*
    e = np.tensordot(lh, R, axes=((0, 1), (0, 1)))  # [vR] [vR*], [vL] [vL*]
    d = np.shape(h)[0]  
    Id = np.reshape(np.eye(d * d), (d, d, d, d))  
    h = h - e * Id  # p1 p2 p1* p2*
    return h


def get_AL_AR(AC, C):
    """From given AC and C, find AL and AR which minimize ||AC - AL C||, ||AC - C AR||.
    
    Left polar decompositions: AC = UL_AC PL_AC, C = UL_C PL_C  ->  AL = UL_AC UL_C^{dagger},
    Right polar decompositions: AC = PR_AC UR_AC, C = PR_C UR_C  ->  AR = UR_C^{dagger} UR_AC.
    """
    D = np.shape(AC)[0]
    d = np.shape(AC)[1]
    UL_AC, _ = polar(np.reshape(AC, (D * d, D)))  # vL.p vR
    UL_C, _ = polar(C)  
    AL = np.matmul(UL_AC, np.conj(UL_C).T)  # vL.p [vR], [vL] vR
    AL = np.reshape(AL, (D, d, D))  # vL p vR
    UR_AC, _ = polar(np.reshape(AC, (D, d * D)), side="left")  # vL p.vR
    UR_C, _ = polar(C, side="left")
    AR = np.matmul(np.conj(UR_C).T, UR_AC)  # vL [vR], [vL] p.vR
    AR = np.reshape(AR, (D, d, D))  # vL p vR
    return AL, AR