from qutip import *
from pylab import *
%pylab inline

def discord(ro):
    ll=[]
    prec=10
    th=linspace(0,pi,prec)
    ph=linspace(0,2.0*pi,prec)
    for i1 in xrange(len(ph)):
        phi=ph[i1]
        for i2 in xrange(len(th)):
            theta=th[i2]
            phi=0.0
            psi=cos(theta/2.0)*basis(2,0)+exp(1J*phi)*sin(theta/2.0)*basis(2,1)
            p0=ket2dm(psi)
            p1=qeye(2)-p0
            P0=tensor(qeye(2),p0)
            P1=tensor(qeye(2),p1)
            ####################
            q0=(P0*ro).tr()
            r20=(P0*ro*P0).ptrace([0])
            ro0=r20.unit()
            ####################
            q1=(P1*ro).tr()
            r21=(P1*ro*P1).ptrace([0])
            ro1=r21.unit()
            ###################
            r1=ro.ptrace([0])
            #r2=ro.ptrace([1])
            JJ=entropy_vn(r1,e)-q0*entropy_vn(ro0,e)-q1*entropy_vn(ro1,e)
            ll=ll+[real(JJ)]
    #
    r2=ro.ptrace([1])
    II=entropy_vn(r1,e)+entropy_vn(r2,e)-entropy_vn(ro,e)
    DD=II-max(ll)
    return DD

B0=(tensor(basis(2,0),basis(2,0))+tensor(basis(2,1),basis(2,1))).unit()

BB0=ket2dm(B0)

discord(BB0)

plist = linspace(0.0, 1.0, 10.0)
clist=zeros(len(plist))

for i1 in xrange(len(plist)):
    p=i1*0.1
    rho=(1-p)*BB0+p/4.0*tensor(qeye(2),qeye(2))
    c=discord(rho)
    clist[i1]=c
                

plot(plist,clist,'-.r',label='', linewidth=4)
xlabel('p',fontsize=20)
ylabel('QD',fontsize=20)
show()

H1=tensor(qeye(2), sigmaz())
H2=tensor(sigmaz(),qeye(2))
Hint=tensor(sigmax(),sigmax())

H=H1+H2+Hint

psi0=BB0

tlist=linspace(0.0, 1.0, 5.0)
qlist=zeros(len(tlist))

qq=mesolve(H, psi0, tlist, [], [])

for i1 in xrange(len(tlist)):
    rho=qq.states[i1]
    c=discord(rho)
    qlist[i1]=c

plot(tlist,qlist,'-.r',label='', linewidth=4)
xlabel('t',fontsize=20)
ylabel('QD',fontsize=20)
show()