#
#  stochastic TMRG program
#     (optimized for systems invariant under spatial reflection)
#
#  (C) 2001 by Tilman Enss (Tilman.Enss@physik.uni-muenchen.de)
#      2002 Tilman Enss: changed buggy dot() to innerproduct()
#

# Set to 1 for debugging output (density matrix spectra).
debug = 0

# Maximum error for the convergence of the power method
# (`power()') and for the selection of nonsingular density matrix
# eigenvalues; should be approximately machine accuracy.
epsilon = 1e-14

# Load Numeric module for matrix functions.
from Numeric import *

# Much time is spent in the diagonalization of the reduced density
# matrix; try to use a fast routine (e.g. Lapack DSYEVR).
# symeigenvec() returns _sorted_ eigenvalues/-vectors.
try:
    from linalg import dsyevr
    def symeigenvec(A): return dsyevr(A,0)
except:
    from LinearAlgebra import eigenvectors
    def symeigenvec(A):
        eval,evec = eigenvectors(A)
        # real eigenvalues, since A is symmetric
        try: eval = eval.real; evec = evec.real
        except: pass
        return eval,evec

# Setup the local transfer matrix tau for a simple reaction-
# diffusion model: monomers aggregate with rate `twoalpha' and
# diffuse with rate `D'; the time step is `dt' units.

def setup_tau(dt=0.1,alpha=0.5,D=0.5):
    n = 2            # number of states per site: 0=empty, 1=A
    tau = zeros((n,n,n,n),Float)    # indices tau[l1,l2,r1,r2]
    expalpha = exp(-2*alpha*dt)
    expD = exp(-2*D*dt)
    tau[0,0,0,0] = 1.0
    tau[1,0,1,0] = 1.0-expalpha
    tau[1,1,1,1] = expalpha
    tau[0,0,1,1] = tau[1,1,0,0] = 0.5*(1.0+expD)
    tau[0,1,1,0] = tau[1,0,0,1] = 0.5*(1.0-expD)
    return n,tau

# Initialize the system (S) and environment (E) blocks for the
# first Trotter step (M=1).  `iweight' is the initial weight
# (bias) of each state; `EA' is the environment block with
# boundary conditions corresponding to the local operator given by
# the expectation values for each final local state, `fweight'.

def init(n,tau,iweight,fweight):
    # S[l2,r2]=sum(l1,r1) iweight[l1]*iweight[r1]*tau[l2,r2,l1,r1]
    S = CS = innerproduct(innerproduct(transpose(tau,(1,3,0,2)),iweight),iweight)
    # E[l1,r1] = 1
    E = ones((n,n),Float)
    CE = ones((1,1),Float)
    # EA[l1,r1] = sum(l2,r2) tau[l1,l2,r1,r2]*fweight[r2]
    EA = sum(innerproduct(tau,fweight),1)
    return S,E,EA,CS,CE

# Normalize the vector v: first compute its norm, if it is nonzero
# then divide by it and return the normalized vector and its norm.

def normalize(v):
    norm = sqrt(innerproduct(v,v))
    if norm > epsilon: v = v/norm
    return v,norm

# Multiply a vector `psi' with the transfer matrix, stored sepa-
# rately as system (`S') and environment (`E').  The order of the
# multiplications with `S' and `E' depends on the exact geometry
# (whether the Trotter number M is odd or even).

def multSE(S,E,odd,psi):
    Sdim,Edim = S.shape[1],E.shape[1]  # dimensions to contract
    if odd:
        # if M is odd, first multiply with S, then with E
        psi = transpose(reshape(psi,(Sdim,-1)),(1,0))
        psi = innerproduct(S,psi)
        psi = innerproduct(reshape(psi,(-1,Edim)),E)
    else:
        # if M is even, first multiply with E, then with S
        psi = innerproduct(reshape(psi,(-1,Edim)),E)
        psi = transpose(reshape(psi,(Sdim,-1)),(1,0))
        psi = innerproduct(S,psi)
    # return flat vector
    return ravel(psi)

# The power method for finding the right eigenvector corresponding
# to the largest eigenvalue of a dense matrix: begin with some
# start vector `psi' and continue to multiply the transfer matrix
# onto it until the eigenvector converges within an error of
# `epsilon'.

MAXITER = 200
def power(S,E,odd,psi):
    for i in range(MAXITER):  # iterate at most MAXITER times
        psinew,lamb = normalize(multSE(S,E,odd,psi))
        err = max(abs((psinew-psi)))/max(abs(psinew))
        if err < epsilon: return lamb,psinew
        psi = psinew
    print "power: largest eigenvector DID NOT converge!"
    return 0,psinew

# Add local tau to S, E, CS, or CE.  `first' comes first in the
# time evolution (system block) and is a square matrix of size
# [a*b,a*b].  `second' comes afterwards (environment block) and is
# a square matrix of size [b*c,b*c].  The order of the multipli-
# cation is determined by the odd or even geometry.  The `b' index
# is contracted.  Finally, the grown matrix is returned as a
# square matrix.

def grow(first,second,odd,a,b,c):
    if odd:
        first = transpose(reshape(first,(a,b,a*b)),(0,2,1))
        second = transpose(reshape(second,(b*c,b,c)),(0,2,1))
    else:
        first = reshape(first,(a*b,a,b))
        second = transpose(reshape(second,(b,c,b*c)),(1,2,0))
    # combine both parts by contracting the `b' index
    combined = innerproduct(first,second)
    # return square matrix
    return reshape(transpose(combined,(0,2,1,3)),(a*b*c,a*b*c))

# Find a truncated basis for renormalization by constructing the
# symmetric reduced density matrix for the system/environment,
# tracing over the resp. other indices and taking the eigenvectors
# corresponding to the largest eigenvalues.  A maximum of `mmax'
# eigenvectors are retained (such that the renormalization reduces
# the state space sufficiently), but only those corresponding to
# _nonsingular_ eigenvalues: `sum(greater(Sval,epsilon))' is the
# number of eigenvalues greaten than `epsilon'.

def renorm_basis(n,m,mmax,psiL,psiR):
    # find reduced system basis
    l = reshape(psiL,(m*n,m))
    r = reshape(psiR,(m*n,m))
    rhosys = innerproduct(l,l) + innerproduct(r,r)
    Sval,Svec = symeigenvec(rhosys/trace(rhosys))
    if debug: print "Sys. eigenvalues:\n",Sval # print eigenvals
    keepsys = min(mmax,sum(greater(Sval,epsilon)))

    # find reduced environment basis
    l = transpose(reshape(psiL,(m,n*m)),(1,0))
    r = transpose(reshape(psiR,(m,n*m)),(1,0))
    rhoenv = innerproduct(l,l) + innerproduct(r,r)
    Eval,Evec = symeigenvec(rhoenv/trace(rhoenv))
    if debug: print "Env. eigenvalues:\n",Eval # print eigenvals
    keepenv = min(mmax,sum(greater(Eval,epsilon)))
    keep = max(keepsys,keepenv)

    # compute approximation error and return truncated basis
    if debug: print "keep m=%i, approx.error %.2e(S), %.2e(E)" % \
       (keep,sum(Sval[keep:]),sum(Eval[keep:]))
    return Svec[:keep],Evec[:keep],keep

# Project a square matrix `M' into a new basis, given by the array
# of basis vectors `basis'.  `l' (`r') dimensions on the left
# (right) are left unchanged.
 
def project(M,basis,l,r):
    newdim,olddim = basis.shape
    # project right side of matrix `M'
    M = transpose(reshape(M,(-1,olddim,r)),(0,2,1))
    M = transpose(innerproduct(M,basis),(0,2,1))
    # project left side of matrix `M'
    M = transpose(reshape(M,(l,olddim,-1)),(0,2,1))
    M = transpose(innerproduct(M,basis),(0,2,1))
    # return projected square matrix
    return reshape(M,(l*newdim*r,l*newdim*r))

# The main Trotter loop: simulate `Mmax' Trotter steps, with
# S(E)-state space dimension of `mmax' states maximum.  The local
# operator to be measured at the end of the time evolution is
# given by `Aweight', its expectation values for the local states.
# `dt' is the time step, the initial weight `iweight' is unbiased
# by default.

def main(Mmax=200,mmax=16,dt=0.1,Aweight=[0.0,1.0],iweight=None):
    # setup the local transfer matrix tau
    n,tau = setup_tau(dt)

    # if not defined, unbiased initial weight
    if iweight is None: iweight = [1.0/n] * n

    # initialize system and environment blocks
    S,E,EA,CS,CE = init(n,tau,iweight,Aweight)
    m = 1

    # write out initial density
    print "M = %3i:  t = %8.4f, Erwartungswert = %.8f" % \
          (0,0,innerproduct(iweight,Aweight))
    outfile = open("density","w")
    outfile.write("%8.4f %.8f\n" % (0,innerproduct(iweight,Aweight)))

    # loop over Trotter number
    for M in range(1,Mmax+1):
        # geometry of S and E alternates with M
        odd = M % 2

        # compute largest right eigenvalue/-vector lmax,psiR of
        # the transfer matrix and multiply psiR by one column of
        # tau's to compute left eigenvector psiL
        psiR = ones((m*n*m,))  # start vector for power method
        lmax,psiR = power(S,E,odd,psiR)
        psiL = multSE(CS,CE,odd,psiR)

        # max. eigenvalue should be approx. 1
        if abs(lmax-1.0) > 0.1:
            print "*** numerical instability ***"
            return

        # compute and print expectation value of local operator
        id = lmax*innerproduct(psiL,psiR) # expect. of identity op
        density = innerproduct(psiL,multSE(S,EA,odd,psiR))/id
        print "M = %3i:  t = %8.4f, Erwartungswert = %.8f" % \
              (M,M*dt,density)
        outfile.write("%8.4f %.8f\n" % (M*dt,density))
        outfile.flush()

        # grow S, E, EA, CS, and CE by one tau each
        S = grow(S,tau,odd,m,n,n)
        E = grow(tau,E,odd,n,n,m)
        EA = grow(tau,EA,odd,n,n,m)
        if odd: CE = grow(tau,CE,odd,n*n,1,m)
        else:   CS = grow(CS,tau,odd,m,1,n*n)

        # if state space is too large, renormalize:
        if m*n > mmax:
            # compute optimal reduced basis
            psiL,norm = normalize(psiL)
            Sbasis,Ebasis,m = renorm_basis(n,m,mmax,psiL,psiR)

            # project down to a smaller state space
            S = project(S,Sbasis,1,n)
            E = project(E,Ebasis,n,1)
            EA = project(EA,Ebasis,n,1)
            CS = project(CS,Sbasis,1,n**(1-odd))
            CE = project(CE,Ebasis,n**odd,1)
        else:
            m = m*n  # grow state space dimension
    outfile.close()