source: inundation/visualiser/vtk_realtime_visualiser.py @ 3461

Last change on this file since 3461 was 3461, checked in by jack, 18 years ago

Rearranging the visualiser files.
File size: 8.4 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, 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        self.running = True
66
67        # Default options
68        for x in self.domain.quantities:
69            self.setup[x] = False
70            self.updating[x] = False
71            self.coloring[x] = False
72        self.start()
73           
74    def run(self):
75        self.initialise_gui()
76        self.add_axes()
77        self.setup_all()
78        self.root.after(100, self.idle.set)
79        self.root.mainloop()
80
81    def initialise_gui(self):
82        """Prepare the GUI for the Visualiser, and set up the
83        renderer.
84
85        """
86        # Using the TK VTK widget allows extendability
87        # should a gui need adding.
88        self.root = Tkinter.Tk()
89
90        # Message handling with after
91        self.root.after(100, self.redraw)
92
93        self.renderWidget = vtkTkRenderWidget(self.root, width=400, height=400)
94        self.renderWidget.pack(expand='true', fill='both')
95        self.renderWindow = self.renderWidget.GetRenderWindow()
96        self.renderer = vtk.vtkRenderer()
97        self.renderWindow.AddRenderer(self.renderer)
98
99        self.quitButton = Tkinter.Button(self.root, text='Quit', command=self.shutdown)
100        self.quitButton.pack(side=Tkinter.BOTTOM)
101
102        self.pauseButton = Tkinter.Button(self.root, text='Pause', command=self.unpaused.clear)
103        self.pauseButton.pack(side=Tkinter.LEFT)
104
105        self.resumeButton = Tkinter.Button(self.root, text='Resume', command=self.unpaused.set)
106        self.resumeButton.pack(side=Tkinter.LEFT)
107       
108        #self.w2if = vtk.vtkWindowToImageFilter()
109        #self.w2if.SetInput(self.renderWindow)
110       
111        #self.pnmWr = vtk.vtkPNMWriter()
112        #self.pnmWr.SetInput(self.w2if.GetOutput())
113
114
115    def add_axes(self):
116        """Add axes to this Visualiser - TODO
117        """
118        pass
119
120    def setup_all(self):
121        """Draw in the data that is specified by setup or update
122        """
123
124        self.N_tri = len(self.domain.triangles)
125        self.N_vert = len(self.vertices)
126        self.cells = vtk.vtkCellArray()
127        self.vertices = self.domain.get_vertex_coordinates()
128        self.vert_index = zeros((self.N_vert,2), Float)
129        for n in range(self.N_vert):
130            for i in range(3):
131                self.vert_index[self.domain.triangles[n][i]] = self.vertices[n][i*2:i*2+2]
132
133        # Prepare the list of cells
134        for t in self.domain.triangles:
135            self.cells.InsertNextCell(3)
136            for i in range(3):
137                self.cells.InsertCellPoint(t[i])
138
139        # Set up the rendering of each quantity
140        for q in self.domain.quantities:
141            if self.setup[q] | self.updating[q]:
142                self.draw_quantity(q)
143
144    def draw_quantity(self, q):
145        if self.scale_z.has_key(q):
146            scale = self.scale_z[q]
147        else:
148            scale = self.default_scale_z
149           
150        if q == 'elevation':
151            shift = 0.001
152        else:
153            shift = 0.0
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                                 shift + 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.SetScalarRange(0.0,0.5)
199            mapper.Update()
200        else:
201            self.mappers[q] = vtk.vtkPolyDataMapper()
202            mapper = self.mappers[q]
203            mapper.SetInput(polydata)
204
205        if not self.actors.has_key(q):
206            self.actors[q] = vtk.vtkActor()
207            actor = self.actors[q]
208            actor.SetMapper(mapper)
209
210            if self.qcolor.has_key(q):
211                actor.GetProperty().SetColor(self.qcolor[q])
212            else:
213                actor.GetProperty().SetColor(0.5, 0.5, 0.5)
214
215            self.renderer.AddActor(actor)
216
217    def redraw(self):
218        if self.redraw_ready.isSet():
219            self.redraw_ready.wait()
220            self.redraw_ready.clear()
221            for q in self.domain.quantities:
222                if self.updating[q]:
223                    self.draw_quantity(q)
224
225            self.renderWindow.Render()
226             
227            self.root.update_idletasks()
228            self.idle.set()
229        self.root.after(100, self.redraw)     
230   
231    def shutdown(self):
232        """Shutdown the visualiser
233        """
234        self.running = False
235        self.idle.set()
236        self.unpaused.set()
237        self.root.withdraw()
238        self.root.quit()
239
240    def update(self):
241        """Update the visualiser's display.
242        Clients are expected to call this in their evolve() loop,
243        to keep the visualiser in sync with the simulation.
244        """
245        if self.running:
246            self.redraw_ready.set()
247            self.idle.wait()
248            self.idle.clear()
249            self.unpaused.wait()
Note: See TracBrowser for help on using the repository browser.