def test_wigner_fock():
"wigner: test wigner function calculation for Fock states"
xvec = np.linspace(-5.0, 5.0, 100)
yvec = xvec
X, Y = np.meshgrid(xvec, yvec)
a = X + 1j * Y # consistent with g=2 option to wigner function
dx = xvec[1] - xvec[0]
dy = yvec[1] - yvec[0]
N = 15
for n in [2, 3, 4, 5, 6]:
psi = fock(N, n)
# calculate the wigner function using qutip and analytic formula
W_qutip = wigner(psi, xvec, yvec, g=2)
W_qutip_cl = wigner(psi, xvec, yvec, g=2, method='clenshaw')
W_analytic = 2 / np.pi * (-1) ** n * \
np.exp(-2 * abs(a) ** 2) * np.polyval(laguerre(n), 4 * abs(a) ** 2)
# check difference
assert_(np.sum(abs(W_qutip - W_analytic)) < 1e-4)
assert_(np.sum(abs(W_qutip_cl - W_analytic)) < 1e-4)
# check normalization
assert_(np.sum(W_qutip) * dx * dy - 1.0 < 1e-8)
assert_(np.sum(W_qutip_cl) * dx * dy - 1.0 < 1e-8)
assert_(np.sum(W_analytic) * dx * dy - 1.0 < 1e-8)
def test_wigner_compare_methods_ket():
"wigner: compare wigner methods for random state vectors"
xvec = np.linspace(-5.0, 5.0, 100)
yvec = xvec
X, Y = np.meshgrid(xvec, yvec)
# a = X + 1j * Y # consistent with g=2 option to wigner function
dx = xvec[1] - xvec[0]
dy = yvec[1] - yvec[0]
N = 15
for n in range(10):
# try ten different random density matrices
psi = rand_ket(N, 0.5 + rand() / 2)
# calculate the wigner function using qutip and analytic formula
W_qutip1 = wigner(psi, xvec, yvec, g=2)
W_qutip2 = wigner(psi, xvec, yvec, g=2, method='laguerre')
# check difference
assert_(np.sum(abs(W_qutip1 - W_qutip1)) < 1e-4)
# check normalization
assert_(np.sum(W_qutip1) * dx * dy - 1.0 < 1e-8)
assert_(np.sum(W_qutip2) * dx * dy - 1.0 < 1e-8)
def test_wigner_compare_methods_ket():
"wigner: compare wigner methods for random state vectors"
xvec = np.linspace(-5.0, 5.0, 100)
yvec = xvec
X, Y = np.meshgrid(xvec, yvec)
# a = X + 1j * Y # consistent with g=2 option to wigner function
dx = xvec[1] - xvec[0]
dy = yvec[1] - yvec[0]
N = 15
for n in range(10):
# try ten different random density matrices
psi = rand_ket(N, 0.5 + rand() / 2)
# calculate the wigner function using qutip and analytic formula
W_qutip1 = wigner(psi, xvec, yvec, g=2)
W_qutip2 = wigner(psi, xvec, yvec, g=2, sparse=True)
# check difference
assert_(np.sum(abs(W_qutip1 - W_qutip2)) < 1e-4)
# check normalization
assert_(np.sum(W_qutip1) * dx * dy - 1.0 < 1e-8)
assert_(np.sum(W_qutip2) * dx * dy - 1.0 < 1e-8)
def test_wigner_coherent():
"wigner: test wigner function calculation for coherent states"
xvec = np.linspace(-5.0, 5.0, 100)
yvec = xvec
X, Y = np.meshgrid(xvec, yvec)
a = X + 1j * Y # consistent with g=2 option to wigner function
dx = xvec[1] - xvec[0]
dy = yvec[1] - yvec[0]
N = 20
beta = rand() + rand() * 1.0j
psi = coherent(N, beta)
# calculate the wigner function using qutip and analytic formula
W_qutip = wigner(psi, xvec, yvec, g=2)
W_qutip_cl = wigner(psi, xvec, yvec, g=2, method='clenshaw')
W_analytic = 2 / np.pi * np.exp(-2 * abs(a - beta) ** 2)
# check difference
assert_(np.sum(abs(W_qutip - W_analytic) ** 2) < 1e-4)
assert_(np.sum(abs(W_qutip_cl - W_analytic) ** 2) < 1e-4)
# check normalization
assert_(np.sum(W_qutip) * dx * dy - 1.0 < 1e-8)
assert_(np.sum(W_qutip_cl) * dx * dy - 1.0 < 1e-8)
assert_(np.sum(W_analytic) * dx * dy - 1.0 < 1e-8)
def test_wigner_coherent():
"wigner: test wigner function calculation for coherent states"
xvec = np.linspace(-5.0, 5.0, 100)
yvec = xvec
X, Y = np.meshgrid(xvec, yvec)
a = X + 1j * Y # consistent with g=2 option to wigner function
dx = xvec[1] - xvec[0]
dy = yvec[1] - yvec[0]
N = 20
beta = rand() + rand() * 1.0j
psi = coherent(N, beta)
# calculate the wigner function using qutip and analytic formula
W_qutip = wigner(psi, xvec, yvec, g=2)
W_qutip_cl = wigner(psi, xvec, yvec, g=2, method='clenshaw')
W_analytic = 2 / np.pi * np.exp(-2 * abs(a - beta)**2)
# check difference
assert_(np.sum(abs(W_qutip - W_analytic)**2) < 1e-4)
assert_(np.sum(abs(W_qutip_cl - W_analytic)**2) < 1e-4)
# check normalization
assert_(np.sum(W_qutip) * dx * dy - 1.0 < 1e-8)
assert_(np.sum(W_qutip_cl) * dx * dy - 1.0 < 1e-8)
assert_(np.sum(W_analytic) * dx * dy - 1.0 < 1e-8)
def test_wigner_fft_comparse_dm():
"Wigner: Compare Wigner fft and iterative for rand. dm"
N = 20
xvec = np.linspace(-10, 10, 128)
for i in range(3):
rho = rand_dm(N)
Wfft, yvec = wigner(rho, xvec, xvec, method='fft')
W = wigner(rho, xvec, yvec, method='iterative')
Wdiff = abs(W - Wfft)
assert_equal(np.sum(abs(Wdiff)) < 1e-7, True)
def test_wigner_clenshaw_sp_iter_dm():
"Wigner: Compare Wigner sparse clenshaw and iterative for rand. dm"
N = 20
xvec = np.linspace(-10, 10, 128)
for i in range(3):
rho = rand_dm(N)
Wclen = wigner(rho, xvec, xvec, method='clenshaw', sparse=True)
W = wigner(rho, xvec, xvec, method='iterative')
Wdiff = abs(W - Wclen)
assert_equal(np.sum(abs(Wdiff)) < 1e-7, True)
def test_wigner_clenshaw_sp_iter_dm():
"Wigner: Compare Wigner sparse clenshaw and iterative for rand. dm"
N = 20
xvec = np.linspace(-10, 10, 128)
for i in range(3):
rho = rand_dm(N)
Wclen = wigner(rho, xvec, xvec, method='clenshaw', sparse=True)
W = wigner(rho, xvec, xvec, method='iterative')
Wdiff = abs(W - Wclen)
assert_equal(np.sum(abs(Wdiff)) < 1e-7, True)
def test_wigner_fft_comparse_dm():
"Wigner: Compare Wigner fft and iterative for rand. dm"
N = 20
xvec = np.linspace(-10, 10, 128)
for i in range(3):
rho = rand_dm(N)
Wfft, yvec = wigner(rho, xvec, xvec, method='fft')
W = wigner(rho, xvec, yvec, method='iterative')
Wdiff = abs(W - Wfft)
assert_equal(np.sum(abs(Wdiff)) < 1e-7, True)
def test_wigner_fock():
"wigner: test wigner function calculation for Fock states"
xvec = np.linspace(-5.0, 5.0, 100)
yvec = xvec
X, Y = np.meshgrid(xvec, yvec)
a = X + 1j * Y # consistent with g=2 option to wigner function
dx = xvec[1] - xvec[0]
dy = yvec[1] - yvec[0]
N = 15
for n in [2, 3, 4, 5, 6]:
psi = fock(N, n)
# calculate the wigner function using qutip and analytic formula
W_qutip = wigner(psi, xvec, yvec, g=2)
W_analytic = 2 / np.pi * (-1) ** n * \
np.exp(-2 * abs(a) ** 2) * np.polyval(laguerre(n), 4 * abs(a) ** 2)
# check difference
assert_(np.sum(abs(W_qutip - W_analytic)) < 1e-4)
# check normalization
assert_(np.sum(W_qutip) * dx * dy - 1.0 < 1e-8)
assert_(np.sum(W_analytic) * dx * dy - 1.0 < 1e-8)
def update(self, rho):
self.data = wigner(rho, self.xvecs[0], self.xvecs[1])
def update(self, rho):
self.data = wigner(rho, self.xvecs[0], self.xvecs[1])
def plot_wigner(rho, fig=None, ax=None, figsize=(8, 4),
cmap=None, alpha_max=7.5, colorbar=False,
method='iterative', projection='2d'):
"""
Plot the the Wigner function for a density matrix (or ket) that describes
an oscillator mode.
Parameters
----------
rho : :class:`qutip.qobj.Qobj`
The density matrix (or ket) of the state to visualize.
fig : a matplotlib Figure instance
The Figure canvas in which the plot will be drawn.
ax : a matplotlib axes instance
The axes context in which the plot will be drawn.
figsize : (width, height)
The size of the matplotlib figure (in inches) if it is to be created
(that is, if no 'fig' and 'ax' arguments are passed).
cmap : a matplotlib cmap instance
The colormap.
alpha_max : float
The span of the x and y coordinates (both [-alpha_max, alpha_max]).
colorbar : bool
Whether (True) or not (False) a colorbar should be attached to the
Wigner function graph.
method : string {'iterative', 'laguerre', 'fft'}
The method used for calculating the wigner function. See the
documentation for qutip.wigner for details.
projection: string {'2d', '3d'}
Specify whether the Wigner function is to be plotted as a
contour graph ('2d') or surface plot ('3d').
Returns
-------
fig, ax : tuple
A tuple of the matplotlib figure and axes instances used to produce
the figure.
"""
if not fig and not ax:
if projection == '2d':
fig, ax = plt.subplots(1, 1, figsize=figsize)
elif projection == '3d':
fig = plt.figure(figsize=figsize)
ax = fig.add_subplot(1, 1, 1, projection='3d')
else:
raise ValueError('Unexpected value of projection keyword argument')
if isket(rho):
rho = ket2dm(rho)
xvec = np.linspace(-alpha_max, alpha_max, 200)
W0 = wigner(rho, xvec, xvec, method=method)
W, yvec = W0 if type(W0) is tuple else (W0, xvec)
wlim = abs(W).max()
if cmap is None:
cmap = cm.get_cmap('RdBu')
if projection == '2d':
cf = ax.contourf(xvec, yvec, W, 100,
norm=mpl.colors.Normalize(-wlim, wlim), cmap=cmap)
elif projection == '3d':
X, Y = np.meshgrid(xvec, xvec)
cf = ax.plot_surface(X, Y, W0, rstride=5, cstride=5, linewidth=0.5,
norm=mpl.colors.Normalize(-wlim, wlim), cmap=cmap)
else:
raise ValueError('Unexpected value of projection keyword argument.')
if xvec is not yvec:
ax.set_ylim(xvec.min(), xvec.max())
ax.set_xlabel(r'$\rm{Re}(\alpha)$', fontsize=12)
ax.set_ylabel(r'$\rm{Im}(\alpha)$', fontsize=12)
if colorbar:
cb = fig.colorbar(cf, ax=ax)
ax.set_title("Wigner function", fontsize=12)
return fig, ax
def wigner_fock_distribution(rho, fig=None, axes=None, figsize=(8, 4), cmap=None, alpha_max=7.5, colorbar=False):
"""
Plot the Fock distribution and the Wigner function for a density matrix
(or ket) that describes an oscillator mode.
Parameters
----------
rho : :class:`qutip.qobj.Qobj`
The density matrix (or ket) of the state to visualize.
fig : a matplotlib Figure instance
The Figure canvas in which the plot will be drawn.
axes : a list of two matplotlib axes instances
The axes context in which the plot will be drawn.
figsize : (width, height)
The size of the matplotlib figure (in inches) if it is to be created
(that is, if no 'fig' and 'ax' arguments are passed).
cmap : a matplotlib cmap instance
The colormap.
alpha_max : float
The span of the x and y coordinates (both [-alpha_max, alpha_max]).
colorbar : bool
Whether (True) or not (False) a colorbar should be attached to the
Wigner function graph.
Returns
-------
fig, ax : tuple
A tuple of the matplotlib figure and axes instances used to produce
the figure.
"""
if not fig and not axes:
fig, axes = plt.subplots(1, 2, figsize=figsize)
if isket(rho):
rho = ket2dm(rho)
fock_distribution(rho, fig=fig, ax=axes[0])
xvec = linspace(-alpha_max, alpha_max, 200)
W = wigner(rho, xvec, xvec)
wlim = abs(W).max()
if cmap is None:
cmap = get_cmap("RdBu")
# cmap = wigner_cmap(W)
cf = axes[1].contourf(xvec, xvec, W, 100, norm=mpl.colors.Normalize(-wlim, wlim), cmap=cmap)
axes[1].set_xlabel(r"$\rm{Re}(\alpha)$", fontsize=12)
axes[1].set_ylabel(r"$\rm{Im}(\alpha)$", fontsize=12)
if colorbar:
cb = fig.colorbar(cf, ax=axes[1])
axes[0].set_title("Fock distribution", fontsize=12)
axes[1].set_title("Wigner function", fontsize=12)
return fig, axes
def wigner_fock_distribution(rho, fig=None, ax=None, figsize=(8, 4),
cmap=None, alpha_max=7.5, colorbar=False):
"""
Plot the Fock distribution and the Wigner function for a density matrix
(or ket) that describes an oscillator mode.
Parameters
----------
rho : :class:`qutip.qobj.Qobj`
The density matrix (or ket) of the state to visualize.
fig : a matplotlib Figure instance
The Figure canvas in which the plot will be drawn.
ax : a matplotlib axes instance
The axes context in which the plot will be drawn.
figsize : (width, height)
The size of the matplotlib figure (in inches) if it is to be created
(that is, if no 'fig' and 'ax' arguments are passed).
cmap : a matplotlib cmap instance
The colormap.
alpha_max : float
The span of the x and y coordinates (both [-alpha_max, alpha_max]).
colorbar : bool
Whether (True) or not (False) a colorbar should be attached to the
Wigner function graph.
Returns
-------
A tuple of matplotlib figure and axes instances.
"""
if not fig and not ax:
fig, axes = plt.subplots(1, 2, figsize=figsize)
if isket(rho):
rho = ket2dm(rho)
fock_distribution(rho, fig=fig, ax=axes[0])
xvec = linspace(-alpha_max, alpha_max, 200)
W = wigner(rho, xvec, xvec)
wlim = abs(W).max()
if cmap is None:
cmap = get_cmap('RdBu')
# cmap = wigner_cmap(W)
cf = axes[1].contourf(xvec, xvec, W, 100,
norm=mpl.colors.Normalize(-wlim, wlim), cmap=cmap)
axes[1].set_xlabel(r'$\rm{Re}(\alpha)$', fontsize=12)
axes[1].set_ylabel(r'$\rm{Im}(\alpha)$', fontsize=12)
if colorbar:
cb = fig.colorbar(cf, ax=axes[1])
axes[0].set_title("Fock distribution", fontsize=12)
axes[1].set_title("Wigner function", fontsize=12)
return fig, ax
def plot_wigner(rho,
fig=None,
ax=None,
figsize=(6, 6),
cmap=None,
alpha_max=7.5,
colorbar=False,
method='clenshaw',
projection='2d'):
"""
Plot the the Wigner function for a density matrix (or ket) that describes
an oscillator mode.
Parameters
----------
rho : :class:`qutip.qobj.Qobj`
The density matrix (or ket) of the state to visualize.
fig : a matplotlib Figure instance
The Figure canvas in which the plot will be drawn.
ax : a matplotlib axes instance
The axes context in which the plot will be drawn.
figsize : (width, height)
The size of the matplotlib figure (in inches) if it is to be created
(that is, if no 'fig' and 'ax' arguments are passed).
cmap : a matplotlib cmap instance
The colormap.
alpha_max : float
The span of the x and y coordinates (both [-alpha_max, alpha_max]).
colorbar : bool
Whether (True) or not (False) a colorbar should be attached to the
Wigner function graph.
method : string {'clenshaw', 'iterative', 'laguerre', 'fft'}
The method used for calculating the wigner function. See the
documentation for qutip.wigner for details.
projection: string {'2d', '3d'}
Specify whether the Wigner function is to be plotted as a
contour graph ('2d') or surface plot ('3d').
Returns
-------
fig, ax : tuple
A tuple of the matplotlib figure and axes instances used to produce
the figure.
"""
if not fig and not ax:
if projection == '2d':
fig, ax = plt.subplots(1, 1, figsize=figsize)
elif projection == '3d':
fig = plt.figure(figsize=figsize)
ax = fig.add_subplot(1, 1, 1, projection='3d')
else:
raise ValueError('Unexpected value of projection keyword argument')
if isket(rho):
rho = ket2dm(rho)
xvec = np.linspace(-alpha_max, alpha_max, 200)
W0 = wigner(rho, xvec, xvec, method=method)
W, yvec = W0 if type(W0) is tuple else (W0, xvec)
wlim = abs(W).max()
if cmap is None:
cmap = cm.get_cmap('RdBu')
if projection == '2d':
cf = ax.contourf(xvec,
yvec,
W,
100,
norm=mpl.colors.Normalize(-wlim, wlim),
cmap=cmap)
elif projection == '3d':
X, Y = np.meshgrid(xvec, xvec)
cf = ax.plot_surface(X,
Y,
W0,
rstride=5,
cstride=5,
linewidth=0.5,
norm=mpl.colors.Normalize(-wlim, wlim),
cmap=cmap)
else:
raise ValueError('Unexpected value of projection keyword argument.')
if xvec is not yvec:
ax.set_ylim(xvec.min(), xvec.max())
ax.set_xlabel(r'$\rm{Re}(\alpha)$', fontsize=12)
ax.set_ylabel(r'$\rm{Im}(\alpha)$', fontsize=12)
if colorbar:
fig.colorbar(cf, ax=ax)
ax.set_title("Wigner function", fontsize=12)
return fig, ax
