source: inundation/pyvolution/vtk_realtime_visualiser.py @ 2221

Last change on this file since 2221 was 2221, checked in by steve, 18 years ago
File size: 7.0 KB
Line 
1# VTK-Based visualiser
2# intended to replace realtime_visualisation_new.py
3
4import threading
5import Tkinter
6import vtk
7from vtk.tk.vtkTkRenderWidget import vtkTkRenderWidget
8
9class Visualiser(threading.Thread):
10
11    """A VTK-powered visualiser, designed to replace the VPython one.
12    Intended to be run in its own thread.
13
14    Customisation options for the visualiser are as follows:
15
16    setup: Dictionary mapping quantity name -> boolean.
17    if setup[q] is true, draw the quantity when the visualiser starts.
18
19    updating: Dictionary mapping quantity name -> boolean.
20    if updating[q] is true, update the rendering of this quantity each
21    timestep.
22
23    qcolor: Dictionary mapping quantity name -> (float, float, float)
24    if the name of a quantity is found in qcolor, the colour specified (in
25    (r, g, b) from 0.0 to 1.0) is used for display of that quantity instead of
26    (0.5, 0.5, 0.5) (the default)
27
28    scale_z: Dictionary mapping quantity name -> float
29    Override the z scaling of this quantity with the float specified.
30
31    default_scale_z: float
32    multiply all z coordinates by this amount unless an overriding value
33    exists in scale_z.
34    """
35
36    def __init__(self, domain, default_scale_z=1.0, rect=None, title='Test'):
37        threading.Thread.__init__(self)
38        # Initialise data structures. setup and updating are maps
39        # quantity name -> boolean, setup means to render it,
40        # update means to update with time.
41        # qcolor maps quantity name -> (float, float, float): the colour (r, g, b)
42        self.setup = {}
43        self.updating = {}
44        self.qcolor = {}
45        self.scale_z = {}
46        self.default_scale_z = default_scale_z
47        self.domain = domain
48        self.vertices = domain.vertex_coordinates
49
50        self.idle = threading.Event()
51        self.redraw_ready = threading.Event()
52
53        # Internal use - storage of vtk objects
54        self.grids = {}
55        self.actors = {}
56        self.polydata = {}
57        self.mappers = {}
58
59        # Default options
60        for x in self.domain.quantities:
61            self.setup[x] = False
62            self.updating[x] = False
63
64        # Bounding box.
65        if rect is None:
66            self.max_x = max(max(self.vertices[:,0]),max(self.vertices[:,2]),max(self.vertices[:,4]))
67            self.min_x = min(min(self.vertices[:,0]),min(self.vertices[:,2]),min(self.vertices[:,4]))
68            self.max_y = max(max(self.vertices[:,1]),max(self.vertices[:,3]),max(self.vertices[:,5]))
69            self.min_y = min(min(self.vertices[:,1]),min(self.vertices[:,3]),min(self.vertices[:,5]))
70        else:
71            self.max_x = rect[2]
72            self.min_x = rect[0]
73            self.max_y = rect[3]
74            self.min_y = rect[1]
75
76        self.range_x = self.max_x - self.min_x
77        self.range_y = self.max_y - self.min_y
78        self.range_xy = max(self.range_x, self.range_y)
79
80    def run(self):
81        self.initialise_gui()
82        self.add_axes()
83        self.setup_all()
84        self.root.after(100, self.idle.set)
85        self.root.mainloop()
86
87    def initialise_gui(self):
88        """Prepare the GUI for the Visualiser, and set up the
89        renderer.
90
91        """
92        # Using the TK VTK widget allows extendability
93        # should a gui need adding.
94        self.root = Tkinter.Tk()
95
96        # Message handling with after
97        self.root.after(100, self.redraw)
98
99        self.renderWidget = vtkTkRenderWidget(self.root, width=400, height=400)
100        self.renderWidget.pack(expand='true', fill='both')
101        self.renderWindow = self.renderWidget.GetRenderWindow()
102        self.renderer = vtk.vtkRenderer()
103        self.renderWindow.AddRenderer(self.renderer)
104
105        self.quitButton = Tkinter.Button(self.root, text='Quit', command=self.shutdown)
106        self.quitButton.pack(side=Tkinter.BOTTOM)
107
108    def add_axes(self):
109        """Add axes to this Visualiser
110        """
111        pass
112
113    def setup_all(self):
114        """Draw in the data that is specified by setup or update
115        """
116
117        self.N_tri = len(self.domain.triangles)
118        self.cells = vtk.vtkCellArray()
119        self.vertices = self.domain.get_vertex_coordinates()
120        self.vert_index = {}
121
122        for n in range(len(self.vertices)):
123            for i in range(3):
124                self.vert_index[self.domain.triangles[n][i]] = self.vertices[n][i*2:i*2+2]
125
126        self.N_vert = len(self.vert_index.keys())
127
128        # Prepare the list of cells
129        for t in self.domain.triangles:
130            self.cells.InsertNextCell(3)
131            for i in range(3):
132                self.cells.InsertCellPoint(t[i])
133
134        # Set up the rendering of each quantity
135        for q in self.domain.quantities:
136            if self.setup[q] | self.updating[q]:
137                self.draw_quantity(q)
138
139    def draw_quantity(self, q):
140        qty_index = {}
141        for n in range(self.N_tri):
142            for v in range(3):
143                qty_index[self.domain.triangles[n][v]] = self.domain.quantities[q].vertex_values[n][v]
144
145        self.grids[q] = vtk.vtkPoints()
146        grid = self.grids[q]
147
148        if self.scale_z.has_key(q):
149            scale = self.scale_z[q]
150        else:
151            scale = self.default_scale_z
152
153        for v in range(self.N_vert):
154            grid.InsertNextPoint(self.vert_index[v][0],
155                                 self.vert_index[v][1],
156                                 qty_index[v] * scale)
157
158        # Can't recycle vtkPolyData objects: Apparently they behave
159        # unusually if the points (i.e. vertex data) is set after
160        # the polys (i.e. triangle data)
161        self.polydata[q] = vtk.vtkPolyData()
162        polydata = self.polydata[q]
163
164        polydata.SetPoints(grid)
165        polydata.SetPolys(self.cells)
166
167        if self.mappers.has_key(q):
168            mapper = self.mappers[q]
169            mapper.SetInput(polydata)
170            mapper.Update()
171        else:
172            self.mappers[q] = vtk.vtkPolyDataMapper()
173            mapper = self.mappers[q]
174            mapper.SetInput(polydata)
175
176        if not self.actors.has_key(q):
177            self.actors[q] = vtk.vtkActor()
178            actor = self.actors[q]
179            actor.SetMapper(mapper)
180
181            if self.qcolor.has_key(q):
182                actor.GetProperty().SetColor(self.qcolor[q])
183            else:
184                actor.GetProperty().SetColor(0.5, 0.5, 0.5)
185
186            self.renderer.AddActor(actor)
187
188    def redraw(self):
189        if self.redraw_ready.isSet():
190            self.redraw_ready.wait()
191            self.redraw_ready.clear()
192            for q in self.domain.quantities:
193                if self.updating[q]:
194                    self.draw_quantity(q)
195
196            self.renderWindow.Render()
197            self.root.update_idletasks()
198            self.idle.set()
199        self.root.after(100, self.redraw)
200
201    def shutdown(self):
202        self.domain.visualise = False
203        self.root.withdraw()
204        self.root.quit()
Note: See TracBrowser for help on using the repository browser.