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

Last change on this file since 7452 was 7452, checked in by steve, 15 years ago

Committing latest parallel code

File size: 11.7 KB
RevLine 
[7452]1#from Numeric import array, Float, ravel, zeros
2import numpy as num
[3493]3from Scientific.IO.NetCDF import NetCDFFile
[4325]4from Tkinter import Button, E, Tk, W, Label, StringVar, Scale, HORIZONTAL
[3493]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.
[3538]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.
[3493]15    """
[4325]16    def __init__(self, source, frameDelay=100, frameStep=1):
[3493]17        """The source parameter is assumed to be a NetCDF sww file.
[3873]18        The frameDelay parameter is the number of milliseconds waited between frames.
[3493]19        """
20        Visualiser.__init__(self, source)
21
[3536]22        self.frameNumber = 0
[3493]23        fin = NetCDFFile(self.source, 'r')
[3536]24        self.maxFrameNumber = fin.variables['time'].shape[0] - 1
[3493]25        fin.close()
[4325]26       
27        #self.frameNumberTkVariable = StringVar()
28        #self.frameNumberTkVariable.set('Frame - %05g'%self.framNumber)
[3493]29
[3873]30        self.frameDelay = frameDelay
31
[3670]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
[4325]39        self.frameStep= frameStep
[4321]40
[3536]41        self.vtk_heightQuantityCache = []
[3537]42        for i in range(self.maxFrameNumber + 1): # maxFrameNumber is zero indexed.
[3536]43            self.vtk_heightQuantityCache.append({})
44
[3548]45        self.paused = False
[4322]46        self.movie = False
[3611]47       
[3493]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()
[3611]57       
[3493]58    def update_height_quantity(self, quantityName, dynamic=True):
[3536]59        polydata = self.vtk_polyData[quantityName] = vtkPolyData()
60        if dynamic is True:
[4325]61            #print ' - Frame',self.frameNumber,'of',self.maxFrameNumber
[3536]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)
[3670]69
70    def get_3d_bounds(self):
71        return [self.xmin, self.xmax, self.ymin, self.ymax, self.zmin, self.zmax]
[3536]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."""
[3493]77        fin = NetCDFFile(self.source, 'r')
[3536]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])
[7452]83        x = num.ravel(num.array(fin.variables['x'], num.float))
84        y = num.ravel(num.array(fin.variables['y'], num.float))
[3536]85        if dynamic is True:
[7452]86            q = num.array(fin.variables[quantityName][frameNumber], num.float)
[3536]87        else:
[7452]88            q = num.ravel(num.array(fin.variables[quantityName], num.float))
[3524]89
[3536]90        q *= self.height_zScales[quantityName]
91        q += self.height_offset[quantityName]
[3524]92
[3536]93        for v in range(N_vert):
94            points.InsertNextPoint(x[v], y[v], q[v])
[3670]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]
[3493]107        fin.close()
[3536]108        return points
[3493]109
[3538]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
[7452]126                quantities[q] = num.ravel(num.array(fin.variables[q], num.float))
[3538]127            else: # Time-varying, get the current timestep data
[7452]128                quantities[q] = num.array(fin.variables[q][self.frameNumber], num.float)
[3538]129        fin.close()
130        return quantities
131
[3493]132    def setup_gui(self):
133        Visualiser.setup_gui(self)
[4325]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)
[3493]140        self.tk_restart.grid(row=1, column=0, sticky=W+E)
[4325]141        self.tk_back10 = Button(self.tk_controlFrame, text="<<", command=self.back10, width=5)
[3493]142        self.tk_back10.grid(row=1, column=1, sticky=W+E)
[4325]143        self.tk_back = Button(self.tk_controlFrame, text="<", command=self.back, width=5)
[3493]144        self.tk_back.grid(row=1, column=2, sticky=W+E)
[3622]145        self.tk_pauseResume = Button(self.tk_controlFrame, text="Pause", command=self.pauseResume, width=15)
[3493]146        self.tk_pauseResume.grid(row=1, column=3, sticky=W+E)
[4325]147        self.tk_forward = Button(self.tk_controlFrame, text=">", command=self.forward, width=5)
[3493]148        self.tk_forward.grid(row=1, column=4, sticky=W+E)
[4325]149        self.tk_forward10 = Button(self.tk_controlFrame, text=">>", command=self.forward10, width=5)
[3493]150        self.tk_forward10.grid(row=1, column=5, sticky=W+E)
[4325]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       
[3493]154
[4325]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       
[3549]164        # Make the buttons stretch to fill all available space
[4322]165        for i in range(7):
[3549]166            self.tk_controlFrame.grid_columnconfigure(i, weight=1)
167
[3611]168    def run(self):
[3873]169        self.alter_tkroot(Tk.after, (self.frameDelay, self.animateForward))
[3611]170        Visualiser.run(self)
171
[3493]172    def restart(self):
[3536]173        self.frameNumber = 0
[3873]174        self.redraw_quantities()
[4325]175        self.update_labels()
[3622]176        self.pause()
[4322]177       
178        if self.movie:
179            self.save_image()
[4325]180 
181    def forwardEnd(self):
182        self.frameNumber = self.maxFrameNumber
183        self.redraw_quantities()
184        self.update_labels()
185        self.pause()
186               
[4322]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   
[3493]229    def back10(self):
[3536]230        if self.frameNumber - 10 >= 0:
231            self.frameNumber -= 10
[3493]232        else:
[3536]233            self.frameNumber = 0
[3873]234        self.redraw_quantities()
[4325]235        self.update_labels()
[3622]236        self.pause()
[3493]237
238    def back(self):
[3536]239        if self.frameNumber > 0:
240            self.frameNumber -= 1
[3873]241            self.redraw_quantities()
[4325]242            self.update_labels()
[3622]243            self.pause()
[3493]244
245    def pauseResume(self):
[3548]246        if self.paused is True:
[3622]247            self.resume()
[3548]248        else:
[3622]249            self.pause()
[3493]250
[3622]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")
[4325]258        self.frameNumber = self.tk_gotoFrame.get()
259        self.frameStep = self.tk_frameStep.get()
[3873]260        self.tk_root.after(self.frameDelay, self.animateForward)
[3622]261
[3493]262    def forward(self):
[3536]263        if self.frameNumber < self.maxFrameNumber:
264            self.frameNumber += 1
[3873]265            self.redraw_quantities()
[4325]266            self.update_labels()
[4321]267            self.pause()
268           
269    def forward_step(self):
[4325]270        if self.frameNumber + self.frameStep <= self.maxFrameNumber:
271            self.frameNumber += self.frameStep
[4321]272            self.redraw_quantities()
[4325]273            self.update_labels()
[3622]274        else:
[4321]275            self.frameNumber = self.maxFrameNumber           
276            self.redraw_quantities()
[4325]277            self.update_labels()   
[3622]278            self.pause()
[4322]279         
280        if self.movie:
281             self.save_image()
282               
[3493]283
284    def forward10(self):
[3536]285        if self.frameNumber + 10 <= self.maxFrameNumber:
286            self.frameNumber += 10
[3493]287        else:
[3536]288            self.frameNumber = self.maxFrameNumber
[3873]289        self.redraw_quantities()
[4325]290        self.update_labels()
[3622]291        self.pause()
[3548]292
293    def animateForward(self):
294        if self.paused is not True:
[3622]295            self.forward_step()
[3873]296            self.tk_root.after(self.frameDelay, self.animateForward)
[4325]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.