%pylab inline

from qutip import *

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)

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

psi0 = tensor(basis(N,0), basis(2,1))

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

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

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

result = mesolve(H, psi0, tlist, c_ops, e_ops)

len(result.expect)

result.expect[0].shape

plot_expectation_values(result);

e_ops = []

result = mesolve(H, psi0, tlist, c_ops, e_ops)

len(result.states)

result.states[0]

# loop over states and calculate the atom entropy
entropy_atom_vec = zeros(len(result.states))

for idx, state in enumerate(result.states):
    entropy_atom_vec[idx] = entropy_linear(ptrace(state, 0))

fig, ax = plt.subplots(1,1)

ax.plot(tlist, entropy_atom_vec)

ax.set_xlabel('time', fontsize=16)
ax.set_ylabel(r'$E(\rho_{\rm atom})$', fontsize=16);

entropy_cavity_vec = zeros(len(tlist))

# calculate the entropy of the cavity in a callback function

def e_ops_func(t, rho):
    # rho is only the data, so do Qobj(rho) to obtain a Qobj representation
    # of rho
    entropy_cavity_vec[e_ops_func.idx] = entropy_linear(ptrace(Qobj(rho), 1))
    e_ops_func.idx += 1
        
e_ops_func.idx = 0

result = mesolve(H, psi0, tlist, c_ops, e_ops_func)

fig, ax = plt.subplots(1,1)

ax.plot(tlist, entropy_atom_vec, label='atom')
ax.plot(tlist, entropy_cavity_vec, label='cavity')

ax.legend()
ax.set_xlabel('time', fontsize=16)
ax.set_ylabel(r'$E$', fontsize=16);

from qutip.ipynbtools import version_table
version_table()