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

from Numeric import array, Float, ravel, zeros
from Scientific.IO.NetCDF import NetCDFFile
from Tkinter import Button, E, Tk, W, Label, StringVar, Scale, HORIZONTAL
from visualiser import Visualiser
from vtk import vtkCellArray, vtkPoints, vtkPolyData

class OfflineVisualiser(Visualiser):
    """A VTK-powered offline visualiser which runs in its own thread.
    In addition to the functions provided by the standard visualiser,
    the following additional functions are provided:

    precache_height_quantities() - Precache all the vtkpoints
    structures for any dynamic height based quantities to render.
    """
    def __init__(self, source, frameDelay=100, frameStep=1):
        """The source parameter is assumed to be a NetCDF sww file.
        The frameDelay parameter is the number of milliseconds waited between frames.
        """
        Visualiser.__init__(self, source)

        self.frameNumber = 0
        fin = NetCDFFile(self.source, 'r')
        self.maxFrameNumber = fin.variables['time'].shape[0] - 1
        fin.close()
        
        #self.frameNumberTkVariable = StringVar()
        #self.frameNumberTkVariable.set('Frame - %05g'%self.framNumber)

        self.frameDelay = frameDelay

        self.xmin = None
        self.xmax = None
        self.ymin = None
        self.ymax = None
        self.zmin = None
        self.zmax = None

        self.frameStep= frameStep 

        self.vtk_heightQuantityCache = []
        for i in range(self.maxFrameNumber + 1): # maxFrameNumber is zero indexed.
            self.vtk_heightQuantityCache.append({})

        self.paused = False
        self.movie = False
        
    def setup_grid(self):
        fin = NetCDFFile(self.source, 'r')
        self.vtk_cells = vtkCellArray()
        N_tri = fin.variables['volumes'].shape[0]
        for v in range(N_tri):
            self.vtk_cells.InsertNextCell(3)
            for i in range(3):
                self.vtk_cells.InsertCellPoint(fin.variables['volumes'][v][i])
        fin.close()
        
    def update_height_quantity(self, quantityName, dynamic=True):
        polydata = self.vtk_polyData[quantityName] = vtkPolyData()
        if dynamic is True:
            #print ' - Frame',self.frameNumber,'of',self.maxFrameNumber
            if not self.vtk_heightQuantityCache[self.frameNumber].has_key(quantityName):
                self.vtk_heightQuantityCache[self.frameNumber][quantityName]\
                    = self.read_height_quantity(quantityName, True, self.frameNumber);
            polydata.SetPoints(self.vtk_heightQuantityCache[self.frameNumber][quantityName])
        else:
            polydata.SetPoints(self.read_height_quantity(quantityName, False))
        polydata.SetPolys(self.vtk_cells)

    def get_3d_bounds(self):
        return [self.xmin, self.xmax, self.ymin, self.ymax, self.zmin, self.zmax]
            
    def read_height_quantity(self, quantityName, dynamic=True, frameNumber=0):
        """Read in a height based quantity from the NetCDF source file
        and return a vtkPoints object. frameNumber is ignored if
        dynamic is false."""
        fin = NetCDFFile(self.source, 'r')
        points = vtkPoints()
        if dynamic is True:
            N_vert = fin.variables[quantityName].shape[1]
        else:
            N_vert = len(fin.variables[quantityName])
        x = ravel(array(fin.variables['x'], Float))
        y = ravel(array(fin.variables['y'], Float))
        if dynamic is True:
            q = array(fin.variables[quantityName][frameNumber], Float)
        else:
            q = ravel(array(fin.variables[quantityName], Float))

        q *= self.height_zScales[quantityName]
        q += self.height_offset[quantityName]

        for v in range(N_vert):
            points.InsertNextPoint(x[v], y[v], q[v])
            if self.xmin == None or self.xmin > x[v]:
                self.xmin = x[v]
            if self.xmax == None or self.xmax < x[v]:
                self.xmax = x[v]
            if self.ymin == None or self.ymin > y[v]:
                self.ymin = y[v]
            if self.ymax == None or self.ymax < y[v]:
                self.ymax = y[v]
            if self.zmin == None or self.zmin > q[v]:
                self.zmin = q[v]
            if self.zmax == None or self.zmax < q[v]:
                self.zmax = q[v]
        fin.close()
        return points

    def precache_height_quantities(self):
        """Precache any height-based quantities. Call before rendering
        beigns."""
        for q in self.height_quantities:
            if self.height_dynamic[q] is True:
                print 'Precaching %s' % q
                for i in range(self.maxFrameNumber + 1): # maxFrameNumber is zero-indexed
                    print ' - Frame %d of %d' % (i, self.maxFrameNumber)
                    self.vtk_heightQuantityCache[i][q]\
                        = self.read_height_quantity(q, True, i)

    def build_quantity_dict(self):
        quantities = {}
        fin = NetCDFFile(self.source, 'r')
        for q in filter(lambda n:n != 'x' and n != 'y' and n != 'z' and n != 'time' and n != 'volumes', fin.variables.keys()):
            if len(fin.variables[q].shape) == 1: # Not a time-varying quantity
                quantities[q] = ravel(array(fin.variables[q], Float))
            else: # Time-varying, get the current timestep data
                quantities[q] = array(fin.variables[q][self.frameNumber], Float)
        fin.close()
        return quantities

    def setup_gui(self):
        Visualiser.setup_gui(self)
        self.tk_quit.grid(row=0, column=0, sticky=W+E)
        self.tk_movie_toggle = Button(self.tk_controlFrame, text="Movie off", command=self.movie_toggle)
        self.tk_movie_toggle.grid(row=0, column=6,  sticky=W+E)
        
                
        self.tk_restart = Button(self.tk_controlFrame, text="<<<", command=self.restart, width=5)
        self.tk_restart.grid(row=1, column=0, sticky=W+E)
        self.tk_back10 = Button(self.tk_controlFrame, text="<<", command=self.back10, width=5)
        self.tk_back10.grid(row=1, column=1, sticky=W+E)
        self.tk_back = Button(self.tk_controlFrame, text="<", command=self.back, width=5)
        self.tk_back.grid(row=1, column=2, sticky=W+E)
        self.tk_pauseResume = Button(self.tk_controlFrame, text="Pause", command=self.pauseResume, width=15)
        self.tk_pauseResume.grid(row=1, column=3, sticky=W+E)
        self.tk_forward = Button(self.tk_controlFrame, text=">", command=self.forward, width=5)
        self.tk_forward.grid(row=1, column=4, sticky=W+E)
        self.tk_forward10 = Button(self.tk_controlFrame, text=">>", command=self.forward10, width=5)
        self.tk_forward10.grid(row=1, column=5, sticky=W+E)
        self.tk_forwardEnd = Button(self.tk_controlFrame, text=">>>", command=self.forwardEnd, width=5)
        self.tk_forwardEnd.grid(row=1, column=6, sticky=W+E)
        

        self.tk_frameNumber = Label(self.tk_controlFrame, text='Frame')
        self.tk_frameNumber.grid(row=2, column=0, sticky=W+E)
        self.tk_gotoFrame = Scale(self.tk_controlFrame, from_=0, to=self.maxFrameNumber, orient=HORIZONTAL)
        self.tk_gotoFrame.grid(row=2, column=1, columnspan=2, sticky=W+E)
        self.tk_stepLabel = Label(self.tk_controlFrame, text='Step')
        self.tk_stepLabel.grid(row=2, column=4, sticky=W+E)        
        self.tk_frameStep = Scale(self.tk_controlFrame, from_=0, to=self.maxFrameNumber, orient=HORIZONTAL)
        self.tk_frameStep.grid(row=2, column=5, columnspan=2, sticky=W+E)
        
        # Make the buttons stretch to fill all available space
        for i in range(7):
            self.tk_controlFrame.grid_columnconfigure(i, weight=1)

    def run(self):
        self.alter_tkroot(Tk.after, (self.frameDelay, self.animateForward))
        Visualiser.run(self)

    def restart(self):
        self.frameNumber = 0
        self.redraw_quantities()
        self.update_labels()
        self.pause()
        
        if self.movie:
            self.save_image()
 
    def forwardEnd(self):
        self.frameNumber = self.maxFrameNumber
        self.redraw_quantities()
        self.update_labels()
        self.pause()
                
    def movie_toggle(self):
        if self.movie == True:
            self.movie = False
            self.tk_movie_toggle.config(text='Movie off')
        else:
            self.movie = True
            self.tk_movie_toggle.config(text='Movie on ')
            
            
        
        
    def save_image(self):
        
        from vtk import vtkJPEGWriter, vtkJPEGWriter, vtkPNGWriter
        from vtk import vtkPNMWriter, vtkWindowToImageFilter
        from os import path
         
        sourcebase, _ = path.splitext(self.source)
        fname = sourcebase+'%05g.png' % self.frameNumber
        #print fname
        
        extmap = {'.jpg' : vtkJPEGWriter,
                  '.jpeg' : vtkJPEGWriter,
                  '.png' : vtkPNGWriter,
                  '.pnm' : vtkPNMWriter,
                  }
        basename, ext = path.splitext(fname)
        try: Writer = extmap[ext.lower()]
        except KeyError:
            error_msg("Don't know how to handle %s files" % ext, parent=self)
            return
    
        renWin = self.vtk_renderer.GetRenderWindow()
        w2i = vtkWindowToImageFilter()
        writer = Writer()
        w2i.SetInput(renWin)
        w2i.Update()
        writer.SetInput(w2i.GetOutput())
        writer.SetFileName(fname)
        renWin.Render()
        writer.Write()        
    
    def back10(self):
        if self.frameNumber - 10 >= 0:
            self.frameNumber -= 10
        else:
            self.frameNumber = 0
        self.redraw_quantities()
        self.update_labels()
        self.pause()

    def back(self):
        if self.frameNumber > 0:
            self.frameNumber -= 1
            self.redraw_quantities()
            self.update_labels()
            self.pause()

    def pauseResume(self):
        if self.paused is True:
            self.resume()
        else:
            self.pause()

    def pause(self):
        self.paused = True
        self.tk_pauseResume.config(text="Resume")

    def resume(self):
        self.paused = False
        self.tk_pauseResume.config(text="Pause")
        self.frameNumber = self.tk_gotoFrame.get()
        self.frameStep = self.tk_frameStep.get()
        self.tk_root.after(self.frameDelay, self.animateForward)

    def forward(self):
        if self.frameNumber < self.maxFrameNumber:
            self.frameNumber += 1
            self.redraw_quantities()
            self.update_labels()
            self.pause()
            
    def forward_step(self):
        if self.frameNumber + self.frameStep <= self.maxFrameNumber:
            self.frameNumber += self.frameStep
            self.redraw_quantities()
            self.update_labels()
        else:
            self.frameNumber = self.maxFrameNumber            
            self.redraw_quantities()
            self.update_labels()    
            self.pause()
         
        if self.movie:
             self.save_image()
                

    def forward10(self):
        if self.frameNumber + 10 <= self.maxFrameNumber:
            self.frameNumber += 10
        else:
            self.frameNumber = self.maxFrameNumber
        self.redraw_quantities()
        self.update_labels()
        self.pause()

    def animateForward(self):
        if self.paused is not True:
            self.forward_step()
            self.tk_root.after(self.frameDelay, self.animateForward)
            
    def update_labels(self): 
        #self.tk_frameNumber.config(text='%05g of %05g'%(self.frameNumber,self.maxFrameNumber))
        self.tk_gotoFrame.set(self.frameNumber)
        self.tk_frameStep.set(self.frameStep)
               
    def shutdown(self):
        #self.pause()
        self.tk_root.withdraw()
        self.tk_root.destroy()
        #Visualiser.shutdown(self)