A Matplotlib plot with bottons and mouse scroll for y-axis zooming

matplotlib_bottons_x_axis_zoom.py

#!/usr/bin/python3
'''
Copyright (C) Saeed Gholami Shahbandi. All rights reserved.
Author: Saeed Gholami Shahbandi

This file is free software: you can redistribute it and/or modify it
under the terms of the GNU Lesser General Public License as published
by the Free Software Foundation, either version 3 of the License, or
(at your option) any later version.  This program is distributed in
the hope that it will be useful, but WITHOUT ANY WARRANTY; without
even the implied warranty of MERCHANTABILITY or FITNESS FOR A
PARTICULAR PURPOSE. See the GNU Lesser General Public License for more
details.  You should have received a copy of the GNU Lesser General
Public License along with this program. If not, see
<http://www.gnu.org/licenses/>
'''

import numpy as np
import matplotlib.pyplot as plt
from matplotlib.widgets import Button, TextBox
from matplotlib.gridspec import GridSpec, GridSpecFromSubplotSpec

class BottonClass(object):
    '''- '''
    ########################################
    def __init__(self, figsize=(18,10), zoom_factor=1.5):
        '''-

        # TODO: add a new argument that takes the (nrow, ncol) for the subplots grid

        Scroll up and down zoom in and out, where the zoom power and
        direction are defined by the zoom_factor parameter.
        
        "Next" and "Previous" bottons iterate trough different colors.
        Left and Right keys from the keyboard also iterate trough different colors.
        
        Paramters
        ---------
        - figsize: tuple (default: (18,10))
        fizgsize parameter for matplotlib.pyplot.figure

        - zoom_factor: float 0< zoom_factor <1 (default: 1.5)
        zoom_factor>1: scroll up is zoom in
        zoom_factor<1: scroll up is zoom out
        The farther the factor is from 1, the more powerful is the zoom.

        '''

        self.zoom_factor = zoom_factor

        self.x = np.arange(0, 2*np.pi, 0.05)
        self.y = np.sin(5* self.x)
        self.idx = 0
        self.clrs = [ 'SteelBlue', 'Tomato', 'MediumPurple',
                      'DarkCyan', 'DarkRed', 'DarkOliveGreen', 'Goldenrod',
                      'Orange', 'MediumVioletRed', 'SandyBrown', 'Chocolate' ]

        self.create_figure(figsize)
        self.plot()
        plt.show()

    ########################################
    def create_figure(self, figsize):
        '''- '''

        self.fig = plt.figure( figsize = figsize )
        self.fig.canvas.mpl_connect('scroll_event', self.zoom_y_axis_event)
        self.fig.canvas.mpl_connect('key_press_event', self.keyboard_ctrl_event)

        gs_global = GridSpec(nrows=3, ncols=1, height_ratios=[4, 4, .5], figure=self.fig)
        gs_widgets = GridSpecFromSubplotSpec(nrows=1, ncols=3, width_ratios=[8,1,1], subplot_spec=gs_global[2])

        ### plotting axes
        self.ax_top = plt.subplot(gs_global[0], label='top_axis')
        self.ax_bot = plt.subplot(gs_global[1], label='bot_axis')

        ### widget axes
        # ax_text = plt.subplot(gs_widgets[0])
        self.ax_prv = plt.subplot(gs_widgets[1])
        self.ax_nxt = plt.subplot(gs_widgets[2])

        ###
        plt.tight_layout()

        ### connecting axes to widgets
        # self.text_box = TextBox(self.ax_txt, 'Evaluate', initial=str(self.vins[0]))
        # self.text_box.on_submit(self.submit_txt)

        self.button_prev = Button(self.ax_prv, 'Previous')
        self.button_prev.on_clicked(self.prev)

        self.button_next = Button(self.ax_nxt, 'Next')
        self.button_next.on_clicked(self.next)
        
        return None

    ########################################
    def keyboard_ctrl_event(self, event):
        '''- '''
        if event.key == 'right':
            self.next()
        elif event.key == 'left':
            self.prev()
        
    
    ########################################
    def zoom_y_axis_event(self, event):
        '''- '''
                
        # print ( axis.get_label() )
        axis = event.inaxes
        xmin, xmax = axis.get_xlim()
        xcursor = event.xdata
        
        if event.button == 'up':
            # zoom out
            xmin_new = xcursor - (xcursor - xmin) * self.zoom_factor
            xmin_new = max( xmin_new, axis.initial_xlim[0] )
            
            xmax_new = xcursor + (xmax - xcursor) * self.zoom_factor
            xmax_new = min( xmax_new, axis.initial_xlim[1] )
            
        elif event.button == 'down':
            # zoom in 
            xmin_new = xcursor - (xcursor - xmin) / self.zoom_factor
            xmin_new = max( xmin_new, axis.initial_xlim[0] )

            xmax_new = xcursor + (xmax - xcursor) / self.zoom_factor
            xmax_new = min( xmax_new, axis.initial_xlim[1] )

        else:
            # double click to reset?
            return None

        axis.set_xlim(xmin_new, xmax_new)
        plt.draw()
        return None


    ########################################
    def plot(self):
        '''- '''

        #################### TOP AXIS
        self.ax_top.clear()
        self.ax_top.plot( self.x, self.y, color=self.clrs[self.idx] )

        # this is for bounding the zoom-out
        self.ax_top.initial_xlim = self.ax_top.get_xlim()
        
        #################### BOTTOM AXES
        self.ax_bot.clear()
        self.ax_bot.plot( self.x, self.y, color=self.clrs[self.idx] )

        # this is for bounding the zoom-out
        self.ax_bot.initial_xlim = self.ax_bot.get_xlim()
        
        #################### wrap up
        plt.draw()
        return None

    ########################################
    def next(self, event=None):
        '''- '''
        self.idx = (self.idx + 1) % len(self.clrs)
        self.plot()
        return None

    ########################################
    def prev(self, event=None):
        '''- '''
        self.idx = (self.idx - 1) % len(self.clrs)
        self.plot()
        return None

bc = BottonClass(figsize=(8,5), zoom_factor=1.5)