# setup the matplotlib graphics library and configure it to show figures inline in the notebook
%pylab inline

# make qutip available in the rest of the notebook
from qutip import *

#qutip.settings.qutip_graphics=False
#qutip.settings.qutip_gui=False

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

tlist = linspace(0,0.6,100)

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

# collapse operator list
c_op_list = []

# decay operator
c_op_list.append(sqrt(kappa * (1 + nth)) * a)

# excitation operator
c_op_list.append(sqrt(kappa * nth) * a.dag())

ntraj = [1, 5, 15, 904] # list of number of trajectories to avg. over

mc = mcsolve(H, psi0, tlist, c_op_list, [a.dag()*a], ntraj)

# get expectation values from mc data (need extra index since ntraj is list)
ex1   = mc.expect[0][0]   # for ntraj=1
ex5   = mc.expect[1][0]   # for ntraj=5
ex15  = mc.expect[2][0]   # for ntraj=15
ex904 = mc.expect[3][0]   # for ntraj=904

# 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) # find steady-state
fexpt = expect(a.dag()*a, final_state)  # find expectation value for particle number

import matplotlib.font_manager
leg_prop = matplotlib.font_manager.FontProperties(size=10)

fig, axes = subplots(4, 1, sharex=True, figsize=(8,12))

fig.subplots_adjust(hspace=0.1) # reduce space between plots

for idx, n in enumerate(ntraj):

    axes[idx].plot(tlist, mc.expect[idx][0], 'b', lw=2)
    axes[idx].plot(tlist, me.expect[0], 'r--', lw=1.5)
    axes[idx].axhline(y=fexpt, color='k', lw=1.5)
    
    axes[idx].set_yticks(linspace(0, 2, 5))
    axes[idx].set_ylim([0, 1.5])
    axes[idx].set_ylabel(r'$\left<N\right>$', fontsize=14)
    
    if idx == 0:
        axes[idx].set_title("Ensemble Averaging of Monte Carlo Trajectories")
        axes[idx].legend(('Single trajectory', 'master equation', 'steady state'), prop=leg_prop)
    else:
        axes[idx].legend(('%d trajectories' % n, 'master equation', 'steady state'), prop=leg_prop)
        

axes[3].xaxis.set_major_locator(MaxNLocator(4))
axes[3].set_xlabel('Time (sec)',fontsize=14)