source: anuga_core/source/anuga/visualiser/visualiser.py @ 7452

Last change on this file since 7452 was 6113, checked in by steve, 15 years ago

Added wireframe visualisation

File size: 11.3 KB
Line 
1from threading import Thread
2from Queue import Queue
3from Tkinter import Tk, Button, Frame, N, E, S, W
4from types import FunctionType, TupleType
5from vtk import vtkActor, vtkCubeAxesActor2D, vtkDelaunay2D, vtkFloatArray, vtkPoints, vtkPolyData, vtkPolyDataMapper, vtkRenderer
6from vtk.tk.vtkTkRenderWidget import vtkTkRenderWidget
7
8class Visualiser(Thread):
9    """Superclass of both the realtime and offline VTK visualisers
10    """
11    def __init__(self, source):
12        Thread.__init__(self)
13
14        self.source = source
15
16        # Structures for Height Based quantities
17        self.height_quantities = []
18        self.height_zScales = {}
19        self.height_dynamic = {}
20        self.height_offset = {}
21        self.height_opacity = {}
22        self.height_wireframe = {}
23
24        # Structures for colouring quantities
25        self.colours_height = {}
26
27        # Structures used for VTK
28        self.vtk_actors = {}
29        self.vtk_axesSet = False
30        self.vtk_drawAxes = False
31        self.vtk_mappers = {}
32        self.vtk_polyData = {}
33
34        # A list of operations to be performed on the cube axes. Type: [(func, (args))]
35        self.conf_axesAlterations = []
36        # A list of all polygons to overlay. Type: [([coords], height, (colour)]
37        self.conf_overlaidPolygons = []
38        # A list of alterations to be performed on the Tk root. Type: [(func, (args))]
39        self.conf_tkAlterations = []
40
41    def run(self):
42        self.vtk_renderer = vtkRenderer()
43        self.setup_gui()
44        self.setup_grid()
45
46        # Handle any deferred configuration
47        # Overlaid polygons
48        for args in self.conf_overlaidPolygons:
49            self.overlay_polygon_internal(*args)
50        # Draw (and maybe alter) the axes
51        if self.vtk_drawAxes:
52            self.vtk_axes = vtkCubeAxesActor2D()
53            # Perform all of the alterations required, by applying func to the vtk_axes instance (with the given args).
54            for func, args in self.conf_axesAlterations:
55                func(*((self.vtk_axes,) + args))
56        # Alter the Tk root as necessary.
57        for func, args in self.conf_tkAlterations:
58            func(*((self.tk_root,) + args))
59        # Finished with deferred configuration.
60
61        # Draw Height Quantities
62        for q in self.height_quantities:
63            self.update_height_quantity(q, self.height_dynamic[q])
64            self.draw_height_quantity(q)
65           
66        self.tk_root.mainloop()
67
68    def redraw_quantities(self):
69        """Redraw all dynamic quantities.
70        """
71        # Height quantities
72        for q in self.height_quantities:
73            if (self.height_dynamic[q]):
74                self.update_height_quantity(q, self.height_dynamic[q])
75                self.draw_height_quantity(q)
76        if self.vtk_drawAxes is True:
77            self.draw_axes()
78
79    # --- Axes --- #
80       
81    def render_axes(self):
82        """Intstruct the visualiser to render cube axes around the render.
83        """
84        self.vtk_drawAxes = True
85
86    def draw_axes(self):
87        """Update the 3D bounds on the axes and add them to the pipeline if not yet connected.
88        """
89        self.vtk_axes.SetBounds(self.get_3d_bounds())
90        if not self.vtk_axesSet:
91            self.vtk_axesSet = True
92            self.vtk_axes.SetCamera(self.vtk_renderer.GetActiveCamera())
93            self.vtk_renderer.AddActor(self.vtk_axes)
94            self.vtk_renderer.ResetCamera(self.get_3d_bounds())
95       
96    def alter_axes(self, func, args):
97        """Attempt to apply the function 'func' with args tuple 'args' to the
98        vtkCubeAxesActor2D instance set up by render_axes. This is done this way to ensure
99        the axes setup is handled in the visualiser thread.
100
101        Example call:
102        from vtk import vtkCubeAxesActor2D
103        alter_axes(vtkCubeAxesActor2D.SetNumberOfPoints, (5,))
104        """
105        self.conf_axesAlterations.append((func, args))
106           
107    # --- Height Based Rendering --- #
108
109    def setup_grid(self):
110        """Create the vtkCellArray instance that represents the
111        triangles. Subclasses are expected to override this function
112        to read from their source as appropriate. The vtkCellArray should
113        be stored to self.vtk_cells.
114        """
115        pass
116
117    def render_quantity_height(self, quantityName, zScale=1.0, offset=0.0, opacity=1.0, dynamic=True, wireframe=False):
118        """Instruct the visualiser to render a quantity using the
119        value at a point as its height.  The value at each point is
120        multiplied by z_scale and is added to offset, and if
121        dynamic=False, the quantity is not recalculated on each
122        update.
123        """
124        self.height_quantities.append(quantityName)
125        self.height_zScales[quantityName] = zScale
126        self.height_offset[quantityName] = offset
127        self.height_dynamic[quantityName] = dynamic
128        self.height_opacity[quantityName] = opacity
129        self.height_wireframe[quantityName] = wireframe
130
131    def update_height_quantity(self, quantityName, dynamic=True):
132        """Create a vtkPolyData object and store it in
133        self.vtk_polyData[quantityName]. Subclasses are expected to override this
134        function.
135        """
136        pass
137
138    def get_3d_bounds(self):
139        """Get the minimum and maximum bounds for the x, y and z directions.
140        Return as a list of double in the order (xmin, xmax, ymin, ymax, zmin, zmax),
141        suitable for passing to vtkCubeAxesActor2D::SetRanges(). Subclasses are expected
142        to override this function.
143        """
144        pass
145
146    def draw_height_quantity(self, quantityName):
147        """Use the vtkPolyData and prepare/update the rest of the VTK
148        rendering pipeline.
149        """
150        if self.vtk_mappers.has_key(quantityName):
151            mapper = self.vtk_mappers[quantityName]
152        else:
153            mapper = self.vtk_mappers[quantityName] = vtkPolyDataMapper()
154        mapper.SetInput(self.vtk_polyData[quantityName])
155        mapper.Update()
156
157        if not self.vtk_actors.has_key(quantityName):
158            actor = self.vtk_actors[quantityName] = vtkActor()
159            actor.GetProperty().SetOpacity(self.height_opacity[quantityName])
160            if self.height_wireframe[quantityName]:
161                actor.GetProperty().SetRepresentationToWireframe()
162            actor.SetMapper(mapper)
163            self.vtk_renderer.AddActor(actor)
164        else:
165            actor = self.vtk_actors[quantityName]
166
167        if self.colours_height.has_key(quantityName):
168            colour = self.colours_height[quantityName]
169            if type(colour) == TupleType:
170                if type(colour[0]) == FunctionType:
171                    # It's a function, so take colour[1] as the
172                    # lower bound on the scalar range and
173                    # colour[2] as the upper bound on the scalar
174                    # range.
175                    scalars = vtkFloatArray()
176                    map(scalars.InsertNextValue, colour[0](self.build_quantity_dict()))
177                    self.vtk_polyData[quantityName].GetPointData().SetScalars(scalars)
178                    mapper.SetScalarRange(colour[1:])
179                    mapper.Update()
180                else:
181                    # It's a 3-tuple representing an RGB value.
182                    actor.GetProperty().SetColor(colour)
183            else:
184                actor.GetProperty().SetColor(0.5, 0.5, 0.5)
185        else:
186            actor.GetProperty().SetColor(0.5, 0.5, 0.5)
187
188    # --- Colour Coding --- #
189
190    def build_quantity_dict(self):
191        """Build and return a dictionary mapping quantity name->Numeric array of vertex
192        values for that quantity. Subclasses are expected to override
193        this function."""
194        pass
195
196    def colour_height_quantity(self, quantityName, colour=(0.5, 0.5, 0.5)):
197        """Add colouring to a height based quantity.
198
199        The colour parameter can be one of the following:
200        - a 3-tuple of values in [0,1] to specify R, G, B values
201        - a 3-tuple of values:
202          - a function that takes a dictionary mapping quantity name->Numeric array of vertex values.
203            This function returns a list of vertex values to be used in the colour coding.
204          - a float for the lower bound on the colouring
205          - a float for the upper bound on the colouring
206        """
207        self.colours_height[quantityName] = colour
208
209    # --- Overlaid Polygons --- #
210
211    def overlay_polygon(self, coords, height=0.0, colour=(1.0, 0.0, 0.0)):
212        """Add a polygon to the output of the visualiser.
213
214        coords is a list of 2-tuples representing x and y coordinates.
215        These are triangulated by vtkDelaunay2D.
216
217        height is the z-value given to all points.
218
219        colour is the colour of the polygon, as a 3-tuple representing
220        r, g, b values between 0 and 1."""
221        self.conf_overlaidPolygons.append((coords, height, colour))
222
223    def overlay_polygon_internal(self, coords, height, colour):
224        """Add a polygon to the output of the visualiser.
225
226        coords is a list of 2-tuples representing x and y coordinates.
227        These are triangulated by vtkDelaunay2D.
228
229        height is the z-value given to all points.
230
231        colour is the colour of the polygon, as a 3-tuple representing
232        r, g, b values between 0 and 1.
233
234        This function should not be called from outside the visualiser thread.
235        Use overlay_polygon instead.
236   
237        """
238        points = vtkPoints()
239        for coord in coords:
240            points.InsertNextPoint(coord[0], coord[1], height)
241        profile = vtkPolyData()
242        profile.SetPoints(points)
243        delny = vtkDelaunay2D()
244        delny.SetInput(profile)
245        mesh = vtkPolyDataMapper()
246        mesh.SetInput(delny.GetOutput())
247        actor = vtkActor()
248        actor.SetMapper(mesh)
249        actor.GetProperty().SetColor(colour)
250        self.vtk_renderer.AddActor(actor)
251       
252    # --- Vector Fields --- #
253
254    # --- GUI Setup --- #
255
256    def setup_gui(self):
257        self.tk_root = Tk()
258        self.tk_root.title("Visualisation")
259        self.tk_root.after(100, self.redraw)
260        self.tk_root.bind("<Destroy>", self.destroyed)
261        self.tk_root.grid_rowconfigure(0, weight=1)
262        self.tk_root.grid_columnconfigure(0, weight=1)
263
264        self.tk_renderWidget = vtkTkRenderWidget(self.tk_root, width=400, height=400)
265        self.tk_renderWidget.grid(row=0, column=0, sticky=N+S+E+W)
266        self.tk_controlFrame = Frame(self.tk_root)
267        self.tk_controlFrame.grid(row=1, column=0, sticky=E+W)
268        self.tk_controlFrame.grid_rowconfigure(0, weight=1)
269        self.tk_controlFrame.grid_columnconfigure(0, weight=1)
270       
271        self.tk_quit = Button(self.tk_controlFrame, text="Quit", command=self.shutdown)
272        self.tk_quit.grid(row=0, column=0, sticky=E+W)
273        self.tk_renderWidget.GetRenderWindow().AddRenderer(self.vtk_renderer)
274
275    def alter_tkroot(self, func, args):
276        """Apply func, with arguments tuple args to the root tk window for this visualiser.
277        """
278        self.conf_tkAlterations.append((func, args))
279
280    # --- GUI Events --- #
281
282    def destroyed(self, event):
283        if event.widget == self.tk_root:
284            self.shutdown()
285
286    def redraw(self):
287        self.tk_renderWidget.GetRenderWindow().Render()
288        self.tk_root.update_idletasks()
289        self.tk_root.after(100, self.redraw)
290
291    def shutdown(self):
292        self.tk_root.withdraw()
293        self.tk_root.destroy()
Note: See TracBrowser for help on using the repository browser.