matplotlib imshow with auto-zoom on the color limits

az_imshow.py

import numpy as np
import matplotlib.pyplot as plt

def iround(x, x0):
    return (np.abs(x-x0)).argmin()


class az_imshow:
    def __init__(self, data, ax=None, **kwarg):
        if ax == None:
            ax = plt.gca()
        self.data = data
        self.ax = ax
        self.image = ax.imshow(data, **kwarg)
        cid = ax.figure.canvas.mpl_connect('draw_event', self.redraw)
        cid = ax.figure.canvas.mpl_connect('resize_event', self.redraw)
        self.lims = np.r_[self.ax.get_xlim(), self.ax.get_ylim()]

    def redraw(self, evt):
        lims = np.r_[self.ax.get_xlim(), self.ax.get_ylim()]
        print lims
        if (self.lims == lims).all():
            return
        self.lims = lims
        xmin, xmax, ymin, ymax = self.image.get_extent()
        x = np.linspace(xmin, xmax, self.data.shape[1])
        y = np.linspace(ymin, ymax, self.data.shape[0])
        xind = [iround(x, ii) for ii in self.ax.get_xlim()]
        yind = [iround(y, ii) for ii in self.ax.get_ylim()]
        local_data = self.data[yind[0]:yind[1],
                               xind[0]:xind[1]]
        self.image.set_clim(local_data.min(), local_data.max())
        self.ax.figure.canvas.draw()

if __name__ == "__main__":
    plt.ioff()
    fig = plt.figure()
    ax1 = fig.add_subplot(211)
    data=np.random.rand(200,100)*np.linspace(0,1,100)
    az_imshow(data, ax1, interpolation='nearest', aspect='auto')
    plt.subplot(212)
    az_imshow(data, aspect='auto')
    plt.show()