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

from qutip import *

w0 = 1.0 * 2 * pi
gamma0 = 0.05

# the temperature of the environment in frequency units
w_th = 0.0 * 2 * pi

# the number of average excitations in the environment mode w0 at temperture w_th
Nth = n_thermal(w0, w_th)

Nth

# squeezing parameter for the environment
r = 1.0
theta = 0.1 * pi

N = Nth * (cosh(r) ** 2 + sinh(r) ** 2) + sinh(r) ** 2

N

M = - cosh(r) * sinh(r) * exp(-1j * theta) * (2 * Nth + 1)

M

# Check, should be zero according to Eq. 3.261 in Breuer and Petruccione
abs(M)**2 - (N * (N + 1) - Nth * (Nth + 1))

sm = sigmam()
sp = sigmap()

H = - 0.5 * w0 * sigmaz()  # by adding the hamiltonian here, so we move back to the schrodinger picture

c_ops = [sqrt(gamma0 * (N + 1)) * sm, sqrt(gamma0 * N) * sp]

L0 = liouvillian(H, c_ops)

L0

Lsq = - gamma0 * M * spre(sp) * spost(sp) - gamma0 * conjugate(M) * spre(sm) * spost(sm)

Lsq

L = L0 + Lsq

L

tlist = np.linspace(0, 50, 1000)

# start in the qubit superposition state
psi0 = (2j * basis(2, 0) + 1 * basis(2, 1)).unit()

e_ops = [sigmax(), sigmay(), sigmaz()]

result1 = mesolve(L, psi0, tlist, [], e_ops)

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

ax.plot(result1.times, result1.expect[0], 'r', label=r'$\langle\sigma_x\rangle$')
ax.plot(result1.times, result1.expect[1], 'g', label=r'$\langle\sigma_y\rangle$')
ax.plot(result1.times, result1.expect[2], 'b', label=r'$\langle\sigma_z\rangle$')

sz_ss_analytical = - 1 / (2 * N + 1)
ax.plot(result1.times, sz_ss_analytical * np.ones(shape(result1.times)), 'k--', 
        label=r'$\langle\sigma_z\rangle_s$ analytical')


ax.set_ylabel(r'$\langle\sigma_z\rangle$', fontsize=16)
ax.set_xlabel("time", fontsize=16)
ax.legend()
ax.set_ylim(-1, 1);

b = Bloch()
b.add_points(result1.expect, meth='l')
b.show()

c_ops = [sqrt(gamma0) * (sm * cosh(r) + sp * sinh(r) * exp(1j*theta))]

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

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

ax.plot(result2.times, result2.expect[0], 'r', label=r'$\langle\sigma_x\rangle$')
ax.plot(result2.times, result2.expect[1], 'g', label=r'$\langle\sigma_y\rangle$')
ax.plot(result2.times, result2.expect[2], 'b', label=r'$\langle\sigma_z\rangle$')

sz_ss_analytical = - 1 / (2 * N + 1)
ax.plot(result2.times, sz_ss_analytical * np.ones(shape(result2.times)), 'k--', 
        label=r'$\langle\sigma_z\rangle_s$ analytical')


ax.set_ylabel(r'$\langle\sigma_z\rangle$', fontsize=16)
ax.set_xlabel("time", fontsize=16)
ax.legend()
ax.set_ylim(-1, 1);

fig, axes = plt.subplots(3, 1, sharex=True, figsize=(12, 9))

axes[0].plot(result1.times, result1.expect[0], 'r', label=r'$\langle\sigma_x\rangle$ - me')
axes[0].plot(result2.times, result2.expect[0], 'b--', label=r'$\langle\sigma_x\rangle$ - me lindblad')
axes[0].legend()
axes[0].set_ylim(-1, 1);

axes[1].plot(result1.times, result1.expect[1], 'r', label=r'$\langle\sigma_y\rangle$ - me')
axes[1].plot(result2.times, result2.expect[1], 'b--', label=r'$\langle\sigma_y\rangle$ - me lindblad')
axes[1].legend()
axes[1].set_ylim(-1, 1);

axes[2].plot(result1.times, result1.expect[2], 'r', label=r'$\langle\sigma_y\rangle$ - me')
axes[2].plot(result2.times, result2.expect[2], 'b--', label=r'$\langle\sigma_y\rangle$ - me lindblad')
axes[2].legend()
axes[2].set_ylim(-1, 1);
axes[2].set_xlabel("time", fontsize=16);

# for vacuum: 
r = 0
theta = 0.0
c_ops = [sqrt(gamma0) * (sm * cosh(r) + sp * sinh(r) * exp(1j*theta))]

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

# for squeezed vacuum: 
r = 1.0
theta = 0.0
c_ops = [sqrt(gamma0) * (sm * cosh(r) + sp * sinh(r) * exp(1j*theta))]

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

fig, axes = plt.subplots(3, 1, sharex=True, figsize=(12, 9))

axes[0].plot(result1.times, result1.expect[0], 'r', label=r'$\langle\sigma_x\rangle$ - vacuum')
axes[0].plot(result2.times, result2.expect[0], 'b', label=r'$\langle\sigma_x\rangle$ - squeezed vacuum')
axes[0].legend()
axes[0].set_ylim(-1, 1);

axes[1].plot(result1.times, result1.expect[1], 'r', label=r'$\langle\sigma_y\rangle$ - vacuum')
axes[1].plot(result2.times, result2.expect[1], 'b', label=r'$\langle\sigma_y\rangle$ - squeezed vacuum')
axes[1].legend()
axes[1].set_ylim(-1, 1);

axes[2].plot(result1.times, result1.expect[2], 'r', label=r'$\langle\sigma_y\rangle$ - vacuum')
axes[2].plot(result2.times, result2.expect[2], 'b', label=r'$\langle\sigma_y\rangle$ - squeezed vacuum')
axes[2].legend()
axes[2].set_ylim(-1, 1);
axes[2].set_xlabel("time", fontsize=16);

from qutip.ipynbtools import version_table; version_table()