| 1 | from Numeric import array, Float |
|---|
| 2 | from anuga.file.netcdf import NetCDFFile |
|---|
| 3 | from visualiser import Visualiser |
|---|
| 4 | from vtk import vtkCellArray |
|---|
| 5 | |
|---|
| 6 | class SWWVisualiser(Visualiser): |
|---|
| 7 | ''' |
|---|
| 8 | A Visualiser that views SWW files. The source for this |
|---|
| 9 | visualiser is a string containing the sww file name. |
|---|
| 10 | ''' |
|---|
| 11 | |
|---|
| 12 | def __init__(self, *args, **kwargs): |
|---|
| 13 | Visualiser.__init__(self, *args, **kwargs) |
|---|
| 14 | fin = NetCDFFile(self.vis_source, 'r') |
|---|
| 15 | self.xPoints = array(fin.variables['x'], Float) |
|---|
| 16 | self.yPoints = array(fin.variables['y'], Float) |
|---|
| 17 | self.quantityCache = {} |
|---|
| 18 | fin.close() |
|---|
| 19 | |
|---|
| 20 | def setupGrid(self): |
|---|
| 21 | fin = NetCDFFile(self.vis_source, 'r') |
|---|
| 22 | nTri = fin.variables['volumes'].shape[0] |
|---|
| 23 | insertNextCell = vtkCellArray.InsertNextCell |
|---|
| 24 | insertCellPoint = vtkCellArray.InsertCellPoint |
|---|
| 25 | |
|---|
| 26 | for v in range(nTri): |
|---|
| 27 | insertNextCell(self.vtk_cells, 3) |
|---|
| 28 | for i in range(3): |
|---|
| 29 | insertCellPoint(self.vtk_cells, fin.variables['volumes'][v][i]) |
|---|
| 30 | |
|---|
| 31 | fin.close() |
|---|
| 32 | |
|---|
| 33 | def getMaxFrameNumber(self): |
|---|
| 34 | fin = NetCDFFile(self.vis_source, 'r') |
|---|
| 35 | rv = fin.variables['time'].shape[0] - 1 |
|---|
| 36 | fin.close() |
|---|
| 37 | return rv |
|---|
| 38 | |
|---|
| 39 | def getQuantityPoints(self, quantityName, dynamic=False): |
|---|
| 40 | try: |
|---|
| 41 | if dynamic: |
|---|
| 42 | q = self.quantityCache[quantityName][self.vis_frame] |
|---|
| 43 | else: |
|---|
| 44 | q = self.quantityCache[quantityName] |
|---|
| 45 | except KeyError: |
|---|
| 46 | fin = NetCDFFile(self.vis_source, 'r') |
|---|
| 47 | if dynamic: |
|---|
| 48 | if not self.quantityCache.has_key(quantityName): |
|---|
| 49 | self.quantityCache[quantityName] = {} |
|---|
| 50 | q = array(fin.variables[quantityName][self.vis_frame], Float) |
|---|
| 51 | self.quantityCache[quantityName][self.vis_frame] = q |
|---|
| 52 | else: |
|---|
| 53 | q = array(fin.variables[quantityName], Float) |
|---|
| 54 | self.quantityCache[quantityName] = q |
|---|
| 55 | fin.close() |
|---|
| 56 | return q |
|---|
| 57 | |
|---|
| 58 | def getQuantityDict(self): |
|---|
| 59 | quantities = {} |
|---|
| 60 | fin = NetCDFFile(self.vis_source, 'r') |
|---|
| 61 | names = [ k for k in fin.variables.keys() |
|---|
| 62 | if k != 'x' |
|---|
| 63 | and k != 'y' |
|---|
| 64 | and k != 'z' |
|---|
| 65 | and k != 'time' |
|---|
| 66 | and k != 'volumes' ] |
|---|
| 67 | argss = [ (name, len(fin.variables[name].shape) != 1) for name in names ] |
|---|
| 68 | fin.close() |
|---|
| 69 | for args in argss: |
|---|
| 70 | quantities[name] = self.getQuantityPoints(*args) |
|---|
| 71 | return quantities |
|---|
