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