source: inundation/pyvolution/vtk_realtime_visualiser.py @ 2378

Last change on this file since 2378 was 2334, checked in by steve, 19 years ago

Adding coloring to vtk visualisation
File size: 8.5 KB
Line 
1# VTK-Based visualiser
2# intended to replace realtime_visualisation_new.py
3
4import threading
5import Tkinter
6import vtk
7from Numeric import *
8from vtk.tk.vtkTkRenderWidget import vtkTkRenderWidget
9from vtk_visualiser_ext import make_vtkpoints
10
11class Visualiser(threading.Thread):
12
13    """A VTK-powered visualiser, designed to replace the VPython one.
14    Intended to be run in its own thread.
15
16    Customisation options for the visualiser are as follows:
17
18    setup: Dictionary mapping quantity name -> boolean.
19    if setup[q] is true, draw the quantity when the visualiser starts.
20
21    updating: Dictionary mapping quantity name -> boolean.
22    if updating[q] is true, update the rendering of this quantity each
23    timestep.
24
25    qcolor: Dictionary mapping quantity name -> (float, float, float)
26    if the name of a quantity is found in qcolor, the colour specified (in
27    (r, g, b) from 0.0 to 1.0) is used for display of that quantity instead of
28    (0.5, 0.5, 0.5) (the default)
29
30    scale_z: Dictionary mapping quantity name -> float
31    Override the z scaling of this quantity with the float specified.
32
33    default_scale_z: float
34    multiply all z coordinates by this amount unless an overriding value
35    exists in scale_z.
36    """
37
38    def __init__(self, domain, default_scale_z=1.0, rect=None, title='Test'):
39        threading.Thread.__init__(self)
40        # Initialise data structures. setup and updating are maps
41        # quantity name -> boolean, setup means to render it,
42        # update means to update with time.
43        # qcolor maps quantity name -> (float, float, float): the colour (r, g, b)
44        self.setup = {}
45        self.updating = {}
46        self.qcolor = {}
47        self.scale_z = {}
48        self.coloring = {}
49        self.default_scale_z = default_scale_z
50        self.domain = domain
51        self.vertices = domain.vertex_coordinates
52
53        self.idle = threading.Event()
54        self.redraw_ready = threading.Event()
55        self.unpaused = threading.Event()
56        self.unpaused.set()
57
58        # Internal use - storage of vtk objects
59        self.grids = {}
60        self.scalars = {}
61        self.actors = {}
62        self.polydata = {}
63        self.mappers = {}
64
65        # Default options
66        for x in self.domain.quantities:
67            self.setup[x] = False
68            self.updating[x] = False
69
70        # Bounding box.
71        if rect is None:
72            self.max_x = max(max(self.vertices[:,0]),max(self.vertices[:,2]),max(self.vertices[:,4]))
73            self.min_x = min(min(self.vertices[:,0]),min(self.vertices[:,2]),min(self.vertices[:,4]))
74            self.max_y = max(max(self.vertices[:,1]),max(self.vertices[:,3]),max(self.vertices[:,5]))
75            self.min_y = min(min(self.vertices[:,1]),min(self.vertices[:,3]),min(self.vertices[:,5]))
76        else:
77            self.max_x = rect[2]
78            self.min_x = rect[0]
79            self.max_y = rect[3]
80            self.min_y = rect[1]
81
82        self.range_x = self.max_x - self.min_x
83        self.range_y = self.max_y - self.min_y
84        self.range_xy = max(self.range_x, self.range_y)
85
86    def run(self):
87        self.initialise_gui()
88        self.add_axes()
89        self.setup_all()
90        self.root.after(100, self.idle.set)
91        self.root.mainloop()
92
93    def initialise_gui(self):
94        """Prepare the GUI for the Visualiser, and set up the
95        renderer.
96
97        """
98        # Using the TK VTK widget allows extendability
99        # should a gui need adding.
100        self.root = Tkinter.Tk()
101
102        # Message handling with after
103        self.root.after(100, self.redraw)
104
105        self.renderWidget = vtkTkRenderWidget(self.root, width=400, height=400)
106        self.renderWidget.pack(expand='true', fill='both')
107        self.renderWindow = self.renderWidget.GetRenderWindow()
108        self.renderer = vtk.vtkRenderer()
109        self.renderWindow.AddRenderer(self.renderer)
110
111        self.quitButton = Tkinter.Button(self.root, text='Quit', command=self.shutdown)
112        self.quitButton.pack(side=Tkinter.BOTTOM)
113
114        self.pauseButton = Tkinter.Button(self.root, text='Pause', command=self.unpaused.clear)
115        self.pauseButton.pack(side=Tkinter.LEFT)
116
117        self.resumeButton = Tkinter.Button(self.root, text='Resume', command=self.unpaused.set)
118        self.resumeButton.pack(side=Tkinter.LEFT)
119
120    def add_axes(self):
121        """Add axes to this Visualiser
122        """
123        pass
124
125    def setup_all(self):
126        """Draw in the data that is specified by setup or update
127        """
128
129        self.N_tri = len(self.domain.triangles)
130        self.N_vert = len(self.vertices)
131        self.cells = vtk.vtkCellArray()
132        self.vertices = self.domain.get_vertex_coordinates()
133        self.vert_index = zeros((self.N_vert,2), Float)
134        for n in range(self.N_vert):
135            for i in range(3):
136                self.vert_index[self.domain.triangles[n][i]] = self.vertices[n][i*2:i*2+2]
137
138        # Prepare the list of cells
139        for t in self.domain.triangles:
140            self.cells.InsertNextCell(3)
141            for i in range(3):
142                self.cells.InsertCellPoint(t[i])
143
144        # Set up the rendering of each quantity
145        for q in self.domain.quantities:
146            if self.setup[q] | self.updating[q]:
147                self.draw_quantity(q)
148
149    def draw_quantity(self, q):
150        if self.scale_z.has_key(q):
151            scale = self.scale_z[q]
152        else:
153            scale = self.default_scale_z
154           
155        #############################################################
156        # Disabled the make_vtkpoints call because the extension is
157        # difficult to get to link under windows
158        #############################################################
159        #self.grids[q] = make_vtkpoints(self.N_tri,
160        #                           self.N_vert,
161        #                           scale,
162        #                           self.domain.quantities[q].vertex_values,
163        #                           self.vert_index,
164        #                           self.domain.triangles)
165        #grid = self.grids[q]
166        #############################################################
167
168        qty_index = zeros(self.N_vert, Float)
169        for n in range(self.N_tri):
170            for v in range(3):
171                qty_index[self.domain.triangles[n][v]] = self.domain.quantities[q].vertex_values[n][v]
172
173        self.grids[q] = vtk.vtkPoints()
174        self.scalars[q] = vtk.vtkFloatArray()
175        grid = self.grids[q]
176        scalars = self.scalars[q]
177
178        for v in range(self.N_vert):
179            grid.InsertNextPoint(self.vert_index[v][0],
180                                 self.vert_index[v][1],
181                                 qty_index[v] * scale)
182            scalars.InsertNextValue(qty_index[v]);
183
184        # Can't recycle vtkPolyData objects: Apparently they behave
185        # unusually if the points (i.e. vertex data) is set after
186        # the polys (i.e. triangle data)
187        self.polydata[q] = vtk.vtkPolyData()
188        polydata = self.polydata[q]
189
190        polydata.SetPoints(grid)
191        if self.coloring[q]:
192            polydata.GetPointData().SetScalars(scalars);
193        polydata.SetPolys(self.cells)
194
195        if self.mappers.has_key(q):
196            mapper = self.mappers[q]
197            mapper.SetInput(polydata)
198            mapper.Update()
199        else:
200            self.mappers[q] = vtk.vtkPolyDataMapper()
201            mapper = self.mappers[q]
202            mapper.SetInput(polydata)
203
204        if not self.actors.has_key(q):
205            self.actors[q] = vtk.vtkActor()
206            actor = self.actors[q]
207            actor.SetMapper(mapper)
208
209            if self.qcolor.has_key(q):
210                actor.GetProperty().SetColor(self.qcolor[q])
211            else:
212                actor.GetProperty().SetColor(0.5, 0.5, 0.5)
213
214            self.renderer.AddActor(actor)
215
216    def redraw(self):
217        if self.redraw_ready.isSet():
218            self.redraw_ready.wait()
219            self.redraw_ready.clear()
220            for q in self.domain.quantities:
221                if self.updating[q]:
222                    self.draw_quantity(q)
223
224            self.renderWindow.Render()
225            self.root.update_idletasks()
226            self.idle.set()
227        self.root.after(100, self.redraw)     
228   
229    def shutdown(self):
230        self.domain.visualise = False
231        self.idle.set()
232        self.unpaused.set()
233        self.root.withdraw()
234        self.root.quit()
Note: See TracBrowser for help on using the repository browser.