~ubuntu-branches/ubuntu/utopic/python-chaco/utopic

« back to all changes in this revision

Viewing changes to examples/xray_plot.py

  • Committer: Bazaar Package Importer
  • Author(s): Varun Hiremath
  • Date: 2011-04-06 19:03:54 UTC
  • mfrom: (7.2.2 sid)
  • Revision ID: james.westby@ubuntu.com-20110406190354-rwd55l2ezjecfo41
Tags: 3.4.0-2
d/rules: fix pyshared directory path (Closes: #621116)

Show diffs side-by-side

added added

removed removed

Lines of Context:
 
1
"""
 
2
Implementation of a plut using a custom overlay and tool
 
3
"""
 
4
 
 
5
from __future__ import with_statement
 
6
 
 
7
import numpy
 
8
 
 
9
from enthought.traits.api import HasTraits, Instance, Enum
 
10
from enthought.traits.ui.api import View, Item
 
11
from enthought.enable.api import ComponentEditor
 
12
from enthought.chaco.api import Plot, ArrayPlotData, AbstractOverlay
 
13
from enthought.enable.api import BaseTool
 
14
from enthought.enable.markers import DOT_MARKER, DotMarker
 
15
 
 
16
class BoxSelectTool(BaseTool):
 
17
    """ Tool for selecting all points within a box
 
18
    
 
19
        There are 2 states for this tool, normal and selecting. While the
 
20
        left mouse button is down the metadata on the datasources will be
 
21
        updated with the current selected bounds.
 
22
        
 
23
        Note that the tool does not actually store the selected point, but the
 
24
        bounds of the box. 
 
25
    """
 
26
 
 
27
    event_state = Enum("normal", "selecting")
 
28
 
 
29
    def normal_left_down(self, event):
 
30
        self.event_state = "selecting"
 
31
        self.selecting_mouse_move(event)
 
32
        
 
33
    def selecting_left_up(self, event):
 
34
        self.event_state = "normal"
 
35
        
 
36
    def selecting_mouse_move(self, event):
 
37
        x1, y1 = self.map_to_data(event.x-25, event.y-25)
 
38
        x2, y2 = self.map_to_data(event.x+25, event.y+25)
 
39
        
 
40
        index_datasource = self.component.index
 
41
        index_datasource.metadata['selections'] = (x1, x2)
 
42
        
 
43
        value_datasource = self.component.value
 
44
        value_datasource.metadata['selections'] = (y1, y2)
 
45
        
 
46
        self.component.request_redraw()
 
47
        
 
48
    def map_to_data(self, x, y):
 
49
        """ Returns the data space coordinates of the given x and y. 
 
50
 
 
51
        Takes into account orientation of the plot and the axis setting.
 
52
        """
 
53
 
 
54
        plot = self.component
 
55
        if plot.orientation == "h":
 
56
            index = plot.x_mapper.map_data(x)
 
57
            value = plot.y_mapper.map_data(y)
 
58
        else:
 
59
            index = plot.y_mapper.map_data(y)
 
60
            value = plot.x_mapper.map_data(x)
 
61
            
 
62
        return index, value
 
63
        
 
64
 
 
65
class XRayOverlay(AbstractOverlay):
 
66
    """ Overlay which draws scatter markers on top of plot data points.
 
67
    
 
68
        This overlay should be combined with a tool which updates the 
 
69
        datasources metadata with selection bounds.
 
70
    """
 
71
 
 
72
    marker = DotMarker()
 
73
    
 
74
    def overlay(self, component, gc, view_bounds=None, mode='normal'):
 
75
        x_range = self._get_selection_index_screen_range()
 
76
        y_range = self._get_selection_value_screen_range()
 
77
        
 
78
        if len(x_range) == 0:
 
79
            return
 
80
        
 
81
        x1, x2 = x_range
 
82
        y1, y2 = y_range
 
83
 
 
84
        with gc:
 
85
            gc.set_alpha(0.8)
 
86
            gc.set_fill_color((1.0,1.0,1.0))
 
87
            gc.rect(x1, y1, x2-x1, y2-y1)
 
88
            gc.draw_path()
 
89
 
 
90
        pts = self._get_selected_points()
 
91
        if len(pts) == 0:
 
92
            return
 
93
        screen_pts = self.component.map_screen(pts)
 
94
        if hasattr(gc, 'draw_marker_at_points'):
 
95
            gc.draw_marker_at_points(screen_pts, 3, DOT_MARKER)
 
96
        else:
 
97
            gc.save_state()
 
98
            for sx,sy in screen_pts:
 
99
                gc.translate_ctm(sx, sy)
 
100
                gc.begin_path()
 
101
                self.marker.add_to_path(gc, 3)
 
102
                gc.draw_path(self.marker.draw_mode)
 
103
                gc.translate_ctm(-sx, -sy)
 
104
            gc.restore_state()
 
105
        
 
106
    def _get_selected_points(self):
 
107
        """ gets all the points within the bounds defined in the datasources
 
108
            metadata
 
109
        """
 
110
        index_datasource = self.component.index
 
111
        index_selection = index_datasource.metadata['selections']
 
112
        index = index_datasource.get_data()
 
113
        
 
114
        value_datasource = self.component.value
 
115
        value_selection = value_datasource.metadata['selections']
 
116
        value = value_datasource.get_data()
 
117
        
 
118
        x_indices = numpy.where((index > index_selection[0]) & (index < index_selection[-1]))
 
119
        y_indices = numpy.where((value > value_selection[0]) & (value < value_selection[-1]))
 
120
        
 
121
        indices = list(set(x_indices[0]) & set(y_indices[0]))
 
122
        
 
123
        sel_index = index[indices]
 
124
        sel_value = value[indices]
 
125
        
 
126
        return zip(sel_index, sel_value)
 
127
        
 
128
    def _get_selection_index_screen_range(self):
 
129
        """ maps the selected bounds which were set by the tool into screen
 
130
            space. The screen space points can be used for drawing the overlay
 
131
        """
 
132
        index_datasource = self.component.index
 
133
        index_mapper = self.component.index_mapper
 
134
        index_selection = index_datasource.metadata['selections']
 
135
        return tuple(index_mapper.map_screen(numpy.array(index_selection)))
 
136
 
 
137
    def _get_selection_value_screen_range(self):
 
138
        """ maps the selected bounds which were set by the tool into screen
 
139
            space. The screen space points can be used for drawing the overlay
 
140
        """
 
141
        value_datasource = self.component.value
 
142
        value_mapper = self.component.value_mapper
 
143
        value_selection = value_datasource.metadata['selections']
 
144
        return tuple(value_mapper.map_screen(numpy.array(value_selection)))
 
145
        
 
146
class PlotExample(HasTraits):
 
147
 
 
148
    plot = Instance(Plot)
 
149
    
 
150
    traits_view = View(Item('plot', editor=ComponentEditor()),
 
151
                       width=600, height=600)
 
152
                       
 
153
    def __init__(self, index, value, *args, **kw):
 
154
        super(PlotExample, self).__init__(*args, **kw)
 
155
        
 
156
        plot_data = ArrayPlotData(index=index)
 
157
        plot_data.set_data('value', value)
 
158
        
 
159
        self.plot = Plot(plot_data)
 
160
        line = self.plot.plot(('index', 'value'))[0]
 
161
        
 
162
        line.overlays.append(XRayOverlay(line))
 
163
        line.tools.append(BoxSelectTool(line))
 
164
        
 
165
index = numpy.arange(0, 25, 0.25)
 
166
value = numpy.sin(index) + numpy.arange(0, 10, 0.1)
 
167
 
 
168
example = PlotExample(index, value)
 
169
example.configure_traits()
 
170