Merge pull request #2057 from KosukeMizuno/plot_wigner_args

add arguments for plot_wigner_fock_distribution

doc/changes/2057.feature

@@ -0,0 +1 @@
+Add arguments of plot_wigner() and plot_wigner_fock_distribution() to specify parameters for wigner(). (#2057)

qutip/animation.py

@@ -9,6 +9,7 @@
                plot_fock_distribution, plot_wigner, plot_spin_distribution,
                plot_qubism, plot_schmidt)
 from .solver import Result
+from numpy import sqrt
 
 
 def _result_state(obj):
@@ -346,9 +347,9 @@ def anim_fock_distribution(rhos, fock_numbers=None, color="green",
     return fig, ani
 
 
-def anim_wigner(rhos, xvec=None, yvec=None, method='clenshaw',
-                projection='2d', *, cmap=None, colorbar=False,
-                fig=None, ax=None):
+def anim_wigner(rhos, xvec=None, yvec=None, method='clenshaw', projection='2d',
+                g=sqrt(2), sparse=False, parfor=False, *,
+                cmap=None, colorbar=False, fig=None, ax=None):
     """
     Animation of the Wigner function for a density matrix (or ket)
     that describes an oscillator mode.
@@ -373,6 +374,18 @@ def anim_wigner(rhos, xvec=None, yvec=None, method='clenshaw',
         Specify whether the Wigner function is to be plotted as a
         contour graph ('2d') or surface plot ('3d').
 
+    g : float
+        Scaling factor for `a = 0.5 * g * (x + iy)`, default `g = sqrt(2)`.
+        See the documentation for qutip.wigner for details.
+
+    sparse : bool {False, True}
+        Flag for sparse format.
+        See the documentation for qutip.wigner for details.
+
+    parfor : bool {False, True}
+        Flag for parallel calculation.
+        See the documentation for qutip.wigner for details.
+
     cmap : a matplotlib cmap instance, optional
         The colormap.
 
@@ -395,7 +408,8 @@ def anim_wigner(rhos, xvec=None, yvec=None, method='clenshaw',
 
     rhos = _result_state(rhos)
 
-    fig, ani = plot_wigner(rhos, xvec, yvec, method, projection,
+    fig, ani = plot_wigner(rhos, xvec, yvec, method=method, g=g, sparse=sparse,
+                           parfor=parfor, projection=projection,
                            cmap=cmap, colorbar=colorbar, fig=fig, ax=ax)
 
     return fig, ani

qutip/visualization.py

@@ -12,7 +12,7 @@
 import warnings
 import itertools as it
 import numpy as np
-from numpy import pi, array, sin, cos, angle, log2
+from numpy import pi, array, sin, cos, angle, log2, sqrt
 
 from packaging.version import parse as parse_version
 
@@ -1126,9 +1126,9 @@ def plot_fock_distribution(rho, fock_numbers=None, color="green",
     return fig, output
 
 
-def plot_wigner(rho, xvec=None, yvec=None, method='clenshaw',
-                projection='2d', *, cmap=None, colorbar=False,
-                fig=None, ax=None):
+def plot_wigner(rho, xvec=None, yvec=None, method='clenshaw', projection='2d',
+                g=sqrt(2), sparse=False, parfor=False, *,
+                cmap=None, colorbar=False, fig=None, ax=None):
     """
     Plot the the Wigner function for a density matrix (or ket) that describes
     an oscillator mode.
@@ -1153,6 +1153,18 @@ def plot_wigner(rho, xvec=None, yvec=None, method='clenshaw',
         Specify whether the Wigner function is to be plotted as a
         contour graph ('2d') or surface plot ('3d').
 
+    g : float
+        Scaling factor for `a = 0.5 * g * (x + iy)`, default `g = sqrt(2)`.
+        See the documentation for qutip.wigner for details.
+
+    sparse : bool {False, True}
+        Flag for sparse format.
+        See the documentation for qutip.wigner for details.
+
+    parfor : bool {False, True}
+        Flag for parallel calculation.
+        See the documentation for qutip.wigner for details.
+
     cmap : a matplotlib cmap instance, optional
         The colormap.
 
@@ -1193,7 +1205,10 @@ def plot_wigner(rho, xvec=None, yvec=None, method='clenshaw',
         if isket(rho):
             rho = ket2dm(rho)
 
-        W0 = wigner(rho, xvec, yvec, method=method)
+        W0 = wigner(
+            rho, xvec, yvec, method=method,
+            g=g, sparse=sparse, parfor=parfor
+        )
 
         W, yvec = W0 if isinstance(W0, tuple) else (W0, yvec)
         Ws.append(W)