#!/usr/bin/env python
# coding: utf-8

# # QuTiP example: Dynamics of an atom-cavity system using three different solvers

# J.R. Johansson and P.D. Nation
# 
# For more information about QuTiP see [http://qutip.org](http://qutip.org)

# In[1]:


get_ipython().run_line_magic('matplotlib', 'inline')
import matplotlib.pyplot as plt


# In[2]:


import numpy as np


# In[3]:


from qutip import *


# ## Model and parameters

# In[4]:


kappa = 2
gamma = 0.2
g = 1 
wc = 0 
w0 = 0 
wl = 0
N = 4
E = 0.5
tlist = np.linspace(0, 10, 200)


# ### mesolve

# In[5]:


def solve(E,kappa,gamma,g,wc,w0,wl,N,tlist):

    ida    = qeye(N)
    idatom = qeye(2)

    # Define cavity field and atomic operators
    a  = tensor(destroy(N),idatom)
    sm = tensor(ida,sigmam())

    # Hamiltonian
    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 = np.sqrt(2*kappa)*a
    C2 = np.sqrt(gamma)*sm
    C1dC1 = C1.dag()*C1
    C2dC2 = C2.dag()*C2

    #intial state
    psi0 = tensor(basis(N,0),basis(2,1))
    rho0 = psi0.dag() * psi0;

    # evolve and calculate expectation values
    output = mesolve(H, psi0, tlist, [C1, C2], [C1dC1, C2dC2, a])  

    return output.expect[0], output.expect[1], output.expect[2]


# In[6]:


count1, count2, infield = solve(E,kappa,gamma,g,wc,w0,wl,N,tlist)


# In[7]:


fig, ax = plt.subplots(figsize=(12,6))
ax.plot(tlist, np.real(count1))
ax.plot(tlist, np.real(count2))
ax.set_xlabel('Time')
ax.set_ylabel('Transmitted Intensity and Spontaneous Emission');


# ### eseries

# In[8]:


def solve(E,kappa,gamma,g,wc,w0,wl,N,tlist):

    # Define cavity field and atomic operators
    a  = tensor(destroy(N),qeye(2))
    sm = tensor(qeye(N),sigmam())

    # Hamiltonian
    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 = np.sqrt(2*kappa)*a
    C2 = np.sqrt(gamma)*sm
    C1dC1 = C1.dag() * C1
    C2dC2 = C2.dag() * C2

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

    # Calculate the Liouvillian
    L = liouvillian(H, [C1, C2])

    # Calculate solution as an exponential series
    rhoES = ode2es(L,rho0);
    
    # Calculate expectation values
    count1  = esval(expect(C1dC1,rhoES),tlist);
    count2  = esval(expect(C2dC2,rhoES),tlist);
    infield = esval(expect(a,rhoES),tlist);

    # alternative
    expt_list = essolve(H, psi0, tlist, [C1, C2], [C1dC1, C2dC2, a]).expect

    return count1, count2, infield, expt_list[0], expt_list[1], expt_list[2]


# In[9]:


count1, count2, infield, count1_2, count2_2, infield_2 = solve(E, kappa, gamma, g, wc, w0, wl, N, tlist)


# In[10]:


fig, ax = plt.subplots(figsize=(12,6))
ax.plot(tlist, np.real(count1), tlist, np.real(count1_2), '.')
ax.plot(tlist, np.real(count2), tlist, np.real(count2_2), '.')
ax.set_xlabel('Time')
ax.set_ylabel('Transmitted Intensity and Spontaneous Emission');


# ### mcsolve

# In[11]:


ntraj = 500 #number of Monte-Carlo trajectories
# 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 = np.sqrt(2*kappa) * a
C2 = np.sqrt(gamma) * sm
C1dC1 = C1.dag() * C1
C2dC2 = C2.dag() * C2
# intial state
psi0=tensor(basis(N,0),basis(2,1))


# In[12]:


avg = mcsolve(H, psi0, tlist, [C1,C2],[C1dC1,C2dC2], ntraj)


# In[13]:


fig, ax = plt.subplots(figsize=(12, 6))
ax.plot(tlist, avg.expect[0],
        tlist, avg.expect[1],'--')
ax.set_xlabel('Time')
ax.set_ylabel('Photocount rates')
ax.legend(('Cavity ouput', 'Spontaneous emission'));


# ## Versions

# In[14]:


from qutip.ipynbtools import version_table

version_table()