source: inundation/pyvolution/vtk_realtime_visualiser.py @ 2950

Last change on this file since 2950 was 2421, checked in by steve, 19 years ago
File size: 8.9 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
9#from 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
54
55
56        self.idle = threading.Event()
57        self.redraw_ready = threading.Event()
58        self.unpaused = threading.Event()
59        self.unpaused.set()
60
61        # Internal use - storage of vtk objects
62        self.grids = {}
63        self.scalars = {}
64        self.actors = {}
65        self.polydata = {}
66        self.mappers = {}
67
68        # Default options
69        for x in self.domain.quantities:
70            self.setup[x] = False
71            self.updating[x] = False
72
73        # Bounding box.
74        if rect is None:
75            self.max_x = max(max(self.vertices[:,0]),max(self.vertices[:,2]),max(self.vertices[:,4]))
76            self.min_x = min(min(self.vertices[:,0]),min(self.vertices[:,2]),min(self.vertices[:,4]))
77            self.max_y = max(max(self.vertices[:,1]),max(self.vertices[:,3]),max(self.vertices[:,5]))
78            self.min_y = min(min(self.vertices[:,1]),min(self.vertices[:,3]),min(self.vertices[:,5]))
79        else:
80            self.max_x = rect[2]
81            self.min_x = rect[0]
82            self.max_y = rect[3]
83            self.min_y = rect[1]
84
85        self.range_x = self.max_x - self.min_x
86        self.range_y = self.max_y - self.min_y
87        self.range_xy = max(self.range_x, self.range_y)
88
89    def run(self):
90        self.initialise_gui()
91        self.add_axes()
92        self.setup_all()
93        self.root.after(100, self.idle.set)
94        self.root.mainloop()
95
96    def initialise_gui(self):
97        """Prepare the GUI for the Visualiser, and set up the
98        renderer.
99
100        """
101        # Using the TK VTK widget allows extendability
102        # should a gui need adding.
103        self.root = Tkinter.Tk()
104
105        # Message handling with after
106        self.root.after(100, self.redraw)
107
108        self.renderWidget = vtkTkRenderWidget(self.root, width=400, height=400)
109        self.renderWidget.pack(expand='true', fill='both')
110        self.renderWindow = self.renderWidget.GetRenderWindow()
111        self.renderer = vtk.vtkRenderer()
112        self.renderWindow.AddRenderer(self.renderer)
113
114        self.quitButton = Tkinter.Button(self.root, text='Quit', command=self.shutdown)
115        self.quitButton.pack(side=Tkinter.BOTTOM)
116
117        self.pauseButton = Tkinter.Button(self.root, text='Pause', command=self.unpaused.clear)
118        self.pauseButton.pack(side=Tkinter.LEFT)
119
120        self.resumeButton = Tkinter.Button(self.root, text='Resume', command=self.unpaused.set)
121        self.resumeButton.pack(side=Tkinter.LEFT)
122       
123        #self.w2if = vtk.vtkWindowToImageFilter()
124        #self.w2if.SetInput(self.renderWindow)
125       
126        #self.pnmWr = vtk.vtkPNMWriter()
127        #self.pnmWr.SetInput(self.w2if.GetOutput())
128
129
130    def add_axes(self):
131        """Add axes to this Visualiser
132        """
133        pass
134
135    def setup_all(self):
136        """Draw in the data that is specified by setup or update
137        """
138
139        self.N_tri = len(self.domain.triangles)
140        self.N_vert = len(self.vertices)
141        self.cells = vtk.vtkCellArray()
142        self.vertices = self.domain.get_vertex_coordinates()
143        self.vert_index = zeros((self.N_vert,2), Float)
144        for n in range(self.N_vert):
145            for i in range(3):
146                self.vert_index[self.domain.triangles[n][i]] = self.vertices[n][i*2:i*2+2]
147
148        # Prepare the list of cells
149        for t in self.domain.triangles:
150            self.cells.InsertNextCell(3)
151            for i in range(3):
152                self.cells.InsertCellPoint(t[i])
153
154        # Set up the rendering of each quantity
155        for q in self.domain.quantities:
156            if self.setup[q] | self.updating[q]:
157                self.draw_quantity(q)
158
159    def draw_quantity(self, q):
160        if self.scale_z.has_key(q):
161            scale = self.scale_z[q]
162        else:
163            scale = self.default_scale_z
164           
165        if q == 'elevation':
166            shift = 0.05
167        else:
168            shift = 0.0
169           
170        #############################################################
171        # Disabled the make_vtkpoints call because the extension is
172        # difficult to get to link under windows
173        #############################################################
174        #self.grids[q] = make_vtkpoints(self.N_tri,
175        #                           self.N_vert,
176        #                           scale,
177        #                           self.domain.quantities[q].vertex_values,
178        #                           self.vert_index,
179        #                           self.domain.triangles)
180        #grid = self.grids[q]
181        #############################################################
182
183        qty_index = zeros(self.N_vert, Float)
184        for n in range(self.N_tri):
185            for v in range(3):
186                qty_index[self.domain.triangles[n][v]] = self.domain.quantities[q].vertex_values[n][v]
187
188        self.grids[q] = vtk.vtkPoints()
189        self.scalars[q] = vtk.vtkFloatArray()
190        grid = self.grids[q]
191        scalars = self.scalars[q]
192
193        for v in range(self.N_vert):
194            grid.InsertNextPoint(self.vert_index[v][0],
195                                 self.vert_index[v][1],
196                                 shift + qty_index[v] * scale)
197            scalars.InsertNextValue(qty_index[v]);
198
199        # Can't recycle vtkPolyData objects: Apparently they behave
200        # unusually if the points (i.e. vertex data) is set after
201        # the polys (i.e. triangle data)
202        self.polydata[q] = vtk.vtkPolyData()
203        polydata = self.polydata[q]
204
205        polydata.SetPoints(grid)
206        if self.coloring[q]:
207            polydata.GetPointData().SetScalars(scalars);
208        polydata.SetPolys(self.cells)
209
210        if self.mappers.has_key(q):
211            mapper = self.mappers[q]
212            mapper.SetInput(polydata)
213            mapper.SetScalarRange(-10.0,10.0)
214            mapper.Update()
215        else:
216            self.mappers[q] = vtk.vtkPolyDataMapper()
217            mapper = self.mappers[q]
218            mapper.SetInput(polydata)
219
220        if not self.actors.has_key(q):
221            self.actors[q] = vtk.vtkActor()
222            actor = self.actors[q]
223            actor.SetMapper(mapper)
224
225            if self.qcolor.has_key(q):
226                actor.GetProperty().SetColor(self.qcolor[q])
227            else:
228                actor.GetProperty().SetColor(0.5, 0.5, 0.5)
229
230            self.renderer.AddActor(actor)
231
232    def redraw(self):
233        if self.redraw_ready.isSet():
234            self.redraw_ready.wait()
235            self.redraw_ready.clear()
236            for q in self.domain.quantities:
237                if self.updating[q]:
238                    self.draw_quantity(q)
239
240            self.renderWindow.Render()
241           
242           
243           
244            self.root.update_idletasks()
245            self.idle.set()
246        self.root.after(100, self.redraw)     
247   
248    def shutdown(self):
249        self.domain.visualise = False
250        self.idle.set()
251        self.unpaused.set()
252        self.root.withdraw()
253        self.root.quit()
Note: See TracBrowser for help on using the repository browser.