qiskit/visualization/state_visualization.py
def plot_state_city(rho, title="", figsize=None, color=None,
                    alpha=1):
    """Plot the cityscape of quantum state.

    Plot two 3d bar graphs (two dimensional) of the real and imaginary
    part of the density matrix rho.

    Args:
        rho (ndarray): Numpy array for state vector or density matrix.
        title (str): a string that represents the plot title
        figsize (tuple): Figure size in inches.
        color (list): A list of len=2 giving colors for real and
        imaginary components of matrix elements.
        alpha (float): Transparency value for bars
    Returns:
         matplotlib.Figure: The matplotlib.Figure of the visualization

    Raises:
        ImportError: Requires matplotlib.
        ValueError: When 'color' is not a list of len=2.
    """
    if not HAS_MATPLOTLIB:
        raise ImportError('Must have Matplotlib installed.')
    rho = _validate_input_state(rho)

    num = int(np.log2(len(rho)))
    # get the real and imag parts of rho
    datareal = np.real(rho)
    dataimag = np.imag(rho)

    # get the labels
    column_names = [bin(i)[2:].zfill(num) for i in range(2**num)]
    row_names = [bin(i)[2:].zfill(num) for i in range(2**num)]

    lx = len(datareal[0])            # Work out matrix dimensions
    ly = len(datareal[:, 0])
    xpos = np.arange(0, lx, 1)    # Set up a mesh of positions
    ypos = np.arange(0, ly, 1)
    xpos, ypos = np.meshgrid(xpos+0.25, ypos+0.25)

    xpos = xpos.flatten()
    ypos = ypos.flatten()
    zpos = np.zeros(lx*ly)

    dx = 0.5 * np.ones_like(zpos)  # width of bars
    dy = dx.copy()
    dzr = datareal.flatten()
    dzi = dataimag.flatten()

    if color is None:
        color = ["#648fff", "#648fff"]
    else:
        if len(color) != 2:
            raise ValueError("'color' must be a list of len=2.")
        if color[0] is None:
            color[0] = "#648fff"
        if color[1] is None:
            color[1] = "#648fff"

    # set default figure size
    if figsize is None:
        figsize = (15, 5)

    fig = plt.figure(figsize=figsize)
    ax1 = fig.add_subplot(1, 2, 1, projection='3d')

    x = [0, max(xpos)+0.5, max(xpos)+0.5, 0]
    y = [0, 0, max(ypos)+0.5, max(ypos)+0.5]
    z = [0, 0, 0, 0]
    verts = [list(zip(x, y, z))]

    fc1 = generate_facecolors(xpos, ypos, zpos, dx, dy, dzr, color[0])
    for idx, cur_zpos in enumerate(zpos):
        if dzr[idx] > 0:
            zorder = 2
        else:
            zorder = 0
        b1 = ax1.bar3d(xpos[idx], ypos[idx], cur_zpos,
                       dx[idx], dy[idx], dzr[idx],
                       alpha=alpha, zorder=zorder)
        b1.set_facecolors(fc1[6*idx:6*idx+6])

    pc1 = Poly3DCollection(verts, alpha=0.15, facecolor='k',
                           linewidths=1, zorder=1)

    if min(dzr) < 0 < max(dzr):
        ax1.add_collection3d(pc1)

    ax2 = fig.add_subplot(1, 2, 2, projection='3d')
    fc2 = generate_facecolors(xpos, ypos, zpos, dx, dy, dzi, color[1])
    for idx, cur_zpos in enumerate(zpos):
        if dzi[idx] > 0:
            zorder = 2
        else:
            zorder = 0
        b2 = ax2.bar3d(xpos[idx], ypos[idx], cur_zpos,
                       dx[idx], dy[idx], dzi[idx],
                       alpha=alpha, zorder=zorder)
        b2.set_facecolors(fc2[6*idx:6*idx+6])

    pc2 = Poly3DCollection(verts, alpha=0.2, facecolor='k',
                           linewidths=1, zorder=1)

    if min(dzi) < 0 < max(dzi):
        ax2.add_collection3d(pc2)
    ax1.set_xticks(np.arange(0.5, lx+0.5, 1))
    ax1.set_yticks(np.arange(0.5, ly+0.5, 1))
    max_dzr = max(dzr)
    min_dzr = min(dzr)
    if max_dzr != min_dzr:
        ax1.axes.set_zlim3d(np.min(dzr), np.max(dzr)+1e-9)
    else:
        if min_dzr == 0:
            ax1.axes.set_zlim3d(np.min(dzr), np.max(dzr)+1e-9)
        else:
            ax1.axes.set_zlim3d(auto=True)
    ax1.zaxis.set_major_locator(MaxNLocator(5))
    ax1.w_xaxis.set_ticklabels(row_names, fontsize=14, rotation=45)
    ax1.w_yaxis.set_ticklabels(column_names, fontsize=14, rotation=-22.5)
    ax1.set_zlabel("Real[rho]", fontsize=14)
    for tick in ax1.zaxis.get_major_ticks():
        tick.label.set_fontsize(14)

    ax2.set_xticks(np.arange(0.5, lx+0.5, 1))
    ax2.set_yticks(np.arange(0.5, ly+0.5, 1))
    min_dzi = np.min(dzi)
    max_dzi = np.max(dzi)
    if min_dzi != max_dzi:
        eps = 0
        ax2.zaxis.set_major_locator(MaxNLocator(5))
        ax2.axes.set_zlim3d(np.min(dzi), np.max(dzi)+eps)
    else:
        if min_dzi == 0:
            ax2.set_zticks([0])
            eps = 1e-9
            ax2.axes.set_zlim3d(np.min(dzi), np.max(dzi)+eps)
        else:
            ax2.axes.set_zlim3d(auto=True)
    ax2.w_xaxis.set_ticklabels(row_names, fontsize=14, rotation=45)
    ax2.w_yaxis.set_ticklabels(column_names, fontsize=14, rotation=-22.5)
    ax2.set_zlabel("Imag[rho]", fontsize=14)
    for tick in ax2.zaxis.get_major_ticks():
        tick.label.set_fontsize(14)
    plt.suptitle(title, fontsize=16)
    plt.tight_layout()
    plt.close(fig)
    return fig
def plot_state_city(rho, title="", figsize=None, color=None,
                    alpha=1):
    """Plot the cityscape of quantum state.

    Plot two 3d bar graphs (two dimensional) of the real and imaginary
    part of the density matrix rho.

    Args:
        rho (ndarray): Numpy array for state vector or density matrix.
        title (str): a string that represents the plot title
        figsize (tuple): Figure size in inches.
        color (list): A list of len=2 giving colors for real and
        imaginary components of matrix elements.
        alpha (float): Transparency value for bars
    Returns:
         matplotlib.Figure: The matplotlib.Figure of the visualization

    Raises:
        ImportError: Requires matplotlib.
        ValueError: When 'color' is not a list of len=2.
    """
    if not HAS_MATPLOTLIB:
        raise ImportError('Must have Matplotlib installed.')
    rho = _validate_input_state(rho)

    num = int(np.log2(len(rho)))
    # get the real and imag parts of rho
    datareal = np.real(rho)
    dataimag = np.imag(rho)

    # get the labels
    column_names = [bin(i)[2:].zfill(num) for i in range(2**num)]
    row_names = [bin(i)[2:].zfill(num) for i in range(2**num)]

    lx = len(datareal[0])            # Work out matrix dimensions
    ly = len(datareal[:, 0])
    xpos = np.arange(0, lx, 1)    # Set up a mesh of positions
    ypos = np.arange(0, ly, 1)
    xpos, ypos = np.meshgrid(xpos+0.25, ypos+0.25)

    xpos = xpos.flatten()
    ypos = ypos.flatten()
    zpos = np.zeros(lx*ly)

    dx = 0.5 * np.ones_like(zpos)  # width of bars
    dy = dx.copy()
    dzr = datareal.flatten()
    dzi = dataimag.flatten()

    if color is None:
        color = ["#648fff", "#648fff"]
    else:
        if len(color) != 2:
            raise ValueError("'color' must be a list of len=2.")
        if color[0] is None:
            color[0] = "#648fff"
        if color[1] is None:
            color[1] = "#648fff"

    # set default figure size
    if figsize is None:
        figsize = (15, 5)

    fig = plt.figure(figsize=figsize)
    ax1 = fig.add_subplot(1, 2, 1, projection='3d')

    x = [0, max(xpos)+0.5, max(xpos)+0.5, 0]
    y = [0, 0, max(ypos)+0.5, max(ypos)+0.5]
    z = [0, 0, 0, 0]
    verts = [list(zip(x, y, z))]

    fc1 = generate_facecolors(xpos, ypos, zpos, dx, dy, dzr, color[0])
    for idx, cur_zpos in enumerate(zpos):
        if dzr[idx] > 0:
            zorder = 2
        else:
            zorder = 0
        b1 = ax1.bar3d(xpos[idx], ypos[idx], cur_zpos,
                       dx[idx], dy[idx], dzr[idx],
                       alpha=alpha, zorder=zorder)
        b1.set_facecolors(fc1[6*idx:6*idx+6])

    pc1 = Poly3DCollection(verts, alpha=0.15, facecolor='k',
                           linewidths=1, zorder=1)

    if min(dzr) < 0 < max(dzr):
        ax1.add_collection3d(pc1)

    ax2 = fig.add_subplot(1, 2, 2, projection='3d')
    fc2 = generate_facecolors(xpos, ypos, zpos, dx, dy, dzi, color[1])
    for idx, cur_zpos in enumerate(zpos):
        if dzi[idx] > 0:
            zorder = 2
        else:
            zorder = 0
        b2 = ax2.bar3d(xpos[idx], ypos[idx], cur_zpos,
                       dx[idx], dy[idx], dzi[idx],
                       alpha=alpha, zorder=zorder)
        b2.set_facecolors(fc2[6*idx:6*idx+6])

    pc2 = Poly3DCollection(verts, alpha=0.2, facecolor='k',
                           linewidths=1, zorder=1)

    if min(dzi) < 0 < max(dzi):
        ax2.add_collection3d(pc2)
    ax1.set_xticks(np.arange(0.5, lx+0.5, 1))
    ax1.set_yticks(np.arange(0.5, ly+0.5, 1))
    max_dzr = max(dzr)
    min_dzr = min(dzr)
    if max_dzr != min_dzr:
        ax1.axes.set_zlim3d(np.min(dzr), np.max(dzr)+1e-9)
    else:
        if min_dzr == 0:
            ax1.axes.set_zlim3d(np.min(dzr), np.max(dzr)+1e-9)
        else:
            ax1.axes.set_zlim3d(auto=True)
    ax1.zaxis.set_major_locator(MaxNLocator(5))
    ax1.w_xaxis.set_ticklabels(row_names, fontsize=14, rotation=45)
    ax1.w_yaxis.set_ticklabels(column_names, fontsize=14, rotation=-22.5)
    ax1.set_zlabel("Real[rho]", fontsize=14)
    for tick in ax1.zaxis.get_major_ticks():
        tick.label.set_fontsize(14)

    ax2.set_xticks(np.arange(0.5, lx+0.5, 1))
    ax2.set_yticks(np.arange(0.5, ly+0.5, 1))
    min_dzi = np.min(dzi)
    max_dzi = np.max(dzi)
    if min_dzi != max_dzi:
        eps = 0
        ax2.zaxis.set_major_locator(MaxNLocator(5))
        ax2.axes.set_zlim3d(np.min(dzi), np.max(dzi)+eps)
    else:
        if min_dzi == 0:
            ax2.set_zticks([0])
            eps = 1e-9
            ax2.axes.set_zlim3d(np.min(dzi), np.max(dzi)+eps)
        else:
            ax2.axes.set_zlim3d(auto=True)
    ax2.w_xaxis.set_ticklabels(row_names, fontsize=14, rotation=45)
    ax2.w_yaxis.set_ticklabels(column_names, fontsize=14, rotation=-22.5)
    ax2.set_zlabel("Imag[rho]", fontsize=14)
    for tick in ax2.zaxis.get_major_ticks():
        tick.label.set_fontsize(14)
    plt.suptitle(title, fontsize=16)
    plt.tight_layout()
    plt.close(fig)
    return fig
Plot the cityscape of quantum state.

    Plot two 3d bar graphs (two dimensional) of the real and imaginary
    part of the density matrix rho.

    Args:
        rho (ndarray): Numpy array for state vector or density matrix.
        title (str): a string that represents the plot title
        figsize (tuple): Figure size in inches.
        color (list): A list of len=2 giving colors for real and
        imaginary components of matrix elements.
        alpha (float): Transparency value for bars
    Returns:
         matplotlib.Figure: The matplotlib.Figure of the visualization

    Raises:
        ImportError: Requires matplotlib.
        ValueError: When 'color' is not a list of len=2.