from qutip import *
from pylab import *
%pylab inline

up=basis(2,0)
dn=basis(2,1)

rpp=up*up.dag()
rpm=up*dn.dag()
rmp=rpm.trans()
rmm=dn*dn.dag()

E=1.0
w0=1.0
w3=1.0
w2=0.0
w1=0.0

H0=(E-w0-w3)*rpp+(E+w0+w3)*rmm+(w1-1j*w2)*rpm+(w1+1j*w2)*rpm

A=1.0
theta=pi/3.0
phi=0.0

HA=A/2.0*((1.0+cos(2.0*theta))*rpp+(1.0-cos(2.0*theta))*rmm+exp(-1j*phi)*sin(2.0*theta)*rpm+exp(1j*phi)*sin(2.0*theta)*rmp) 

H=H0+HA

psi0=cos(theta)**2.0*rpp+sin(theta)**2.0*rmm+cos(theta)*sin(theta)*(exp(-1j*phi)*rpm+exp(1j*phi)*rmp)
psi_mu=qeye(2)-psi0

D=sqrt(0.1)*sigmaz()

tlist=linspace(0,5.0,10.0)
outlist=zeros(len(tlist))

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

for i1 in range(len(tlist)):
    rho=qq.states[i1]
    outlist[i1]=abs((rho*psi_mu).tr())

plot(tlist,outlist,'-',label='', linewidth=4)
xlabel('t',fontsize=20)
ylabel('P(nu->mu)',fontsize=20)
show()