from qutip import *
from mpl_toolkits.mplot3d import Axes3D
from matplotlib import *
from pylab import *
%pylab inline

N=10

z=3

kot_p=(coherent(N,z)+coherent(N,-z)).unit()
kot_np=(coherent(N,z)-coherent(N,-z)).unit()

rho_p=ket2dm(kot_p)
rho_np=ket2dm(kot_np)

rho=rho_p
xvec = linspace(-5,5,200)
X, Y = meshgrid(xvec, xvec)
W = wigner(rho, xvec, xvec)
fig = figure()
ax = Axes3D(fig, azim=-62, elev=25)
ax.plot_surface(X, Y, W, rstride=2, cstride=2, cmap=cm.jet, lw=.1)
ax.set_xlim3d(-6, 6)
ax.set_ylim3d(-6, 6)
ax.set_zlim3d(0, 0.4)
title('Funkcja Wignera')
show()

rho=rho_p
Q = qfunc(rho, xvec, xvec)
fig = figure()
ax = Axes3D(fig, azim=-62, elev=25)
ax.plot_surface(X, Y, Q, rstride=2, cstride=2, cmap=cm.jet, lw=.1)
ax.set_xlim3d(-6, 6)
ax.set_ylim3d(-6, 6)
#ax.set_zlim3d(0, 0.4)
title('Funkcja Q')
show()

rho=rho_p
fock_distribution(rho)#, fig=None, ax=None, figsize=(8, 6), title=None)
show()

rho=rho_np
fock_distribution(rho)#, fig=None, ax=None, figsize=(8, 6), title=None)
show()