main.py

import numpy as np
import scipy.linalg as linalg
from random import random
from math import sqrt, exp, isclose

from scipy.special import lmbda

import pim


class SpinChain:
    def __init__(self, N):
        self.N = N
        self.model = None
        self.full_state = None
        self.MPS = CanMPS(self.N)

    def set_Ising_model(self, J, g):
        self.model = Ising(J, g)

    def rand_State(self):
        self.full_state = np.random.rand(2**self.N, 1)
        norm = np.sqrt(np.dot(self.full_state.T, self.full_state))
        self.full_state = self.full_state/norm

    def rand_MPS(self, chi):
        self.MPS.rand_MPS(chi)

    def Theta(self, i):
        """ Construct the two-site state tensor Theta for sites (i, i+1)
        Returns:
            Theta (np.ndarray) : shape (2, 2, χ_{i-1}, χ_{i+1}), indices (s_i, s_{i+1}, a_{i-1}, a_{i+1})
        """

        if i < 1 or i > self.N - 1:
            print("ERROR: Application of Theta out of index.")
            return

        L0 = self.MPS.Lambdas[i - 1]  # shape (χ_{i-1}, χ_{i-1})
        G1 = self.MPS.Gammas[i]  # shape (2, χ_{i-1}, χ_i)
        L1 = self.MPS.Lambdas[i]  # shape (χ_i, χ_i)
        G2 = self.MPS.Gammas[i + 1]  # shape (2, χ_i, χ_{i+1})
        L2 = self.MPS.Lambdas[i + 1]  # shape (χ_{i+1}, χ_{i+1})


        # Form two-site theta
        # we want: theta indices s_i, s_{i+1}, a_{i-1}, a_{i+1}
        # = sum over a_i:
        #       L0_(a_{i-1}, a_{i-1}) * G1^[s_i]_(a_{i-1}, a_i) *
        #     * L1_{a_i, a_i} *
        #     * G2^[s_{i+1}]_(a_i, a_{i+1}) *
        #     * L2_(a_{i+1}, a_{i+1})

        #print("L0: ", L0.shape, "  G1: ",  G1.shape)

        # L0 * G1
        # Contract axis 1 (a_{i-1}) of L0 with axis 1 (a_{i-1})  of G1
        # (a_{i-1}, a_{i-1}) * (s_i, a_{i-1}, a_i) -> (a_{i-1}, s_i, a_i)
        # L0G1(a_{i-1}, s_i, a_i) = sum over k: L0(a_{i-1}, k) G1(s_i, k, a_i) = (L0 is diag) =
        #                              = L0(a_{i-1}, a_{i-1}) G1(s_i, a_{i-1}, a_i) = (3.36)
        L0G1 = np.tensordot(L0, G1, axes=(1, 1))    # (x_{i-1}, 2, x_i)

        #print("L0G1: ", L0G1.shape)

        # (L0 * G1) * L1
        # Contract axis 2 (a_i) of L0G1 with axis 0 (a_i) of L1
        # (a_{i-1}, s_i, a_i) * (a_i, a_i) -> (a_{i-1}, s_i, a_i)
        # L0G1L1(a_{i-1}, s_i, a_i) = sum over k: L0G1(a_{i-1}, s_i, k) L1(k, a_i) = (L1 is diag) =
        #                              = L0G1(a_{i-1}, s_i, a_i) L1(a_i, a_i) = (3.36)
        L0G1L1 = np.tensordot(L0G1, L1, axes=(2, 0))    # (χ_{i-1}, 2, χ_i)

        #print("L0G1L1: ", L0G1L1.shape)

        # G2 * L2
        # Contract axis 2 (a_{i+1}) of G2 with axis 0 (a_{i+1}) of L2
        # (s_{i+1}, a_i, a_{i+1}) * (a_{i+1}, a_{i+1}) -> (s_{i+1}, a_i, a_{i+1})
        # G2L2(s_{i+1}, a_i, a_{i+1}) = sum over k: G2(s_{i+1}, a_i, k) * L2(k, a_{i+1})  = (L2 is diag) =
        #                              = G2(s_{i+1}, a_i, a_{i+1}) * L2(a_{i+1}, a_{i+1})  = (3.36)
        G2L2 = np.tensordot(G2, L2, axes=(2, 1))    # (2, x_i, x_{i+1})

        #print("G2L2: ", G2L2.shape)

        # Theta = ((L0 * G1) * L1) * (G2 * L2)
        # Contract axis 2 (a_i) of L0G1L1 with axis 1 (a_i) of G2L2
        # (a_{i-1}, s_i, a_i) *  (s_{i+1}, a_i, a_{i+1}) ->  (a_{i-1}, s_i, s_{i+1}, a_{i+1})
        # Theta(a_{i-1}, s_i, s_{i+1}, a_{i+1}) = sum over k: L0G1L1(a_{i-1}, s_i, k) * G2L2(s_{i+1}, k, a_{i+1})  = (3.36)
        Theta = np.tensordot(L0G1L1, G2L2, axes=(2, 1))     # (χ_{i-1}, 2, 2, χ_{i+1})



        Theta = np.moveaxis(Theta, 0, 2)


        return Theta  # shape: (2, 2, χ_{i-1},  χ_{i+1})

    def normalize_Theta(self, Theta):
        norm = np.tensordot(Theta.conj(), Theta, axes = ([0, 1, 2, 3], [0,1,2,3]))
        norm = sqrt(norm)
        Theta_normal = Theta / norm

        return Theta_normal

    def update_Theta(self, i, Op):
        """ Applies a two-site operator Op to the two-site state Theta at sites (i, i+1).
        Args:
            i (int): Site index
            Op (np.ndarray): Two-site operator
                            indices: (s'_i, s'_{i+1}, s_i, s_{i+1})
        Returns:
            Theta_new (np.ndarray): updated Theta
                            indices: (s'_i, s'_{i+1}, a_{i-1}, a_{i+1})
        """
        # Construct two-site state Theta
        Theta = self.Theta(i)   # shape (s_i, s_{i+1}, a_{i-1}, a_{i+1})


        # Contract axis 2, 3 (s_i, s_{i+1}) of Op with axis 0,1 (s_i, s_{i+1}) of Theta
        # (s'_i, s'_{i+1}, s_i, s_{i+1}) * (s_i, s_{i+1}, a_{i-1}, a_{i+1}) -> (s'_i, s'_{i+1}, a_{i-1}, a_{i+1})
        # Gamma_i(s'_i, s'_{i+1}, a_{i-1}, a_{i+1}) =
        #                   = sum over k, l: Op(s'_i, s'_{i+1}, k, l) * Theta(k, l, a_{i-1}, a_{i+1}) =
        #                   = sum over k, l: Theta(k, l, a_{i-1}, a_{i+1}) * Op(s'_i, s'_{i+1}, k, l) = (3.56)
        #print(Op.shape, Theta.shape)
        Theta_new = np.tensordot(Op, Theta, axes=([2, 3], [0, 1]))

        #Theta_new = self.normalize_Theta(Theta_new)

        return Theta_new

    def trunc(self, M, n):
        p = M.shape[1] # = min(2*x_{i-1}, 2*x_{i+1})
        k = min(chi, p)

        # Truncate so that U is (2*x_{i-1}, x)
        U_trunc = np.zeros((2 * chiL, chi)) # pad in case p < x
        U_trunc[:, :k] = U[:, :k]
        # Truncate so that S is (x)
        S_trunc = np.zeros((chi)) # pad in case p < x
        S_trunc[:k] = S[:k]
        # Truncate so that V_dagger is (x, 2*x_{i+1})
        Vh_trunc = np.zeros((chi, 2*chiR)) # pad in case p < x
        Vh_trunc[:k, :] = Vh[:k, :]

    def apply_two_site_op(self, i, Op):
        """
        Applies a two-site operator Op to the MPS at sites (i, i+1), updating Gammas and Lambdas.

        Args:
            i (int): Site index
            Op (np.ndarray): Two-site operator
                             indices: (s'_i, s'_{i+1}, s_i, s_{i+1})
        """
        #print("Gamma i is left-canonical: ", self.MPS.site_is_left_canonical(i))
        #print("Gamma i + 1 is left-canonical: ", self.MPS.site_is_left_canonical(i+1))

        #print("---------------   i = ", i, "   ----------------------")

        # (s'_i, s'_{i+1}, a_{i-1}, a_{i+1})
        Theta_new = self.update_Theta(i, Op)    # apply Op to Theta
        Theta_new = np.transpose(Theta_new, axes=(0, 2, 1, 3))


        # Reshape Theta_new to a matrix for SVD
        #  (s'_i, s'_{i+1}, a_{i-1}, a_{i+1}) -> ( [s'_i, a_{i-1}] , [s'_{i+1}, a_{i+1}] )
        chiL = Theta_new.shape[1]
        chiR = Theta_new.shape[3]
        Theta_mat = Theta_new.reshape(2 * chiL, 2 * chiR)

        # SVD decomposition Theta_mat = U*S*Vh
        # U: (2*x_{i-1}, p)
        # S: (p)
        # Vh: (p, 2*x_{i+1})
        U, S, Vh = np.linalg.svd(Theta_mat, full_matrices=False)
        #Vh = Vh.T

        #print("------------ ", i, " ------------------")
        #print(Vh)
        #print("-----------------------")

        chi = self.MPS.chi
        p = U.shape[1] # = min(2*x_{i-1}, 2*x_{i+1})
        k = min(chi, p)

        # Truncate so that U is (2*x_{i-1}, x)
        U_trunc = np.zeros((2 * chiL, chi)) # pad in case p < x
        U_trunc[:, :k] = U[:, :k]
        # Truncate so that S is (x)
        S_trunc = np.zeros((chi)) # pad in case p < x
        S_trunc[:k] = S[:k]
        # Truncate so that V_dagger is (x, 2*x_{i+1})
        Vh_trunc = np.zeros((chi, 2*chiR)) # pad in case p < x
        Vh_trunc[:k, :] = Vh[:k, :]




        # Reshape U_trunc: (2*x_{i-1} , x) → (2, x_{i-1}, x)
        # ( [s'_i * a_{i-1}] , a_i') → (s'_i, a_{i-1}, a_i')
        U_trunc = U_trunc.reshape((2, chiL, chi))

        #print("Canonicality of U_trunc: ", self.MPS.tensor_is_left_canonical(U_trunc))
        #print("Uh_trunc norm : ", linalg.norm(Vh_trunc))

        # Reshape Vh_trunc: (x, 2*x_{i+1}) → (x, 2, x_{i+1})
        # (a_i', [s'_{i+1} * a_{i+1}] ) → (a_i', s'_{i+1}, a_{i+1})
        Vh_trunc = Vh_trunc.reshape((chi, 2, chiR))

        #print("Canonicality of Vh_trunc: ", self.MPS.tensor_is_left_canonical(Vh_trunc))

        #print(f"S Not normalized S norm = {np.linalg.norm(S_trunc)}")
        norm = np.linalg.norm(S_trunc)
        #print("Vh_trunc norm : ", linalg.norm(Vh_trunc))
        #S_trunc = S_trunc / norm   # Normalize S_trunc
        Lambda_i = np.diag(S_trunc)   # construct new Lambda^(i) based on S_trunc

        #print(f"S Norm = {np.linalg.norm(S_trunc)} | U norm = {np.linalg.norm(U_trunc)} | Vh norm = {np.linalg.norm(Vh_trunc)} ")

        Lambda_prev = self.MPS.Lambdas[i - 1]
        chi_prev = Lambda_prev.shape[0]
        #Lambda_prev[Lambda_prev < 10**(-6)] = 0
        #print("Lambda prev")
        #print(Lambda_prev)

        #Lambda_prev_inv = np.array(np.diag(Lambda_prev))
        #Lambda_prev_inv[Lambda_prev_inv > 0] = Lambda_prev_inv[Lambda_prev_inv > 0]**(-1)
        #Lambda_prev_inv = np.diag(Lambda_prev_inv)

        #print("Lambda prev inv")
        #print(Lambda_prev_inv)
        Lambda_prev_inv = np.linalg.pinv(Lambda_prev) # (pseudo) inverse of Lambda^(i-1)

        # New Gamma^(i) = inv(Lambda^(i-1)) * U
        # Contract axis 1 (a_{i-1}) of inv(Lambda^(i-1)) with axis 1 (a_{i-1}) of U_trunc
        # (a_{i-1}, a_{i-1}) * (s'_i, a_{i-1}, a_i') -> (a_{i-1}, s'_i, a_i')
        # Gamma_i(a_{i-1}, s'_i, a_i') = sum over k: L(a_{i-1}, k) U(s'_i, k, a_i') = (L is diag) =
        #                              = L(a_{i-1}, a_{i-1}) U(s'_i, a_{i-1}, a_i') = (3.58)
        Gamma_i = np.tensordot(Lambda_prev_inv, U_trunc, axes=(1, 1))
        Gamma_i = np.moveaxis(Gamma_i, 1, 0)  # (s'_i, a_{i-1}, a_i')



        Lambda_next = self.MPS.Lambdas[i +1]
        #Lambda_next[Lambda_next < 10**(-6)] = 0
        #print("Lambda next")
        #print(Lambda_next)

        #Lambda_next_inv = np.array(np.diag(Lambda_next))
        #Lambda_next_inv[Lambda_next_inv > 0] = Lambda_next_inv[Lambda_next_inv > 0]**(-1)
        #Lambda_next_inv = np.diag(Lambda_next_inv)

        #Lambda_next_inv /= np.linalg.norm(Lambda_next_inv)*sqrt(2)

        #print("Lambda next inv")
        #print(Lambda_next_inv)

        Lambda_next_inv = np.linalg.pinv(self.MPS.Lambdas[i + 1])  # (pseudo) inverse of Lambda^(i+1)

        # New Gamma^(i+1)= Vh_trunc * inv(Lambda^(i+1))
        # Contract axis 2 (a_{i+1}) of Vh_trunc with axis 0 (a_{i+1}) of inv(Lambda^(i+1))
        # (a_i', s'_i, a_{i+1}) * (a_{i+1}, a_{i+1}) -> (a_i', s'_i, a_{i+1})
        # Gamma_{i+1}(a_i', s'_i, a_{i+1}) = sum over k: V(a_i', s'_i, k) * L(k, a_{i+1})  = (L is diag) =
        #                              = V(a_i', s'_i, a_{i+1}) * L(a_{i+1}, a_{i+1})  = (3.58)
        Gamma_next = np.tensordot(Vh_trunc, Lambda_next_inv, axes=(2, 0))
        Gamma_next = np.moveaxis(Gamma_next, 1, 0) # (s'_i, a_i', a_{i+1})

        #print("Lambda[i-1]: ")
        #print(self.MPS.Lambdas[i - 1])
        #print("Lambda[i-1] norm : ", linalg.norm(self.MPS.Lambdas[i - 1]))
        #print("inv(Lambda[i-1]) norm : ", linalg.norm(Lambda_prev_inv))
        #print("Vh_trunc norm : ", linalg.norm(Vh_trunc))
        #print("inv(Lambda[i-1])^t * inv(Lambda[i-1]) : i = ", i-1 )
        #print(np.dot(Lambda_prev_inv.conj().T, Lambda_prev_inv))

        #print("Gamma i: ", Gamma_i.shape)
        #print("Gamma i+1: ", Gamma_next.shape)
        #print("Lambda i: ", Lambda_i.shape)


        # Update MPS
        self.MPS.Gammas[i] = Gamma_i
        self.MPS.Lambdas[i] = Lambda_i
        self.MPS.Gammas[i + 1] = Gamma_next
        #print("Lambda i: = ", i)
        #print(Lambda_i)



        #print("Gamma ", i, " is left canonical: ", self.MPS.tensor_is_left_canonical(Gamma_i))
        #print("Gamma ", i + 1, " is left canonical: ", self.MPS.tensor_is_left_canonical(Gamma_next))
        #print("MPS is left canonical after one operation: ", chain.MPS.is_left_canonical())


        #print("Gamma i is left-canonical: ", self.MPS.site_is_left_canonical(i))
        #print("Gamma i + 1 is left-canonical: ", self.MPS.site_is_left_canonical(i+1))

        return norm

    def ground_energy(self, dt):
        norm = self.MPS.norm()
        E1 = - np.log(norm**2) / (2 * dt)
        return E1

    def itebd_evolve(self, dt, steps):
        """ Perform iTEBD: imaginary time evolution on the MPS.

        Args:
            dt (float): Time step
            steps (int): Number of time steps to evolve
        """


        N = self.MPS.N

        U_even = self.model.U_local(dt/2) # e^(H_i * dt/2)
        U_odd = self.model.U_local(dt) # e^(H_i * dt/2)
        H, H_left, H_right = pim.Hamiltonian_Ising_model(g, J)
        #U = pim.two_site_Hamiltonian(H, dt)  # e^(H_i * dt/2)
        #U_end = pim.two_site_Hamiltonian(H_right, dt)  # e^(H_i * dt/2)

        # For N:th site, H_local is a one-site operator
        U_even_end = self.model.U_local(dt/2, is_end = True) # e^(H_N * dt/2)
        U_odd_end = self.model.U_local(dt, is_end=True)  # e^(H_N * dt/2)

        #lambdas, gammas, loc_size = pim.Initializing_State(N, chi, 2)
        #print(loc_size)
        #self.MPS = pim_to_my(gammas, lambdas, N, chi)


        ground_energies = []    # Track numerical calculations for E1 after each iteration
        for step in range(steps):
            #print("MPS satisfies canonical form: ", chain.MPS.is_left_canonical())
            # S_even sweep: apply H_local to even sites: i = 2, 4, 6, ..., N - 1

            #lambdas, gammas = my_to_pim(self.MPS)


            #print("Gammas : ",  gammas.shape)
            #print("Lambdas : ",  lambdas.shape)

            print("%%%%%%%%%%%%%%%%%%%%%%% t = ", step, "%%%%%%%%%%%%%%%%%%%%%%%%%" )

            for i in range(2, N, 2):  # even sites
                U_step = U_even if i != N-1 else U_even_end     # if i = N, U_local is different
                self.apply_two_site_op(i, U_step)
            #for i in range(0, N-1, 2):  # even sites
                #gammas, lambdas = pim.Two_site_Operator (i , gammas , lambdas, pim.O_arr ,N , 2, chi , loc_size )

            # S_odd sweep: apply H_local to odd sites: i = 1, 3, 5, ..., N - 1
            for i in range(1, N, 2):  # odd sites
                U_step = U_odd if i != N-1 else U_odd_end     # if i = N, U_local is different
                self.apply_two_site_op(i, U_step)
            #for i in range(1, N - 1, 2):  # Odd bonds
                #gammas, lambdas = pim.Two_site_Operator (i , gammas , lambdas, pim.O_arr ,N , 2, chi , loc_size )

            # S_even sweep: apply H_local again to even sites: i = 2, 4, 6, ..., N - 1
            for i in range(2, N, 2):  # even sites
                U_step = U_even if i != N-1 else U_even_end     # if i = N, U_local is different
                self.apply_two_site_op(i, U_step)
            #for i in range(0, N - 1, 2):  # even sites
                #gammas, lambdas = pim.Two_site_Operator(i, gammas, lambdas, pim.O_arr, N, 2, chi, loc_size)

            #self.MPS = pim_to_my(gammas, lambdas, N, chi)

            print("MPS norm: ", self.MPS.norm())

            pim_norm = sqrt(self.MPS.pim_norm())
            print("Pim norm: ", pim_norm)

            #self.MPS.normalize()
            #print("Norm: ", self.MPS.norm())
            print("My energy: ", chain.ground_energy(dt))
            print("Pim energy: ", - np.log(pim_norm**2) / (2 * dt))
            #chain.MPS.normalize()

            #norm = pim_norm ** (1 / self.N)
            norm = self.MPS.norm() ** (1 / (self.N-1))
            # normalize before next step
            for i in range(1, self.N ):
                #norm = linalg.norm(self.MPS.Lambdas[i])
                self.MPS.Lambdas[i] /= norm


            print("MPS norm: ", self.MPS.norm())
            pim_norm = sqrt(self.MPS.pim_norm())
            print("Pim norm: ", sqrt(pim_norm))

            print("-----------------------------------------------------------")




        return ground_energies



class Ising:
    def __init__(self, J, g):
        self.J = J
        self.g = g
    def H_local(self, is_end = False): # checked
        """ Local two-site Hamiltonian.
            OBS: same for all sites i = 1, ..., N-1
                 but for i = N, H_local is a single site operator
        Returns:
            H (np.ndarray) : (4, 4) matrix rep. of local Hamiltonian
        """
        Sx = 1/2*np.array([[0., 1.], [1., 0.]])
        Sz = 1/2*np.array([[1., 0.], [0., -1.]] )
        H =  - self.J * np.kron(Sz, Sz) + self.g * np.kron(Sx, np.eye(2, 2))
        if is_end:
            H = H + self.g * np.kron(np.eye(2, 2), Sx)

        return H

    def U_local(self, time_step, is_end=False): # checked
        # Two-site gate for imaginary time evolution
        # U is (4,4), indices (x, y) with x = 2s'_i + s'_{i+1}, y = 2s_i + s_{i+1}
        # s_i, s_{i+1}, s'_i, s'_{i+1} in {0, 1}
        # U_(x,y) = < s'_i, s'_{i+1} | U | s_i, s_{i+1} >
        U = linalg.expm(-time_step* self.H_local(is_end))

        # Reshape U to have indices s'_i, s'_{i+1}, s_i, s_{i+1}
        U = U.reshape(2, 2, 2, 2)
        return U



class CanMPS:
    def __init__(self, N):
        self.N = N
        self.Gammas = []
        self.Lambdas = []
        self.chi = None

    def rand_MPS(self, chi):
        """Set self.state as randomly generated state in truncated canonical MPS form
            with fixed bond dimension.
        Args:
            chi (int) : bond dimension for each site,
                 (OBS: chi is a constant, so the MPS is a truncated approximation)
        Returns:
            Gammas (list of np.ndarray): len = N + 1
                Canonical Gamma tensors, shapes (2, chi, chi).
                Dummy 0:th Gamma is 1x1 np.ndarray [[1.0]].
                1:st Gammas has shape (2, 1, chi), N:th shape (2, chi, 1)
            Lambdas (list of np.ndarray): len = N+1
                Normalized Schmidt diagonal matrices, shapes (chi, chi).
                Dummy 0:th and N:th lambdas are 1x1 np.ndarrays [[1.0]].
        """
        self.chi = chi
        Gammas = []
        Lambdas = []

        Lambdas.append(np.array([[1.0]])) # dummy 0:th Lambda = 1 for cleaner tensor operations
        Gammas.append(np.array([[1.0]]))  # dummy 0:th Gamma = 1 to align indices of Lambda and Gamma
        np.random.seed()
        chi_prev = 1  # start at boundary with bond dim 1
        for i in range(1, self.N+1):
            # for i = 1, ..., N

            # Set bond dimension for this site to chi
            chi_curr = chi if i < self.N else 1  # N:th site ends in bond dim 1

            # Generate random matrix A = QR: (chi_prev * 2) x chi_curr
            A = np.random.rand(chi_prev * 2, chi_curr)
            if chi_curr > chi_prev*2:
                A = A.T
                Q, _ = np.linalg.qr(A)  # Q is orthonormal
                Q = Q.T
            else:
                Q, _ = np.linalg.qr(A)  # Q is orthonormal

            Gamma_i = Q.reshape(2, chi_prev, chi_curr)
            Gammas.append(Gamma_i)

            Gi_is_can = False
            S = np.dot(Gamma_i[0].conj().T, Gamma_i[0]) + np.dot(Gamma_i[1].conj().T, Gamma_i[1])
            if np.allclose(S, np.eye(Gamma_i[0].shape[1])):
                Gi_is_can = True
            else:
                print(S)
            print("Gamma i is left canonical at initial generation: ", Gi_is_can)

            # Generate normalized random Lambda (Schmidt vector)
            if i == self.N:
                Lambdas.append(np.array([[1.0]]))  # dummy N:th lambda = 1 for cleaner tensor operations
            else:
                Lambda_i = np.random.rand(chi_curr)
                Lambda_i /= np.linalg.norm(Lambda_i)
                Lambdas.append(np.diag(Lambda_i))
                chi_prev = chi_curr

        self.Gammas = Gammas
        self. Lambdas = Lambdas
        return

    def tensor_is_left_canonical(self, T):
        S = np.dot(T[0].conj().T, T[0]) + np.dot(T[1].conj().T, T[1])
        n = S.shape[1]
        p = n - 1
        while isclose(S[p, p], 0, abs_tol=10**(-5)) and p >= 0:
            p = p - 1
        #print(S[:p+1, :p+1])
        #print("p : ", p)
        #S[S < 10 ** (-9)] = 0
        #print(S)
        if np.allclose(S[:p+1, :p+1], np.eye(p+1)):
            return True
        #print("Canonical value: ", linalg.norm(S))
        #print("p : ", p)
        return False

    def is_left_canonical(self):
        x = True
        #print("-----------------")
        for i in range(1, self.N+1):
            #print("site : ", i)
            if not self.tensor_is_left_canonical(self.Gammas[i]):
                x = False
        #print("-----------------")
        return x

    def norm(self):
        E = np.array([[1.0]])
        #print(self.Lambdas)
        #print(self.Gammas)
        for i in range(1, self.N+1):
            # get MPS matrix A at site i
            if i == 1:
                # A = G^(i)
                A = self.Gammas[i]      # (s_i, a_{i-1}, a_i)
            else:
                # A = L^(i-1)* G^(i)
                # (a_{i-1}, a_{i-1}) * (s_i, a_{i-1}, a_i) -> (a_{i-1}, s_i, a_i)

                A = np.tensordot(self.Lambdas[i-1], self.Gammas[i], axes=(1,1))     # (a_{i-1}, s_i, a_i)
                A = np.moveaxis(A, 1, 0)    # (s_i, a_{i-1}, a_i)
            # if i = 0:     B = G^(i)^t
            # else:         B = G^(i)^t L^(i-1)^t
            B = np.transpose(A.conj(), (0,2,1))

            # E'_(a_i', a_i) = sum over s_i, a_{i.1}, a_{i-1}':
            #                        B_{s_i, a_i', a_{i-1}') * E_(a_{i-1}', a_{i-1}) * A_(s_i, a_{i-1}, a_i)

            # Contract E with A and B
            # (s_i, a_i', a_{i-1}') * (a_{i-1}', a_{i-1}) -> (s_i, a_i', a_{i-1})
            BE = np.tensordot(B, E, axes=(2, 0))
            # (s_i, a_i', a_{i-1})  * (s_i, a_{i-1}, a_i) -> (a_i', a_i)
            E = np.tensordot(BE, A, axes=([0, 2], [0, 1]))

        #print("----------------------------------------------------------------------- Norm: ", sqrt(np.squeeze(E)))

        return sqrt(np.squeeze(E))

    def pim_norm(self):
        lambdas, gammas = my_to_pim(self)
        return pim.calc_norm(gammas, lambdas, self.N, self.chi)

    def normalize(self):
        norm = self.norm()
        for i in range(1,self.N+1):
            self.Gammas[i] /= (norm**(1/self.N))
        return

    def init_fm_mps(self, chi): #checked
        """ Returns FM Ising MPS"""
        d = 2
        B = []
        s = []

        self.chi = chi
        Gammas = []
        Lambdas = []

        Lambdas.append(np.array([[1.0]])) # dummy 0:th Lambda = 1 for cleaner tensor operations
        Gammas.append(np.array([[[1.0]], [[1.0]]]))  # dummy 0:th Gamma = 1 to align indices of Lambda and Gamma

        chi_prev = 1  # start at boundary with bond dim 1
        chi_next = 5
        for i in range(1, self.N +1 ):
            if i == self.N:
                chi_next = 1
            Gammas.append(np.zeros([2, chi_prev, chi_next]))
            Gammas[-1][0, 0, 0] = 1
            chi_prev = 5
        for i in range(1, self.N ):
            L = np.zeros((chi,chi))
            L[0,0] = 1.0
            Lambdas.append(L)
        Lambdas.append(np.array([[1.0]]))
        #Lambdas.append(np.array([[1.0]]))
        #Gammas.append(np.array([[[1.0]], [[1.0]]]))  # dummy N+1:th Gamma = 1 for cleaner tensor operations

        self.Gammas = Gammas
        self.Lambdas = Lambdas

def my_to_pim(MPS):
    chi = MPS.chi
    N = MPS.N
    lambdas = np. zeros ((N + 1 , chi ))
    gammas = np.zeros((N ,chi , chi , 2))
    for i in range(N):
        G = np.transpose(MPS.Gammas[i+1], (1, 2, 0))
        G = np.pad(G, ((0,chi), (0, chi), (0,0)))
        gammas[i] = G[:chi, :chi, :]

        L = np.diag(MPS.Lambdas[i])
        L = np.pad(L, ((0, chi)))
        lambdas[i] = L[:chi]

    L = np.diag(MPS.Lambdas[N])
    L = np.pad(L, ((0, chi)))
    lambdas[N] = L[:chi]

    return lambdas, gammas

def pim_to_my(gammas, lambdas, N, chi):
    MPS_my = CanMPS(N)
    MPS_my.chi = chi

    MPS_my.Lambdas.append(np.array([[1.0]]))  # i = 0
    for i in range(1, N):   # i = 1, ..., N-1
        MPS_my.Lambdas.append(np.diag(lambdas[i, :]))
    MPS_my.Lambdas.append(np.array([[1.0]]))  # i = N

    MPS_my.Gammas.append(np.array([[1.0]]))  # i = 0
    gammas = np.transpose(gammas, axes=(0, 3, 1, 2))
    MPS_my.Gammas.append(gammas[0, :, :1, :])    # i = 1
    for i in range(1, N-1):   # i = 2, ..., N-1
        MPS_my.Gammas.append(gammas[i, :, :, :])
    MPS_my.Gammas.append(gammas[N-1, :, :, :1])   # i = N

    return MPS_my


np.set_printoptions(3, suppress=True)

J = 0.5
g = 1
chi = 5
N = 10
dt = 0.05
steps = 200

chain = SpinChain(N)
chain.set_Ising_model(J, g)
#chain.rand_MPS(chi)

chain.MPS.init_fm_mps(chi)


"""
for i in range(len(chain.MPS.Gammas)):
    print("Gamma ", i, " :")
    for s in [0, 1]:
        print("     s = ", s, " :")
        print("         ", chain.MPS.Gammas[i][s])
for i in range(len(chain.MPS.Lambdas)):
    print("Lambda ", i, " :")
    print("     ", chain.MPS.Lambdas[i])
"""

is_canonical = chain.MPS.is_left_canonical()
print(f"MPS satisfies canonical form: {is_canonical}")
print("Initial norm: ", chain.MPS.norm(), chain.MPS.pim_norm())

#chain.MPS.normalize()
print("Norm: ", chain.MPS.norm())

chain.itebd_evolve(dt, steps)

#print("Post evolution MPS is left canonical: ", chain.MPS.is_left_canonical())

#print(chain.model.H_local(True))
#print(chain.model.H_local())
#print("U")
#print(chain.model.U_local(dt))
#print(pim.np. array ([[ -3* g /4 ,0 ,0 , - J /4] ,[0 , g /4 , - J /4 ,0] ,[0 , - J /4 , - g /4 ,0] ,[ - J /4 ,0 ,0 ,3* g /4]]))
#print(np. array ([[ - g /2 ,0 ,0 , - J /4] ,[0 ,0 , - J /4 ,0] ,[0 , - J /4 ,0 ,0] ,[ - J /4 ,0 ,0 , g /2]]))

E1_Teitsma = -1.11696246824
E1_Pim  = -3.0953724999136853
print("Sought norm Teitsma: ", exp(-E1_Teitsma*2*dt), " | Sought norm Pim: ", exp(-E1_Pim*2*dt))
print("Sought energy Teitsma: ", - np.log(exp(-E1_Teitsma*2*dt)) / (2 * dt), " | Sought energy Pim: ", - np.log(exp(-E1_Pim*2*dt)) / (2 * dt))

"""
for i in range(len(chain.MPS.Gammas)):
    print("Gamma ", i, ":  ", chain.MPS.Gammas[i].shape)
for i in range(len(chain.MPS.Lambdas)):
    print("Lambda ", i, ":  ", chain.MPS.Lambdas[i].shape)
"""

#print("Norm: ", chain.MPS.norm())
#print(chain.ground_energy(dt))

"""

# Below is for debugging. MPS.norm() correctly calculates a simple normalized states norm.

chain2 = SpinChain(3)
chain2.chi = 2

Lambda_1 = np.array([[1.0, 0.0], [0.0, 1.0]])  # Diagonal matrix for the first bond (2x2)
Lambda_2 = np.array([[1.0, 0.0], [0.0, 1.0]])  # Diagonal matrix for the second bond (2x2)

# Gammas (Local spin states)
Gamma_1 = np.array([[[1, 0]], [[1, 0]]])

Gamma_2 = np.array([np.eye(2), np.eye(2)])

Gamma_3 = np.array([[[1], [0]], [[1], [0]]])

chain2.MPS.Gammas.append(np.array([[1.0]]))
chain2.MPS.Gammas.append(Gamma_1)
chain2.MPS.Gammas.append(Gamma_2)
chain2.MPS.Gammas.append(Gamma_3)
chain2.MPS.Lambdas.append(np.array([[1.0]]))
chain2.MPS.Lambdas.append(Lambda_1)
chain2.MPS.Lambdas.append(Lambda_2)
chain2.MPS.Lambdas.append(np.array([[1.0]]))
for i in range(len(chain2.MPS.Gammas)):
    print("Gamma ", i, ":  ", chain2.MPS.Gammas[i].shape)
for i in range(len(chain2.MPS.Lambdas)):
    print("Lambda ", i, ":  ", chain2.MPS.Lambdas[i].shape)
print(chain2.MPS.norm()**2)


"""