utils.py 文件源码

def plot_jet_image(
        ax, image, vmin=1e-9, vmax=1e-2, cmap="jet", title="Intensity",
        label_axes=True, visible_axes=False, show_colorbar=True, colorbar_inside=False):
    """Display jet image.

    Args:
        ax: matplotlib axes to plot on.
        image: array representing image to plot.
        vmin, vmax: min, max intensity values to plot.
    """
    width, height = image.T.shape
    dw, dh = 1./width, 1./height
    if not (vmin is None) and not (vmax is None):
        if vmin < 0:
            norm = MidPointNorm(vmin=vmin, vmax=vmax)
            ticks = None
        else:
            norm = LogNorm(vmin=vmin, vmax=vmax)
            ticks = np.logspace(
                np.log10(vmin), np.log10(vmax),
                1 + np.log10(vmax) - np.log10(vmin))
    else:
        norm = None
        ticks = None

    p = ax.imshow(
        image.T, extent=(-(1+dw), 1+dw, -(1+dh), 1+dh), origin='low',
        interpolation='nearest', norm=norm, cmap=cmap)

    if show_colorbar:
        if colorbar_inside:
            cax = ax.figure.add_axes([0.85, 0.08, 0.03, 0.82])
        else:
            divider = make_axes_locatable(ax)
            cax = divider.append_axes("right", size="5%", pad=0.05)
        cbar = plt.colorbar(p, cax=cax, ticks=ticks)
        cbar.set_label(title, rotation=90, fontsize=18)
        cbar.ax.tick_params(labelsize=12)
        if colorbar_inside:
            cbar.ax.yaxis.set_ticks_position('left')

    if label_axes:
        ax.set_xlabel(r'$x_1$', fontsize=18)
        ax.set_ylabel(r'$x_2$', fontsize=18)
        ax.tick_params(axis='both', which='major', labelsize=12)
    else:
        ax.axes.get_xaxis().set_ticks([])
        ax.axes.get_yaxis().set_ticks([])

    ax.axes.get_xaxis().set_visible(visible_axes)
    ax.axes.get_yaxis().set_visible(visible_axes)
    if visible_axes:
        for spine in ['top', 'bottom', 'left', 'right']:
            ax.spines[spine].set_linewidth(1)
            ax.spines[spine].set_color('k')
    else:
        for spine in ['top', 'bottom', 'left', 'right']:
            ax.spines[spine].set_visible(False)
    ax.axes.grid(False)