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

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

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

File size: 11.1 KB
RevLine 
[3540]1from threading import Thread
[3873]2from Queue import Queue
[3549]3from Tkinter import Tk, Button, Frame, N, E, S, W
[3538]4from types import FunctionType, TupleType
[3813]5from vtk import vtkActor, vtkCubeAxesActor2D, vtkDelaunay2D, vtkFloatArray, vtkPoints, vtkPolyData, vtkPolyDataMapper, vtkRenderer
[3465]6from vtk.tk.vtkTkRenderWidget import vtkTkRenderWidget
7
8class Visualiser(Thread):
9    """Superclass of both the realtime and offline VTK visualisers
10    """
[3493]11    def __init__(self, source):
[3465]12        Thread.__init__(self)
13
[3493]14        self.source = source
[3465]15
16        # Structures for Height Based quantities
17        self.height_quantities = []
18        self.height_zScales = {}
19        self.height_dynamic = {}
[3493]20        self.height_offset = {}
[4267]21        self.height_opacity = {}
[3465]22
[3538]23        # Structures for colouring quantities
24        self.colours_height = {}
25
[3493]26        # Structures used for VTK
27        self.vtk_actors = {}
[3670]28        self.vtk_axesSet = False
29        self.vtk_drawAxes = False
[3493]30        self.vtk_mappers = {}
31        self.vtk_polyData = {}
32
[3958]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 = []
[3548]39
[3465]40    def run(self):
[3873]41        self.vtk_renderer = vtkRenderer()
42        self.setup_gui()
[3493]43        self.setup_grid()
[3873]44
[3958]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.
[3873]59
[3493]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)
[3873]64           
[3465]65        self.tk_root.mainloop()
66
[3958]67    def redraw_quantities(self):
[3873]68        """Redraw all dynamic quantities.
[3493]69        """
70        # Height quantities
71        for q in self.height_quantities:
[3873]72            if (self.height_dynamic[q]):
[3493]73                self.update_height_quantity(q, self.height_dynamic[q])
74                self.draw_height_quantity(q)
[3670]75        if self.vtk_drawAxes is True:
[3873]76            self.draw_axes()
77
78    # --- Axes --- #
[3670]79       
80    def render_axes(self):
81        """Intstruct the visualiser to render cube axes around the render.
82        """
83        self.vtk_drawAxes = True
[3873]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)
[3958]93            self.vtk_renderer.ResetCamera(self.get_3d_bounds())
[3670]94       
[3873]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,))
[3670]103        """
[3958]104        self.conf_axesAlterations.append((func, args))
[3873]105           
106    # --- Height Based Rendering --- #
107
[3493]108    def setup_grid(self):
109        """Create the vtkCellArray instance that represents the
110        triangles. Subclasses are expected to override this function
[3625]111        to read from their source as appropriate. The vtkCellArray should
112        be stored to self.vtk_cells.
[3465]113        """
[3493]114        pass
[3465]115
[4267]116    def render_quantity_height(self, quantityName, zScale=1.0, offset=0.0, opacity=1.0, dynamic=True):
[3493]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
[4267]127        self.height_opacity[quantityName] = opacity
[3493]128
129    def update_height_quantity(self, quantityName, dynamic=True):
130        """Create a vtkPolyData object and store it in
[3625]131        self.vtk_polyData[quantityName]. Subclasses are expected to override this
[3493]132        function.
133        """
134        pass
135
[3670]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
[3493]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()
[4267]157            actor.GetProperty().SetOpacity(self.height_opacity[quantityName])
[3493]158            actor.SetMapper(mapper)
159            self.vtk_renderer.AddActor(actor)
[3873]160        else:
161            actor = self.vtk_actors[quantityName]
[3493]162
[3538]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
[3465]184    # --- Colour Coding --- #
185
[3538]186    def build_quantity_dict(self):
[3625]187        """Build and return a dictionary mapping quantity name->Numeric array of vertex
[3538]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
[3813]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."""
[3958]217        self.conf_overlaidPolygons.append((coords, height, colour))
[3873]218
[3958]219    def overlay_polygon_internal(self, coords, height, colour):
[3873]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        """
[3813]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       
[3465]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)
[3549]257        self.tk_root.grid_rowconfigure(0, weight=1)
258        self.tk_root.grid_columnconfigure(0, weight=1)
[3465]259
260        self.tk_renderWidget = vtkTkRenderWidget(self.tk_root, width=400, height=400)
[3493]261        self.tk_renderWidget.grid(row=0, column=0, sticky=N+S+E+W)
[3549]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)
[3465]269        self.tk_renderWidget.GetRenderWindow().AddRenderer(self.vtk_renderer)
270
[3873]271    def alter_tkroot(self, func, args):
272        """Apply func, with arguments tuple args to the root tk window for this visualiser.
273        """
[3958]274        self.conf_tkAlterations.append((func, args))
[3873]275
[3465]276    # --- GUI Events --- #
277
278    def destroyed(self, event):
279        if event.widget == self.tk_root:
280            self.shutdown()
281
282    def redraw(self):
[3493]283        self.tk_renderWidget.GetRenderWindow().Render()
284        self.tk_root.update_idletasks()
285        self.tk_root.after(100, self.redraw)
[3465]286
287    def shutdown(self):
288        self.tk_root.withdraw()
[3547]289        self.tk_root.destroy()
Note: See TracBrowser for help on using the repository browser.