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

« back to all changes in this revision

Viewing changes to examples/demo/vtk/cmap_scatter.py

  • Committer: Package Import Robot
  • Author(s): Andrew Starr-Bochicchio
  • Date: 2014-06-01 17:04:08 UTC
  • mfrom: (7.2.5 sid)
  • Revision ID: package-import@ubuntu.com-20140601170408-m86xvdjd83a4qon0
Tags: 4.4.1-1ubuntu1
* Merge from Debian unstable. Remaining Ubuntu changes:
 - Let the binary-predeb target work on the usr/lib/python* directory
   as we don't have usr/share/pyshared anymore.

Show diffs side-by-side

added added

removed removed

Lines of Context:
1
 
#!/usr/bin/env python
2
 
"""
3
 
Draws a colormapped scatterplot of some random data.
4
 
 
5
 
Interactions are the same as simple_line, and additionally, range selection
6
 
is available on the colorbar.  Right-click-drag will select a range of
7
 
colors on the colormap.  This range can be dragged around, and the main
8
 
plot will respond accordingly.  Left-click anywhere on the colorbar to
9
 
cancel the range selection.
10
 
"""
11
 
 
12
 
# Major library imports
13
 
from numpy import exp, sort
14
 
from numpy.random import random
15
 
 
16
 
# VTK-related stuff
17
 
from tvtk.api import tvtk
18
 
from mayavi import mlab
19
 
from enable.vtk_backend.vtk_window import EnableVTKWindow
20
 
 
21
 
 
22
 
# Chaco imports
23
 
from chaco.api import ArrayPlotData, ColorBar, \
24
 
    ColormappedSelectionOverlay, OverlayPlotContainer, \
25
 
    jet, LinearMapper, Plot
26
 
from chaco.tools.api import PanTool, ZoomTool, RangeSelection, \
27
 
    RangeSelectionOverlay, MoveTool
28
 
 
29
 
#===============================================================================
30
 
# # Create the Chaco plot.
31
 
#===============================================================================
32
 
def create_plot():
33
 
 
34
 
    # Create some data
35
 
    numpts = 200
36
 
    x = sort(random(numpts))
37
 
    y = random(numpts)
38
 
    color = exp(-(x**2 + y**2))
39
 
 
40
 
    # Create a plot data obect and give it this data
41
 
    pd = ArrayPlotData()
42
 
    pd.set_data("index", x)
43
 
    pd.set_data("value", y)
44
 
    pd.set_data("color", color)
45
 
 
46
 
    # Create the plot
47
 
    plot = Plot(pd)
48
 
    plot.plot(("index", "value", "color"),
49
 
              type="cmap_scatter",
50
 
              name="my_plot",
51
 
              color_mapper=jet,
52
 
              marker = "square",
53
 
              fill_alpha = 0.5,
54
 
              marker_size = 6,
55
 
              outline_color = "black",
56
 
              border_visible = True,
57
 
              bgcolor = "white")
58
 
 
59
 
    # Tweak some of the plot properties
60
 
    plot.title = "Colormapped Scatter Plot"
61
 
    plot.padding = 50
62
 
    plot.x_grid.visible = False
63
 
    plot.y_grid.visible = False
64
 
    plot.x_axis.font = "modern 16"
65
 
    plot.y_axis.font = "modern 16"
66
 
 
67
 
    # Set colors
68
 
    #plot.title_color = "white"
69
 
    #for axis in plot.x_axis, plot.y_axis:
70
 
    #    axis.set(title_color="white", tick_label_color="white")
71
 
 
72
 
    # Right now, some of the tools are a little invasive, and we need the
73
 
    # actual ColomappedScatterPlot object to give to them
74
 
    cmap_renderer = plot.plots["my_plot"][0]
75
 
 
76
 
    # Attach some tools to the plot
77
 
    plot.tools.append(PanTool(plot, constrain_key="shift"))
78
 
    zoom = ZoomTool(component=plot, tool_mode="box", always_on=False)
79
 
    plot.overlays.append(zoom)
80
 
    selection = ColormappedSelectionOverlay(cmap_renderer, fade_alpha=0.35,
81
 
                                            selection_type="mask")
82
 
    cmap_renderer.overlays.append(selection)
83
 
    plot.tools.append(MoveTool(plot, drag_button="right"))
84
 
    return plot
85
 
 
86
 
def create_colorbar(colormap):
87
 
    colorbar = ColorBar(index_mapper=LinearMapper(range=colormap.range),
88
 
                    color_mapper=colormap, orientation='v', resizable='',
89
 
                    height=400, width=30, padding=20)
90
 
    colorbar.tools.append(RangeSelection(component=colorbar))
91
 
    colorbar.overlays.append(RangeSelectionOverlay(component=colorbar,
92
 
           border_color="white", alpha=0.8, fill_color="lightgray"))
93
 
    colorbar.tools.append(MoveTool(colorbar, drag_button="left"))
94
 
    return colorbar
95
 
 
96
 
def start_vtk(component):
97
 
    f = mlab.figure(size=(700,500))
98
 
    m = mlab.test_mesh()
99
 
    scene = mlab.gcf().scene
100
 
    render_window = scene.render_window
101
 
    renderer = scene.renderer
102
 
    rwi = scene.interactor
103
 
    window = EnableVTKWindow(rwi, renderer,
104
 
            component = component,
105
 
            istyle_class = tvtk.InteractorStyleTrackballCamera,
106
 
            bgcolor = "transparent",
107
 
            event_passthrough = True,
108
 
            )
109
 
    mlab.show()
110
 
 
111
 
def main():
112
 
    plot = create_plot()
113
 
    plot.bounds = [400,300]
114
 
    plot.outer_position = [30,30]
115
 
    plot.resizable = ""
116
 
    cmap_renderer = plot.plots["my_plot"][0]
117
 
 
118
 
    # Create the colorbar, handing in the appropriate range and colormap
119
 
    colorbar = create_colorbar(plot.color_mapper)
120
 
    colorbar.outer_position = [450,30]
121
 
    colorbar.plot = cmap_renderer
122
 
    colorbar.padding_top = plot.padding_top
123
 
    colorbar.padding_bottom = plot.padding_bottom
124
 
 
125
 
    container = OverlayPlotContainer(bgcolor = "transparent",
126
 
                    fit_window = True)
127
 
    container.add(plot)
128
 
    container.add(colorbar)
129
 
 
130
 
    start_vtk(container)
131
 
 
132
 
 
133
 
if __name__ == "__main__":
134
 
    main()
135