| 1 | from feature import Feature |
|---|
| 2 | from Tkinter import Button |
|---|
| 3 | from vtk import vtkFloatArray, vtkPoints, vtkPolyData, vtkPolyDataMapper |
|---|
| 4 | from time import time |
|---|
| 5 | class HeightQuantity(Feature): |
|---|
| 6 | ''' |
|---|
| 7 | A height quantity, such as stage, elevation, ... |
|---|
| 8 | ''' |
|---|
| 9 | def __init__(self, quantityName, zScale=1.0, |
|---|
| 10 | offset=0.0, **kwargs): |
|---|
| 11 | ''' |
|---|
| 12 | Parameters: |
|---|
| 13 | quantityName: string - name of a quantity |
|---|
| 14 | zScale: float - multiply point z-values by this |
|---|
| 15 | offset: float - add this to point z-values |
|---|
| 16 | ''' |
|---|
| 17 | Feature.__init__(self, **kwargs) |
|---|
| 18 | |
|---|
| 19 | self.quantityName = quantityName |
|---|
| 20 | self.zScale = zScale |
|---|
| 21 | self.offset = offset |
|---|
| 22 | |
|---|
| 23 | def button(self, tk_component): |
|---|
| 24 | return Button(tk_component, |
|---|
| 25 | text='H: ' + self.quantityName) |
|---|
| 26 | |
|---|
| 27 | def draw(self, renderer): |
|---|
| 28 | ### FIXME this should be made faster (C++ Module? How to deal with C++ linkage problems?) |
|---|
| 29 | ### Sticking the vtkPoints objects in a cache would help somewhat but not on the first view. |
|---|
| 30 | ### - Jack |
|---|
| 31 | if not self.drawn: |
|---|
| 32 | vtk_points = vtkPoints() |
|---|
| 33 | points = self.visualiser.getQuantityPoints(self.quantityName, dynamic=self.dynamic) |
|---|
| 34 | nPoints = len(points) |
|---|
| 35 | vtk_points.SetNumberOfPoints(nPoints) |
|---|
| 36 | setPoint = vtkPoints.SetPoint |
|---|
| 37 | for i in xrange(nPoints): |
|---|
| 38 | z = points[i] * self.zScale + self.offset |
|---|
| 39 | setPoint(vtk_points, i, self.visualiser.xPoints[i], self.visualiser.yPoints[i], z) |
|---|
| 40 | |
|---|
| 41 | polyData = vtkPolyData() |
|---|
| 42 | polyData.SetPoints(vtk_points) |
|---|
| 43 | polyData.SetPolys(self.visualiser.vtk_cells) |
|---|
| 44 | mapper = vtkPolyDataMapper() |
|---|
| 45 | mapper.SetInput(polyData) |
|---|
| 46 | setValue = vtkFloatArray.SetValue |
|---|
| 47 | if hasattr(self.colour[0], '__call__'): |
|---|
| 48 | scalars = self.colour[0](self.visualiser.getQuantityDict()) |
|---|
| 49 | nScalars = len(scalars) |
|---|
| 50 | vtk_scalars = vtkFloatArray() |
|---|
| 51 | vtk_scalars.SetNumberOfValues(nScalars) |
|---|
| 52 | for i in xrange(nScalars): |
|---|
| 53 | setValue(vtk_scalars, i, scalars[i]) |
|---|
| 54 | polyData.GetPointData().SetScalars(vtk_scalars) |
|---|
| 55 | mapper.SetScalarRange(self.colour[1:3]) |
|---|
| 56 | mapper.Update() |
|---|
| 57 | self.actor.SetMapper(mapper) |
|---|
| 58 | Feature.draw(self, renderer) |
|---|
