Matplotlib plot zooming with scroll wheel

Thanks guys, the examples were very helpful. I had to make a few changes to work with a scatter plot and I added panning with a left button drag. Hopefully someone will find this useful.

from matplotlib.pyplot import figure, show
import numpy

class ZoomPan:
    def __init__(self):
        self.press = None
        self.cur_xlim = None
        self.cur_ylim = None
        self.x0 = None
        self.y0 = None
        self.x1 = None
        self.y1 = None
        self.xpress = None
        self.ypress = None


    def zoom_factory(self, ax, base_scale = 2.):
        def zoom(event):
            cur_xlim = ax.get_xlim()
            cur_ylim = ax.get_ylim()

            xdata = event.xdata # get event x location
            ydata = event.ydata # get event y location

            if event.button == 'down':
                # deal with zoom in
                scale_factor = 1 / base_scale
            elif event.button == 'up':
                # deal with zoom out
                scale_factor = base_scale
            else:
                # deal with something that should never happen
                scale_factor = 1
                print event.button

            new_width = (cur_xlim[1] - cur_xlim[0]) * scale_factor
            new_height = (cur_ylim[1] - cur_ylim[0]) * scale_factor

            relx = (cur_xlim[1] - xdata)/(cur_xlim[1] - cur_xlim[0])
            rely = (cur_ylim[1] - ydata)/(cur_ylim[1] - cur_ylim[0])

            ax.set_xlim([xdata - new_width * (1-relx), xdata + new_width * (relx)])
            ax.set_ylim([ydata - new_height * (1-rely), ydata + new_height * (rely)])
            ax.figure.canvas.draw()

        fig = ax.get_figure() # get the figure of interest
        fig.canvas.mpl_connect('scroll_event', zoom)

        return zoom

    def pan_factory(self, ax):
        def onPress(event):
            if event.inaxes != ax: return
            self.cur_xlim = ax.get_xlim()
            self.cur_ylim = ax.get_ylim()
            self.press = self.x0, self.y0, event.xdata, event.ydata
            self.x0, self.y0, self.xpress, self.ypress = self.press

        def onRelease(event):
            self.press = None
            ax.figure.canvas.draw()

        def onMotion(event):
            if self.press is None: return
            if event.inaxes != ax: return
            dx = event.xdata - self.xpress
            dy = event.ydata - self.ypress
            self.cur_xlim -= dx
            self.cur_ylim -= dy
            ax.set_xlim(self.cur_xlim)
            ax.set_ylim(self.cur_ylim)

            ax.figure.canvas.draw()

        fig = ax.get_figure() # get the figure of interest

        # attach the call back
        fig.canvas.mpl_connect('button_press_event',onPress)
        fig.canvas.mpl_connect('button_release_event',onRelease)
        fig.canvas.mpl_connect('motion_notify_event',onMotion)

        #return the function
        return onMotion


fig = figure()

ax = fig.add_subplot(111, xlim=(0,1), ylim=(0,1), autoscale_on=False)

ax.set_title('Click to zoom')
x,y,s,c = numpy.random.rand(4,200)
s *= 200

ax.scatter(x,y,s,c)
scale = 1.1
zp = ZoomPan()
figZoom = zp.zoom_factory(ax, base_scale = scale)
figPan = zp.pan_factory(ax)
show()

The other answers using ax.set_xlim() and ax.set_ylim() did not give a satisfactory user experience for figures where setting the axes is slow. (for me this was an axes with a pcolormesh) The method ax.drag_pan() is much faster, and I believe it is more suited for most cases:

def mousewheel_move( event):
    ax=event.inaxes
    ax._pan_start = types.SimpleNamespace(
            lim=ax.viewLim.frozen(),
            trans=ax.transData.frozen(),
            trans_inverse=ax.transData.inverted().frozen(),
            bbox=ax.bbox.frozen(),
            x=event.x,
            y=event.y)
    if event.button == 'up':
        ax.drag_pan(3, event.key, event.x+10, event.y+10)
    else: #event.button == 'down':
        ax.drag_pan(3, event.key, event.x-10, event.y-10)
    fig=ax.get_figure()
    fig.canvas.draw_idle()

Then connect your figure with:

fig.canvas.mpl_connect('scroll_event',mousewheel_move)

Tested with matplotlib 3.0.2 using the the TkAgg backend and python 3.6

There is another way of doing it as far as I know. By chance I came across the Axis.zoom method. I don't know if this is faster or a good way in general, but it works and is certainly less code:

    def __init(self):
        ...

        self.cid_zoom = self.canvas.mpl_connect('scroll_event', self.zoom)

    def zoom(self, event):
        if event.inaxes == self.ax:
            scale_factor = np.power(self.zoom_factor, -event.step)*event.step
            self.ax.get_xaxis().zoom(scale_factor)
            self.ax.get_yaxis().zoom(scale_factor)
            self.ax.invert_yaxis()
            self.canvas.draw_idle()

If you plot an image though, for some reason, you have to invert the y-axis again.

You can also sort of implement panning that way, but it doesn't really work that well. I'm not sure why:

    def __init(self):
        ...

        self.cid_motion = self.canvas.mpl_connect(
            'motion_notify_event', self.pan_move
        )
        self.cid_button = self.canvas.mpl_connect(
            'button_press_event', self.pan_press
        )


    def pan_press(self, event):
        if event.inaxes == self.ax:
            self.x_press = event.xdata
            self.y_press = event.ydata

    def pan_move(self, event):
        if event.button == 1 and event.inaxes == self.ax:
            xdata = event.xdata
            ydata = event.ydata
            dx = (xdata - self.x_press)/np.diff(self.ax.get_xlim())
            dy = (ydata - self.y_press)/np.diff(self.ax.get_ylim())
            self.ax.get_xaxis().pan(-dx)
            self.ax.get_yaxis().pan(-dy)
            self.ax.drag_pan(event.button, event.key, dx, dy)
            self.canvas.draw()

make tacaswell 's answer 'smooth'

def zoom_factory(ax, base_scale=2.):
    prex = 0
    prey = 0
    prexdata = 0
    preydata = 0

    def zoom_fun(event):
        nonlocal prex, prey, prexdata, preydata
        curx = event.x
        cury = event.y

        # if not changed mouse position(or changed so little)
        # remain the pre scale center
        if abs(curx - prex) < 10 and abs(cury - prey) < 10:
            # remain same
            xdata = prexdata
            ydata = preydata
        # if changed mouse position ,also change the cur scale center
        else:
            # change
            xdata = event.xdata  # get event x location
            ydata = event.ydata  # get event y location

            # update previous location data
            prex = event.x
            prey = event.y
            prexdata = xdata
            preydata = ydata

        # get the current x and y limits
        cur_xlim = ax.get_xlim()
        cur_ylim = ax.get_ylim()

        cur_xrange = (cur_xlim[1] - cur_xlim[0]) * .5
        cur_yrange = (cur_ylim[1] - cur_ylim[0]) * .5

        # log.debug((xdata, ydata))
        if event.button == 'up':
            # deal with zoom in
            scale_factor = 1 / base_scale
        elif event.button == 'down':
            # deal with zoom out
            scale_factor = base_scale
        else:
            # deal with something that should never happen
            scale_factor = 1
            print(event.button)
        # set new limits
        ax.set_xlim([
            xdata - cur_xrange * scale_factor,
            xdata + cur_xrange * scale_factor
        ])
        ax.set_ylim([
            ydata - cur_yrange * scale_factor,
            ydata + cur_yrange * scale_factor
        ])
        plt.draw()  # force re-draw

    fig = ax.get_figure()  # get the figure of interest
    # attach the call back
    fig.canvas.mpl_connect('scroll_event', zoom_fun)

    # return the function
    return zoom_fun