%pylab inline

from qutip import *

N = 20
alpha = 2.5

a1 = tensor(destroy(N), qeye(N))
a2 = tensor(qeye(N), destroy(N))

def visualize_matrices(a1, a2, rho):
    fig, ax = subplots(2, 3, figsize=(10,5), subplot_kw={'projection':'3d'})

    # field operators
    C = correlation_matrix_field(a1, a2, rho)
    lbls = [r'$a_1$', r'$a_1^\dagger$', r'$a_2$', r'$a_2^\dagger$']
    matrix_histogram(real(C), xlabels=lbls, ylabels=lbls,
                     fig=fig, ax=ax[0,0], colorbar=False, title="field operator corr.mat.")
    
    basis = [a1, a1.dag(), a2, a2.dag()]
    C = covariance_matrix(basis, rho)
    lbls = [r'$a_1$', r'$a_1^\dagger$', r'$a_2$', r'$a_2^\dagger$']
    matrix_histogram(real(C), xlabels=lbls, ylabels=lbls,
                     fig=fig, ax=ax[1,0], colorbar=False, title="field operator cov.mat.")
    
    # quadratures
    C = correlation_matrix_quadrature(a1, a2, rho)
    lbls = [r'$x_1$', r'$p_1$', r'$x_2$', r'$p_2$']
    matrix_histogram(real(C), xlabels=lbls, ylabels=lbls,
                     fig=fig, ax=ax[0,1], colorbar=False, title="quadrature operator corr.mat.")

    x1 = (a1 + a1.dag()) / np.sqrt(2)
    p1 = -1j * (a1 - a1.dag()) / np.sqrt(2)
    x2 = (a2 + a2.dag()) / np.sqrt(2)
    p2 = -1j * (a2 - a2.dag()) / np.sqrt(2)
    basis = [x1, p1, x2, p2]
    C = covariance_matrix(basis, rho)
    lbls = [r'$a_1$', r'$a_1^\dagger$', r'$a_2$', r'$a_2^\dagger$']
    matrix_histogram(real(C), xlabels=lbls, ylabels=lbls,
                     fig=fig, ax=ax[1,1], colorbar=False, title="quadrature operator cov.mat.")
    
    
    # wigner covariance matrix
    C = wigner_covariance_matrix(a1=a1, a2=a2, rho=rho)
    lbls = [r'$x_1$', r'$p_1$', r'$x_2$', r'$p_2$']
    matrix_histogram(real(C), xlabels=lbls, ylabels=lbls,
                     fig=fig, ax=ax[0,2], colorbar=False, title="wigner cov.mat. #1")

    R = covariance_matrix(basis, rho)
    C = wigner_covariance_matrix(R=R)
    lbls = [r'$x_1$', r'$p_1$', r'$x_2$', r'$p_2$']
    matrix_histogram(real(C), xlabels=lbls, ylabels=lbls,
                     fig=fig, ax=ax[1,2], colorbar=False, title="wigner cov.mat. #2")

    fig.tight_layout()
    
    print "logarithmic negativity =", logarithmic_negativity(C)

rho = tensor(basis(N, 0), basis(N, 0))

visualize_matrices(a1, a2, rho)

# correlated modes
rho = tensor(coherent_dm(N, alpha), coherent_dm(N, alpha/2))

visualize_matrices(a1, a2, rho)

# mixture of coherent states
rho = (tensor(coherent_dm(N, alpha), coherent_dm(N, -0.5j*alpha)) +
       tensor(coherent_dm(N, -0.5j*alpha), coherent_dm(N, alpha))).unit()

visualize_matrices(a1, a2, rho)

# superposition of coherent states
rho = ket2dm((tensor(coherent(N, alpha), coherent(N, -0.5j*alpha)) +
             tensor(coherent(N, -0.5j*alpha), coherent(N, alpha))).unit())

visualize_matrices(a1, a2, rho)

rho = tensor(thermal_dm(N, alpha**2), thermal_dm(N, (alpha/2)**2))

visualize_matrices(a1, a2, rho)

# mixture of thermal states and the vacuum state
rho = (tensor(thermal_dm(N, (alpha)**2), fock_dm(N, 0)) + 
       tensor(fock_dm(N, 0), thermal_dm(N, (alpha/2)**2))).unit()

visualize_matrices(a1, a2, rho)

# Squeezed vacuum

S = squeezing(a1, a2, 0.35)

vac = tensor(fock(N, 0), fock(N, 0))

rho = ket2dm(S * vac)

visualize_matrices(a1, a2, rho)

# Squeezed two-mode coherent state

S = squeezing(a1, a2, 0.35)

vac = tensor(coherent(N, alpha), coherent(N, alpha))

rho = ket2dm(S * vac)

visualize_matrices(a1, a2, rho)

from qutip.ipynbtools import version_table
version_table()