mps.py

#!/usr/bin/env python

import tensorflow as tf

from tensorflow.python import debug as tf_debug

D = 2
t = -0.001

with tf.name_scope("constant"):
    H = tf.reshape(tf.constant([[0.25,0,0,0],[0,-0.25,0.5,0],[0,0.5,-0.25,0],[0,0,0,0.25]]),[2,2,2,2],name="Hamiltonian")
    I = tf.reshape(tf.constant([[1.,0,0,0],[0,1,0,0],[0,0,1,0],[0,0,0,1]]),[2,2,2,2],"Identity")
    expH = tf.exp(t*H,name="expH")
    Z = tf.zeros(D,name="Zeros")

with tf.name_scope("variables"):
    A = tf.Variable(tf.random_normal([D, D, 2], stddev=1),name="A")
    B = tf.Variable(tf.random_normal([D, D, 2], stddev=1),name="B")
    EAB = tf.Variable(tf.ones(D), name="EAB")
    EBA = tf.Variable(tf.ones(D), name="EBA")
    L = tf.Variable(tf.random_normal([D, D], stddev=1),name="L")
    R = tf.Variable(tf.random_normal([D, D], stddev=1),name="R")
    init = tf.global_variables_initializer()

def update(A,B,EAB,EBA,s):
    label = "%s_H_%s"%(s[0],s[1])
    with tf.name_scope(label):
        Sqrt_EBA = tf.sqrt(EBA,name="sqrt_%s"%s[3])
        EBA_A = tf.multiply(tf.reshape(Sqrt_EBA,[D,1,1]),A,name="%s_%s"%(s[3],s[0]))
        EBA_A_EAB = tf.multiply(EBA_A,tf.reshape(EAB,[1,D,1]),name="%s_%s_%s"%(s[3],s[0],s[2]))
        B_EBA = tf.multiply(B,tf.reshape(Sqrt_EBA,[1,D,1]),name="%s_%s"%(s[1],s[3]))
        A_H = tf.tensordot(EBA_A_EAB,expH,[[2],[0]],name="%s_H"%s[0])
        BIG = tf.tensordot(A_H,B_EBA,[[1,2],[0,2]],name=label) # D 2 2 D
        S, U, V = tf.svd(tf.reshape(BIG,[2*D,2*D]),name="%s_SVD"%label) # D 2 2D, 2D, 2D 2 D
        for p in xrange(4):
            tf.summary.scalar("s%d"%p,S[p])
        nEAB = tf.slice(S,[0],[D],name="n%s"%s[2]) #s[:D]
        pre_reci = tf.reciprocal(Sqrt_EBA,name="pre_reci") # D
        reci = tf.where(tf.is_inf(pre_reci),Z,pre_reci,name="%s_Eviron_Reci"%label) # D
        nA = tf.transpose(tf.multiply(tf.reshape(U[:,:D],[D,2,D]),tf.reshape(reci,[D,1,1])),[0,2,1],name="n%s"%s[0])
        nB = tf.transpose(tf.multiply(tf.reshape(V[:D,:],[D,2,D]),tf.reshape(reci,[1,1,D])),[0,2,1],name="n%s"%s[1])
        return tf.assign(A,nA),tf.assign(B,nB),tf.assign(EAB,nEAB)

def normer():
    with tf.name_scope("Normer"):
        nA = A/tf.norm(A)
        nB = B/tf.norm(B)
        nEAB = EAB/tf.norm(EAB)
        nEBA = EBA/tf.norm(EBA)
        return tf.assign(A,nA),tf.assign(B,nB),tf.assign(EAB,nEAB),tf.assign(EBA,nEBA)

updateAB = update(A,B,EAB,EBA,["A","B","EAB","EBA"])
updateBA = update(B,A,EBA,EAB,["B","A","EBA","EAB"])
normerize = normer()

with tf.name_scope("updateR"):
    RB = tf.tensordot(tf.tensordot(R,B,[[1],[1]]),B,[[0,2],[1,2]],name="RB")
    RBE = tf.multiply(tf.multiply(RB,tf.reshape(EAB,[D,1])),tf.reshape(EAB,[1,D]),name="RBE")
    RBEA = tf.tensordot(tf.tensordot(RBE/tf.norm(RBE),A,[[1],[1]]),A,[[0,2],[1,2]],name="RBEA")
    RBEAE = tf.multiply(tf.multiply(RBEA,tf.reshape(EBA,[D,1])),tf.reshape(EBA,[1,D]),name="RBEAE")
    updateR = tf.assign(R,RBEAE/tf.norm(RBEAE))

with tf.name_scope("updateL"):
    LA = tf.tensordot(tf.tensordot(L,A,[[1],[0]]),A,[[0,2],[0,2]],name="RA")
    LAE = tf.multiply(tf.multiply(LA,tf.reshape(EAB,[D,1])),tf.reshape(EAB,[1,D]),name="RAE")
    LAEB = tf.tensordot(tf.tensordot(LAE/tf.norm(LAE),B,[[1],[0]]),B,[[0,2],[0,2]],name="RAEB")
    LAEBE = tf.multiply(tf.multiply(LAEB,tf.reshape(EBA,[D,1])),tf.reshape(EBA,[1,D]),name="RAEBE")
    updateL = tf.assign(L,LAEBE/tf.norm(LAEBE))

with tf.name_scope("Energy"):
    E_A = tf.multiply(A,tf.reshape(EAB,[1,D,1]),name="E_A")
    H_LA = tf.tensordot(tf.tensordot(L,E_A,[[0],[0]]),E_A,[[0],[0]],name="H_LA")
    H_LAH = tf.tensordot(H_LA,H,[[1,3],[0,2]],name="H_LAH")
    H_LAHB = tf.tensordot(tf.tensordot(H_LAH,B,[[0,2],[0,2]]),B,[[0,1],[0,2]],name="H_LAHB")
    H_LAHBR = tf.tensordot(H_LAHB,R,[[0,1],[0,1]],name="H_LAHBR")
    I_LA = tf.tensordot(tf.tensordot(L,E_A,[[0],[0]]),E_A,[[0,2],[0,2]],name="I_LA")
    I_LAB = tf.tensordot(tf.tensordot(I_LA,B,[[0],[0]]),B,[[0,2],[0,2]],name="I_LAB")
    I_LABR = tf.tensordot(I_LAB,R,[[0,1],[0,1]],name="I_LAHBR")

E = H_LAHBR/I_LABR
tf.summary.scalar("E",E)
merged = tf.summary.merge_all()

sess = tf.Session()
sess.run(init)
logger = tf.summary.FileWriter("/tmp/data/",sess.graph)

sess = tf_debug.LocalCLIDebugWrapperSession(sess)

for i in xrange(1000):
    sess.run(updateAB)
    sess.run(normerize)
    sess.run(updateBA)
    sess.run(normerize)
    sess.run(updateBA)
    sess.run(normerize)
    sess.run(updateAB)
    sess.run(normerize)
    sess.run(updateL)
    sess.run(updateR)
    logger.add_summary(sess.run(merged),i)

logger.flush()

print sess.run(E)