%pylab inline

from qutip import *

import qutip.settings
qutip.settings.debug = False

def plot_correlation(taulist, d_list):
    
    fig, axes = subplots(len(d_list), 2, sharex=True, sharey=False, figsize=(12, 4*len(d_list)))
    
    for idx, d in enumerate(d_list):
        axes[idx,0].plot(taulist, real(d), label="data %d" % idx)
        axes[idx,1].plot(taulist, imag(d), label="data %d" % idx)
        
    axes[0,0].set_title("real")
    axes[0,1].set_title("imag")
        
    axes[idx,0].set_xlabel(r'$\tau$')
    axes[idx,1].set_xlabel(r'$\tau$');    

def plot_correlation_2t(t1_list, t2_list, d_list):
    
    fig, axes = subplots(len(d_list), 2, sharex=True, figsize=(8, 4*len(d_list)))
    
    T1, T2 = meshgrid(t1_list, t2_list)
    
    for idx, d in enumerate(d_list):
        axes[idx,0].pcolor(T1, T2, real(d), cmap=cm.RdBu, vmin=-abs(d).max(), vmax=abs(d).max())
        axes[idx,0].autoscale(tight=True)
        axes[idx,1].pcolor(T1, T2, imag(d), cmap=cm.RdBu, vmin=-abs(d).max(), vmax=abs(d).max())
        axes[idx,1].autoscale(tight=True)

N = 20
w0 = 2 * pi
G1 = 0.75
n_th = 2.00
taulist = linspace(0, 10.0, 500)

# operators and hamiltonian 
a = destroy(N)
H = w0 * a.dag() * a
c_ops = [sqrt(G1*(1+n_th)) * a, sqrt(G1*n_th) * a.dag()]

# initial state
rho0 = thermal_dm(N, 4.0)
rho0 = coherent_dm(N, sqrt(4.0))

A = a.dag()
B = a
n = mesolve(H, rho0, taulist, c_ops, [A * B]).expect[0]
corr = correlation(H, rho0, None, taulist, c_ops, A, B)
plot_correlation(taulist, [n, corr])

corr = correlation(H, rho0, None, taulist, c_ops, A, B, reverse=True)
plot_correlation(taulist, [n, corr])

A = a + a.dag()
B = A
n = mesolve(H, rho0, taulist, c_ops, [A * B]).expect[0]
corr1 = correlation_ss(H, taulist, c_ops, A, B, rho0=rho0)
corr2 = correlation_ss(H, taulist, c_ops, A, B, rho0=rho0, reverse=True)
plot_correlation(taulist, [n, corr1, corr2])

A = -1j*(a + a.dag())
B = A.dag()
n = mesolve(H, rho0, taulist, c_ops, [A * B]).expect[0]
corr1 = correlation_ss(H, taulist, c_ops, A, B, rho0=rho0)
corr2 = correlation_ss(H, taulist, c_ops, A, B, rho0=rho0, reverse=True)
plot_correlation(taulist, [n, corr1, corr2])

A = 1j * a
B = a.dag()
n = mesolve(H, rho0, taulist, c_ops, [A * B]).expect[0]
corr1 = correlation_ss(H, taulist, c_ops, A, B, rho0=rho0)
corr2 = correlation_ss(H, taulist, c_ops, A, B, rho0=rho0, reverse=True)
plot_correlation(taulist, [n, corr1, corr2])

A = a.dag()
B = a

A = 1j * a
B = a.dag()

n = mesolve(H, rho0, taulist, c_ops, [A * B]).expect[0]

corr1 = correlation(H, rho0, None, taulist, c_ops, A, B, solver="me", reverse=True)
corr2 = correlation(H, rho0, None, taulist, c_ops, A, B, solver="es", reverse=True)

plot_correlation(taulist, [corr1, corr2])

A = a.dag()
B = a
psi0 = coherent(N, 3.0)

n = mesolve(H, rho0, taulist, c_ops, [A * B]).expect[0]

corr1 = correlation(H, psi0, None, taulist, c_ops, A, B, solver="me", reverse=False)
corr2 = correlation(H, psi0, None, taulist, c_ops, A, B, solver="mc", reverse=False)

plot_correlation(taulist, [corr1, corr2])

# use a small delay time range for 2t correlations
taulist = linspace(0, 2.5, 50)

A = a.dag()
B = a
psi0 = coherent(N, 3.0)

corr1 = correlation(H, psi0, taulist, taulist, c_ops, A, B, solver="me")
corr2 = correlation(H, psi0, taulist, taulist, c_ops, A, B, solver="me", reverse=True)

plot_correlation_2t(taulist, taulist, [corr1, corr2])

corr2 = correlation(H, psi0, taulist, taulist, c_ops, A, B, solver="es")
plot_correlation_2t(taulist, taulist, [corr1, corr2])

A = a.dag()
B = a

taulist = linspace(0, 2.5, 50)

corr1 = correlation(H, None, taulist, taulist, c_ops, A, B, solver="me")
corr2 = correlation(H, None, taulist, taulist, c_ops, A, B, solver="es")

plot_correlation_2t(taulist, taulist, [corr1, corr2])

N = 20
w0 = 2 * pi
kappa = 0.75
n_th = 2.00
taulist = linspace(0, 10.0, 500)
a = destroy(N)
H = w0 * a.dag() * a
c_ops = [sqrt(kappa*(1+n_th)) * a, sqrt(kappa*n_th) * a.dag()]

# steady state initial condition
g1, G1 = coherence_function_g1(H, None, taulist, c_ops, a)
g2, G2 = coherence_function_g2(H, None, taulist, c_ops, a)
plot_correlation(taulist, [g1, G1, g2, G2])

rho0 = coherent_dm(N, sqrt(4.0))
g1, G1 = coherence_function_g1(H, rho0, taulist, c_ops, a)
g2, G2 = coherence_function_g2(H, rho0, taulist, c_ops, a)
plot_correlation(taulist, [g1, g2])

rho0 = thermal_dm(N, 4.0)
g1, G1 = coherence_function_g1(H, rho0, taulist, c_ops, a)
g2, G2 = coherence_function_g2(H, rho0, taulist, c_ops, a)
plot_correlation(taulist, [g1, g2])

from qutip.ipynbtools import version_table
version_table()