def write_to_vtk(conf, ar, filename): """ This function writes a numpy array into vtk format file. It reads configuration into DispalyParams object. The given array is cropped to the configured size if the crop parameter is defined in configuration file. Based on configuration parameters the grid needed for the tvtk StructuredGrid is calculated. Other StructuredGrid attributes are set. The initialized StructuredGrid object is then written into vtk type file. If configuration parameter save_two_files is set to True, one file will be saved for amplitudes with "_Amp.vtk" ending, and one for phases with "_Phase.vtk" ending. Otherwise, a single file will be saved that contains double array with amplitudes a nd phases. Parameters ---------- conf : str configuration file name ar : numpy array a complex array thatwill be saved filename : str a prefix to an output filename Returns ------- none """ params = DispalyParams(conf) if params.crop is None: dims = ar.shape arr_cropped = ar.ravel() else: dims = params.crop arr_cropped = crop_array_center(ar, params.crop).ravel() amps = np.abs(arr_cropped) phases = np.arctan2(arr_cropped.imag, arr_cropped.real) geometry = params.get_geometry() coordinates = get_coords(dims, geometry) sg = tvtk.StructuredGrid() sg.points = coordinates sg.dimensions = (dims[2], dims[1], dims[0]) sg.extent = 0, dims[2] - 1, 0, dims[1] - 1, 0, dims[0] - 1 if params.save_two_files: sg.point_data.scalars = amps sg.point_data.scalars.name = "Amp" write_array(sg, filename + "_Amp.vtk") sg.point_data.scalars = phases sg.point_data.scalars.name = "Phase" write_array(sg, filename + "_Phase.vtk") else: sg.point_data.scalars = amps sg.point_data.scalars.name = "Amp" ph = tvtk.FloatArray() ph.from_array(phases) ph.name = "Phase" sg.point_data.add_array(ph) write_array(sg, filename + ".vtk")
def make_data(): points = N.array([[0,0,0], [1,0,0], [0,1,0], [0,0,1], # tets [1,0,0], [2,0,0], [1,1,0], [1,0,1], [2,0,0], [3,0,0], [2,1,0], [2,0,1], ], 'f') tets = N.array([[0, 1, 2, 3], [4, 5, 6, 7], [8, 9, 10, 11]]) tet_type = tvtk.Tetra().cell_type ug = tvtk.UnstructuredGrid(points=points) ug.set_cells(tet_type, tets) # Setup the point attributes. temp = N.random.random(12) v = N.random.randn(12, 3) ten = N.random.randn(12, 9) a = tvtk.FloatArray(name='p') a.from_array(N.random.randn(12)) ug.point_data.add_array(a) ug.point_data.scalars = temp ug.point_data.scalars.name = 't' ug.point_data.vectors = v ug.point_data.vectors.name = 'v' ug.point_data.tensors = ten ug.point_data.tensors.name = 'ten' # Setup the cell attributes. temp = N.random.random(3) v = N.random.randn(3, 3) ten = N.random.randn(3, 9) ug.cell_data.scalars = temp ug.cell_data.scalars.name = 't' ug.cell_data.vectors = v ug.cell_data.vectors.name = 'v' ug.cell_data.tensors = ten ug.cell_data.tensors.name = 'ten' return ug
def update_data(path, polydata, mask, variables, t=0): """ Add arrays for each variable to polydata """ ds = netCDF4.Dataset(path) for key in variables: arr = ds.variables[key][t, ...] array = tvtk.FloatArray() array.from_array(arr[~mask]) array.name = key polydata.cell_data.add_array(array) polydata.modified() ds.close()
def test_array_conversion(self): """Test if Numeric/VTK array conversion works.""" # This is only a simple test. data = numpy.array( [[0, 0, 0, 10], [1, 0, 0, 20], [0, 1, 0, 20], [0, 0, 1, 30]], 'f') triangles = numpy.array([[0, 1, 3], [0, 3, 2], [1, 2, 3], [0, 2, 1]]) points = data[:, :3] temperature = data[:, -1] mesh = tvtk.PolyData() mesh.points = points mesh.polys = triangles mesh.point_data.scalars = temperature # Test if a normal float array also works. temp = tvtk.FloatArray() temp.from_array(temperature) mesh.point_data.scalars = temp
def test_data_array(self): """Test if vtkDataArrays behave in a Pythonic fashion.""" # Check a 3D array. f = tvtk.FloatArray() a = numpy.array([[0., 0, 0], [1, 1, 1]]) f.from_array(a) self.assertEqual(f.number_of_components, 3) self.assertEqual(f.number_of_tuples, 2) self.assertEqual(mysum(f.to_array() - a), 0) for i, j in zip(a, f): self.assertEqual(mysum(i-j), 0.0) self.assertEqual(f[0], (0.0, 0.0, 0.0)) self.assertEqual(f[-1], (1., 1., 1.)) self.assertEqual(repr(f), '[(0.0, 0.0, 0.0), (1.0, 1.0, 1.0)]') f.append((2, 2, 2)) f.extend([[3, 3, 3], [4, 4, 4]]) self.assertEqual(len(f), 5) self.assertEqual(f.number_of_components, 3) self.assertEqual(f.number_of_tuples, 5) f[1] = [-1, -1, -1] self.assertEqual(f[1], (-1.0, -1.0, -1.0)) self.assertRaises(IndexError, f.__getitem__, 100) self.assertRaises(IndexError, f.__setitem__, 100, 100) # Check a 5D arrray a = numpy.array([[0., 0, 0, 0, 0], [1, 1, 1, 1, 1]]) f.from_array(a) self.assertEqual(mysum(f.to_array() - a), 0.0) for i, j in zip(a, f): self.assertEqual(mysum(i - j), 0.0) self.assertEqual(f[0], (0.0, 0.0, 0.0, 0.0, 0.0)) self.assertEqual(f[-1], (1., 1., 1., 1., 1.)) self.assertEqual( repr(f), '[(0.0, 0.0, 0.0, 0.0, 0.0), (1.0, 1.0, 1.0, 1.0, 1.0)]') f.append((2, 2, 2, 2, 2)) f.extend([[3, 3, 3, 3, 3], [4, 4, 4, 4, 4]]) self.assertEqual(len(f), 5) self.assertEqual(f.number_of_components, 5) self.assertEqual(f.number_of_tuples, 5) self.assertEqual(f[-1], (4., 4., 4., 4., 4.)) f[1] = [-1, -1, -1, -1, -1] self.assertEqual(f[1], (-1.0, -1.0, -1.0, -1.0, -1.0)) self.assertRaises(IndexError, f.__getitem__, 100) self.assertRaises(IndexError, f.__setitem__, 100, 100)
def screenshot(figure=None, mode='rgb', antialiased=False): """ Return the current figure pixmap as an array. **Parameters** :figure: a figure instance or None, optional If specified, the figure instance to capture the view of. :mode: {'rgb', 'rgba'} The color mode of the array captured. :antialiased: {True, False} Use anti-aliasing for rendering the screenshot. Uses the number of aa frames set by figure.scene.anti_aliasing_frames **Notes** On most systems, this works similarly to taking a screenshot of the rendering window. Thus if it is hidden by another window, you will capture the other window. This limitation is due to the heavy use of the hardware graphics system. **Examples** This function can be useful for integrating 3D plotting with Mayavi in a 2D plot created by matplotlib. >>> from mayavi import mlab >>> mlab.test_plot3d() >>> arr = mlab.screenshot() >>> import pylab as pl >>> pl.imshow(arr) >>> pl.axis('off') >>> pl.show() """ if figure is None: figure = gcf() # don't use figure.scene.get_size() here because this will be incorrect # for HiDPI systems x, y = tuple(figure.scene.render_window.size) # Try to lift the window figure.scene._lift() if mode == 'rgb': out = tvtk.UnsignedCharArray() shape = (y, x, 3) pixel_getter = figure.scene.render_window.get_pixel_data if vtk_major_version > 7: pg_args = (0, 0, x - 1, y - 1, 1, out, 0) else: pg_args = (0, 0, x - 1, y - 1, 1, out) elif mode == 'rgba': out = tvtk.FloatArray() shape = (y, x, 4) pixel_getter = figure.scene.render_window.get_rgba_pixel_data if vtk_major_version > 7: pg_args = (0, 0, x - 1, y - 1, 1, out, 0) else: pg_args = (0, 0, x - 1, y - 1, 1, out) else: raise ValueError('mode type not understood') if antialiased: # save the current aa value to restore it later render_window = figure.scene.render_window if hasattr(render_window, 'aa_frames'): old_aa = render_window.aa_frames render_window.aa_frames = figure.scene.anti_aliasing_frames else: old_aa = render_window.multi_samples render_window.multi_samples = figure.scene.anti_aliasing_frames figure.scene.render() pixel_getter(*pg_args) if hasattr(render_window, 'aa_frames'): render_window.aa_frames = old_aa else: render_window.multi_samples = old_aa figure.scene.render() else: pixel_getter(*pg_args) # Return the array in a way that pylab.imshow plots it right: out = out.to_array() out.shape = shape out = np.flipud(out) return out
def test_set_scalars(self): """Test if SetScalars works without a segfault.""" mesh = tvtk.PolyData() sc = tvtk.FloatArray() # If this does not segfault, we are OK. mesh.point_data.scalars = sc
# self.start() def add_actor(self, actor): self.renderer.add_actor(actor) # update window here. ### DATA points = array([[0, 0, 0, 10], [1, 0, 0, 20], [0, 1, 0, 20], [0, 0, 1, 30]]) triangles = array([[0, 1, 3], [0, 3, 2], [1, 2, 3], [0, 2, 1]]) ### PIPELINE # Convert list of points to VTK structure verts = tvtk.Points() temperature = tvtk.FloatArray() for p in points: verts.insert_next_point(p[0], p[1], p[2]) temperature.insert_next_value(p[3]) # Define triangular cells from the vertex # "ids" (index) and append to polygon list. polygons = tvtk.CellArray() for tri in triangles: cell = tvtk.IdList() cell.insert_next_id(tri[0]) cell.insert_next_id(tri[1]) cell.insert_next_id(tri[2]) polygons.insert_next_cell(cell) # Create a mesh from these lists