1 | # VTK-Based visualiser |
---|
2 | # intended to replace realtime_visualisation_new.py |
---|
3 | |
---|
4 | import threading |
---|
5 | from Queue import Queue, Empty |
---|
6 | import Tkinter |
---|
7 | import vtk |
---|
8 | from vtk.tk.vtkTkRenderWidget import vtkTkRenderWidget |
---|
9 | |
---|
10 | class Visualiser(threading.Thread): |
---|
11 | |
---|
12 | """A VTK-powered visualiser, designed to replace the VPython one. |
---|
13 | Intended to be run in its own thread |
---|
14 | """ |
---|
15 | |
---|
16 | def __init__(self, domain, default_scale_z=1.0, rect=None, title='Test'): |
---|
17 | threading.Thread.__init__(self) |
---|
18 | # Initialise data structures. coloring, setup and updating are maps |
---|
19 | # quantity name -> boolean, setup means to render it, |
---|
20 | # update means to update with time. |
---|
21 | # qcolor maps quantity name -> (float, float, float): the colour (r, g, b) |
---|
22 | # to render a quantity. |
---|
23 | self.setup = {} |
---|
24 | self.updating = {} |
---|
25 | self.coloring = {} |
---|
26 | self.qcolor = {} |
---|
27 | self.scale_z = {} |
---|
28 | self.default_scale_z = default_scale_z |
---|
29 | self.domain = domain |
---|
30 | self.vertices = domain.vertex_coordinates |
---|
31 | |
---|
32 | self.messages = Queue() |
---|
33 | |
---|
34 | # Internal use - storage of vtk objects |
---|
35 | self.grids = {} |
---|
36 | self.actors = {} |
---|
37 | self.polydata = {} |
---|
38 | self.mappers = {} |
---|
39 | |
---|
40 | # Default options |
---|
41 | for x in self.domain.quantities: |
---|
42 | self.setup[x] = False |
---|
43 | self.updating[x] = False |
---|
44 | self.qcolor[x] = (0.0, 0.0, 0.0) |
---|
45 | self.coloring[x] = False |
---|
46 | |
---|
47 | # Bounding box. |
---|
48 | if rect is None: |
---|
49 | self.max_x = max(max(self.vertices[:,0]),max(self.vertices[:,2]),max(self.vertices[:,4])) |
---|
50 | self.min_x = min(min(self.vertices[:,0]),min(self.vertices[:,2]),min(self.vertices[:,4])) |
---|
51 | self.max_y = max(max(self.vertices[:,1]),max(self.vertices[:,3]),max(self.vertices[:,5])) |
---|
52 | self.min_y = min(min(self.vertices[:,1]),min(self.vertices[:,3]),min(self.vertices[:,5])) |
---|
53 | else: |
---|
54 | self.max_x = rect[2] |
---|
55 | self.min_x = rect[0] |
---|
56 | self.max_y = rect[3] |
---|
57 | self.min_y = rect[1] |
---|
58 | |
---|
59 | self.range_x = self.max_x - self.min_x |
---|
60 | self.range_y = self.max_y - self.min_y |
---|
61 | self.range_xy = max(self.range_x, self.range_y) |
---|
62 | |
---|
63 | def run(self): |
---|
64 | self.initialise_gui() |
---|
65 | self.add_axes() |
---|
66 | self.setup_all() |
---|
67 | self.root.mainloop() |
---|
68 | |
---|
69 | def initialise_gui(self): |
---|
70 | """Prepare the GUI for the Visualiser, and set up the |
---|
71 | renderer. |
---|
72 | |
---|
73 | """ |
---|
74 | # Using the TK VTK widget allows extendability |
---|
75 | # should a gui need adding. |
---|
76 | self.root = Tkinter.Tk() |
---|
77 | |
---|
78 | # Message handling with after |
---|
79 | self.root.after(100, self.handle_messages) |
---|
80 | |
---|
81 | self.renderWidget = vtkTkRenderWidget(self.root, width=400, height=400) |
---|
82 | self.renderWidget.pack(expand='true', fill='both') |
---|
83 | self.renderWindow = self.renderWidget.GetRenderWindow() |
---|
84 | self.renderer = vtk.vtkRenderer() |
---|
85 | self.renderWindow.AddRenderer(self.renderer) |
---|
86 | |
---|
87 | def add_axes(self): |
---|
88 | """Add axes to this Visualiser |
---|
89 | """ |
---|
90 | pass |
---|
91 | |
---|
92 | def setup_all(self): |
---|
93 | """Draw in the data that is specified by setup or update |
---|
94 | """ |
---|
95 | |
---|
96 | self.N_tri = len(self.domain.triangles) |
---|
97 | self.cells = vtk.vtkCellArray() |
---|
98 | self.vertices = self.domain.get_vertex_coordinates() |
---|
99 | self.vert_index = {} |
---|
100 | |
---|
101 | for n in range(len(self.vertices)): |
---|
102 | for i in range(3): |
---|
103 | self.vert_index[self.domain.triangles[n][i]] = self.vertices[n][i*2:i*2+2] |
---|
104 | |
---|
105 | self.N_vert = len(self.vert_index.keys()) |
---|
106 | |
---|
107 | # Prepare the list of cells |
---|
108 | for t in self.domain.triangles: |
---|
109 | self.cells.InsertNextCell(3) |
---|
110 | for i in range(3): |
---|
111 | self.cells.InsertCellPoint(t[i]) |
---|
112 | |
---|
113 | # Set up the rendering of each quantity |
---|
114 | for q in self.domain.quantities: |
---|
115 | if self.setup[q] | self.updating[q]: |
---|
116 | self.draw_quantity(q) |
---|
117 | |
---|
118 | def draw_quantity(self, q): |
---|
119 | qty_index = {} |
---|
120 | for n in range(self.N_tri): |
---|
121 | for v in range(3): |
---|
122 | qty_index[self.domain.triangles[n][v]] = self.domain.quantities[q].vertex_values[n][v] |
---|
123 | |
---|
124 | self.grids[q] = vtk.vtkPoints() |
---|
125 | grid = self.grids[q] |
---|
126 | |
---|
127 | if self.scale_z.has_key(q): |
---|
128 | scale = self.scale_z[q] |
---|
129 | else: |
---|
130 | scale = self.default_scale_z |
---|
131 | |
---|
132 | for v in range(self.N_vert): |
---|
133 | grid.InsertNextPoint(self.vert_index[v][0], |
---|
134 | self.vert_index[v][1], |
---|
135 | qty_index[v] * scale) |
---|
136 | |
---|
137 | # Can't recycle vtkPolyData objects: Apparently they behave |
---|
138 | # unusually if the points (i.e. vertex data) is set after |
---|
139 | # the polys (i.e. triangle data) |
---|
140 | self.polydata[q] = vtk.vtkPolyData() |
---|
141 | polydata = self.polydata[q] |
---|
142 | |
---|
143 | polydata.SetPoints(grid) |
---|
144 | polydata.SetPolys(self.cells) |
---|
145 | |
---|
146 | if self.mappers.has_key(q): |
---|
147 | mapper = self.mappers[q] |
---|
148 | mapper.SetInput(polydata) |
---|
149 | mapper.Update() |
---|
150 | else: |
---|
151 | self.mappers[q] = vtk.vtkPolyDataMapper() |
---|
152 | mapper = self.mappers[q] |
---|
153 | mapper.SetInput(polydata) |
---|
154 | |
---|
155 | if not self.actors.has_key(q): |
---|
156 | self.actors[q] = vtk.vtkActor() |
---|
157 | actor = self.actors[q] |
---|
158 | actor.SetMapper(mapper) |
---|
159 | |
---|
160 | if self.coloring[q] is True: |
---|
161 | actor.GetProperty().SetColor(self.qcolor[q]) |
---|
162 | else: |
---|
163 | actor.GetProperty().SetColor(0.5, 0.5, 0.5) |
---|
164 | |
---|
165 | self.renderer.AddActor(actor) |
---|
166 | |
---|
167 | def handle_messages(self): |
---|
168 | # Expandable to include more kinds of messages if need be |
---|
169 | # This lookup table is copied because it is modified later |
---|
170 | handlers = {'redraw': self.message_redraw}.copy() |
---|
171 | try: |
---|
172 | while True: |
---|
173 | message = self.messages.get_nowait() |
---|
174 | try: |
---|
175 | handlers[message]() |
---|
176 | # Guard against redrawing more than once per |
---|
177 | # call to handle_messages. |
---|
178 | if message == 'redraw': |
---|
179 | del handlers[message] |
---|
180 | except KeyError: |
---|
181 | pass |
---|
182 | except Empty: |
---|
183 | pass |
---|
184 | self.root.after(100, self.handle_messages) |
---|
185 | |
---|
186 | def message_redraw(self): |
---|
187 | """Re-render and repaint the rendered image. |
---|
188 | This function should only be called from the message handler. |
---|
189 | Use request_redraw to ask for a redraw""" |
---|
190 | |
---|
191 | for q in self.domain.quantities: |
---|
192 | if self.updating[q]: |
---|
193 | self.draw_quantity(q) |
---|
194 | |
---|
195 | self.renderWindow.Render() |
---|
196 | self.root.update_idletasks() |
---|
197 | |
---|
198 | def request_redraw(self): |
---|
199 | self.messages.put('redraw') |
---|
200 | |
---|