%matplotlib inline
import matplotlib.pyplot as plt
import numpy as np

from qutip import *

Omega = 1.0 * 2 * pi

gamma0 = 0.05
w_th = 0.0 
N = n_thermal(Omega, w_th)

def system_spec(Omega, gamma0, N):
    HL = -0.5 * Omega * (sigmap() + sigmam())
    c_ops = [sqrt(gamma0 * (N + 1)) * sigmam(), sqrt(gamma0 * N) * sigmap()]
    return HL, c_ops

HL, c_ops = system_spec(Omega, gamma0, N)

e_ops = [sigmax(), sigmay(), sigmaz(), sigmam(), sigmap(), num(2)]

psi0 = basis(2, 0)

tlist = np.linspace(0, 20/(2*pi), 200)
result = mesolve(HL, psi0, tlist, c_ops, e_ops)

fig, axes = plt.subplots(2, 1, figsize=(12, 6), sharex=True)

axes[0].plot(result.times, result.expect[0], 'r', label=r'$\langle\sigma_x\rangle$')
axes[0].plot(result.times, result.expect[1], 'g', label=r'$\langle\sigma_y\rangle$')
axes[0].plot(result.times, result.expect[2], 'b', label=r'$\langle\sigma_z\rangle$')
axes[0].legend()
axes[0].set_ylim(-1, 1);


axes[1].plot(result.times, result.expect[5], 'b', label=r'$P_e$')

#axes[1].set_ylabel(r'$\langle\sigma_z\rangle$', fontsize=16)
axes[1].set_xlabel("time", fontsize=16)
axes[1].legend()
axes[1].set_ylim(0, 1);

fig, ax = plt.subplots(1, 1, figsize=(12, 6), sharex=True)


for idx, gamma0 in enumerate([0.1 * Omega, 0.5 * Omega, 1.0 * Omega]):

    HL, c_ops = system_spec(Omega, gamma0, N)
    result = mesolve(HL, psi0, tlist, c_ops, e_ops)

    ax.plot(result.times, result.expect[5], 'b', label=r'$\langle\sigma_z\rangle$')

ax.set_ylim(0, 1);

fig, ax = plt.subplots(1, 1, figsize=(12, 6), sharex=True)


for idx, gamma0 in enumerate([0.1 * Omega, 0.5 * Omega, 1.0 * Omega]):

    HL, c_ops = system_spec(Omega, gamma0, N)
    result = mesolve(HL, psi0, tlist, c_ops, e_ops)

    ax.plot(result.times, imag(result.expect[4]), label=r'im $\langle\sigma_+\rangle$')

ax.set_ylim(-.5, 0.5);

fig, axes = plt.subplots(1, 2, figsize=(12, 6))

taulist = np.linspace(0, 100, 10000)

for idx, gamma0 in enumerate([2 * Omega, 0.5 * Omega, 0.25 * Omega]):

    HL, c_ops = system_spec(Omega, gamma0, N)
    corr_vec = correlation_2op_1t(HL, None, taulist, c_ops, sigmap(), sigmam())
    w, S = spectrum_correlation_fft(taulist, corr_vec)

    axes[0].plot(taulist, corr_vec, label=r'$<\sigma_+(\tau)\sigma_-(0)>$')
    axes[1].plot(-w / (gamma0), S, 'b', label=r'$S(\omega)$')
    axes[1].plot( w / (gamma0), S, 'b', label=r'$S(\omega)$')

axes[0].set_xlim(0, 10)
axes[1].set_xlim(-5, 5);

from qutip.ipynbtools import version_table; version_table()