quick_imshow.py

from pylab import *

def quick_imshow(nrows, ncols, images, titles=None, colorbar=True, vmax=None,
                 vmin=None, figsize=None, figtitle=None):
    """-------------------------------------------------------------------------
    Desc.:      convenience function that make subplots of imshow
    Args.:      nrows - number of rows
                ncols - number of cols
                images - list of images
                titles - list of titles
                vmax - tuple of vmax for the colormap. If scalar,
                        the same value is used for all subplots. If one
                        of the entries is None, no colormap for that
                        subplot will be drawn.
                 vmin - tuple of vmin
    Returns:    f - the figure handle
                axes - axes or array of axes objects
                caxes - tuple of axes image
    -------------------------------------------------------------------------"""
    if figsize == None:
        # 1.0 translates to 100 pixels of the figure
        s = 3.5
        if figtitle:
            figsize = (s * ncols, s * nrows + 0.5)
        else:
            figsize = (s * ncols, s * nrows)

    if nrows == ncols == 1:
        f, ax = subplots(figsize=figsize)
        cax = ax.imshow(images, cmap='jet', vmax=vmax, vmin=vmin)
        if colorbar:
            f.colorbar(cax, ax=ax)
        if titles != None:
            ax.set_title(titles)
        if figtitle != None:
            f.suptitle(figtitle)
        cax.axes.get_xaxis().set_visible(False)
        cax.axes.get_yaxis().set_visible(False)
        return f, ax, cax

    f, axes = subplots(nrows, ncols, figsize=figsize)
    caxes = []
    i = 0
    for ax, img in zip(axes.flat, images):
        if isinstance(vmax, tuple) and isinstance(vmin, tuple):
            if vmax[i] is not None and vmin[i] is not None:
                cax = ax.imshow(img, cmap='jet', vmax=vmax[i], vmin=vmin[i])
            else:
                cax = ax.imshow(img, cmap='jet')
        elif isinstance(vmax, tuple) and vmin is None:
            if vmax[i] is not None:
                cax = ax.imshow(img, cmap='jet', vmax=vmax[i], vmin=0)
            else:
                cax = ax.imshow(img, cmap='jet')
        elif vmax is None and vmin is None:
            cax = ax.imshow(img, cmap='jet')
        else:
            cax = ax.imshow(img, cmap='jet', vmax=vmax, vmin=vmin)
        if titles != None:
            ax.set_title(titles[i])
        if colorbar:
            f.colorbar(cax, ax=ax)
        caxes.append(cax)
        cax.axes.get_xaxis().set_visible(False)
        cax.axes.get_yaxis().set_visible(False)
        i = i + 1
    if figtitle != None:
        f.suptitle(figtitle)
    return f, axes, tuple(caxes)


def update_subplots(images, caxes, f=None, axes=None, indices=(), vmax=None,
                    vmin=None):
    """-------------------------------------------------------------------------
    Desc.:      update subplots in a figure
    Args.:      images  - new images to plot
                caxes   - caxes returned at figure creation
                indices - specific indices of subplots to be updated
    Returns:
    -------------------------------------------------------------------------"""
    for i in range(len(images)):
        if len(indices) > 0:
            ind = indices[i]
        else:
            ind = i
        img = images[i]
        caxes[ind].set_data(img)
        cbar = caxes[ind].colorbar
        if isinstance(vmax, tuple) and isinstance(vmin, tuple):
            if vmax[i] is not None and vmin[i] is not None:
                cbar.set_clim([vmin[i], vmax[i]])
            else:
                cbar.set_clim([img.min(), img.max()])
        elif isinstance(vmax, tuple) and vmin is None:
            if vmax[i] is not None:
                cbar.set_clim([0, vmax[i]])
            else:
                cbar.set_clim([img.min(), img.max()])
        elif vmax is None and vmin is None:
            cbar.set_clim([img.min(), img.max()])
        else:
            cbar.set_clim([vmin, vmax])
        cbar.update_normal(caxes[ind])
    pause(0.01)
    tight_layout()