Skip to content

visualization

Visualization utils.

plot_qp(state, pts, ax=None, contour=True, qp_type=WIGNER)

Plot quasi-probability distribution.

TODO: decouple this from qutip.

Parameters:

Name Type Description Default
state

statevector

 required
pts

points to evaluate quasi-probability distribution on

 required
dim

dimensions of state

 required
ax

matplotlib axis to plot on

 None
contour

make the plot use contouring

 True
qp_type

type of quasi probability distribution ("wigner", "qfunc")

 WIGNER

Returns:

Type Description

axis on which the plot was plotted.
Source code in jaxquantum/core/visualization.py 
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
def plot_qp(state, pts, ax=None, contour=True, qp_type=WIGNER):
    """Plot quasi-probability distribution.

    TODO: decouple this from qutip.

    Args:
        state: statevector
        pts: points to evaluate quasi-probability distribution on
        dim: dimensions of state
        ax: matplotlib axis to plot on
        contour: make the plot use contouring
        qp_type: type of quasi probability distribution ("wigner", "qfunc")

    Returns:
        axis on which the plot was plotted.
    """
    pts = np.array(pts)
    state = jqt2qt(state)
    if ax is None:
        _, ax = plt.subplots(1, figsize=(4, 3), dpi=200)
    # fig = ax.get_figure()

    if qp_type == WIGNER:
        vmin = -1
        vmax = 1
        scale = np.pi / 2
        cmap = "seismic"
    elif qp_type == QFUNC:
        vmin = 0
        vmax = 1
        scale = np.pi
        cmap = "jet"

    QP = scale * getattr(qt, qp_type)(state, pts, pts, g=2)

    if contour:
        im = ax.contourf(
            pts,
            pts,
            QP,
            cmap=cmap,
            vmin=vmin,
            vmax=vmax,
            levels=np.linspace(vmin, vmax, 101),
        )
    else:
        im = ax.pcolormesh(
            pts,
            pts,
            QP,
            cmap=cmap,
            vmin=vmin,
            vmax=vmax,
        )
    ax.axhline(0, linestyle="-", color="black", alpha=0.7)
    ax.axvline(0, linestyle="-", color="black", alpha=0.7)
    ax.grid()
    ax.set_aspect("equal", adjustable="box")
    return im