%pylab inline

from sympy import *
import numpy

init_printing()

TP = 2 * numpy.pi

ai = Symbol('a_in')
ao = Symbol('a_out')
bi = Symbol('b_in')
bo = Symbol('b_out')

# fundamental parameters
wR = Symbol('omega_R', real=True) # probe field frequency
w0 = Symbol('omega_0', real=True) # cavity frequency

k1 = Symbol('kappa_1', real=True)
k2 = Symbol('kappa_2', real=True)
ki = Symbol('kappa_i', real=True)
sigma_e = Symbol('sigma_e', real=True)

gamma  = Symbol('gamma', real=True)

tc = Symbol('t_c', real=True)
de = Symbol('delta_epsilon', real=True)
e  = Symbol('epsilon', real=True)
gc = Symbol('g_c', real=True)

X, Y = symbols("X, Y")

param_fundamental     = {tc: 3.5 * TP, 
                         w0: 6.1948 * TP,
                         gc: 0.030 * TP, 
                         k1: 0.015, 
                         k2: 0.015, 
                         ki: 0.00, 
                         sigma_e: 5.1 * TP, 
                         gamma: 0.066} # 15ns

param_fundamental_all = {wR: 6.1948 * TP,
                         e: 0}

param_fundamental_all.update(param_fundamental)

# derived parameters
delta0 = Symbol('Delta_0', real=True)
delta  = Symbol('Delta',   real=True)
W      = Symbol('Omega', real=True)

geff   = Symbol('g_eff', real=True)
chi    = Symbol('chi', real=True)
k      = Symbol('kappa', real=True)

ao = -I * sqrt(k1*k2) * bi / (delta0 - I*k/2 + geff * chi)

ao

# the transmission coefficient
T = ao.subs(ai,0) / (bi)

T

W_num = sqrt((2*tc)**2 + e**2)

W_num

# probe field detuning from atom
delta_num = W - wR

delta_num

# cavity frequency detuning from probe frequency
delta0_num = w0 - wR

delta0_num

geff_num = gc * 2 * tc / W

geff_num

chi_num = geff / (-delta + I * gamma / 2)

chi_num

delta_phi = - arg(T) * 180/pi

param_derived = {W: W_num,
                 geff: geff_num,
                 chi:  chi_num,
                 delta0: delta0_num,
                 delta: delta_num,
                 k: k1 + k2 + ki}

chi.subs(param_derived).subs(param_derived)

T.subs(param_derived).subs(param_derived)

delta_phi.subs(param_derived).subs(param_derived)

(2*tc > wR).subs(param_derived).subs(param_fundamental_all)

(delta > gamma/2).subs(param_derived).subs(param_fundamental_all)

(2*tc/wR).subs(param_derived).subs(param_fundamental_all).evalf()

(delta/(gamma/2)).subs(param_derived).subs(param_fundamental_all).evalf()

def GHz_to_meV(w_GHz):
    # 1 GHz = 4.1357e-6 eV = 4.1357e-3 meV
    w_meV = w_GHz * 4.1357e-3
    return w_meV

def meV_to_GHz(w_meV):
    # 1 meV = 1.0/4.1357e-3 GHz
    w_GHz = w_meV / 4.1357e-3
    return w_GHz

wR_vec = linspace(6.190, 6.20, 1000) * TP

T0_lambda = lambdify(X, T.subs(param_derived).subs(param_derived).subs({e: 0.0, wR: X}).subs(param_fundamental), 'numpy')
T1_lambda = lambdify(X, T.subs(param_derived).subs(param_derived).subs({e: 3.0 * TP, wR: X}).subs(param_fundamental), 'numpy')

T0_vec = T0_lambda(wR_vec)
T1_vec = T1_lambda(wR_vec)

fig, axes = subplots(1,2, figsize=(12,4))

# according to phase definition in Petersson et al
axes[0].plot(wR_vec/TP, -numpy.angle(T0_vec) * 180/pi, color="green")
axes[0].plot(wR_vec/TP, -numpy.angle(T1_vec) * 180/pi, 'b:')
axes[0].set_ylabel("phase", fontsize=18)
axes[0].set_xlabel(r'$\omega_R$', fontsize=18)

axes[1].plot(wR_vec/TP, numpy.abs(T0_vec)**2, color="green")
axes[1].plot(wR_vec/TP, numpy.abs(T1_vec)**2, 'b:')
axes[1].set_ylabel(r'$|T|^2$', fontsize=18)
axes[1].set_xlabel(r'$\omega_R$', fontsize=18);

fig.tight_layout()

e_vec = linspace(-20.0, 20.0, 500) * TP

T_lambda = lambdify(X, T.subs(param_derived).subs(param_derived).subs({tc: (3.5/2) * TP, e: X}).subs(param_fundamental_all), 'numpy')

T_vec = T_lambda(e_vec)

fig, axes = subplots(1,2, figsize=(12,4))

# according to phase definition in Petersson et al
axes[0].plot(e_vec/TP, -numpy.angle(T_vec) * 180/pi)
axes[1].plot(e_vec/TP, numpy.abs(T_vec)**2)

axes[0].set_ylabel("phase", fontsize=18)
axes[0].set_xlabel(r'$\epsilon$ (GHz)', fontsize=18)
axes[1].set_ylabel(r'$|T|^2$', fontsize=18)
axes[1].set_xlabel(r'$\epsilon$ (GHz)', fontsize=18);

fig.tight_layout()

tc_vec = array([1.8, 5.1, 5.8, 6.2, 7.0]) / 2.0 * TP

e_vec = np.linspace(meV_to_GHz(-0.12), meV_to_GHz(0.12), 1000) * TP

W_lambda = lambdify((X, Y), W.subs(param_derived).subs(param_derived).subs(e, X).subs(tc, Y), 'numpy')
W_mat = [W_lambda(e_vec, tc_num) for tc_num in tc_vec]

tc_num = tc_vec[0]
chi_lambda = lambdify((X, Y), chi.subs(param_derived).subs(param_derived).subs({e: X, tc: Y}).subs(param_fundamental_all), 'numpy')
chi_vec = chi_lambda(e_vec, tc_num)

w0_num = (w0 / (2*pi)).subs(param_fundamental_all).evalf()

fig, axes = subplots(1,2, figsize=(10, 5))

axes[0].plot(GHz_to_meV(e_vec/TP), +0.5 * real(W_mat[0])/TP)
axes[0].plot(GHz_to_meV(e_vec/TP), -0.5 * real(W_mat[0])/TP)
axes[0].set_ylabel(r'$E_\pm$', fontsize=18)
axes[0].set_xlabel(r'$\epsilon$ (meV)', fontsize=18)

axes[1].plot(GHz_to_meV(e_vec/TP), real(chi_vec))
axes[1].set_ylabel(r'$\chi$', fontsize=18)
axes[1].set_xlabel(r'$\epsilon$ (meV)', fontsize=18)

fig.tight_layout()

fig, axes = subplots(1,2, figsize=(10,5))

# GHz
for W_vec in W_mat:
    axes[0].plot(e_vec/TP, real(W_vec)/TP)

axes[0].plot(e_vec/TP, w0_num * np.ones(shape(e_vec)), 'k--')

axes[0].set_ylabel(r'$\Omega$', fontsize=18)
axes[0].set_xlabel(r'$\epsilon$ (GHz)', fontsize=18)
axes[0].set_ylim(0, 11)
axes[1].set_xlim(meV_to_GHz(-0.05), meV_to_GHz(0.05))

# meV
for W_vec in W_mat:
    axes[1].plot(GHz_to_meV(e_vec/TP), real(W_vec)/TP)

axes[1].plot(GHz_to_meV(e_vec/TP), w0_num * np.ones(shape(e_vec)), 'k--')

axes[1].set_ylabel(r'$\Omega$', fontsize=18)
axes[1].set_xlabel(r'$\epsilon$ (meV)', fontsize=18)
axes[1].set_ylim(0, 11)
axes[1].set_xlim(-0.05, 0.05)

fig.tight_layout()

#delta_phi_mat = array([[delta_phi.subs(param_derived).subs({e: e_num, tc: tc_num}).subs(param_fundamental_all).evalf()
#                        for e_num in e_vec]
#                       for tc_num in tc_vec], dtype=complex)

delta_phi_lambda = lambdify((X, Y), delta_phi.subs(param_derived).subs(param_derived).subs({e: X, tc: Y}).subs(param_fundamental_all), 'numpy')
delta_phi_mat = array([delta_phi_lambda(e_vec, tc_num) for tc_num in tc_vec], dtype=complex)

fig, axes = subplots(1,2, figsize=(10,5))

# GHz
for idx, delta_phi_vec in enumerate(delta_phi_mat):
    axes[0].plot(e_vec/TP, - idx * 0.0 + real(delta_phi_vec))

axes[0].set_ylabel(r'$\Delta\varphi$', fontsize=18)
axes[0].set_xlabel(r'$\epsilon$ (GHz)', fontsize=18)
axes[1].set_xlim(meV_to_GHz(-0.05), meV_to_GHz(0.05))

# meV
for idx, delta_phi_vec in enumerate(delta_phi_mat):
    axes[1].plot(GHz_to_meV(e_vec/TP), real(delta_phi_vec))

axes[1].set_ylabel(r'$\Delta\varphi$', fontsize=18)
axes[1].set_xlabel(r'$\epsilon$ (meV)', fontsize=18)
axes[1].set_xlim(-0.05, 0.05)

fig.tight_layout()

def cavity_convolution_func(w_vec, w0):
    g = numpy.exp(-(w_vec-w0)**2/param_fundamental[sigma_e]**2)

    return g / numpy.sum(g)

def cavity_convolution(w_vec, f_vec):
    return array([numpy.sum(f_vec * cavity_convolution_func(w_vec, w)) for w in w_vec])

fig, ax = subplots(1, 1, figsize=(10,5))

f_vec = abs(e_vec) < 5*TP
w0 = 0 * TP
ax.plot(e_vec/TP, cavity_convolution_func(e_vec, w0))
ax.plot(e_vec/TP, f_vec)
ax.plot(e_vec/TP, cavity_convolution(e_vec, f_vec));

fig, axes = subplots(1,2, figsize=(10,5))

# GHz
for idx, delta_phi_vec in enumerate(delta_phi_mat):
    axes[0].plot(e_vec/TP, - idx * 7.5 + real(cavity_convolution(e_vec, delta_phi_vec)))

axes[0].set_ylabel(r'$\Delta\varphi$', fontsize=18)
axes[0].set_xlabel(r'$\epsilon$ (GHz)', fontsize=18)
axes[1].set_xlim(meV_to_GHz(-0.05), meV_to_GHz(0.05))

# meV
for idx, delta_phi_vec in enumerate(delta_phi_mat):
    axes[1].plot(GHz_to_meV(e_vec/TP), - idx * 7.5 + real(cavity_convolution(e_vec, delta_phi_vec)))

axes[1].set_ylabel(r'$\Delta\varphi$', fontsize=18)
axes[1].set_xlabel(r'$\epsilon$ (meV)', fontsize=18)
axes[1].set_xlim(-0.05, 0.05)

fig.tight_layout()

%load_ext version_information

%version_information numpy, sympy, matplotlib