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

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