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

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

Animation support added to the offline visualiser.

File size: 6.1 KB
Line 
1from threading import Thread
2from Tkinter import Tk, Button, N, E, S, W
3from types import FunctionType, TupleType
4from vtk import vtkActor, vtkFloatArray, vtkPolyDataMapper, vtkRenderer
5from vtk.tk.vtkTkRenderWidget import vtkTkRenderWidget
6
7class Visualiser(Thread):
8    """Superclass of both the realtime and offline VTK visualisers
9    """
10    def __init__(self, source):
11        Thread.__init__(self)
12
13        self.source = source
14
15        # Structures for Height Based quantities
16        self.height_quantities = []
17        self.height_zScales = {}
18        self.height_dynamic = {}
19        self.height_offset = {}
20
21        # Structures for colouring quantities
22        self.colours_height = {}
23
24        # Structures used for VTK
25        self.vtk_actors = {}
26        self.vtk_mappers = {}
27        self.vtk_polyData = {}
28
29        self.setup_gui()
30
31    def run(self):
32        self.setup_grid()
33        # Draw Height Quantities
34        for q in self.height_quantities:
35            self.update_height_quantity(q, self.height_dynamic[q])
36            self.draw_height_quantity(q)
37        self.tk_root.mainloop()
38
39    def redraw_quantities(self, dynamic_only=False):
40        """Redraw all dynamic quantities, unless dynamic_only is True.
41        """
42        # Height quantities
43        for q in self.height_quantities:
44            if (dynamic_only is False) or (self.height_dynamic[q]):
45                self.update_height_quantity(q, self.height_dynamic[q])
46                self.draw_height_quantity(q)
47
48    # --- Height Based Rendering --- #
49
50    def setup_grid(self):
51        """Create the vtkCellArray instance that represents the
52        triangles. Subclasses are expected to override this function
53        to read from their source as appropriate.
54        """
55        pass
56
57    def render_quantity_height(self, quantityName, zScale=1.0, offset=0.0, dynamic=True):
58        """Instruct the visualiser to render a quantity using the
59        value at a point as its height.  The value at each point is
60        multiplied by z_scale and is added to offset, and if
61        dynamic=False, the quantity is not recalculated on each
62        update.
63        """
64        self.height_quantities.append(quantityName)
65        self.height_zScales[quantityName] = zScale
66        self.height_offset[quantityName] = offset
67        self.height_dynamic[quantityName] = dynamic
68
69    def update_height_quantity(self, quantityName, dynamic=True):
70        """Create a vtkPolyData object and store it in
71        self.vtk_polyData[q]. Subclasses are expected to override this
72        function.
73        """
74        pass
75
76
77    def draw_height_quantity(self, quantityName):
78        """Use the vtkPolyData and prepare/update the rest of the VTK
79        rendering pipeline.
80        """
81        if self.vtk_mappers.has_key(quantityName):
82            mapper = self.vtk_mappers[quantityName]
83        else:
84            mapper = self.vtk_mappers[quantityName] = vtkPolyDataMapper()
85        mapper.SetInput(self.vtk_polyData[quantityName])
86        mapper.Update()
87
88        if not self.vtk_actors.has_key(quantityName):
89            actor = self.vtk_actors[quantityName] = vtkActor()
90            actor.SetMapper(mapper)
91            self.vtk_renderer.AddActor(actor)
92
93        if self.colours_height.has_key(quantityName):
94            colour = self.colours_height[quantityName]
95            if type(colour) == TupleType:
96                if type(colour[0]) == FunctionType:
97                    # It's a function, so take colour[1] as the
98                    # lower bound on the scalar range and
99                    # colour[2] as the upper bound on the scalar
100                    # range.
101                    scalars = vtkFloatArray()
102                    map(scalars.InsertNextValue, colour[0](self.build_quantity_dict()))
103                    self.vtk_polyData[quantityName].GetPointData().SetScalars(scalars)
104                    mapper.SetScalarRange(colour[1:])
105                    mapper.Update()
106                else:
107                    # It's a 3-tuple representing an RGB value.
108                    actor.GetProperty().SetColor(colour)
109            else:
110                actor.GetProperty().SetColor(0.5, 0.5, 0.5)
111        else:
112            actor.GetProperty().SetColor(0.5, 0.5, 0.5)
113
114    # --- Colour Coding --- #
115
116    def build_quantity_dict(self):
117        """Build a dictionary mapping quantity name->Numeric array of vertex
118        values for that quantity. Subclasses are expected to override
119        this function."""
120        pass
121
122    def colour_height_quantity(self, quantityName, colour=(0.5, 0.5, 0.5)):
123        """Add colouring to a height based quantity.
124
125        The colour parameter can be one of the following:
126        - a 3-tuple of values in [0,1] to specify R, G, B values
127        - a 3-tuple of values:
128          - a function that takes a dictionary mapping quantity name->Numeric array of vertex values.
129            This function returns a list of vertex values to be used in the colour coding.
130          - a float for the lower bound on the colouring
131          - a float for the upper bound on the colouring
132        """
133        self.colours_height[quantityName] = colour
134           
135    # --- Vector Fields --- #
136
137    # --- GUI Setup --- #
138
139    def setup_gui(self):
140        self.tk_root = Tk()
141        self.tk_root.title("Visualisation")
142        self.tk_root.after(100, self.redraw)
143        self.tk_root.bind("<Destroy>", self.destroyed)
144
145        self.tk_renderWidget = vtkTkRenderWidget(self.tk_root, width=400, height=400)
146        self.tk_renderWidget.grid(row=0, column=0, sticky=N+S+E+W)
147        self.tk_quit = Button(self.tk_root, text="Quit", command=self.shutdown)
148        self.tk_quit.grid(row=2, column=0, sticky=E+W)
149        self.vtk_renderer = vtkRenderer()
150        self.tk_renderWidget.GetRenderWindow().AddRenderer(self.vtk_renderer)
151
152    # --- GUI Events --- #
153
154    def destroyed(self, event):
155        if event.widget == self.tk_root:
156            self.shutdown()
157
158    def redraw(self):
159        self.tk_renderWidget.GetRenderWindow().Render()
160        self.tk_root.update_idletasks()
161        self.tk_root.after(100, self.redraw)
162
163    def shutdown(self):
164        self.tk_root.withdraw()
165        self.tk_root.destroy()
Note: See TracBrowser for help on using the repository browser.