source: inundation/pyvolution/vtk_realtime_visualiser.py @ 2195

Last change on this file since 2195 was 2195, checked in by jack, 19 years ago

Fixed problems with an alternate way of accessing the visualiser.

File size: 6.6 KB
Line 
1# VTK-Based visualiser
2# intended to replace realtime_visualisation_new.py
3
4import threading
5from Queue import Queue, Empty
6import Tkinter
7import vtk
8from vtk.tk.vtkTkRenderWidget import vtkTkRenderWidget
9
10class 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
Note: See TracBrowser for help on using the repository browser.