visualization

Visualization utils.

plot_cf(state, pts_x, pts_y=None, axs=None, contour=True, qp_type=WIGNER, cbar_label='', axis_scale_factor=1, plot_cbar=True, plot_grid=True, x_ticks=None, y_ticks=None, z_ticks=None, subtitles=None, figtitle=None, gif=False, gif_params=None)

Plot a characteristic function as paired real/imag subplots.

Each batch element produces two adjacent subplots — real part followed by imaginary part — so the rendered grid has shape (rows, 2 * cols).

Parameters:

Name	Type	Description	Default
state		state with arbitrary number of batch dimensions, result will be flattened to a 2d grid to allow for plotting	required
pts_x		x points to evaluate the characteristic function at	required
pts_y		y points to evaluate the characteristic function at	None
axs		matplotlib axes to plot on	None
contour		make the plot use contouring	True
qp_type		type of characteristic function. Currently only "wigner" is supported.	WIGNER
cbar_label		labels for the real and imaginary cbar (overridden internally based on qp_type)	''
axis_scale_factor		scale of the axes labels relative	1
plot_cbar		whether to plot cbar	True
plot_grid		whether to draw gridlines on each subplot	True
x_ticks		tick position for the x-axis	None
y_ticks		tick position for the y-axis	None
z_ticks		tick position for the z-axis	None
subtitles		subtitles for the subplots (shape must match state.bdims)	None
figtitle		figure title	None
gif		if True, render an animation over one batch axis instead of a tiled grid. Returns a matplotlib.animation.FuncAnimation that auto-renders inline in Jupyter.	False
gif_params		dict of options for the gif path. Recognized keys: save_path (default None) — if set, save the animation here via PillowWriter; interval_ms (default 200) — milliseconds per frame; ts (default None) — optional 1D array of timestamps matching the animation-axis length; when set, each frame's suptitle gets a t = … label; batch_animation_axis (default 0) — index into state.bdims selecting which axis becomes the animation/time axis (the remaining batch dims form the per-frame subplot grid).	None

Returns:

Type	Description
	(axs, im) in the static case, or a FuncAnimation when
	gif=True.

Source code in jaxquantum/core/visualization.py

def plot_cf(
        state,
        pts_x,
        pts_y=None,
        axs=None,
        contour=True,
        qp_type=WIGNER,
        cbar_label="",
        axis_scale_factor=1,
        plot_cbar=True,
        plot_grid=True,
        x_ticks=None,
        y_ticks=None,
        z_ticks=None,
        subtitles=None,
        figtitle=None,
        gif=False,
        gif_params=None,
):
    """Plot a characteristic function as paired real/imag subplots.

    Each batch element produces two adjacent subplots — real part followed
    by imaginary part — so the rendered grid has shape ``(rows, 2 * cols)``.

    Args:
        state: state with arbitrary number of batch dimensions, result will
            be flattened to a 2d grid to allow for plotting
        pts_x: x points to evaluate the characteristic function at
        pts_y: y points to evaluate the characteristic function at
        axs: matplotlib axes to plot on
        contour: make the plot use contouring
        qp_type: type of characteristic function. Currently only
            ``"wigner"`` is supported.
        cbar_label: labels for the real and imaginary cbar (overridden
            internally based on ``qp_type``)
        axis_scale_factor: scale of the axes labels relative
        plot_cbar: whether to plot cbar
        plot_grid: whether to draw gridlines on each subplot
        x_ticks: tick position for the x-axis
        y_ticks: tick position for the y-axis
        z_ticks: tick position for the z-axis
        subtitles: subtitles for the subplots (shape must match ``state.bdims``)
        figtitle: figure title
        gif: if True, render an animation over one batch axis instead of a
            tiled grid. Returns a ``matplotlib.animation.FuncAnimation``
            that auto-renders inline in Jupyter.
        gif_params: dict of options for the gif path. Recognized keys:
            ``save_path`` (default None) — if set, save the animation here
            via PillowWriter; ``interval_ms`` (default 200) — milliseconds
            per frame; ``ts`` (default None) — optional 1D array of
            timestamps matching the animation-axis length; when set, each
            frame's suptitle gets a ``t = …`` label;
            ``batch_animation_axis`` (default 0) — index into
            ``state.bdims`` selecting which axis becomes the animation/time
            axis (the remaining batch dims form the per-frame subplot grid).

    Returns:
        ``(axs, im)`` in the static case, or a ``FuncAnimation`` when
        ``gif=True``.
    """
    if pts_y is None:
        pts_y = pts_x
    pts_x = jnp.array(pts_x)
    pts_y = jnp.array(pts_y)

    if gif:
        return _plot_cf_gif(
            state=state,
            pts_x=pts_x,
            pts_y=pts_y,
            axs=axs,
            contour=contour,
            qp_type=qp_type,
            axis_scale_factor=axis_scale_factor,
            plot_cbar=plot_cbar,
            plot_grid=plot_grid,
            x_ticks=x_ticks,
            y_ticks=y_ticks,
            z_ticks=z_ticks,
            subtitles=subtitles,
            figtitle=figtitle,
            gif_params=gif_params or {},
        )

    bdims = state.bdims
    added_baxes = 0

    if subtitles is not None:
        if subtitles.shape != bdims:
            raise ValueError(
                f"labels must have same shape as bdims, "
                f"got shapes {subtitles.shape} and {bdims}"
            )

    if len(bdims) == 0:
        bdims = (1,)
        added_baxes += 1
    if len(bdims) == 1:
        bdims = (1, bdims[0])
        added_baxes += 1

    extra_dims = bdims[2:]
    if extra_dims != ():
        state = state.reshape_bdims(
            bdims[0] * int(jnp.prod(jnp.array(extra_dims))), bdims[1]
        )
        if subtitles is not None:
            subtitles = subtitles.reshape(
                bdims[0] * int(jnp.prod(jnp.array(extra_dims))), bdims[1]
            )
        bdims = state.bdims

    if axs is None:
        _, axs = plt.subplots(
            bdims[0],
            bdims[1]*2,
            figsize=(4 * bdims[1]*2, 3 * bdims[0]),
            dpi=200,
        )


    if qp_type == WIGNER:
        vmin = -1
        vmax = 1
        scale = 1
        cmap = "seismic"
        cbar_label = [r"$\mathcal{Re}(\chi_W(\alpha))$", r"$\mathcal{"
                                                         r"Im}(\chi_W("
                                                         r"\alpha))$"]
        QP = scale * cf_wigner(state, pts_x, pts_y)

    for _ in range(added_baxes):
        QP = jnp.array([QP])
        axs = np.array([axs])
        if subtitles is not None:
            subtitles = np.array([subtitles])

    if added_baxes==2:
        axs = axs[0] # When the input state is zero-dimensional, remove an
                     # axis that is automatically added due to the subcolumns


    pts_x = pts_x * axis_scale_factor
    pts_y = pts_y * axis_scale_factor

    x_ticks = (
        jnp.linspace(jnp.min(pts_x), jnp.max(pts_x),
                     5) if x_ticks is None else x_ticks
    )
    y_ticks = (
        jnp.linspace(jnp.min(pts_y), jnp.max(pts_y),
                     5) if y_ticks is None else y_ticks
    )
    z_ticks = jnp.linspace(vmin, vmax, 11) if z_ticks is None else z_ticks

    im = _render_cf_grid(
        axs,
        QP,
        pts_x,
        pts_y,
        contour=contour,
        cmap=cmap,
        vmin=vmin,
        vmax=vmax,
        x_ticks=x_ticks,
        y_ticks=y_ticks,
        z_ticks=z_ticks,
        cbar_label=cbar_label,
        plot_cbar=plot_cbar,
        plot_grid=plot_grid,
        subtitles=subtitles,
        decorate=True,
    )

    fig = axs[0, 0].get_figure()
    fig.tight_layout()
    if figtitle is not None:
        fig.suptitle(figtitle, y=1.04)
    return axs, im

plot_cf_wigner(state, pts_x, pts_y=None, axs=None, contour=True, cbar_label='', axis_scale_factor=1, plot_cbar=True, plot_grid=True, x_ticks=None, y_ticks=None, z_ticks=None, subtitles=None, figtitle=None, gif=False, gif_params=None)

Plot the Wigner characteristic function of the state.

Thin wrapper around :func:plot_cf with qp_type='wigner'. Each batch element is rendered as two subplots side-by-side: real then imaginary part of the characteristic function.

Parameters:

Name	Type	Description	Default
state		state with arbitrary number of batch dimensions, result will be flattened to a 2d grid to allow for plotting	required
pts_x		x points to evaluate the characteristic function at	required
pts_y		y points to evaluate the characteristic function at	None
axs		matplotlib axes to plot on	None
contour		make the plot use contouring	True
cbar_label		label for the cbar	''
axis_scale_factor		scale of the axes labels relative	1
plot_cbar		whether to plot cbar	True
plot_grid		whether to draw gridlines on each subplot	True
x_ticks		tick position for the x-axis	None
y_ticks		tick position for the y-axis	None
z_ticks		tick position for the z-axis	None
subtitles		subtitles for the subplots	None
figtitle		figure title	None
gif		if True, render an animation over one batch axis instead of a tiled subplot grid. See :func:plot_cf for details.	False
gif_params		dict of options for the gif path. Recognized keys: save_path (default None), interval_ms (default 200), ts (default None — adds a t = … label per frame), batch_animation_axis (default 0).	None

Returns:

Type	Description
	(axs, im) in the static case, or a matplotlib.animation.FuncAnimation
	when gif=True.

Source code in jaxquantum/core/visualization.py

def plot_cf_wigner(
    state,
    pts_x,
    pts_y=None,
    axs=None,
    contour=True,
    cbar_label="",
    axis_scale_factor=1,
    plot_cbar=True,
    plot_grid=True,
    x_ticks=None,
    y_ticks=None,
    z_ticks=None,
    subtitles=None,
    figtitle=None,
    gif=False,
    gif_params=None,
):
    """Plot the Wigner characteristic function of the state.

    Thin wrapper around :func:`plot_cf` with ``qp_type='wigner'``. Each batch
    element is rendered as two subplots side-by-side: real then imaginary
    part of the characteristic function.

    Args:
        state: state with arbitrary number of batch dimensions, result will
            be flattened to a 2d grid to allow for plotting
        pts_x: x points to evaluate the characteristic function at
        pts_y: y points to evaluate the characteristic function at
        axs: matplotlib axes to plot on
        contour: make the plot use contouring
        cbar_label: label for the cbar
        axis_scale_factor: scale of the axes labels relative
        plot_cbar: whether to plot cbar
        plot_grid: whether to draw gridlines on each subplot
        x_ticks: tick position for the x-axis
        y_ticks: tick position for the y-axis
        z_ticks: tick position for the z-axis
        subtitles: subtitles for the subplots
        figtitle: figure title
        gif: if True, render an animation over one batch axis instead of a
            tiled subplot grid. See :func:`plot_cf` for details.
        gif_params: dict of options for the gif path. Recognized keys:
            ``save_path`` (default None), ``interval_ms`` (default 200),
            ``ts`` (default None — adds a ``t = …`` label per frame),
            ``batch_animation_axis`` (default 0).

    Returns:
        ``(axs, im)`` in the static case, or a ``matplotlib.animation.FuncAnimation``
        when ``gif=True``.
    """
    return plot_cf(
        state=state,
        pts_x=pts_x,
        pts_y=pts_y,
        axs=axs,
        contour=contour,
        qp_type=WIGNER,
        cbar_label=cbar_label,
        axis_scale_factor=axis_scale_factor,
        plot_cbar=plot_cbar,
        plot_grid=plot_grid,
        x_ticks=x_ticks,
        y_ticks=y_ticks,
        z_ticks=z_ticks,
        subtitles=subtitles,
        figtitle=figtitle,
        gif=gif,
        gif_params=gif_params,
    )

plot_qfunc(state, pts_x, pts_y=None, g=2, axs=None, contour=True, cbar_label='', axis_scale_factor=1, plot_cbar=True, x_ticks=None, y_ticks=None, z_ticks=None, subtitles=None, figtitle=None, gif=False, gif_params=None)

Plot the husimi (Q) function of the state.

Thin wrapper around :func:plot_qp with qp_type='husimi'.

Parameters:

Name	Type	Description	Default
state		state with arbitrary number of batch dimensions, result will be flattened to a 2d grid to allow for plotting	required
pts_x		x points to evaluate quasi-probability distribution at	required
pts_y		y points to evaluate quasi-probability distribution at	None
g		float, default 2. Scaling factor for a = 0.5 * g * (x + iy). The value of g is related to the value of :math:\hbar in the commutation relation :math:[x,\,y] = i\hbar via :math:\hbar=2/g^2.	2
axs		matplotlib axes to plot on	None
contour		make the plot use contouring	True
cbar_label		label for the cbar	''
axis_scale_factor		scale of the axes labels relative	1
plot_cbar		whether to plot cbar	True
x_ticks		tick position for the x-axis	None
y_ticks		tick position for the y-axis	None
z_ticks		tick position for the z-axis	None
subtitles		subtitles for the subplots	None
figtitle		figure title	None
gif		if True, render an animation over one batch axis instead of a tiled subplot grid. See :func:plot_qp for details.	False
gif_params		dict of options for the gif path. Recognized keys: save_path (default None), interval_ms (default 200), ts (default None — adds a t = … label per frame), batch_animation_axis (default 0).	None

Returns:

Type	Description
	(axs, im) in the static case, or a matplotlib.animation.FuncAnimation
	when gif=True.

Source code in jaxquantum/core/visualization.py

def plot_qfunc(
    state,
    pts_x,
    pts_y=None,
    g=2,
    axs=None,
    contour=True,
    cbar_label="",
    axis_scale_factor=1,
    plot_cbar=True,
    x_ticks=None,
    y_ticks=None,
    z_ticks=None,
    subtitles=None,
    figtitle=None,
    gif=False,
    gif_params=None,
):
    """Plot the husimi (Q) function of the state.

    Thin wrapper around :func:`plot_qp` with ``qp_type='husimi'``.

    Args:
        state: state with arbitrary number of batch dimensions, result will
            be flattened to a 2d grid to allow for plotting
        pts_x: x points to evaluate quasi-probability distribution at
        pts_y: y points to evaluate quasi-probability distribution at
        g: float, default 2. Scaling factor for ``a = 0.5 * g * (x + iy)``.
            The value of ``g`` is related to the value of :math:`\\hbar` in
            the commutation relation :math:`[x,\,y] = i\\hbar` via
            :math:`\\hbar=2/g^2`.
        axs: matplotlib axes to plot on
        contour: make the plot use contouring
        cbar_label: label for the cbar
        axis_scale_factor: scale of the axes labels relative
        plot_cbar: whether to plot cbar
        x_ticks: tick position for the x-axis
        y_ticks: tick position for the y-axis
        z_ticks: tick position for the z-axis
        subtitles: subtitles for the subplots
        figtitle: figure title
        gif: if True, render an animation over one batch axis instead of a
            tiled subplot grid. See :func:`plot_qp` for details.
        gif_params: dict of options for the gif path. Recognized keys:
            ``save_path`` (default None), ``interval_ms`` (default 200),
            ``ts`` (default None — adds a ``t = …`` label per frame),
            ``batch_animation_axis`` (default 0).

    Returns:
        ``(axs, im)`` in the static case, or a ``matplotlib.animation.FuncAnimation``
        when ``gif=True``.
    """
    return plot_qp(
        state=state,
        pts_x=pts_x,
        pts_y=pts_y,
        g=g,
        axs=axs,
        contour=contour,
        qp_type=HUSIMI,
        cbar_label=cbar_label,
        axis_scale_factor=axis_scale_factor,
        plot_cbar=plot_cbar,
        x_ticks=x_ticks,
        y_ticks=y_ticks,
        z_ticks=z_ticks,
        subtitles=subtitles,
        figtitle=figtitle,
        gif=gif,
        gif_params=gif_params,
    )

plot_qp(state, pts_x, pts_y=None, g=2, axs=None, contour=True, qp_type=WIGNER, cbar_label='', axis_scale_factor=1, plot_cbar=True, x_ticks=None, y_ticks=None, z_ticks=None, subtitles=None, figtitle=None, gif=False, gif_params=None)

Plot a quasi-probability distribution (Wigner or Husimi-Q).

The state may carry an arbitrary number of batch dimensions; they are flattened to a 2D (rows, cols) grid of subplots. With gif=True, one batch axis is animated instead and the remaining batch dims form the per-frame subplot grid.

Parameters:

Name	Type	Description	Default
state		state with arbitrary number of batch dimensions; result will be flattened to a 2d grid to allow for plotting	required
pts_x		x points to evaluate the quasi-probability distribution at	required
pts_y		y points to evaluate the quasi-probability distribution at; defaults to pts_x	None
g		float, default 2. Scaling factor for a = 0.5 * g * (x + iy). The value of g is related to the value of :math:\hbar in the commutation relation :math:[x,\,y] = i\hbar via :math:\hbar=2/g^2.	2
axs		matplotlib axes to plot on (created if None)	None
contour		use contourf if True, otherwise pcolormesh	True
qp_type		type of quasi-probability distribution ("wigner" or "husimi")	WIGNER
cbar_label		label for the cbar (overridden internally based on qp_type)	''
axis_scale_factor		multiplicative scale applied to the axis tick positions and labels	1
plot_cbar		whether to draw a colorbar on each subplot	True
x_ticks		tick positions for the x-axis (auto if None)	None
y_ticks		tick positions for the y-axis (auto if None)	None
z_ticks		tick positions for the colorbar (auto if None)	None
subtitles		subtitles for the subplots; shape must match state.bdims (or the per-frame batch dims when gif=True)	None
figtitle		figure title	None
gif		if True, render an animation over one batch axis instead of a tiled subplot grid. Returns a matplotlib.animation.FuncAnimation that auto-renders inline in Jupyter (its _repr_html_ is patched to to_jshtml).	False
gif_params		dict of options for the gif path (ignored if gif=False). Recognized keys: save_path (default None) — if set, save the animation to this path via matplotlib.animation.PillowWriter. interval_ms (default 200) — milliseconds per frame; also derives fps = round(1000 / interval_ms) for the writer. ts (default None) — optional 1D array of timestamps matching the animation-axis length; when set, each frame's suptitle gets a t = … label. batch_animation_axis (default 0) — index into state.bdims selecting which axis becomes the animation axis. The remaining batch dims form the per-frame subplot grid.	None

Returns:

Type	Description
	(axs, im) in the static case, or a FuncAnimation when
	gif=True.

Source code in jaxquantum/core/visualization.py

def plot_qp(
    state,
    pts_x,
    pts_y=None,
    g=2,
    axs=None,
    contour=True,
    qp_type=WIGNER,
    cbar_label="",
    axis_scale_factor=1,
    plot_cbar=True,
    x_ticks=None,
    y_ticks=None,
    z_ticks=None,
    subtitles=None,
    figtitle=None,
    gif=False,
    gif_params=None,
):
    """Plot a quasi-probability distribution (Wigner or Husimi-Q).

    The state may carry an arbitrary number of batch dimensions; they are
    flattened to a 2D ``(rows, cols)`` grid of subplots. With ``gif=True``,
    one batch axis is animated instead and the remaining batch dims form
    the per-frame subplot grid.

    Args:
        state: state with arbitrary number of batch dimensions; result will
            be flattened to a 2d grid to allow for plotting
        pts_x: x points to evaluate the quasi-probability distribution at
        pts_y: y points to evaluate the quasi-probability distribution at;
            defaults to ``pts_x``
        g: float, default 2. Scaling factor for ``a = 0.5 * g * (x + iy)``.
            The value of ``g`` is related to the value of :math:`\\hbar` in
            the commutation relation :math:`[x,\,y] = i\\hbar` via
            :math:`\\hbar=2/g^2`.
        axs: matplotlib axes to plot on (created if None)
        contour: use ``contourf`` if True, otherwise ``pcolormesh``
        qp_type: type of quasi-probability distribution
            (``"wigner"`` or ``"husimi"``)
        cbar_label: label for the cbar (overridden internally based on
            ``qp_type``)
        axis_scale_factor: multiplicative scale applied to the axis tick
            positions and labels
        plot_cbar: whether to draw a colorbar on each subplot
        x_ticks: tick positions for the x-axis (auto if None)
        y_ticks: tick positions for the y-axis (auto if None)
        z_ticks: tick positions for the colorbar (auto if None)
        subtitles: subtitles for the subplots; shape must match
            ``state.bdims`` (or the per-frame batch dims when ``gif=True``)
        figtitle: figure title
        gif: if True, render an animation over one batch axis instead of a
            tiled subplot grid. Returns a
            ``matplotlib.animation.FuncAnimation`` that auto-renders inline
            in Jupyter (its ``_repr_html_`` is patched to ``to_jshtml``).
        gif_params: dict of options for the gif path (ignored if
            ``gif=False``). Recognized keys:

            - ``save_path`` (default ``None``) — if set, save the animation
              to this path via ``matplotlib.animation.PillowWriter``.
            - ``interval_ms`` (default ``200``) — milliseconds per frame;
              also derives ``fps = round(1000 / interval_ms)`` for the writer.
            - ``ts`` (default ``None``) — optional 1D array of timestamps
              matching the animation-axis length; when set, each frame's
              suptitle gets a ``t = …`` label.
            - ``batch_animation_axis`` (default ``0``) — index into
              ``state.bdims`` selecting which axis becomes the animation
              axis. The remaining batch dims form the per-frame subplot grid.

    Returns:
        ``(axs, im)`` in the static case, or a ``FuncAnimation`` when
        ``gif=True``.
    """
    if pts_y is None:
        pts_y = pts_x
    pts_x = jnp.array(pts_x)
    pts_y = jnp.array(pts_y)

    if len(state.bdims)==1 and state.bdims[0]==1:
        state = state[0]

    if gif:
        return _plot_qp_gif(
            state=state,
            pts_x=pts_x,
            pts_y=pts_y,
            g=g,
            axs=axs,
            contour=contour,
            qp_type=qp_type,
            axis_scale_factor=axis_scale_factor,
            plot_cbar=plot_cbar,
            x_ticks=x_ticks,
            y_ticks=y_ticks,
            z_ticks=z_ticks,
            subtitles=subtitles,
            figtitle=figtitle,
            gif_params=gif_params or {},
        )

    bdims = state.bdims
    added_baxes = 0

    if subtitles is not None:
        if subtitles.shape != bdims:
            raise ValueError(
                f"labels must have same shape as bdims, "
                f"got shapes {subtitles.shape} and {bdims}"
            )

    if len(bdims) == 0:
        bdims = (1,)
        added_baxes += 1
    if len(bdims) == 1:
        bdims = (1, bdims[0])
        added_baxes += 1

    extra_dims = bdims[2:]
    if extra_dims != ():
        state = state.reshape_bdims(
            bdims[0] * int(jnp.prod(jnp.array(extra_dims))), bdims[1]
        )
        if subtitles is not None:
            subtitles = subtitles.reshape(
                bdims[0] * int(jnp.prod(jnp.array(extra_dims))), bdims[1]
            )
        bdims = state.bdims

    if axs is None:
        _, axs = plt.subplots(
            bdims[0],
            bdims[1],
            figsize=(4 * bdims[1], 3 * bdims[0]),
            dpi=200,
        )

    if qp_type == WIGNER:
        vmin = -1
        vmax = 1
        scale = np.pi / 2
        cmap = "seismic"
        cbar_label = r"$\mathcal{W}(\alpha)$"
        QP = scale * wigner(state, pts_x, pts_y, g=g)

    elif qp_type == HUSIMI:
        vmin = 0
        vmax = 1
        scale = np.pi
        cmap = "jet"
        cbar_label = r"$\mathcal{Q}(\alpha)$"
        QP = scale * qfunc(state, pts_x, pts_y, g=g)



    for _ in range(added_baxes):
        QP = jnp.array([QP])
        axs = np.array([axs])
        if subtitles is not None:
            subtitles = np.array([subtitles])




    pts_x = pts_x * axis_scale_factor
    pts_y = pts_y * axis_scale_factor

    x_ticks = (
        jnp.linspace(jnp.min(pts_x), jnp.max(pts_x), 5) if x_ticks is None else x_ticks
    )
    y_ticks = (
        jnp.linspace(jnp.min(pts_y), jnp.max(pts_y), 5) if y_ticks is None else y_ticks
    )
    z_ticks = jnp.linspace(vmin, vmax, 3) if z_ticks is None else z_ticks

    im = _render_qp_grid(
        axs,
        QP,
        pts_x,
        pts_y,
        contour=contour,
        cmap=cmap,
        vmin=vmin,
        vmax=vmax,
        x_ticks=x_ticks,
        y_ticks=y_ticks,
        z_ticks=z_ticks,
        cbar_label=cbar_label,
        plot_cbar=plot_cbar,
        subtitles=subtitles,
        decorate=True,
    )

    fig = axs[bdims[0] - 1, bdims[1] - 1].get_figure()
    fig.tight_layout()
    if figtitle is not None:
        fig.suptitle(figtitle, y=1.04)
    return axs, im

plot_wigner(state, pts_x, pts_y=None, g=2, axs=None, contour=True, cbar_label='', axis_scale_factor=1, plot_cbar=True, x_ticks=None, y_ticks=None, z_ticks=None, subtitles=None, figtitle=None, gif=False, gif_params=None)

Plot the wigner function of the state.

Thin wrapper around :func:plot_qp with qp_type='wigner'.

Parameters:

Name	Type	Description	Default
state		state with arbitrary number of batch dimensions, result will be flattened to a 2d grid to allow for plotting	required
pts_x		x points to evaluate quasi-probability distribution at	required
pts_y		y points to evaluate quasi-probability distribution at	None
g		float, default 2. Scaling factor for a = 0.5 * g * (x + iy). The value of g is related to the value of :math:\hbar in the commutation relation :math:[x,\,y] = i\hbar via :math:\hbar=2/g^2.	2
axs		matplotlib axes to plot on	None
contour		make the plot use contouring	True
cbar_label		label for the cbar	''
axis_scale_factor		scale of the axes labels relative	1
plot_cbar		whether to plot cbar	True
x_ticks		tick position for the x-axis	None
y_ticks		tick position for the y-axis	None
z_ticks		tick position for the z-axis	None
subtitles		subtitles for the subplots	None
figtitle		figure title	None
gif		if True, render an animation over one batch axis instead of a tiled subplot grid. See :func:plot_qp for details.	False
gif_params		dict of options for the gif path. Recognized keys: save_path (default None), interval_ms (default 200), ts (default None — adds a t = … label per frame), batch_animation_axis (default 0).	None

Returns:

Type	Description
	(axs, im) in the static case, or a matplotlib.animation.FuncAnimation
	when gif=True.

Source code in jaxquantum/core/visualization.py

def plot_wigner(
    state,
    pts_x,
    pts_y=None,
    g=2,
    axs=None,
    contour=True,
    cbar_label="",
    axis_scale_factor=1,
    plot_cbar=True,
    x_ticks=None,
    y_ticks=None,
    z_ticks=None,
    subtitles=None,
    figtitle=None,
    gif=False,
    gif_params=None,
):
    """Plot the wigner function of the state.

    Thin wrapper around :func:`plot_qp` with ``qp_type='wigner'``.

    Args:
        state: state with arbitrary number of batch dimensions, result will
            be flattened to a 2d grid to allow for plotting
        pts_x: x points to evaluate quasi-probability distribution at
        pts_y: y points to evaluate quasi-probability distribution at
        g: float, default 2. Scaling factor for ``a = 0.5 * g * (x + iy)``.
            The value of ``g`` is related to the value of :math:`\\hbar` in
            the commutation relation :math:`[x,\,y] = i\\hbar` via
            :math:`\\hbar=2/g^2`.
        axs: matplotlib axes to plot on
        contour: make the plot use contouring
        cbar_label: label for the cbar
        axis_scale_factor: scale of the axes labels relative
        plot_cbar: whether to plot cbar
        x_ticks: tick position for the x-axis
        y_ticks: tick position for the y-axis
        z_ticks: tick position for the z-axis
        subtitles: subtitles for the subplots
        figtitle: figure title
        gif: if True, render an animation over one batch axis instead of a
            tiled subplot grid. See :func:`plot_qp` for details.
        gif_params: dict of options for the gif path. Recognized keys:
            ``save_path`` (default None), ``interval_ms`` (default 200),
            ``ts`` (default None — adds a ``t = …`` label per frame),
            ``batch_animation_axis`` (default 0).

    Returns:
        ``(axs, im)`` in the static case, or a ``matplotlib.animation.FuncAnimation``
        when ``gif=True``.
    """
    return plot_qp(
        state=state,
        pts_x=pts_x,
        pts_y=pts_y,
        g=g,
        axs=axs,
        contour=contour,
        qp_type=WIGNER,
        cbar_label=cbar_label,
        axis_scale_factor=axis_scale_factor,
        plot_cbar=plot_cbar,
        x_ticks=x_ticks,
        y_ticks=y_ticks,
        z_ticks=z_ticks,
        subtitles=subtitles,
        figtitle=figtitle,
        gif=gif,
        gif_params=gif_params,
    )