%autosave 0
%matplotlib inline
from __future__ import division
from sglib import *
from sympy import sqrt, E, Rational, latex, Symbol
L = latex # Shorthand

# Zero and One states (they're also + and - z)
s0 = col(1,0); s1 = col(0,1)
pz = s0; mz = s1
# Create a formatter for stuff like |00>
fmt = sg_format_state(['0','1'], '').format
zfmt = sg_format_state(['+z','-z'], '').format
xfmt = sg_format_state(['+x','-x'], '').format
# Create two-bit states and basis
s00=TP(s0,s0); s01=TP(s0,s1); s10=TP(s1,s0); s11=TP(s1,s1)
b2bit = [s00, s01, s10, s11]
# Some misc definitions
half = Rational(1,2) # Define 1/2
alpha = Symbol('alpha') # Define some symbols
beta = Symbol('beta')
# The Bell states
bell_1 = 1/sqrt(2)*s00 + 1/sqrt(2)*s11 # phi+
bell_2 = 1/sqrt(2)*s00 - 1/sqrt(2)*s11 # phi-
bell_3 = 1/sqrt(2)*s01 + 1/sqrt(2)*s10 # psi+ (Triplet)
bell_4 = 1/sqrt(2)*s01 - 1/sqrt(2)*s10 # psi- (Singlet)
bell = [bell_1, bell_2, bell_3, bell_4]

from sympy import var
a_11, a_12, a_21, a_22 = var('a_11, a_12, a_21, a_22')
A = mat(a_11, a_12, a_21, a_22)
MOP = matrix_as_outer_product(A) # Just to get access to the "demo" Tr()
from sglib import pZ, mZ, pX, mX, pY, mY
trace = MOP.Tr(A, [pZ, mZ], V=True)
trace = MOP.Tr(A, [pX, mX], V=True)
trace = MOP.Tr(A, [pY, mY], V=True)

psi = 1/sqrt(2)*s0 + 1/sqrt(2)*s1
rho = OP(psi,psi)
examine_dm(rho)

half = Rational(1,2)
rho = half*OP(s0,s0) + half*OP(s1,s1)
examine_dm(rho)

rho = .15*OP(s0,s0) + .85*OP(s1,s1)
examine_dm(rho)

rho = OP(bell_1, bell_1)
examine_dm(rho)

rho = .25*OP(bell_1, bell_1) + .25*OP(bell_2, bell_2)\
      + .25*OP(bell_3, bell_3) + .25*OP(bell_4, bell_4)
examine_dm(rho)

rho = .25*OP(s00,s00)+.15*OP(s01,s01)+.05*OP(s10,s10)+.55*OP(s11,s11)
print 'trace is %s'%rho.trace() # Make sure it adds up!
examine_dm(rho)

rho = .50*OP(s00,s00)+.50*OP(s10,s10)
print 'trace is %s'%rho.trace() # Make sure it adds up!
examine_dm(rho)

# The state we are given
psi = 1/sqrt(2)*pz + 1/sqrt(2)*mz
rho = OP(psi,psi)
phi = col(1/sqrt(2), 1/sqrt(2)) # This is +x
phi = col(1,0) # This is +z
Print(r'Given: $%s=%s$' %( zfmt(psi), L(psi) ))
Print(r'Find the probability of measuring $%s=%s$' %(zfmt(phi),L(phi) ))
Print(r'$\langle\phi|\rho|\phi\rangle=%s%s%s=%s$'
      %(L(phi.adjoint()),L(rho),L(phi), (phi.adjoint()*rho*phi)[0] ))
foo = phi*phi.adjoint()*rho
string = r'$\mathrm{Tr}(|\phi\rangle\langle\phi|\rho)'
string += r'=\mathrm{Tr}\left(%s%s%s\right)=\mathrm{Tr}\left(%s\right)=%s$'
Print(string
      %(L(phi),L(phi.adjoint()),L(rho),L(foo),foo.trace()),font_size=3)

one = 1/sqrt(2)*s0 + 1/sqrt(2)*s1
two = sqrt(3)/2*s0 + Rational(1,2)*s1
state = TP(two, one)
Print(r'State = $%s$' %fmt(state, which_ltx='sympy'))
rho = OP(state, state)
Print(r'Density matrix is: $\rho=%s$' %latex(rho))
MOP = matrix_as_outer_product(rho)
Print(r'In "outer product form":')
Print(r'$\rho=%s$' %MOP.latex())
MOP.partial_trace(1)
Print('After tracing out the second bit:')
Print(r'$\rho=%s$' %MOP.latex())
Print('And in matrix form:')
Print(r'$\rho=%s$' %latex(MOP.M))

state = 1/sqrt(2)*s00 + half*s01 + half*s11
Print(r'State = $%s$' %fmt(state, which_ltx='mine'))
rho = OP(state, state)
Print(r'Density matrix is: $\rho=%s$' %latex(rho))
MOP = matrix_as_outer_product(rho)
Print(r'In "outer product form":')
Print(r'$\rho=%s$' %MOP.latex())
MOP.partial_trace(1)
Print('After tracing out the second bit:')
Print(r'$\rho=%s$' %MOP.latex())
Print('And in matrix form:')
Print(r'$\rho=%s$' %latex(MOP.M))

Print(r'$\psi = %s$'%fmt(bell_1))
rho_maxE = OP(bell_1, bell_1)
examine_dm(rho_maxE)

rho = matrix_as_outer_product(rho_maxE)
rho.partial_trace(1)
examine_dm(rho.M)

rho = matrix_as_outer_product(rho_maxE)
rho.partial_trace(0)
examine_dm(rho.M)

rho = mat(.5,0,0,0, 0,0,0,0, 0,0,0,0, 0,0,0,.5)
S = Entropy(rho, base=E)
Print(r'$\rho = %s, \;\; S = %s$' %(latex(rho), float(S)))

r3 = 1/sqrt(3)
psi = r3*s00 + r3*s01 + r3*s10
Print(r'$\psi = %s$'%fmt(psi))
rho = OP(psi,psi)
examine_dm(rho)
# Then do the reduced density for both bit. Should be the same.
rho_op = matrix_as_outer_product(rho)
rho_op.partial_trace(1)
examine_dm(rho_op.M)
rho_op = matrix_as_outer_product(rho)
rho_op.partial_trace(0)
examine_dm(rho_op.M)

psi = half*s00 + half*s01 + half*s10 - half*s11
Print(r'$\psi = %s$'%fmt(psi))
rho = OP(psi,psi)
S = Entropy(rho)
Print(r'$\rho = %s, \;\; S = %s$' %(latex(rho), latex(S)))
# Then do the reduced density for bit zero
rho_op = matrix_as_outer_product(rho)
rho_op.partial_trace(1)
examine_dm(rho_op.M)

psi = half*s00 + half*s01 + half*s10 + half*s11
Print(r'$\psi = %s$'%fmt(psi))
rho = OP(psi,psi)
S = Entropy(rho)
Print(r'$\rho = %s, \;\; S = %s$' %(latex(rho), latex(S)))
# Then do the reduced density for bit one
rho_op = matrix_as_outer_product(rho)
rho_op.partial_trace(1)
examine_dm(rho_op.M)

P = [] # Purity values
S = [] # Entropy values
M = [] # Percentage of -z mixed in
for percent_mZ in range(51):
    frac_mZ = percent_mZ/100.0
    rho = (1-frac_mZ)*OP(pZ, pZ) + frac_mZ*OP(mZ,mZ)
    M += [ percent_mZ, ]
    P += [ Purity(rho), ]
    S += [ Entropy(rho), ]
pyplot.plot(M, P, label="Purity", color="green");
pyplot.plot(M, S, label="Entropy", color="red");
pyplot.legend(loc='best');
pyplot.title('One bit purity/entropy vs percent mixed');

psi = alpha*s0 - beta*s1
Print(r'$\psi=%s$' %myltx(psi))
rho = OP(psi,psi)
Print(r'$\rho=%s$' %myltx(rho))
MOP = matrix_as_outer_product(rho)
Print(r'$%s$' %MOP.latex())
Print(r'$\rho=%s$' %latex(MOP.M))

M = mat(1,2,3,4,5,6,7,8,
        1,2,3,4,5,6,7,8,
        1,2,3,4,5,6,7,8,
        1,2,3,4,5,6,7,8,
        1,2,3,4,5,6,7,8,
        1,2,3,4,5,6,7,8,
        1,2,3,4,5,6,7,8,
        1,2,3,4,5,6,7,8)
#M = mat(-1,2,3,4,5,6,7,8,9,10,11,12,13,14,15,-16)
MOP = matrix_as_outer_product(M)
Print('The original matrix:')
Print(r'$%s$' %MOP.latex())
Print('Partial Trace of bit 0:')
MOP.partial_trace(0)
Print(r'$%s$' %MOP.latex())
Print('Another partial trace of bit 0:')
MOP.partial_trace(0)
Print(r'$%s$' %MOP.latex())

#
# How about matrices with symbolic values?
#

psi = TP(1/sqrt(2)*s0 + 1/sqrt(2)*s1, col(alpha, beta))
Print(r'$\psi=%s$'%latex(psi))
Print(r'$\psi=%s$'%fmt(psi))
M = OP(psi,psi)
Print(r'$\rho = %s$' %latex(M))
MOP = matrix_as_outer_product(M)
Print(r'$%s$'%MOP.latex())
MOP.partial_trace(0)
Print('Tracing out bit zero:')
Print(r'$%s$'%MOP.latex())

#
# Now the teleportation state
#

psi = half*TP(s00, col(alpha,beta))\
    + half*TP(s01, col(beta,alpha))\
    + half*TP(s10, col(alpha,-beta))\
    + half*TP(s11, col(-beta,alpha))
Print(r'$\psi=%s$'%fmt(psi)) # THIS IS A PROBLEM !!!
Print(r'$\psi=%s$'%latex(psi))
M = OP(psi,psi)
Print(r'$\rho = %s$' %latex(M))
MOP = matrix_as_outer_product(M)
Print(r'$%s$'%MOP.latex())
MOP.partial_trace(0)
Print('Tracing out Alice\'s first bit:')
Print(r'$%s$'%MOP.latex())
MOP.partial_trace(0)
Print('Tracing out Alice\'s second bit:')
Print(r'$%s$'%MOP.latex())
Print('And the reduced density matrix for Bob is:')
Print(r'$\rho_B = %s$' %latex(MOP.M))