%pylab inline

from qutip import *
from pylab import *

def integrate(N, h, Jx, Jy, Jz, psi0, tlist, gamma, solver):

    si = qeye(2)
    sx = sigmax()
    sy = sigmay()
    sz = sigmaz()

    sx_list = []
    sy_list = []
    sz_list = []
    for n in range(N):
        op_list = []
        for m in range(N):
            op_list.append(si)
        op_list[n] = sx
        sx_list.append(tensor(op_list))
        op_list[n] = sy
        sy_list.append(tensor(op_list))
        op_list[n] = sz
        sz_list.append(tensor(op_list))
        
    # construct the hamiltonian
    H = 0    
    
    # energy splitting terms
    for n in range(N):
        H += - 0.5 * h[n] * sz_list[n]
        
    # interaction terms
    for n in range(N-1):
        H += - 0.5 * Jx[n] * sx_list[n] * sx_list[n+1]
        H += - 0.5 * Jy[n] * sy_list[n] * sy_list[n+1]
        H += - 0.5 * Jz[n] * sz_list[n] * sz_list[n+1] 

    # collapse operators
    c_op_list = []

    # spin dephasing
    for n in range(N):
        if gamma[n] > 0.0:
            c_op_list.append(sqrt(gamma[n]) * sz_list[n])
            
    # evolve and calculate expectation values
    if solver == "me":
        output = mesolve(H, psi0, tlist, c_op_list, sz_list)
    elif solver == "mc":
        output = mcsolve(H, psi0, tlist, c_op_list, sz_list)
    elif solver == "mcf90":
        output = mcsolve_f90(H, psi0, tlist, c_op_list, sz_list)

    return output.expect

N = 4           # number of spins

# array of spin energy splittings and coupling strengths. here we use
# uniform parameters, but in general we don't have too
h  = 1.0 * 2 * pi * ones(N) 
Jz = 0.1 * 2 * pi * ones(N)
Jx = 0.1 * 2 * pi * ones(N)
Jy = 0.1 * 2 * pi * ones(N)

# dephasing rate
gamma = 0.01 * ones(N)

# intial state, first spin in state |1>, the rest in state |0>
psi_list = []
psi_list.append(basis(2,1))
for n in range(N-1):
    psi_list.append(basis(2,0))
psi0 = tensor(psi_list)

tlist = linspace(0, 50, 300)

sz_expt = integrate(N, h, Jx, Jy, Jz, psi0, tlist, gamma, "mcf90")

rc('font', family='Bitstream Vera Sans')
for n in range(N):
    plot(tlist, real(sz_expt[n]), label=r'$\langle\sigma_z($'+str(n)+r'$)\rangle$',lw=2)
xlabel(r'Time [ns]',fontsize=14)
ylabel(r'$\langle\sigma_{z}\rangle$',fontsize=14)
title(r'Dynamics of a Heisenberg spin chain')
legend(loc = "lower right")

show()

# set system parameters
kappa=2.0 # mirror coupling
gamma=0.2 # spontaneous emission rate
g=1       # atom/cavity coupling strength
wc=0      # cavity frequency
w0=0      # atom frequency
wl=0      # driving frequency
E=0.5     # driving amplitude
N=4 # number of cavity energy levels (0->3 Fock states)
tlist=linspace(0,10,101) #times for expectation values

# construct Hamiltonian
ida=qeye(N)
idatom=qeye(2)
a=tensor(destroy(N),idatom)
sm=tensor(ida,sigmam())
H=(w0-wl)*sm.dag()*sm+(wc-wl)*a.dag()*a+1j*g*(a.dag()*sm-sm.dag()*a)+E*(a.dag()+a)

# collapse operators
C1=sqrt(2*kappa)*a
C2=sqrt(gamma)*sm
C1dC1=C1.dag()*C1
C2dC2=C2.dag()*C2

# intial state
psi0=tensor(basis(N,0),basis(2,1))

# run monte-carlo solver with default 500 trajectories
data = mcsolve_f90(H, psi0, tlist, [C1,C2], [C1dC1,C2dC2])

# plot expectation values
plot(tlist,data.expect[0],tlist,data.expect[1],lw=2)
legend(('Transmitted Cavity Intensity','Spontaneous Emission'))
ylabel('Counts')
xlabel('Time')
show()

wa  = 1.0 * 2 * pi   # frequency of system a
wb  = 1.0 * 2 * pi   # frequency of system a
wab = 0.2 * 2 * pi   # coupling frequency
ga = 0.2 * 2 * pi    # dissipation rate of system a
gb = 0.1 * 2 * pi    # dissipation rate of system b
Na = 10              # number of states in system a
Nb = 10              # number of states in system b
E = 1.0 * 2 * pi     # Oscillator A driving strength 

a = tensor(destroy(Na), qeye(Nb))
b = tensor(qeye(Na), destroy(Nb))
na = a.dag() * a
nb = b.dag() * b
H = wa*na + wb*nb + wab*(a.dag()*b+a*b.dag()) + E*(a.dag()+a)

# start with both oscillators in ground state
psi0 = tensor(basis(Na), basis(Nb))

c_op_list = []
c_op_list.append(sqrt(ga) * a)
c_op_list.append(sqrt(gb) * b)

tlist = linspace(0, 5, 101)

data = mcsolve_f90(H, psi0, tlist, c_op_list, [na,nb])

# plot results
figure()
plot(tlist,data.expect[0],'b',tlist,data.expect[1],'r',lw=2)
xlabel('Time',fontsize=14)
ylabel('Excitations',fontsize=14)
legend(('Oscillator A', 'Oscillator B'))
show()

N = 4             # number of basis states to consider
kappa = 1.0/0.129 # coupling to heat bath
nth = 0.063       # temperature with <n>=0.063

# create operators and initial |1> state
a = destroy(N)    # cavity destruction operator
H = a.dag()*a     # harmonic oscillator Hamiltonian
psi0 = basis(N,1) # initial Fock state with one photon

# collapse operators
c_op_list = []
c_op_list.append(sqrt(kappa * (1 + nth)) * a)
c_op_list.append(sqrt(kappa * nth) * a.dag())

# run monte carlo simulation
tlist = linspace(0,0.6,100)
ex1 = mcsolve_f90(H,psi0,tlist,c_op_list,[a.dag()*a], ntraj=1).expect[0]
ex5 = mcsolve_f90(H,psi0,tlist,c_op_list,[a.dag()*a], ntraj=5).expect[0]
ex15 = mcsolve_f90(H,psi0,tlist,c_op_list,[a.dag()*a], ntraj=15).expect[0]
ex904 = mcsolve_f90(H,psi0,tlist,c_op_list,[a.dag()*a], ntraj=904).expect[0]

## run master equation to get ensemble average expectation values ## 
me = mesolve(H,psi0,tlist,c_op_list, [a.dag()*a])

#  calulate final state using steadystate solver
final_state = steadystate(H,c_op_list) 
fexpt = expect(a.dag()*a,final_state) 

#
# plot results using vertically stacked plots
#

# set legend fontsize
import matplotlib.font_manager
leg_prop = matplotlib.font_manager.FontProperties(size=10)

figure(figsize=(6,9))
subplots_adjust(hspace=0.001)

# subplot 1 (top)
ax1 = subplot(411)
ax1.plot(tlist,ex1,'b',lw=2)
ax1.axhline(y=fexpt,color='k',lw=1.5)
yticks(linspace(0,2,5))
ylim([-0.1,1.5])
ylabel('$\left< N \\right>$',fontsize=14)
title("Ensemble Averaging of Monte Carlo Trajectories")
legend(('Single trajectory','steady state'),prop=leg_prop)

# subplot 2
ax2=subplot(412,sharex=ax1)
ax2.plot(tlist,ex5,'b',lw=2)
ax2.axhline(y=fexpt,color='k',lw=1.5)
yticks(linspace(0,2,5))
ylim([-0.1,1.5])
ylabel('$\left< N \\right>$',fontsize=14)
legend(('5 trajectories','steadystate'),prop=leg_prop)

# subplot 3
ax3=subplot(413,sharex=ax1)
ax3.plot(tlist,ex15,'b',lw=2)
ax3.plot(tlist,me.expect[0],'r--',lw=1.5)
ax3.axhline(y=fexpt,color='k',lw=1.5)
yticks(linspace(0,2,5))
ylim([-0.1,1.5])
ylabel('$\left< N \\right>$',fontsize=14)
legend(('15 trajectories','master equation','steady state'),prop=leg_prop)

# subplot 4 (bottom)
ax4=subplot(414,sharex=ax1)
ax4.plot(tlist,ex904,'b',lw=2)
ax4.plot(tlist,me.expect[0],'r--',lw=1.5)
ax4.axhline(y=fexpt,color='k',lw=1.5)
yticks(linspace(0,2,5))
ylim([-0.1,1.5])
ylabel('$\left< N \\right>$',fontsize=14)
legend(('904 trajectories','master equation','steady state'),prop=leg_prop)

# remove x-axis tick marks from top 3 subplots
xticklabels = ax1.get_xticklabels()+ax2.get_xticklabels()+ax3.get_xticklabels()
setp(xticklabels, visible=False)

ax1.xaxis.set_major_locator(MaxNLocator(4))
xlabel('Time (sec)',fontsize=14)
show()

# number of states for each mode
N0=15
N1=15
N2=15

# define operators
a0=tensor(destroy(N0),qeye(N1),qeye(N2))
a1=tensor(qeye(N0),destroy(N1),qeye(N2))
a2=tensor(qeye(N0),qeye(N1),destroy(N2))

# number operators for each mode
num0=a0.dag()*a0
num1=a1.dag()*a1
num2=a2.dag()*a2

# initial state: coherent mode 0 & vacuum for modes #1 & #2
alpha=sqrt(7)#initial coherent state param for mode 0
psi0=tensor(coherent(N0,alpha),basis(N1,0),basis(N2,0))

# trilinear Hamiltonian
H=1.0j*(a0*a1.dag()*a2.dag()-a0.dag()*a1*a2)

# run Monte-Carlo
tlist = linspace(0,2.5,50)
output = mcsolve_f90(H,psi0,tlist,[],[],ntraj=1, serial=True)
#output=mcsolve(H,psi0,tlist,[],[],ntraj=1)

# extrace mode 1 using ptrace
mode1=[psi.ptrace(1) for psi in output.states]
# get diagonal elements
diags1=[k.diag() for k in mode1]
# calculate num of particles in mode 1
num1=[expect(num1,k) for k in output.states]
# generate thermal state with same # of particles
thermal=[thermal_dm(N1,k).diag() for k in num1]

# plot results
from mpl_toolkits.mplot3d import Axes3D
from matplotlib import cm
colors=['m', 'g','orange','b', 'y','pink']
x=arange(N1)
params = {'axes.labelsize': 14,'text.fontsize': 14,'legend.fontsize': 12,'xtick.labelsize': 14,'ytick.labelsize': 14}
rcParams.update(params)

fig = figure()

ax = Axes3D(fig)
for j in range(5):
    ax.bar(x, diags1[10*j], zs=tlist[10*j], zdir='y', color=colors[j], linewidth=1.0, alpha=0.6, align='center')
    ax.plot(x, thermal[10*j], zs=tlist[10*j], zdir='y', color='r', linewidth=3, alpha=1)
ax.set_zlabel(r'Probability')
ax.set_xlabel(r'Number State')
ax.set_ylabel(r'Time')
ax.set_zlim3d(0,1)
show()

N0 = N1 = N2 = 6
gamma0 = 0.1
gamma1 = 0.4
gamma2 = 0.1
alpha = sqrt(2)
tlist = linspace(0,4,200)
ntraj = 500

# define operators
a0=tensor(destroy(N0),qeye(N1),qeye(N2))
a1=tensor(qeye(N0),destroy(N1),qeye(N2))
a2=tensor(qeye(N0),qeye(N1),destroy(N2))

# number operators for each mode
num0=a0.dag()*a0
num1=a1.dag()*a1
num2=a2.dag()*a2

# dissipative operators for zero-temp. baths
C0=sqrt(2.0*gamma0)*a0
C1=sqrt(2.0*gamma1)*a1
C2=sqrt(2.0*gamma2)*a2

# initial state: coherent mode 0 & vacuum for modes #1 & #2
psi0=tensor(coherent(N0,alpha),basis(N1,0),basis(N2,0))

# trilinear Hamiltonian
H=1j*(a0*a1.dag()*a2.dag()-a0.dag()*a1*a2)

# run Monte-Carlo
data = mcsolve_f90(H,psi0,tlist,[C0,C1,C2],[num0,num1,num2])

# plot results
fig = figure()
ax = fig.add_subplot(111)
cs=['b','r','g']
for k in range(ntraj):
    if len(data.col_times[k])>0:
        colors=[cs[j] for j in data.col_which[k]]
        xdat=[k for x in range(len(data.col_times[k]))]
        ax.scatter(xdat,data.col_times[k],marker='o',c=colors)
ax.set_xlim([-1,ntraj+1])
ax.set_ylim([0,tlist[-1]])
ax.set_xlabel('Trajectory',fontsize=14)
ax.set_ylabel('Collpase Time',fontsize=14)
ax.set_title('Blue = C0, Red = C1, Green= C2')
show()

from qutip.ipynbtools import version_table
version_table()