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

from scipy import *

from qutip import *

N = 10
gamma = 1
eta = 0.9

def solve(N, gamma, kappa, eta):
    
    E = kappa * 0.25
    
    # create operators
    a = tensor(destroy(N), identity(2))
    sm = tensor(identity(N), destroy(2))

    # Hamiltonian
    H = 0.5j * (E * a.dag() ** 2 - conjugate(E) * a ** 2)

    # master equation superoperators
    L0 = liouvillian(H, [sqrt(kappa) * a, sqrt(gamma) * sm])
    L1 = - sqrt(kappa * gamma * eta) * (
        spre(sm.dag() * a) - spre(a) * spost(sm.dag()) + 
        spost(a.dag() * sm) - spre(sm) * spost(a.dag()))
    
    L = L0 + L1
    
    # steady state
    rhoss = steadystate(L)
    
    # correlation function and spectrum
    taulist = linspace(0, 500, 2500)
    c = correlation_2op_1t(L, rhoss, taulist, [], sm.dag(), sm)
    w, S = spectrum_correlation_fft(taulist, c)
    
    ww = hstack([fliplr(-array([w])).squeeze(), w])
    SS = hstack([fliplr(array([S])).squeeze(), S])

    return rhoss, ww, SS

rhoss2, w2, S2 = solve(N, gamma, 2, eta)
rhoss4, w4, S4 = solve(N, gamma, 4, eta)
rhoss8, w8, S8 = solve(N, gamma, 8, eta)

wigner_fock_distribution(rhoss2.ptrace(0));
wigner_fock_distribution(rhoss4.ptrace(0));
wigner_fock_distribution(rhoss8.ptrace(0));

fig, ax = plt.subplots()
ax.plot(w2, S2 / S2.max(), label=r'$\kappa = 2$')
ax.plot(w4, S4 / S4.max(), label=r'$\kappa = 4$')
ax.plot(w8, S8 / S8.max(), label=r'$\kappa = 8$')
ax.plot(w8, 0.25/((0.5 * gamma)**2 + w8**2), 'k:', label=r'Lorentian')
ax.legend()
ax.set_ylabel(r'Flouresence spectrum', fontsize=16)
ax.set_xlabel(r'$\omega$', fontsize=18)
ax.set_xlim(-2, 2);

e1 = 0.5
e2 = 0.5

gamma1 = 2
gamma2 = 2

E = 2 / sqrt(e1 * gamma1)

sm1 = tensor(destroy(2), identity(2))
sp1 = sm1.dag()
sm2 = tensor(identity(2), destroy(2))
sp2 = sm2.dag()

H = -1j * sqrt(e1 * gamma1) * (E * sp1 - conjugate(E) * sm1)

L0 = liouvillian(H, [sqrt(gamma1) * sm1, sqrt(gamma2) * sm2])

L1 = - sqrt((1 - e1) * (1 - e2) * gamma1 * gamma2) * \
        (spre(sp2 * sm1) - spre(sm1) * spost(sp2) + 
         spost(sp1 * sm2) - spre(sm2) * spost(sp1))

L = L0 + L1

# steady state
rhoss = steadystate(L)

# correlation function and spectrum
taulist = linspace(0, 4, 250)

G2_11 = correlation_4op_1t(L, rhoss, taulist, [], sp1, sp1, sm1, sm1)
g2_11 = G2_11 / (expect(sp1*sm1, rhoss) * expect(sp1*sm1, rhoss))

G2_22 = correlation_4op_1t(L, rhoss, taulist, [], sp2, sp2, sm2, sm2)
g2_22 = G2_22 / (expect(sp2*sm2, rhoss) * expect(sp2*sm2, rhoss))

G2_12 = correlation_4op_1t(L, rhoss, taulist, [], sp2, sp1, sm1, sm2)
g2_12 = G2_12 / (expect(sp1*sm1, rhoss) * expect(sp2*sm2, rhoss))

G2_21 = correlation_4op_1t(L, rhoss, taulist, [], sp1, sp2, sm2, sm1)
g2_21 = G2_21 / (expect(sp2*sm2, rhoss) * expect(sp1*sm1, rhoss))

fig, ax = plt.subplots()

ax.plot(taulist, g2_11, label=r'$g^{(2)}_{11}(\tau)$')
ax.plot(taulist, g2_22, label=r'$g^{(2)}_{22}(\tau)$')
ax.plot(taulist, g2_12, label=r'$g^{(2)}_{12}(\tau)$')
ax.plot(taulist, g2_21, label=r'$g^{(2)}_{21}(\tau)$')

ax.legend(loc=4)
ax.set_xlabel(r'$\tau$');

N = 10

e1 = 0.5
e2 = 0.5
ek = 0.5

n_th = 1
kappa = 0.1
gamma1 = 1
gamma2 = 1

E = 0.025

taulist = linspace(0, 5, 250)

a   = tensor(destroy(N), identity(2), identity(2))
sm1 = tensor(identity(N), destroy(2), identity(2))
sp1 = sm1.dag()
sm2 = tensor(identity(N), identity(2), destroy(2))
sp2 = sm2.dag()

def solve(ek, e1, e2, gamma1, gamma2, kappa, n_th, E):
    
    eta1 = (1 - ek) * e1
    eta2 = (1 - e1) * (1 - e2)
    
    H = 1j * E * (a - a.dag())
    
    L0 = liouvillian(H, [sqrt(kappa * (1 + n_th)) * a, sqrt(kappa * n_th) * a.dag(),
                     sqrt(gamma1) * sm1, sqrt(gamma2) * sm2])
    
    L1 = - sqrt(2 * kappa * eta1 * gamma1) * \
            (spre(sp1 * a) - spre(a) * spost(sp1) + 
             spost(a.dag() * sm1) - spre(sm1) * spost(a.dag())) + \
         - sqrt(eta2 * gamma1 * gamma2) * \
            (spre(sp2 * sm1) - spre(sm1) * spost(sp2) + 
             spost(sp1 * sm2) - spre(sm2) * spost(sp1))
    
    L = L0 + L1
    
    rhoss = steadystate(L)
    
    G2_11 = correlation_4op_1t(L, rhoss, taulist, [], sp1, sp1, sm1, sm1)
    g2_11 = G2_11 / (expect(sp1*sm1, rhoss) * expect(sp1*sm1, rhoss))

    G2_22 = correlation_4op_1t(L, rhoss, taulist, [], sp2, sp2, sm2, sm2)
    g2_22 = G2_22 / (expect(sp2*sm2, rhoss) * expect(sp2*sm2, rhoss))

    G2_12 = correlation_4op_1t(L, rhoss, taulist, [], sp2, sp1, sm1, sm2)
    g2_12 = G2_12 / (expect(sp1*sm1, rhoss) * expect(sp2*sm2, rhoss))
    
    G2_21 = correlation_4op_1t(L, rhoss, taulist, [], sp1, sp2, sm2, sm1)
    g2_21 = G2_21 / (expect(sp2*sm2, rhoss) * expect(sp1*sm1, rhoss))
    
    return rhoss, g2_11, g2_12, g2_21, g2_22

# thermal
rhoss_t, g2_11_t, g2_12_t, g2_21_t, g2_22_t = solve(ek, e1, e2, gamma1, gamma2, 
                                                    kappa, n_th, 0.0)

# visualize the cavity state
wigner_fock_distribution(rhoss_t.ptrace(0));

fig, ax = plt.subplots(figsize=(8,4))

ax.plot(taulist, g2_11_t, label=r'$g^{(2)}_{11}(\tau)$')
ax.plot(taulist, g2_22_t, label=r'$g^{(2)}_{22}(\tau)$')
ax.plot(taulist, g2_12_t, label=r'$g^{(2)}_{12}(\tau)$')
ax.plot(taulist, g2_21_t, label=r'$g^{(2)}_{21}(\tau)$')

ax.legend(loc=4)
ax.set_xlabel(r'$\tau$', fontsize=16);

from qutip.ipynbtools import version_table; version_table()