#
#  Light Cone Renormalization Group (LCRG) algorithm
#
#  (C) 2011 by Tilman Enss (Tilman.Enss@gmail.com)
#
#  You are welcome to use this code but if you do so please cite
#  the following article:
#     Tilman Enss and Jesko Sirker,
#     "Lightcone renormalization and quantum quenches in
#      one-dimensional Hubbard models",
#     New J. Phys. 14, 023008 (2012).
#
#  This program is written in Python (http://www.python.org) and uses
#  the 'numpy' matrix package (http://numpy.scipy.org) to perform
#  tensor operations.
#
#  To run, call
#
#      python lcrg.py
#
#  from the command line.
#
#
#  This implementation of the LCRG algorithm is optimized for product
#  initial states with period two, like the Neel state, and a
#  translationally invariant Hamiltonian.  It operates directly on the
#  reduced left density matrix rho which grows to the left with every
#  time step.
#

from numpy import *

def LCRG(Delta=0,dt=.1,tmax=5.,chimax=64):
    """
    LCRG algorithm:
    time evolution of staggered magnetization from initial Neel state
    within the XXZ model (Jz=Delta*Jx);
    Trotter step 'dt', maximum simulation time 'tmax', and
    maximum dimension of reduced Hilbert space 'chimax'.
    """
    M = 2                             # local states: 0=down, 1=up
    tau = XXZ(Delta,dt)               # local update for XXZ model
    initstate = array([[0.,0],[1,0]]) # initial Neel state up-down
    op = array([[-.5,0],[0,.5]])      # local magnetization operator 0.5*sigma_z

    L = initstate.reshape(1,M,M,1)    # initial L[ml,sl,sr,mr]=initstate[sl,sr]
    rho = array([[1.]])               # initial rho[m',m]=1, block dimension dim(m)=chi=1

    print("# LCRG algorithm by Tilman Enss and Jesko Sirker (2011)")
    print("#    [version without conserved quantum numbers]")
    print("# Time evolution of staggered magnetization from Neel state in XXZ model")
    print("# columns: time, staggered magnetization, truncation error")

    for t in arange(0.,tmax,dt):      # loop over Trotter steps
        # first half time step
        rho = grow_rho(rho,L)         # grow density matrix rho
        magn,tr_rho = eval_op(rho,op) # evaluate local staggered magnetization
        print("%.2f %.10g %.10g" % (t,magn,1-tr_rho))
        if tr_rho < 0.99: break       # stop if truncation error becomes too large
        rho,U = renorm(rho,chimax)    # find reduced Hilbert space basis
        L = grow_L(L,U,tau)           # project to new basis and grow left transfer matrix L

        # second half time step
        rho = grow_rho(rho,L)         # grow density matrix rho
        rho,U = renorm(rho,chimax)    # find reduced Hilbert space basis
        L = grow_L(L,U,tau)           # project and grow left transfer matrix L
    return

def XXZ(Delta,dt):
    """local update tau=exp(-i*H*dt) for XXZ model with Jz=Delta*Jx, time in units 1/Jx"""
    tau = zeros((2,2,2,2),complex)
    expD = exp(-.25j*Delta*dt)
    tau[0,0,0,0] = tau[1,1,1,1] = expD
    tau[0,1,0,1] = tau[1,0,1,0] = cosh(-.5j*dt)/expD
    tau[0,1,1,0] = tau[1,0,0,1] = sinh(-.5j*dt)/expD
    return tau

def grow_rho(rho,L):
    """grow density matrix rho to the left"""
    # multiply L onto lower block index: rho[m,s,sr,mr']=sum(mr) L[m,s,sr,mr] rho[mr',mr]
    rho = tensordot(L,rho,axes=([3],[1]))
    # L* onto upper block index: rho[m',s',m,s]=sum(sr,mr') L*[m',s',sr,mr'] rho[m,s,sr,mr']
    rho = tensordot(L.conjugate(),rho,axes=([2,3],[2,3]))
    return rho

def eval_op(rho,op):
    """evaluate local operator"""
    # local density matrix rho_local[s',s]=sum(m) rho[m,s',m,s]
    rho_local = trace(rho,axis1=0,axis2=2) # trace over block index
    observ = trace(dot(rho_local,op)).real # evaluate local operator
    tr_rho = trace(rho_local).real         # trace of local density matrix
    return observ,tr_rho

def renorm(rho,chimax):
    """find reduced Hilbert space basis"""
    chi,M = rho.shape[:2]             # block and site dimensions of rho
    # diagonalize hermitean matrix rho[(m's'),(ms)]
    eigenval,eigenvec = linalg.eigh(rho.reshape(chi*M,chi*M))
    # eigenvectors form unitary transformation from old indices (m,s) to new index m'
    U = eigenvec.transpose().reshape(chi*M,chi,M)
    if chi*M > chimax:                # choose chimax states with largest eigenvalues
        eigenval = eigenval[-chimax:] # discarded weight: sum(eigenval[:-chimax])
        U = U[-chimax:]
    rho = diagflat(eigenval)          # rho is diagonal in the new basis
    return rho,U

def grow_L(L,U,tau):
    """project to new basis and grow left transfer matrix L"""
    # project left block index: L[ml',sr,mr]=sum(ml,sl) U[ml',ml,sl] L[ml,sl,sr,mr]
    L = tensordot(U,L,axes=([1,2],[0,1]))
    # project right block index: L[ml',sl,mr',sr]=sum(mr) L[ml',sl,mr] U*[mr',mr,sr]
    L = tensordot(L,U.conjugate(),axes=([2],[1]))
    # add local update: L[ml',sl',sr',mr']=sum(sl,sr) tau[sl',sr',sl,sr] L[ml',sl,mr',sr]
    L = tensordot(tau,L,axes=([2,3],[1,3])).transpose(2,0,1,3)
    return L

# run the algorithm for the Delta=0.5 XXZ model
LCRG(Delta=0.5,dt=.1)