laydi/system/plots.py

import pygtk
import gtk
import matplotlib
import scipy
from matplotlib.backends.backend_gtkagg import FigureCanvasGTKAgg as FigureCanvas
from matplotlib.backends.backend_gtkagg import NavigationToolbar2GTK as NavigationToolbar2
from matplotlib.axes import Subplot
from matplotlib.figure import Figure
from matplotlib.numerix import arange, sin, pi
from matplotlib.widgets import RectangleSelector
from system import logger


class MainView (gtk.Notebook):
    def __init__(self, view_listener):
        gtk.Notebook.__init__(self)
        self.set_show_tabs(False)
        self.set_show_border(False)

        # Add a multiple pane view and a single pane view.
        self.small_view = SmallView()
        self.small_view.set_view_listener(view_listener)
        self.small_view.show()
        self.large_view = LargeView()
        self.large_view.show()

        self.append_page(self.small_view)
        self.append_page(self.large_view)
        self.set_current_page(0)
    
    def goto_small(self):
        if self.get_current_page() == 0:
            return None
        self.set_current_page(0)
        view = self.large_view.remove_child()
        self.small_view.return_current(view)

    def goto_large(self):
        if self.get_current_page() == 1:
            return None
        self.set_current_page(1)
        view = self.small_view.borrow_current()
        self.large_view.set_child(view)

    def show(self):
        gtk.Notebook.show(self)

    def insert_view(self, view):
        self.small_view.insert_view(view)


class Plot (gtk.Frame):

    def __init__(self, title):
        gtk.Frame.__init__(self)
        self.mark_active(False)
        self.connect('button_press_event', self.on_button_press)
        self.sel_obj = None
        self.active = False
        self.title = title

    def get_title(self):
        return self.title

    def on_button_press(self, *rest):
#        logger.log('debug', 'button pressed in plot')
        self.mark_active(True)

    def mark_active(self, active):
        if active:
            self.set_shadow_type(gtk.SHADOW_IN)
        else:
            self.set_shadow_type(gtk.SHADOW_OUT)
        self.active = active

    def update(self, key):
        pass

    def set_project(self, project):
        self.project = project

    def get_toolbar(self):
        return None
        

class SmallView (gtk.Table):
        
    def __init__(self):
        gtk.Table.__init__(self, 2, 2, True)

        self.child_views = [[EmptyView(), EmptyView()],
                            [EmptyView(), EmptyView()]]
        self.cols = 2
        self.rows = 2

        self.active_x = 0
        self.active_y = 0

        self.set_row_spacings(3)
        self.set_col_spacings(3)
        self._listener = None
        
        for x in range(self.cols):
            for y in range(self.rows):
                child = self.child_views[x][y]
                child.parent_signalling = child.connect('button_press_event', self.__view_button_event__)
                self.attach(child, x, x+1, y, y+1)

    def set_view_listener(self, listener):
        self._listener = listener

    def set_child(self, child, col, row):
        cur_widget = self.child_views[col][row]
        cur_widget.disconnect(cur_widget.parent_signalling)
        self.remove(cur_widget)
        self.attach(child, col, col+1, row, row+1)
        child.parent_signalling = child.connect('button_press_event', self.__view_button_event__)
        self.child_views[col][row] = child
        if cur_widget.active:
            child.mark_active(True)
            cur_widget.mark_active(False)
        child.show()

    def borrow_current(self):
        self.borrowed = self.child_views[self.active_x][self.active_y]
        self.remove_child(self.active_x, self.active_y)
        return self.borrowed

    def return_current(self, view):
        self.set_child(view, self.active_x, self.active_y)

    def remove_child(self, col, row):
        self.remove(self.child_views[col][row])
        self.attach(EmptyView(), col, col+1, row, row+1)
       
    def insert_view(self, child):
        if not self.find_child(child):
            self.set_child(child, self.active_x, self.active_y)

    def show(self):
        for x in self.child_views:
            for y in x:
                y.show()
        gtk.Table.show(self)

    def hide(self):
        for x in self.child_views:
            for y in x:
                y.hide()
        gtk.Table.hide(self)

    def set_active(self, x, y):
        old_focus = self.child_views[self.active_x][self.active_y]
        new_focus = self.child_views[x][y]
        old_focus.mark_active(False)
        new_focus.mark_active(True)
        self.active_x = x
        self.active_y = y
        if self._listener:
            logger.log("debug", "emitting")
            self._listener(new_focus)

    def find_child(self, child):
        for i, row in enumerate(self.child_views):
            for j, v in enumerate(row):
                if v == child:
                    return (i, j)
        return None

    def __view_button_event__(self, view, *rest):
        loc = self.find_child(view)
        if loc:
            self.set_active(loc[0], loc[1])

    def get_view(self, x, y):
        return self.child_views[x][y]

class LargeView (gtk.Frame):
    def __init__(self):
        gtk.Frame.__init__(self)
        self.child_view = EmptyView()
        self.add(self.child_view)
    
    def set_child(self, child):
        self.remove(self.child_view)
        self.child_view.hide()
        self.add(child)
        self.child_view = child
        child.show()

    def hide(self):
        self.child_view.hide()
        gtk.Frame.hide(self)

    def show(self):
        self.child_view.show()
        gtk.Frame.show(self)

    def remove_child(self):
        child = self.child_view
        child.hide()
        self.remove(child)
        return child
        
    
class EmptyView (Plot):
    def __init__(self):
        Plot.__init__(self, 'Empty view')
        
        label = gtk.Label('No view')
        ebox = gtk.EventBox()
        ebox.add(label)
        self.add(ebox)
        label.connect('button_press_event', self.on_button_press)

        self.label = label
        self.ebox = ebox
        self.show()

    def show(self):
        self.ebox.show()
        self.label.show()
        Plot.show(self)
        
    def hide(self):
        self.label.hide()
        self.ebox.hide()
        Plot.hide(self)


class NavToolbar(NavigationToolbar2):
    toolitems = (('Select', 'Select within rectangle', 'zoom_to_rect.png',
                  'select'),) + NavigationToolbar2.toolitems

    def __init__(self, *args):
        NavigationToolbar2.__init__(self, *args)
        self._select_callback = None

    def select(self, *args):
        """Selection mode selected handler."""
        if self._active == 'SELECT':
            self._active = None
        else:
            self._active = 'SELECT'

        if self._idPress is not None:
            self._idPress = self.canvas.mpl_disconnect(self._idPress)
            self.mode = ''

        if self._idRelease is not None:
            self._idRelease = self.canvas.mpl_disconnect(self._idRelease)
            self.mode = ''

        if self._active:
            self._idPress = self.canvas.mpl_connect('button_press_event', self.press_select)
            self._idRelease = self.canvas.mpl_connect('button_release_event', self.release_select)
            self.mode = 'Select rectangle mode'

        self.set_message(self.mode)

    def set_message(self, s):
        """Set status in toolbar to string s.

        Overrided to make sure message can be updated even when
        drawing rubberband.
        """
        self.message.set_label(s)

    def press_select(self, event):
        """Mouse button pressed handler for selection mode."""
        x, y = event.x, event.y

        for i, a in enumerate(self.canvas.figure.get_axes()):
            if event.inaxes==a and event.inaxes.get_navigate():
                xmin, xmax = a.get_xlim()
                ymin, ymax = a.get_ylim()
                lim = xmin, xmax, ymin, ymax
                self._xypress = x, y, a, i, lim, a.transData.deepcopy()
                break
        
        self.press(event)

    def release_select(self, event):
        """Mouse button released handler for selection mode."""
        # only release if button was pressed inside first?
        if self._xypress:
            x, y = event.x, event.y
            lastx, lasty, a, ind, lim, trans = self._xypress
            lastx, lasty = a.transData.inverse_xy_tup( (lastx, lasty) )
            x, y = a.transData.inverse_xy_tup( (x, y) )

            if self._select_callback:
                self._select_callback(lastx, lasty, x, y)
        
        self._xypress = None
        self.draw()
        self.release(event)

    def mouse_move(self, event):
        """Extend NavigationToolbar2.mouse_move to provide selection support."""
        from matplotlib.backend_bases import cursors

        # Only update the rubberband for selection mode when mouse is
        # within the plotting area.
        if event.inaxes and self._active=='SELECT':
            if self._lastCursor != cursors.SELECT_REGION:
                self.set_cursor(cursors.SELECT_REGION)
                self._lastCursor = cursors.SELECT_REGION
            if self._xypress is not None:
                x, y = event.x, event.y
                lastx, lasty, a, ind, lim, trans= self._xypress
                self.draw_rubberband(event, x, y, lastx, lasty)

        NavigationToolbar2.mouse_move(self, event)

    def set_select_callback(self, listener):
        """Allow plots to register a callback for selection events.

        The callback will be called as listener(x1, y1, x2, y2).  All
        coordinates are in the plot coordinate system, not pixels or
        widget coordinates.
        """
        self._select_callback = listener
        

class SinePlot(Plot):
    def __init__(self):
        Plot.__init__(self, 'Sine plot')

        fig = Figure(figsize=(5,4), dpi=72)
        ax = fig.add_subplot(111)
        t = arange(0.0,3.0,0.01)
        s = sin(2*pi*t)
        ax.plot(t,s)

        self.canvas = FigureCanvas(fig)
        self._toolbar = NavToolbar(self.canvas, None)
        self._toolbar.set_property('show-arrow', False)
        self.add(self.canvas)
        self.canvas.show()

    def get_toolbar(self):
        return self._toolbar


class ScatterPlot(Plot):
    def __init__(self, dataset,id_dim, sel_dim,id_1,id_2):
        Plot.__init__(self, 'Scatter plot')
        fig = Figure(figsize=(5,4), dpi=72)
        self.ax = ax = fig.add_subplot(111)
        self.current_dim = id_dim
        # testing testing
        self.dataset = dataset
        x_index = dataset[sel_dim][id_1]
        y_index = dataset[sel_dim][id_2]
        
        self.xaxis_data = dataset._array[:,x_index]
        self.yaxis_data = dataset._array[:,y_index]
        ax.plot(self.xaxis_data,self.yaxis_data,'og')

        ###

        self.canvas = FigureCanvas(fig)
        self.add(self.canvas)
        self.canvas.show()

        self._toolbar = NavToolbar(self.canvas, None)
        self._toolbar.set_property('show-arrow', False)
        self._toolbar.set_select_callback(self.rectangle_select_callback)

    def get_toolbar(self):
        return self._toolbar

    def rectangle_select_callback(self, x1, y1, x2, y2):
        'event1 and event2 are the press and release events'
        ydata = self.yaxis_data
        xdata = self.xaxis_data

        # find indices of selected area
        if x1>x2:
            if y1<y2:
                index =scipy.nonzero((xdata<x1) & (xdata>x2) & (ydata>y1) & (ydata<y2))
            else:
                index =scipy.nonzero((xdata<x1) & (xdata>x2) & (ydata<y1) & (ydata>y2))
        else:
            #logger.log('debug','Selection x_start less than x_end')
            if y1<y2:
                #logger.log('debug','Selection y_start less than y_end')
                index =scipy.nonzero((xdata>x1) & (xdata<x2) & (ydata>y1) & (ydata<y2))
            else:
                #logger.log('debug','Selection y_start bigger than y_end')
                index =scipy.nonzero((xdata>x1) & (xdata<x2) & (ydata<y1) & (ydata>y2))

        # generate ids for selected indices
        reverse = {}
        for key, value in self.dataset[self.current_dim].items():
            reverse[value] = key
        ids = []
        for ind in index:
            ids.append(reverse[ind])
            
        self.project.set_selection(self.current_dim,ids)

    def update(self, key):
        curr_sel = self.project.get_selection() # get selection object
        ids = curr_sel[self.current_dim] # current identifiers
        
        index = [ind for id,ind in self.dataset[self.current_dim].items() if id in ids] #conversion to index
        xdata_new = scipy.take(self.xaxis_data,index) #take data
        ydata_new = scipy.take(self.yaxis_data,index)
        self.ax.clear()
        self.ax.plot(self.xaxis_data,self.yaxis_data,'og')
        self.ax.plot(xdata_new,ydata_new,'or')
        self.canvas.draw()