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

Last change on this file since 5453 was 4267, checked in by jack, 17 years ago

Added support for custom opacity for height quantities. See the updated visualise_rectangle.py

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