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] |
---|
39 | z *= self.zScale |
---|
40 | z += self.offset |
---|
41 | setPoint(vtk_points, i, self.visualiser.xPoints[i], self.visualiser.yPoints[i], z) |
---|
42 | |
---|
43 | polyData = vtkPolyData() |
---|
44 | polyData.SetPoints(vtk_points) |
---|
45 | polyData.SetPolys(self.visualiser.vtk_cells) |
---|
46 | mapper = vtkPolyDataMapper() |
---|
47 | mapper.SetInput(polyData) |
---|
48 | setValue = vtkFloatArray.SetValue |
---|
49 | if hasattr(self.colour[0], '__call__'): |
---|
50 | scalars = self.colour[0](self.visualiser.getQuantityDict()) |
---|
51 | nScalars = len(scalars) |
---|
52 | vtk_scalars = vtkFloatArray() |
---|
53 | vtk_scalars.SetNumberOfValues(nScalars) |
---|
54 | for i in xrange(nScalars): |
---|
55 | setValue(vtk_scalars, i, scalars[i]) |
---|
56 | polyData.GetPointData().SetScalars(vtk_scalars) |
---|
57 | mapper.SetScalarRange(self.colour[1:3]) |
---|
58 | mapper.Update() |
---|
59 | self.actor.SetMapper(mapper) |
---|
60 | Feature.draw(self, renderer) |
---|