source: trunk/anuga_core/source/anuga/visualiser/offline.py @ 8968

Last change on this file since 8968 was 8780, checked in by steve, 12 years ago

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