%pylab inline

from qutip import *

from qutip.ipynbtools import HTMLProgressBar
from qutip.gui.progressbar import TextProgressBar

N = 10

w = 1.0 * 2 * pi
g = 0.1 * 2 * pi
kappa = 0.005
W21 = 0.01
W12 = 0.05

tlist = linspace(0, 250, 250)

# cavity operators
a  = tensor(destroy(N), identity(2))
ad = tensor(create(N), identity(2))

# atomic operators
sz = tensor(identity(N), sigmaz())
sx = tensor(identity(N), sigmax())
sm = tensor(identity(N), destroy(2))
sp = tensor(identity(N), create(2))

#psi0 = tensor(fock(N, 0), fock(2, 0)) # start with the vacuum + ground state
psi0 = tensor(coherent(N, 0.5), fock(2, 0)) # start a small coherent sate + ground state

theta = 0.0

H = w * ad * a - 0.5 * w * (cos(2*theta) * sz + sin(2*theta) * sx)

H

e_ops = [ad * a, sm.dag() * sm, sm, sp]

c_ops = [sqrt(2*kappa) * a, sqrt(W21) * sm, sqrt(W12) * sp]

result = mesolve(H, psi0, tlist, c_ops, e_ops, options=Odeoptions(store_final_state=True)) 

fig, axes = subplots(2,1)

axes[0].plot(result.times, result.expect[0], label=r'$a^\dagger a$')
axes[0].plot(result.times, result.expect[1], label=r'$\sigma_+\sigma_-$')
axes[0].set_ylim(-0.1, 1.1)
axes[0].legend();

axes[1].plot(result.times, abs(result.expect[2]), label=r'$\sigma_-$')
axes[1].plot(result.times, abs(result.expect[3]), label=r'$\sigma_+$')
axes[1].set_ylim(-0.1, 1.1)
axes[1].legend();

wigner_fock_distribution(ptrace(result.final_state, 0));

d = (W12 - W21) / (W12 + W21)

d

gamma = 0.5 * (W12 + W21)

gamma

C = d * (g/(2*pi)) ** 2 / (gamma * kappa)

C

def Heff(t, rho, args):
    """ not an optimized implementation ... """
    q_rho = Qobj(vec2mat(rho))
    Heff = 1.0j * g * (expect(sm, q_rho) * ad - expect(sp, q_rho) * a) \
         + 1.0j * g * (expect(ad, q_rho) * sm - expect(a, q_rho) * sp)
    
    return args[0] + liouvillian_fast(Heff, []).data

sop_a_data = liouvillian_fast(a, []).data
sop_ad_data = liouvillian_fast(ad, []).data
sop_sm_data = liouvillian_fast(sm, []).data
sop_sp_data = liouvillian_fast(sp, []).data

sop_a_L_data = spre(a).data
sop_ad_L_data = spre(ad).data
sop_sm_L_data = spre(sm).data
sop_sp_L_data = spre(sp).data

def Heff_fast(t, rho, args):
    """ a somewhat more optimized version ... """
    
    phi_sm = expect_rho_vec1d(sop_sm_L_data, rho)
    phi_a = expect_rho_vec1d(sop_a_L_data, rho)
    
    Heff = 1.0j * g * (phi_sm * sop_ad_data  - conjugate(phi_sm) * sop_a_data) \
         + 1.0j * g * (conjugate(phi_a) * sop_sm_data - phi_a * sop_sp_data)
    
    return args[0] + Heff

result = mesolve(Heff_fast, psi0, tlist, c_ops, e_ops, args=[H],
                 options=Odeoptions(rhs_with_state=True, store_final_state=True),
                 progress_bar=TextProgressBar())

fig, ax = subplots()

ax.plot(result.times, result.expect[0], label=r'$a^\dagger a$')
ax.plot(result.times, result.expect[1], label=r'$\sigma_+\sigma_-$')
ax.set_ylim(-0.1, 3)
ax.legend();

wigner_fock_distribution(ptrace(result.final_state, 0));

def a_coeff(t, rho, args):
    return - 1.0j * g * expect_rho_vec1d(sop_sp_L_data, rho)

def ad_coeff(t, rho, args):
    return + 1.0j * g * expect_rho_vec1d(sop_sm_L_data, rho)

def sm_coeff(t, rho, args):
    return + 1.0j * g * expect_rho_vec1d(sop_ad_L_data, rho)

def sp_coeff(t, rho, args):
    return - 1.0j * g * expect_rho_vec1d(sop_a_L_data, rho)

result = mesolve([H, [a, a_coeff], [ad, ad_coeff], [sm, sm_coeff], [sp, sp_coeff]],
                 psi0, tlist, c_ops, e_ops, args={}, 
                 options=Odeoptions(rhs_with_state=True, store_final_state=True),
                 progress_bar=TextProgressBar())

fig, ax = subplots()

ax.plot(result.times, result.expect[0], label=r'$a^\dagger a$')
ax.plot(result.times, result.expect[1], label=r'$\sigma_+\sigma_-$')
ax.set_ylim(-0.1, 3)
ax.legend();

tlist = linspace(0, 50, 150)

H_data = H.data
a_data = a.data
ad_data = ad.data
sm_data = sm.data
sp_data = sp.data

def Heff_func(t, psi, args):
    
    phi_sm = expect_psi(sm_data, psi[:, newaxis])
    phi_a  = expect_psi(a_data, psi[:, newaxis])
    phi_sp = expect_psi(sp_data, psi[:, newaxis])
    phi_ad = expect_psi(ad_data, psi[:, newaxis])
    
    Heff = 1.0j * g * (phi_sm * ad_data - phi_sp * a_data) \
         + 1.0j * g * (phi_ad * sm_data - phi_a * sp_data)
    
    return H_data + Heff

result = sesolve(Heff_func, psi0, tlist, e_ops,
                 options=Odeoptions(rhs_with_state=True, store_final_state=True),
                 progress_bar=TextProgressBar())

fig, ax = subplots()

ax.plot(result.times, real(result.expect[0]), label=r'$a^\dagger a$')
ax.plot(result.times, real(result.expect[1]), label=r'$\sigma_+\sigma_-$')
ax.set_ylim(-0.1, 1)
ax.legend();

def a_coeff(t, psi, args):
    return - 1.0j * g * expect_psi(sp_data, psi[:, newaxis])

def ad_coeff(t, psi, args):
    return + 1.0j * g * expect_psi(sm_data, psi[:, newaxis])

def sm_coeff(t, psi, args):
    return + 1.0j * g * expect_psi(ad_data, psi[:, newaxis])

def sp_coeff(t, psi, args):
    return - 1.0j * g * expect_psi(a_data, psi[:, newaxis])

result = sesolve([H, [a, a_coeff], [a.dag(), ad_coeff], [sm, sm_coeff], [sp, sp_coeff]],
                 psi0, tlist, e_ops, 
                 options=Odeoptions(rhs_with_state=True, store_final_state=True),
                 progress_bar=TextProgressBar())

fig, ax = subplots()

ax.plot(result.times, real(result.expect[0]), label=r'$a^\dagger a$')
ax.plot(result.times, real(result.expect[1]), label=r'$\sigma_+\sigma_-$')
ax.set_ylim(-0.1, 1)
ax.legend();

tlist = linspace(0, 250, 250)

result = mcsolve(Heff_func, psi0, tlist, c_ops, e_ops, ntraj=64,
                 options=Odeoptions(rhs_with_state=True, store_final_state=True, gui=False, num_cpus=4),
                 progress_bar=TextProgressBar())

fig, ax = subplots()

ax.plot(result.times, real(result.expect[0]), label=r'$a^\dagger a$')
ax.plot(result.times, real(result.expect[1]), label=r'$\sigma_+\sigma_-$')
ax.set_ylim(-0.1, 3)
ax.legend();

result = mcsolve([H, [a, a_coeff], [ad, ad_coeff], [sm, sm_coeff], [sp, sp_coeff]],
                 psi0, tlist, c_ops, e_ops, args={}, ntraj=512,
                 options=Odeoptions(rhs_with_state=True, store_final_state=True, gui=False, num_cpus=4),
                 progress_bar=TextProgressBar())

fig, ax = subplots()

ax.plot(result.times, real(result.expect[0]), label=r'$a^\dagger a$')
ax.plot(result.times, real(result.expect[1]), label=r'$\sigma_+\sigma_-$')
ax.set_ylim(-0.1, 3)
ax.legend();

from qutip.ipynbtools import version_table

version_table()