1import sys
2
3try:
4    import numpy
5except ImportError:
6    print "Numpy (http://numpy.scipy.org) not found.",
7    print "This test requires numpy!"
8    sys.exit(0)
9
10import vtk
11import vtk.numpy_interface.dataset_adapter as dsa
12import vtk.numpy_interface.algorithms as algs
13
14w = vtk.vtkRTAnalyticSource()
15
16bp = vtk.vtkBrownianPoints()
17bp.SetInputConnection(w.GetOutputPort())
18bp.Update()
19
20elev = vtk.vtkElevationFilter()
21elev.SetInputConnection(bp.GetOutputPort())
22elev.SetLowPoint(-10, 0, 0)
23elev.SetHighPoint(10, 0, 0)
24elev.SetScalarRange(0, 20)
25
26g = vtk.vtkMultiBlockDataGroupFilter()
27g.AddInputConnection(elev.GetOutputPort())
28g.AddInputConnection(elev.GetOutputPort())
29
30g.Update()
31
32elev2 = vtk.vtkElevationFilter()
33elev2.SetInputConnection(bp.GetOutputPort())
34elev2.SetLowPoint(0, -10, 0)
35elev2.SetHighPoint(0, 10, 0)
36elev2.SetScalarRange(0, 20)
37
38g2 = vtk.vtkMultiBlockDataGroupFilter()
39g2.AddInputConnection(elev2.GetOutputPort())
40g2.AddInputConnection(elev2.GetOutputPort())
41
42g2.Update()
43
44elev3 = vtk.vtkElevationFilter()
45elev3.SetInputConnection(bp.GetOutputPort())
46elev3.SetLowPoint(0, 0, -10)
47elev3.SetHighPoint(0, 0, 10)
48elev3.SetScalarRange(0, 20)
49
50g3 = vtk.vtkMultiBlockDataGroupFilter()
51g3.AddInputConnection(elev3.GetOutputPort())
52g3.AddInputConnection(elev3.GetOutputPort())
53
54g3.Update()
55
56cd = dsa.CompositeDataSet(g.GetOutput())
57randomVec = cd.PointData['BrownianVectors']
58elev = cd.PointData['Elevation']
59
60cd2 = dsa.CompositeDataSet(g2.GetOutput())
61elev2 = cd2.PointData['Elevation']
62
63cd3 = dsa.CompositeDataSet(g3.GetOutput())
64elev3 = cd3.PointData['Elevation']
65
66npa = randomVec.Arrays[0]
67
68# Test operators
69assert algs.all(1 + randomVec - 1 - randomVec < 1E-4)
70
71assert (1 + randomVec).DataSet is randomVec.DataSet
72
73# Test slicing and indexing
74assert algs.all(randomVec[randomVec[:,0] > 0.2].Arrays[0] - npa[npa[:,0] > 0.2] < 1E-7)
75assert algs.all(randomVec[algs.where(randomVec[:,0] > 0.2)].Arrays[0] - npa[numpy.where(npa[:,0] > 0.2)] < 1E-7)
76assert algs.all(randomVec[dsa.VTKCompositeDataArray([(slice(None, None, None), slice(0,2,None)), 2])].Arrays[0] - npa[:, 0:2] < 1E-6)
77
78# Test ufunc
79assert algs.all(algs.cos(randomVec) - numpy.cos(npa) < 1E-7)
80assert algs.cos(randomVec).DataSet is randomVec.DataSet
81
82# Various numerical ops implemented in VTK
83g = algs.gradient(elev)
84assert algs.all(g[0] == (1, 0, 0))
85
86v = algs.make_vector(elev, g[:,0], elev)
87assert algs.all(algs.gradient(v) == [[1, 0, 0], [0, 0, 0], [1, 0, 0]])
88
89v = algs.make_vector(elev, g[:,0], elev2)
90assert algs.all(algs.curl(v) == [1, 0, 0])
91
92v = algs.make_vector(elev, elev2, 2*elev3)
93g = algs.gradient(v)
94assert g.DataSet is v.DataSet
95assert algs.all(algs.det(g) == 2)
96
97assert algs.all(algs.eigenvalue(g) == [2, 1, 1])
98
99assert algs.all(randomVec[:,0] == randomVec[:,0])
100
101ssource = vtk.vtkSphereSource()
102ssource.Update()
103
104output = ssource.GetOutput()
105
106fd = vtk.vtkFloatArray()
107fd.SetNumberOfTuples(11)
108fd.FillComponent(0, 5)
109fd.SetName("field array")
110
111output.GetFieldData().AddArray(fd)
112
113g2 = vtk.vtkMultiBlockDataGroupFilter()
114g2.AddInputData(output)
115g2.AddInputData(output)
116
117g2.Update()
118
119sphere = dsa.CompositeDataSet(g2.GetOutput())
120
121vn = algs.vertex_normal(sphere)
122assert algs.all(algs.mag(vn) - 1 < 1E-6)
123
124sn = algs.surface_normal(sphere)
125assert algs.all(algs.mag(sn) - 1 < 1E-6)
126
127dot = algs.dot(vn, vn)
128assert dot.DataSet is sphere
129assert algs.all(dot == 1)
130assert algs.all(algs.cross(vn, vn) == [0, 0, 0])
131
132fd = sphere.FieldData['field array']
133assert algs.all(fd == 5)
134assert algs.shape(fd) == (22,)
135
136assert vn.DataSet is sphere
137
138# --------------------------------------
139
140na = dsa.NoneArray
141
142# Test operators
143assert (1 + na - 1 - randomVec) is na
144
145# Test slicing and indexing
146assert na[:, 0] is na
147assert algs.where(na[:, 0] > 0) is na
148assert (na > 0) is na
149
150# Test ufunc
151assert algs.cos(na) is na
152
153# Various numerical ops implemented in VTK
154assert algs.gradient(na) is na
155assert algs.cross(na, na) is na
156assert algs.cross(v.Arrays[0], na) is na
157assert algs.cross(na, v.Arrays[0]) is na
158
159assert algs.make_vector(na, g[:,0], elev) is na
160
161pd = vtk.vtkPolyData()
162pdw = dsa.WrapDataObject(pd)
163pdw.PointData.append(na, 'foo')
164assert pdw.PointData.GetNumberOfArrays() == 0
165
166# --------------------------------------
167
168na2 = dsa.VTKCompositeDataArray([randomVec.Arrays[0], na])
169
170# Test operators
171assert (1 + na2 - 1 - randomVec).Arrays[1] is na
172
173# Test slicing and indexing
174assert na2[:, 0].Arrays[1] is na
175assert algs.where(na2[:, 0] > 0).Arrays[1] is na
176assert (na2 > 0).Arrays[1] is na
177
178# Test ufunc
179assert algs.cos(na2).Arrays[1] is na
180
181# Various numerical ops implemented in VTK
182assert algs.gradient(na2).Arrays[1] is na
183assert algs.cross(na2, na2).Arrays[1] is na
184assert algs.cross(v, na2).Arrays[1] is na
185assert algs.cross(na2, v).Arrays[1] is na
186
187assert algs.make_vector(na2[:, 0], elev, elev).Arrays[1] is na
188assert algs.make_vector(elev, elev, na2[:, 0]).Arrays[1] is na
189assert algs.make_vector(elev, na2[:, 0], elev).Arrays[1] is na
190
191mb = vtk.vtkMultiBlockDataSet()
192mb.SetBlock(0, pd)
193pd2 = vtk.vtkPolyData()
194mb.SetBlock(1, pd2)
195mbw = dsa.WrapDataObject(mb)
196
197mbw.PointData.append(dsa.NoneArray, 'foo')
198assert mbw.GetBlock(0).GetPointData().GetNumberOfArrays() == 0
199assert mbw.GetBlock(1).GetPointData().GetNumberOfArrays() == 0
200
201mbw.PointData.append(na2, 'foo')
202assert mbw.GetBlock(0).GetPointData().GetNumberOfArrays() == 1
203assert mbw.GetBlock(1).GetPointData().GetNumberOfArrays() == 0
204assert mbw.GetBlock(0).GetPointData().GetArray(0).GetName() == 'foo'
205
206mbw.PointData.append(algs.max(na2), "maxfoo")
207assert mbw.GetBlock(0).GetPointData().GetNumberOfArrays() == 2
208assert mbw.GetBlock(1).GetPointData().GetNumberOfArrays() == 1
209assert mbw.GetBlock(0).GetPointData().GetArray(1).GetName() == 'maxfoo'
210
211# --------------------------------------
212
213mb = vtk.vtkMultiBlockDataSet()
214mb.SetBlock(0, vtk.vtkImageData())
215mb.SetBlock(1, vtk.vtkImageData())
216assert dsa.WrapDataObject(mb).Points is na
217
218mb = vtk.vtkMultiBlockDataSet()
219mb.SetBlock(0, vtk.vtkStructuredGrid())
220mb.SetBlock(1, vtk.vtkImageData())
221assert dsa.WrapDataObject(mb).Points is na
222
223mb = vtk.vtkMultiBlockDataSet()
224sg = vtk.vtkStructuredGrid()
225sg.SetPoints(vtk.vtkPoints())
226mb.SetBlock(0, sg)
227mb.SetBlock(1, vtk.vtkImageData())
228assert dsa.WrapDataObject(mb).Points.Arrays[0] is not na
229assert dsa.WrapDataObject(mb).Points.Arrays[1] is na
230