1#!/usr/bin/env python
2
3# This script tests picking props, points and cells rendered with a vtkCompositePolyDataMapper2.
4#
5import vtk
6import math
7import time
8from vtk.util.misc import vtkGetDataRoot
9VTK_DATA_ROOT = vtkGetDataRoot()
10
11# Create a vtkMultiBlockDataSet with 3
12# vtkPolyData blocks, rendered using a vtkCompositePolyDataMapper2
13# and picked using vtkPicker, vtkPointPicker and vtkCellPicker.
14
15# points for two 'stacks' of 3 cells, each belonging to a different block
16#     first stack                                 second stack
17xyz=[[[ 4., 4.],[10., 4.],[10., 6.],[ 4., 6.],    [20., 0.],[25.,0.],[30., 0.],[30.,1.], [25.,1.],[20.,1.]],
18     [[ 2., 2.],[12., 2.],[12., 8.],[ 2., 8.],    [22.,-1.],         [28.,-1.],[28.,2.],          [22.,2.]],
19     [[ 0., 0.],[14., 0.],[14.,10.],[ 0.,10.],    [24.,-2.],         [26.,-2.],[26.,3.],          [24.,3.]]]
20
21polyconn=[[[0,1,2,3],                             [4,5,8,9],[5,6,7,8]],
22          [[0,1,2],[2,3,0],                       [4,5,6],[6,7,4]],
23          [[0,1,2,3],                             [4,5,6,7]]]
24
25# These define the tests
26# pick-coordinate, prop-pick-expectations, point-pick-expectations, cell-pick-expectations
27# prop-pick-expectations: pick-result, block#
28# point-pick-expectations: pick-result, block#, point-id
29# cell-pick-expectations: pick-result, block#, cell-id, point-id
30pickdata=[
31  # Pick the first stack
32  ([ 0.,0.],(True,3),(True,3,0),(True,3,0,0)),
33  ([ 4.,4.],(True,1),(True,1,0),(True,1,0,0)),
34  ([ 5.,5.],(True,1),(False,-1,-1),(True,1,0,0)),
35
36  # Pick in between the stacks
37  # note: the ray passes between the data sets but the block's bounding box encompasses both data sets
38  ([18.,-1.],(True,2),(False,-1,-1),(False,-1,-1,-1)),
39  ([18.,-3.],(False,-1),(False,-1,-1),(False,-1,-1,-1)),
40
41  # Pick the second stack
42  ([25.,0.],(True,1),(True,1,5),(True,1,1,5)),
43  # note: the prop pick hits block #1 because the bounding box of block #1 is picked. The cell picker does not hit a cell of block #1
44  ([28.,2.],(True,1),(True,2,6),(True,2,2,6)) ]
45
46
47
48# Construct the data-set, and set up color mapping of the blocks
49mbd=vtk.vtkMultiBlockDataSet()
50mbd.SetNumberOfBlocks(3)
51cda=vtk.vtkCompositeDataDisplayAttributes()
52m=vtk.vtkCompositePolyDataMapper2()
53m.SetInputDataObject(mbd)
54m.SetCompositeDataDisplayAttributes(cda)
55
56for blk in range(0,3):
57  mbd.SetBlock(blk,vtk.vtkPolyData())
58  poly=mbd.GetBlock(blk)
59  coords=xyz[blk]
60  pts=vtk.vtkPoints()
61  for coord in coords:
62    pts.InsertNextPoint(coord[0],coord[1],float(blk))
63
64  polys=vtk.vtkCellArray()
65  for cell in polyconn[blk]:
66     polys.InsertNextCell(len(cell))
67     for pid in cell:
68       polys.InsertCellPoint(pid)
69
70  poly.SetPoints(pts)
71  poly.SetPolys(polys)
72  m.SetBlockColor(blk,(blk%3)==0,(blk+1)%3==0,(blk+2)%3==0)
73
74# Set up the actor
75a=vtk.vtkActor()
76a.SetMapper(m)
77a.GetProperty().EdgeVisibilityOn()
78a.GetProperty().SetEdgeColor(1,1,1)
79
80# Render the actor
81r = vtk.vtkRenderer()
82r.AddViewProp(a)
83r.SetBackground(0,0,0)
84rw = vtk.vtkRenderWindow()
85rw.AddRenderer(r)
86rw.Render()
87
88# Define the pickers
89propPicker = vtk.vtkPicker()
90cellPicker = vtk.vtkCellPicker()
91pointPicker = vtk.vtkPointPicker()
92
93# this can be switched on to print some debug output
94debug=False
95
96def propPick(p,xy,data):
97    """
98    Pick a prop. camera has been positioned such that a ray pick at
99    screen-coordinates p hits points with x-y coordinates xy.
100    data contains: expected pick result, block#
101    """
102    errors = 0
103    result=propPicker.Pick(p[0],p[1],0,r)==1
104    if result != data[0]:
105        print('prop pick at {} result: expected {} actual {}'
106                .format(xy,data[0],result))
107        errors+=1
108    if propPicker.GetFlatBlockIndex() != data[1]:
109        print('prop pick at {} flat block index expected {} actual {}'
110                .format(xy,data[1],propPicker.GetFlatBlockIndex()))
111        errors+=1
112
113    if debug and errors:
114        if result>0:
115           if propPicker.GetFlatBlockIndex():
116                print('prop pick at {} picked block {}'
117                        .format(xy,propPicker.GetFlatBlockIndex()))
118        else:
119            print('prop pick at {} result:{}'.format(xy,result))
120    return errors
121
122def pointPick(p,xy,data):
123    """
124    Pick a prop. camera has been positioned such that a ray pick at
125    screen-coordinates p hits points with x-y coordinates xy.
126    data contains: expected pick result, block#, point id
127    """
128    errors = 0
129    result=pointPicker.Pick(p[0],p[1],0,r)==1
130    if result != data[0]:
131        print('point pick at {} result: expected {} actual {}'
132                .format(xy,data[0],result))
133        errors+=1
134    if pointPicker.GetFlatBlockIndex() != data[1]:
135        print('point pick at {} flat block index expected {} actual {}'
136                .format(xy,data[1],pointPicker.GetFlatBlockIndex()))
137        errors+=1
138    if pointPicker.GetPointId() != data[2]:
139        print('point pick at {} point id expected {} actual {}'
140                .format(xy,data[2],pointPicker.GetPointId()))
141        errors+=1
142    if debug and errors:
143        if result>0:
144            if pointPicker.GetFlatBlockIndex():
145               print('point pick at {} picked point {} of block {}'
146                       .format(xy,pointPicker.GetPointId(),pointPicker.GetFlatBlockIndex()))
147            else:
148                print('point pick at {} picked point {}'
149                        .format(xy,pointPicker.GetPointId()))
150        else:
151            print('point pick at {} result:{}'.format(xy,result))
152    return errors
153
154def cellPick(p,xy,data):
155    """
156    Pick a prop. camera has been positioned such that a ray pick at
157    screen-coordinates p hits points with x-y coordinates xy.
158    data contains: expected pick result, block#, cell id, point id
159    """
160    errors=0
161    result=cellPicker.Pick(p[0],p[1],0,r)==1
162    if result != data[0]:
163        print('cell pick at {} result: expected {} actual {}'
164                .format(xy,data[0],result))
165        errors+=1
166    if cellPicker.GetFlatBlockIndex() != data[1]:
167        print('cell pick at {} flat block index expected {} actual {}'
168                .format(xy,data[1],cellPicker.GetFlatBlockIndex()))
169        errors+=1
170    if cellPicker.GetCellId() != data[2]:
171        print('cell pick at {} cell id expected {} actual {}'
172                .format(xy,data[2],cellPicker.GetCellId()))
173        errors+=1
174    if cellPicker.GetPointId() != data[3]:
175        print('cell pick at {} point id expected {} actual {}'
176                .format(xy,data[3],cellPicker.GetPointId()))
177        errors+=1
178
179    if debug and errors:
180        if result>0:
181           if cellPicker.GetFlatBlockIndex():
182                print('cell pick at {} picked cell {} point {} of block {}'
183                        .format(xy,cellPicker.GetCellId(),cellPicker.GetPointId(),cellPicker.GetFlatBlockIndex()))
184           else:
185                print('cell pick at {} picked cell {} point {}'
186                        .format(xy,cellPicker.GetCellId(),cellPicker.GetPointId()))
187        else:
188            print('cell pick at {} result:{}'.format(xy,result))
189    return errors
190
191
192def pick(xy,propPickData,pointPickData,cellPickData):
193    """
194    Position the camera such that a ray through the screen center
195    points in z direction at position xy, and pick with vtkPicker,
196    vtkPointPicker and vtkCellPicker.
197    Returns the number of errors
198    """
199    size=rw.GetSize()
200    p=[size[0]/2,size[1]/2,0]
201
202    cam=r.GetActiveCamera()
203    cam.SetPosition(xy[0],xy[1],-10.0)
204    cam.SetFocalPoint(xy[0],xy[1],0.0)
205    cam.SetViewUp(0,1,0)
206    cam.SetViewAngle(90)
207    cam.ParallelProjectionOff()
208    r.ResetCameraClippingRange()
209    rw.Render()
210
211    errors = 0
212
213    if debug:
214      print('Picking at {}'.format(xy))
215    errors += propPick(p,xy,propPickData)
216    errors += pointPick(p,xy,pointPickData)
217    errors += cellPick(p,xy,cellPickData)
218
219    if debug:
220      time.sleep(3)
221
222    return errors
223
224
225# Run all of the tests defined by pickdata
226errors = 0
227for data in pickdata:
228  errors += pick( data[0], data[1], data[2], data[3] )
229
230if errors:
231  print("Encountered {} errors".format(errors))
232  sys.exit(1)
233
234# --- end of script --
235