%pylab inline

from qutip import *

from qutip.distributions import *

N = 20
alpha = 2.0 + 2j
rho = (coherent(N, alpha) + coherent(N, -alpha)).unit()

w = WignerDistribution(rho, extent=[[-10, 10], [-10, 10]])

w.visualize();

q = QDistribution(rho, extent=[[-10, 10], [-10, 10]])

q.visualize(style="surface");

alpha = 1.0

psi = (tensor(coherent(N, alpha), basis(N, 0)) + tensor(basis(N, 0), coherent(N, -alpha))).unit()  # superposition
rho = (ket2dm(tensor(coherent(N, alpha), basis(N, 0))) + ket2dm(tensor(basis(N, 0), coherent(N, -alpha)))).unit()  # mixture

p = TwoModeQuadratureCorrelation(psi)

p.visualize();

p = TwoModeQuadratureCorrelation(rho)

p.visualize();

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

w0 = WignerDistribution(ptrace(rho, 0))
w1 = WignerDistribution(ptrace(rho, 1))

w0.visualize(fig=fig, ax=axes[0]);
w1.visualize(fig=fig, ax=axes[1]);

alpha = 2.0
rho = (coherent(N, alpha) + coherent(N, -alpha)).unit()

w = WignerDistribution(rho)

w.visualize();

wx = w.marginal(dim=0)

fig, ax = wx.visualize();

wy = w.marginal(dim=1)

fig, ax = wy.visualize();

wx = w.project(dim=1)

fig, ax = wx.visualize();

M=8

fig, ax = subplots(M, 1, figsize=(10, 12), sharex=True)

for n in range(M):
    psi = fock(N, n)
    wf = HarmonicOscillatorWaveFunction(psi, 1.0, extent=[-10, 10])
    wf.visualize(fig=fig, ax=ax[M-n-1], show_ylabel=False, show_xlabel=(n == 0))  

psi = coherent(50, 4.0)
wf = HarmonicOscillatorWaveFunction(psi, 1.0, extent=[-15, 15])
wf.visualize();

from qutip.ipynbtools import version_table
version_table()