source: anuga_core/source/anuga/visualiser/offline.py @ 4325

Last change on this file since 4325 was 4325, checked in by steve, 17 years ago
File size: 11.7 KB
Line 
1from Numeric import array, Float, ravel, zeros
2from Scientific.IO.NetCDF import NetCDFFile
3from Tkinter import Button, E, Tk, W, Label, StringVar, Scale, HORIZONTAL
4from visualiser import Visualiser
5from vtk import vtkCellArray, vtkPoints, vtkPolyData
6
7class OfflineVisualiser(Visualiser):
8    """A VTK-powered offline visualiser which runs in its own thread.
9    In addition to the functions provided by the standard visualiser,
10    the following additional functions are provided:
11
12    precache_height_quantities() - Precache all the vtkpoints
13    structures for any dynamic height based quantities to render.
14    """
15    def __init__(self, source, frameDelay=100, frameStep=1):
16        """The source parameter is assumed to be a NetCDF sww file.
17        The frameDelay parameter is the number of milliseconds waited between frames.
18        """
19        Visualiser.__init__(self, source)
20
21        self.frameNumber = 0
22        fin = NetCDFFile(self.source, 'r')
23        self.maxFrameNumber = fin.variables['time'].shape[0] - 1
24        fin.close()
25       
26        #self.frameNumberTkVariable = StringVar()
27        #self.frameNumberTkVariable.set('Frame - %05g'%self.framNumber)
28
29        self.frameDelay = frameDelay
30
31        self.xmin = None
32        self.xmax = None
33        self.ymin = None
34        self.ymax = None
35        self.zmin = None
36        self.zmax = None
37
38        self.frameStep= frameStep
39
40        self.vtk_heightQuantityCache = []
41        for i in range(self.maxFrameNumber + 1): # maxFrameNumber is zero indexed.
42            self.vtk_heightQuantityCache.append({})
43
44        self.paused = False
45        self.movie = False
46       
47    def setup_grid(self):
48        fin = NetCDFFile(self.source, 'r')
49        self.vtk_cells = vtkCellArray()
50        N_tri = fin.variables['volumes'].shape[0]
51        for v in range(N_tri):
52            self.vtk_cells.InsertNextCell(3)
53            for i in range(3):
54                self.vtk_cells.InsertCellPoint(fin.variables['volumes'][v][i])
55        fin.close()
56       
57    def update_height_quantity(self, quantityName, dynamic=True):
58        polydata = self.vtk_polyData[quantityName] = vtkPolyData()
59        if dynamic is True:
60            #print ' - Frame',self.frameNumber,'of',self.maxFrameNumber
61            if not self.vtk_heightQuantityCache[self.frameNumber].has_key(quantityName):
62                self.vtk_heightQuantityCache[self.frameNumber][quantityName]\
63                    = self.read_height_quantity(quantityName, True, self.frameNumber);
64            polydata.SetPoints(self.vtk_heightQuantityCache[self.frameNumber][quantityName])
65        else:
66            polydata.SetPoints(self.read_height_quantity(quantityName, False))
67        polydata.SetPolys(self.vtk_cells)
68
69    def get_3d_bounds(self):
70        return [self.xmin, self.xmax, self.ymin, self.ymax, self.zmin, self.zmax]
71           
72    def read_height_quantity(self, quantityName, dynamic=True, frameNumber=0):
73        """Read in a height based quantity from the NetCDF source file
74        and return a vtkPoints object. frameNumber is ignored if
75        dynamic is false."""
76        fin = NetCDFFile(self.source, 'r')
77        points = vtkPoints()
78        if dynamic is True:
79            N_vert = fin.variables[quantityName].shape[1]
80        else:
81            N_vert = len(fin.variables[quantityName])
82        x = ravel(array(fin.variables['x'], Float))
83        y = ravel(array(fin.variables['y'], Float))
84        if dynamic is True:
85            q = array(fin.variables[quantityName][frameNumber], Float)
86        else:
87            q = ravel(array(fin.variables[quantityName], Float))
88
89        q *= self.height_zScales[quantityName]
90        q += self.height_offset[quantityName]
91
92        for v in range(N_vert):
93            points.InsertNextPoint(x[v], y[v], q[v])
94            if self.xmin == None or self.xmin > x[v]:
95                self.xmin = x[v]
96            if self.xmax == None or self.xmax < x[v]:
97                self.xmax = x[v]
98            if self.ymin == None or self.ymin > y[v]:
99                self.ymin = y[v]
100            if self.ymax == None or self.ymax < y[v]:
101                self.ymax = y[v]
102            if self.zmin == None or self.zmin > q[v]:
103                self.zmin = q[v]
104            if self.zmax == None or self.zmax < q[v]:
105                self.zmax = q[v]
106        fin.close()
107        return points
108
109    def precache_height_quantities(self):
110        """Precache any height-based quantities. Call before rendering
111        beigns."""
112        for q in self.height_quantities:
113            if self.height_dynamic[q] is True:
114                print 'Precaching %s' % q
115                for i in range(self.maxFrameNumber + 1): # maxFrameNumber is zero-indexed
116                    print ' - Frame %d of %d' % (i, self.maxFrameNumber)
117                    self.vtk_heightQuantityCache[i][q]\
118                        = self.read_height_quantity(q, True, i)
119
120    def build_quantity_dict(self):
121        quantities = {}
122        fin = NetCDFFile(self.source, 'r')
123        for q in filter(lambda n:n != 'x' and n != 'y' and n != 'z' and n != 'time' and n != 'volumes', fin.variables.keys()):
124            if len(fin.variables[q].shape) == 1: # Not a time-varying quantity
125                quantities[q] = ravel(array(fin.variables[q], Float))
126            else: # Time-varying, get the current timestep data
127                quantities[q] = array(fin.variables[q][self.frameNumber], Float)
128        fin.close()
129        return quantities
130
131    def setup_gui(self):
132        Visualiser.setup_gui(self)
133        self.tk_quit.grid(row=0, column=0, sticky=W+E)
134        self.tk_movie_toggle = Button(self.tk_controlFrame, text="Movie off", command=self.movie_toggle)
135        self.tk_movie_toggle.grid(row=0, column=6,  sticky=W+E)
136       
137               
138        self.tk_restart = Button(self.tk_controlFrame, text="<<<", command=self.restart, width=5)
139        self.tk_restart.grid(row=1, column=0, sticky=W+E)
140        self.tk_back10 = Button(self.tk_controlFrame, text="<<", command=self.back10, width=5)
141        self.tk_back10.grid(row=1, column=1, sticky=W+E)
142        self.tk_back = Button(self.tk_controlFrame, text="<", command=self.back, width=5)
143        self.tk_back.grid(row=1, column=2, sticky=W+E)
144        self.tk_pauseResume = Button(self.tk_controlFrame, text="Pause", command=self.pauseResume, width=15)
145        self.tk_pauseResume.grid(row=1, column=3, sticky=W+E)
146        self.tk_forward = Button(self.tk_controlFrame, text=">", command=self.forward, width=5)
147        self.tk_forward.grid(row=1, column=4, sticky=W+E)
148        self.tk_forward10 = Button(self.tk_controlFrame, text=">>", command=self.forward10, width=5)
149        self.tk_forward10.grid(row=1, column=5, sticky=W+E)
150        self.tk_forwardEnd = Button(self.tk_controlFrame, text=">>>", command=self.forwardEnd, width=5)
151        self.tk_forwardEnd.grid(row=1, column=6, sticky=W+E)
152       
153
154        self.tk_frameNumber = Label(self.tk_controlFrame, text='Frame')
155        self.tk_frameNumber.grid(row=2, column=0, sticky=W+E)
156        self.tk_gotoFrame = Scale(self.tk_controlFrame, from_=0, to=self.maxFrameNumber, orient=HORIZONTAL)
157        self.tk_gotoFrame.grid(row=2, column=1, columnspan=2, sticky=W+E)
158        self.tk_stepLabel = Label(self.tk_controlFrame, text='Step')
159        self.tk_stepLabel.grid(row=2, column=4, sticky=W+E)       
160        self.tk_frameStep = Scale(self.tk_controlFrame, from_=0, to=self.maxFrameNumber, orient=HORIZONTAL)
161        self.tk_frameStep.grid(row=2, column=5, columnspan=2, sticky=W+E)
162       
163        # Make the buttons stretch to fill all available space
164        for i in range(7):
165            self.tk_controlFrame.grid_columnconfigure(i, weight=1)
166
167    def run(self):
168        self.alter_tkroot(Tk.after, (self.frameDelay, self.animateForward))
169        Visualiser.run(self)
170
171    def restart(self):
172        self.frameNumber = 0
173        self.redraw_quantities()
174        self.update_labels()
175        self.pause()
176       
177        if self.movie:
178            self.save_image()
179 
180    def forwardEnd(self):
181        self.frameNumber = self.maxFrameNumber
182        self.redraw_quantities()
183        self.update_labels()
184        self.pause()
185               
186    def movie_toggle(self):
187        if self.movie == True:
188            self.movie = False
189            self.tk_movie_toggle.config(text='Movie off')
190        else:
191            self.movie = True
192            self.tk_movie_toggle.config(text='Movie on ')
193           
194           
195       
196       
197    def save_image(self):
198       
199        from vtk import vtkJPEGWriter, vtkJPEGWriter, vtkPNGWriter
200        from vtk import vtkPNMWriter, vtkWindowToImageFilter
201        from os import path
202         
203        sourcebase, _ = path.splitext(self.source)
204        fname = sourcebase+'%05g.png' % self.frameNumber
205        #print fname
206       
207        extmap = {'.jpg' : vtkJPEGWriter,
208                  '.jpeg' : vtkJPEGWriter,
209                  '.png' : vtkPNGWriter,
210                  '.pnm' : vtkPNMWriter,
211                  }
212        basename, ext = path.splitext(fname)
213        try: Writer = extmap[ext.lower()]
214        except KeyError:
215            error_msg("Don't know how to handle %s files" % ext, parent=self)
216            return
217   
218        renWin = self.vtk_renderer.GetRenderWindow()
219        w2i = vtkWindowToImageFilter()
220        writer = Writer()
221        w2i.SetInput(renWin)
222        w2i.Update()
223        writer.SetInput(w2i.GetOutput())
224        writer.SetFileName(fname)
225        renWin.Render()
226        writer.Write()       
227   
228    def back10(self):
229        if self.frameNumber - 10 >= 0:
230            self.frameNumber -= 10
231        else:
232            self.frameNumber = 0
233        self.redraw_quantities()
234        self.update_labels()
235        self.pause()
236
237    def back(self):
238        if self.frameNumber > 0:
239            self.frameNumber -= 1
240            self.redraw_quantities()
241            self.update_labels()
242            self.pause()
243
244    def pauseResume(self):
245        if self.paused is True:
246            self.resume()
247        else:
248            self.pause()
249
250    def pause(self):
251        self.paused = True
252        self.tk_pauseResume.config(text="Resume")
253
254    def resume(self):
255        self.paused = False
256        self.tk_pauseResume.config(text="Pause")
257        self.frameNumber = self.tk_gotoFrame.get()
258        self.frameStep = self.tk_frameStep.get()
259        self.tk_root.after(self.frameDelay, self.animateForward)
260
261    def forward(self):
262        if self.frameNumber < self.maxFrameNumber:
263            self.frameNumber += 1
264            self.redraw_quantities()
265            self.update_labels()
266            self.pause()
267           
268    def forward_step(self):
269        if self.frameNumber + self.frameStep <= self.maxFrameNumber:
270            self.frameNumber += self.frameStep
271            self.redraw_quantities()
272            self.update_labels()
273        else:
274            self.frameNumber = self.maxFrameNumber           
275            self.redraw_quantities()
276            self.update_labels()   
277            self.pause()
278         
279        if self.movie:
280             self.save_image()
281               
282
283    def forward10(self):
284        if self.frameNumber + 10 <= self.maxFrameNumber:
285            self.frameNumber += 10
286        else:
287            self.frameNumber = self.maxFrameNumber
288        self.redraw_quantities()
289        self.update_labels()
290        self.pause()
291
292    def animateForward(self):
293        if self.paused is not True:
294            self.forward_step()
295            self.tk_root.after(self.frameDelay, self.animateForward)
296           
297    def update_labels(self): 
298        #self.tk_frameNumber.config(text='%05g of %05g'%(self.frameNumber,self.maxFrameNumber))
299        self.tk_gotoFrame.set(self.frameNumber)
300        self.tk_frameStep.set(self.frameStep)
301               
302    def shutdown(self):
303        #self.pause()
304        self.tk_root.withdraw()
305        self.tk_root.destroy()
306        #Visualiser.shutdown(self)
Note: See TracBrowser for help on using the repository browser.