%pylab inline

from qutip import *

def visualize(result):
    fig, ax = subplots(1, 1, figsize=(8,5))
    
    for e in result.expect:
        ax.plot(result.times, e)

    ax.set_xlabel('Time')
    ax.set_ylabel('Occupation probability')
    ax.set_ylim(0, 1)

wc = 1.0  * 2 * pi  # cavity frequency
wa = 1.0  * 2 * pi  # atom frequency
g  = 0.05 * 2 * pi  # coupling strength
kappa = 0.005       # cavity dissipation rate
gamma = 0.05        # atom dissipation rate
N = 15              # number of cavity fock states
tlist = linspace(0,25,100)

psi0 = tensor(basis(N,0), basis(2,1))
a  = tensor(destroy(N), qeye(2))
sm = tensor(qeye(N), destroy(2))

H = wc * a.dag() * a + wa * sm.dag() * sm + g * (a.dag() * sm + a * sm.dag())

c_ops = [sqrt(kappa) * a, sqrt(gamma) * sm]

result = mesolve(H, psi0, tlist, c_ops, [a.dag() * a, sm.dag() * sm])

visualize(result)

H = wc * a.dag() * a + wa * sm.dag() * sm + g * (a.dag() * sm + a * sm.dag())
c_ops = [sqrt(kappa) * a, sqrt(gamma) * sm]

L = liouvillian(H, c_ops)

result = mesolve(L, psi0, tlist, [], [a.dag() * a, sm.dag() * sm])

visualize(result)

H = wc * a.dag() * a + wa * sm.dag() * sm + g * (a.dag() * sm + a * sm.dag())
c_ops = [sqrt(kappa) * a, sqrt(gamma) * sm]

L = liouvillian(None, c_ops)

result = mesolve(H, psi0, tlist, [[L, '1.0']], [a.dag() * a, sm.dag() * sm])

visualize(result)

H = wc * a.dag() * a + wa * sm.dag() * sm + g * (a.dag() * sm + a * sm.dag())
c_ops = [sqrt(kappa) * a, sqrt(gamma) * sm]

L = liouvillian(None, c_ops)

result = mesolve(H, psi0, tlist, [[L, lambda t, args: 1.0]], [a.dag() * a, sm.dag() * sm])

visualize(result)

H = wc * a.dag() * a + wa * sm.dag() * sm + g * (a.dag() * sm + a * sm.dag())
c_ops = [sqrt(kappa) * a, sqrt(gamma) * sm]

L = liouvillian(None, c_ops)

result = mesolve(H, psi0, tlist, L, [a.dag() * a, sm.dag() * sm])

visualize(result)

H = wc * a.dag() * a + wa * sm.dag() * sm + g * (a.dag() * sm + a * sm.dag())
c_ops = [sqrt(kappa) * a, liouvillian(None, [sqrt(gamma) * sm])]

result = mesolve(H, psi0, tlist, c_ops, [a.dag() * a, sm.dag() * sm])

visualize(result)

psi0 = tensor(basis(N,0), basis(2,1))    # start with an excited atom

a  = tensor(destroy(N), qeye(2))
sm = tensor(qeye(N), destroy(2))

H = wc * a.dag() * a + wa * sm.dag() * sm + g * (a.dag() * sm + a * sm.dag())
c_ops = [sqrt(kappa) * a, sqrt(gamma) * sm]

L = sum([spre(c)*spost(c.dag()) - 0.25 * spre(c.dag()*c) - 0.5 * spost(c.dag()*c) for c in c_ops])

result = mesolve(H, psi0, tlist, L, [a.dag() * a, sm.dag() * sm])

visualize(result)

from qutip.ipynbtools import version_table
version_table()