def add_attribute(self, array, name, category='point'): """Add an attribute to the dataset to specified category ('point' or 'cell'). One may add a scalar, vector (3/4 components) or a tensor (9 components). Note that it is the user's responsibility to set the correct size of the arrays. Also no automatic transposing of the data is done. Parameters ---------- array: numpy array/list : array data to add. name: str: name of the array. category: 'point'/'cell': the category of the attribute data. """ array = np.asarray(array) assert len(array.shape) <= 2, "Only 2D arrays can be added." data = getattr(self.image_data, '%s_data' % category) if len(array.shape) == 2: assert array.shape[1] in [1, 3, 4, 9], \ "Only N x m arrays where (m in [1,3,4,9]) are supported" va = tvtk.to_tvtk(array2vtk(array)) va.name = name data.add_array(va)
def test_reference_to_array(self): """Does to_array return an existing array instead of a new copy.""" arr = numpy.arange(0.0, 10.0, 0.1) arr = numpy.reshape(arr, (25, 4)) vtk_arr = array_handler.array2vtk(arr) arr1 = array_handler.vtk2array(vtk_arr) # Now make sure these are using the same memory. arr[0][0] = 100.0 self.assertEqual(arr[0][0], arr1[0][0]) self.assertEqual(arr.shape, arr1.shape)
def set_arrays(dataset, particle_array): """ Code to add all the arrays to a dataset given a particle array.""" props = set(particle_array.properties.keys()) # Add the vector data. vec = numpy.empty((len(particle_array.x), 3), dtype=float) vec[:, 0] = particle_array.u vec[:, 1] = particle_array.v vec[:, 2] = particle_array.w va = tvtk.to_tvtk(array2vtk(vec)) va.name = 'velocity' dataset.data.point_data.add_array(vec) # Now add the scalar data. scalars = props - set(('u', 'v', 'w')) for sc in scalars: arr = particle_array.get(sc) va = tvtk.to_tvtk(array2vtk(arr)) va.name = sc dataset.data.point_data.add_array(va) dataset._update_data()
def set_arrays(dataset, particle_array): """ Code to add all the arrays to a dataset given a particle array.""" props = set(particle_array.properties.keys()) # Add the vector data. vec = numpy.empty((len(particle_array.x), 3), dtype=float) vec[:,0] = particle_array.u vec[:,1] = particle_array.v vec[:,2] = particle_array.w va = tvtk.to_tvtk(array2vtk(vec)) va.name = 'velocity' dataset.data.point_data.add_array(vec) # Now add the scalar data. scalars = props - set(('u', 'v', 'w')) for sc in scalars: arr = particle_array.get(sc) va = tvtk.to_tvtk(array2vtk(arr)) va.name = sc dataset.data.point_data.add_array(va) dataset._update_data()
def write_vtk(self, array_name, props): # ceck if it is possible if not TVTK: raise RuntimeError("Cannot generate VTK output!") # check if the array is legal if array_name not in self.array_dict.keys(): raise RuntimeError("Array %s not defined" % array_name) # create a list of props if type(props) != list: props = [props] # create an output folder for the vtk files dirname = path.join(self.dirname, "vtk") utils.mkdir(dirname) array = self.array_dict[array_name] num_particles = array.num_real_particles # save the points points = np.zeros(shape=(num_particles, 3)) points[:, 0] = array.z points[:, 1] = array.y points[:, 2] = array.x mesh = tvtk.PolyData(points=points) # add the scalar props for prop in props: if prop == "vmag": u, v, w = array.get("u", "v", "w") numpy_array = np.sqrt(u ** 2 + v ** 2 + w ** 2) else: numpy_array = array.get(prop) vtkarray = array2vtk(numpy_array) vtkarray.SetName(prop) # add the array as point data mesh.point_data.add_array(vtkarray) # set the last prop as the active scalar mesh.point_data.set_active_scalars(props[-1]) # spit it out if self.fileno is None: _fname = "%s" % (array_name) else: _fname = "%s_%03d" % (array_name, self.fileno) self._write_vtk_snapshot(mesh, dirname, _fname)
def write_vtk(self, array_name, props): # ceck if it is possible if not TVTK: raise RuntimeError('Cannot generate VTK output!') # check if the array is legal if array_name not in list(self.array_dict.keys()): raise RuntimeError('Array %s not defined' % array_name) # create a list of props if type(props) != list: props = [props] # create an output folder for the vtk files dirname = path.join(self.dirname, 'vtk') utils.mkdir(dirname) array = self.array_dict[array_name] num_particles = array.num_real_particles # save the points points = np.zeros(shape=(num_particles, 3)) points[:, 0] = array.z points[:, 1] = array.y points[:, 2] = array.x mesh = tvtk.PolyData(points=points) # add the scalar props for prop in props: if prop == 'vmag': u, v, w = array.get('u', 'v', 'w') numpy_array = np.sqrt(u**2 + v**2 + w**2) else: numpy_array = array.get(prop) vtkarray = array2vtk(numpy_array) vtkarray.SetName(prop) # add the array as point data mesh.point_data.add_array(vtkarray) # set the last prop as the active scalar mesh.point_data.set_active_scalars(props[-1]) # spit it out if self.fileno is None: _fname = '%s' % (array_name) else: _fname = '%s_%03d' % (array_name, self.fileno) self._write_vtk_snapshot(mesh, dirname, _fname)
def add_array(self, array, name, category='point'): """ Add an array to the dataset to specified category ('point' or 'cell'). """ assert len(array.shape) <= 2, "Only 2D arrays can be added." data = getattr(self.dataset, '%s_data' % category) if len(array.shape) == 2: assert array.shape[1] in [1, 3, 4, 9], \ "Only Nxm arrays where (m in [1,3,4,9]) are supported" va = tvtk.to_tvtk(array2vtk(array)) va.name = name data.add_array(va) mapping = {1: 'scalars', 3: 'vectors', 4: 'scalars', 9: 'tensors'} dict = getattr(self, '%s_%s' % (category, mapping[array.shape[1]])) dict[name] = array else: va = tvtk.to_tvtk(array2vtk(array)) va.name = name data.add_array(va) dict = getattr(self, '%s_scalars' % (category)) dict[name] = array
def add_attribute(self, array, name, category='point'): """Add an attribute to the dataset to specified category ('point' or 'cell'). One may add a scalar, vector (3/4 components) or a tensor (9 components). Note that it is the user's responsibility to set the correct size of the arrays. Parameters ---------- array: numpy array/list : array data to add. name: str: name of the array. category: 'point'/'cell': the category of the attribute data. """ array = np.asarray(array) assert len(array.shape) <= 2, "Only 2D arrays can be added." data = getattr(self.data, '%s_data' % category) if len(array.shape) == 2: assert array.shape[1] in [1, 3, 4, 9], \ "Only Nxm arrays where (m in [1,3,4,9]) are supported" va = tvtk.to_tvtk(array2vtk(array)) va.name = name data.add_array(va) mapping = {1: 'scalars', 3: 'vectors', 4: 'scalars', 9: 'tensors'} attribute = '_%s_%s_list' % (category, mapping[array.shape[1]]) else: va = tvtk.to_tvtk(array2vtk(array)) va.name = name data.add_array(va) attribute = '_%s_scalars_list' % category names = set(getattr(self, attribute) + [name]) setattr(self, attribute, sorted(names))
def write_vtk(self, array_name, props): if not TVTK: return # create a list of props if type(props) != list: props = [props] # create an output folder for the vtk files dirname = path.join(self.dirname, 'vtk') utils.mkdir(dirname) nfiles = self.nfiles for i in range(self.start, nfiles): f = self.files[i] data = utils.load(f) array = data['arrays'][array_name] num_particles = array.num_real_particles # save the points points = np.zeros(shape=(num_particles, 3)) points[:, 0] = array.z points[:, 1] = array.y points[:, 2] = array.x mesh = tvtk.PolyData(points=points) # add the scalar props for prop in props: if prop == 'vmag': u, v, w = array.get('u', 'v', 'w') numpy_array = np.sqrt(u**2 + v**2 + w**2) else: numpy_array = array.get(prop) vtkarray = array2vtk(numpy_array) vtkarray.SetName(prop) # add the array as point data mesh.point_data.add_array(vtkarray) # set the last prop as the active scalar mesh.point_data.set_active_scalars(props[-1]) # spit it out fileno = data['solver_data']['count'] _fname = self.fname + '_%s_%s' % (array_name, fileno) self._write_vtk_snapshot(mesh, dirname, _fname)
def write_vtk(self, array_name, props): if not TVTK: return # create a list of props if type(props) != list: props = [props] # create an output folder for the vtk files dirname = path.join(self.dirname, "vtk") utils.mkdir(dirname) nfiles = self.nfiles for i in range(self.start, nfiles): f = self.files[i] data = utils.load(f) array = data["arrays"][array_name] num_particles = array.num_real_particles # save the points points = np.zeros(shape=(num_particles, 3)) points[:, 0] = array.z points[:, 1] = array.y points[:, 2] = array.x mesh = tvtk.PolyData(points=points) # add the scalar props for prop in props: if prop == "vmag": u, v, w = array.get("u", "v", "w") numpy_array = np.sqrt(u ** 2 + v ** 2 + w ** 2) else: numpy_array = array.get(prop) vtkarray = array2vtk(numpy_array) vtkarray.SetName(prop) # add the array as point data mesh.point_data.add_array(vtkarray) # set the last prop as the active scalar mesh.point_data.set_active_scalars(props[-1]) # spit it out fileno = data["solver_data"]["count"] _fname = self.fname + "_%s_%s" % (array_name, fileno) self._write_vtk_snapshot(mesh, dirname, _fname)
def add_array(self, array, name, category='point'): """ Add an array to the dataset to specified category ('point' or 'cell'). """ assert len(array.shape) <= 2, "Only 2D arrays can be added." data = getattr(self.dataset, '%s_data'%category) if len(array.shape) == 2: assert array.shape[1] in [1, 3, 4, 9], \ "Only Nxm arrays where (m in [1,3,4,9]) are supported" va = tvtk.to_tvtk(array2vtk(array)) va.name = name data.add_array(va) mapping = {1:'scalars', 3: 'vectors', 4: 'scalars', 9: 'tensors'} dict = getattr(self, '%s_%s'%(category, mapping[array.shape[1]])) dict[name] = array else: va = tvtk.to_tvtk(array2vtk(array)) va.name = name data.add_array(va) dict = getattr(self, '%s_scalars'%(category)) dict[name] = array
def test_vtk2array_appended_array(self): """Test the vtk2array can tolerate appending a cached array.""" # array is cached upon array2vtk is called arr = numpy.arange(8).reshape(2, 4) vtk_arr = array_handler.array2vtk(arr) arr1 = array_handler.vtk2array(vtk_arr) # the vtk array is appended, shapes don't match cached array anymore extra_row = (1, 2, 3, 4) vtk_arr.InsertTuple4(2, *extra_row) # arr2 has a different shape arr2 = array_handler.vtk2array(vtk_arr) self.assertEqual(arr2.shape, (3, 4)) # check values expected = numpy.vstack((arr, numpy.array(extra_row))) self.assertEqual(numpy.sum(arr2 - expected), 0)
def vtk_add_field(grid, fun): V = fun.function_space() family = V.ufl_element().family() degree = V.ufl_element().degree() if fun.value_rank() > 0: idx = np.column_stack( [V.sub(i).dofmap().dofs() for i in range(0, V.num_sub_spaces())]) fval = fun.vector().array()[idx] else: if family in ["Discontinuous Lagrange"]: fval = fun.vector().array() elif family in ["Real"]: fval = fun.vector().array()[0] * np.ones( int(grid.GetNumberOfPoints())) elif family in ["Quadrature"]: # Take the average over the quadrature points within # a given cell. Visualize as pointwise cell averages fval_cell = np.zeros((V.mesh().num_cells(), )) for c in dolfin.cells(V.mesh()): idx = c.index() fval_cell[idx] = fun.vector()[V.dofmap().cell_dofs(idx)].mean() # Visualize at the vertices by reducing the value at the # nearby quadrature points to one value at the vertex # Compute connvectivity between dimension 1 and 3 V.mesh().init(1, 3) fval_vert = np.zeros((V.mesh().num_vertices(), )) for v in dolfin.vertices(V.mesh()): idx = v.index() # Find the cells who have v as a vertex c = V.mesh().topology()(1, 3)(idx) # Find the dofs for the quadrature points within these cells dofs = reduce(np.union1d, [V.dofmap().cell_dofs(ci) for ci in c]) # Take the average of the value at these dofs fval_vert[idx] = fun.vector()[dofs].mean() else: vtd = dolfin.vertex_to_dof_map(V) fval_tmp = fun.vector().array() fval = np.zeros(len(fval_tmp)) fval = fval_tmp[vtd] if fun.name() == "displacement": # add zero columns if necessary gdim = V.num_sub_spaces() fval = np.hstack([fval, np.zeros((fval.shape[0], 3 - gdim))]) from tvtk.array_handler import array2vtk if family in ["Discontinuous Lagrange"]: funvtk = array2vtk(fval) funvtk.SetName(fun.name()) grid.GetCellData().AddArray(funvtk) elif family in ["Quadrature"]: funvtk_cell = array2vtk(fval_cell) funvtk_cell.SetName(fun.name()) grid.GetCellData().AddArray(funvtk_cell) funvtk_vert = array2vtk(fval_vert) funvtk_vert.SetName(fun.name()) grid.GetPointData().AddArray(funvtk_vert) else: funvtk = array2vtk(fval) funvtk.SetName(fun.name()) grid.GetPointData().AddArray(funvtk)
def test_array2vtk(self): """Test Numeric array to VTK array conversion and vice-versa.""" # Put all the test arrays here. t_z = [] # Test the different types of arrays. t_z.append(numpy.array([-128, 0, 127], numpy.int8)) # FIXME: character arrays are a problem since there is no # unique mapping to a VTK data type and back. #t_z.append(numpy.array([-128, 0, 127], numpy.character)) t_z.append(numpy.array([-32768, 0, 32767], numpy.int16)) t_z.append(numpy.array([-2147483648, 0, 2147483647], numpy.int32)) t_z.append(numpy.array([0, 255], numpy.uint8)) t_z.append(numpy.array([0, 65535], numpy.uint16)) t_z.append(numpy.array([0, 4294967295L], numpy.uint32)) t_z.append(numpy.array([-1.0e38, 0, 1.0e38], 'f')) t_z.append(numpy.array([-1.0e299, 0, 1.0e299], 'd')) # Check multi-component arrays. t_z.append(numpy.array([[1], [2], [300]], 'd')) t_z.append(numpy.array([[1, 20], [300, 4000]], 'd')) t_z.append(numpy.array([[1, 2, 3], [4, 5, 6]], 'f')) t_z.append(numpy.array([[1, 2, 3],[4, 5, 6]], 'd')) t_z.append(numpy.array([[1, 2, 3, 400],[4, 5, 6, 700]], 'd')) t_z.append(numpy.array([range(9),range(10,19)], 'f')) # Test if a Python list also works. t_z.append(numpy.array([[1., 2., 3., 400.],[4, 5, 6, 700]], 'd')) # Test if arrays with number of components not in [1,2,3,4,9] work. t_z.append(numpy.array([[1, 2, 3, 400, 5000], [4, 5, 6, 700, 8000]], 'd')) t_z.append(numpy.array([range(10), range(10,20)], 'd')) for z in t_z: vtk_arr = array_handler.array2vtk(z) # Test for memory leaks. self.assertEqual(vtk_arr.GetReferenceCount(), array_handler.BASE_REFERENCE_COUNT) self._check_arrays(z, vtk_arr) z1 = array_handler.vtk2array(vtk_arr) if len(z.shape) == 1: self.assertEqual(len(z1.shape), 1) if z.dtype.char != 'c': #print z1 self.assertEqual(sum(numpy.ravel(z) - numpy.ravel(z1)), 0) else: #print z1.astype('c') self.assertEqual(z, z1.astype('c')) # Check if type conversion works correctly. z = numpy.array([-128, 0, 127], numpy.int8) vtk_arr = vtk.vtkDoubleArray() ident = id(vtk_arr) vtk_arr = array_handler.array2vtk(z, vtk_arr) # Make sure this is the same array! self.assertEqual(ident, id(vtk_arr)) self._check_arrays(z, vtk_arr) # Check the vtkBitArray. vtk_arr = vtk.vtkBitArray() vtk_arr.InsertNextValue(0) vtk_arr.InsertNextValue(1) vtk_arr.InsertNextValue(0) vtk_arr.InsertNextValue(1) arr = array_handler.vtk2array(vtk_arr) self.assertEqual(numpy.sum(arr - [0,1,0,1]), 0) vtk_arr = array_handler.array2vtk(arr, vtk_arr) self.assertEqual(vtk_arr.GetValue(0), 0) self.assertEqual(vtk_arr.GetValue(1), 1) self.assertEqual(vtk_arr.GetValue(2), 0) self.assertEqual(vtk_arr.GetValue(3), 1) # ---------------------------------------- # Test if the array is copied or not. a = numpy.array([[1, 2, 3],[4, 5, 6]], 'd') vtk_arr = array_handler.array2vtk(a) # Change the numpy array and see if the changes are # reflected in the VTK array. a[0] = [10.0, 20.0, 30.0] self.assertEqual(vtk_arr.GetTuple3(0), (10., 20., 30.)) # Make sure the cache is doing its job. key = vtk_arr.__this__ z = array_handler._array_cache.get(vtk_arr) self.assertEqual(numpy.sum(z - numpy.ravel(a)), 0.0) l1 = len(array_handler._array_cache) # del the Numeric array and see if this still works. del a self.assertEqual(vtk_arr.GetTuple3(0), (10., 20., 30.)) # Check the cache -- just making sure. self.assertEqual(len(array_handler._array_cache), l1) # Delete the VTK array and see if the cache is cleared. del vtk_arr self.assertEqual(len(array_handler._array_cache), l1-1) self.assertEqual(array_handler._array_cache._cache.has_key(key), False) # Make sure bit arrays are copied. vtk_arr = vtk.vtkBitArray() a = numpy.array([0,1,0,1], numpy.int32) vtk_arr = array_handler.array2vtk(a, vtk_arr) del a self.assertEqual(vtk_arr.GetValue(0), 0) self.assertEqual(vtk_arr.GetValue(1), 1) self.assertEqual(vtk_arr.GetValue(2), 0) self.assertEqual(vtk_arr.GetValue(3), 1) # Make sure the code at least runs for all the non-complex # numerical dtypes in numpy. for dtype in (numpy.sctypes['int'] + numpy.sctypes['uint'] + numpy.sctypes['float']): array_handler.array2vtk(numpy.zeros((1,), dtype=dtype))
def test_array2vtk(self): """Test Numeric array to VTK array conversion and vice-versa.""" # Put all the test arrays here. t_z = [] # Test the different types of arrays. t_z.append(numpy.array([-128, 0, 127], numpy.int8)) # FIXME: character arrays are a problem since there is no # unique mapping to a VTK data type and back. #t_z.append(numpy.array([-128, 0, 127], numpy.character)) t_z.append(numpy.array([-32768, 0, 32767], numpy.int16)) t_z.append(numpy.array([-2147483648, 0, 2147483647], numpy.int32)) t_z.append(numpy.array([0, 255], numpy.uint8)) t_z.append(numpy.array([0, 65535], numpy.uint16)) t_z.append(numpy.array([0, 4294967295], numpy.uint32)) t_z.append(numpy.array([-1.0e38, 0, 1.0e38], 'f')) t_z.append(numpy.array([-1.0e299, 0, 1.0e299], 'd')) # Check multi-component arrays. t_z.append(numpy.array([[1], [2], [300]], 'd')) t_z.append(numpy.array([[1, 20], [300, 4000]], 'd')) t_z.append(numpy.array([[1, 2, 3], [4, 5, 6]], 'f')) t_z.append(numpy.array([[1, 2, 3], [4, 5, 6]], 'd')) t_z.append(numpy.array([[1, 2, 3, 400], [4, 5, 6, 700]], 'd')) t_z.append(numpy.array([list(range(9)), list(range(10, 19))], 'f')) # Test if a Python list also works. t_z.append(numpy.array([[1., 2., 3., 400.], [4, 5, 6, 700]], 'd')) # Test if arrays with number of components not in [1,2,3,4,9] work. t_z.append( numpy.array([[1, 2, 3, 400, 5000], [4, 5, 6, 700, 8000]], 'd')) t_z.append(numpy.array([list(range(10)), list(range(10, 20))], 'd')) for z in t_z: vtk_arr = array_handler.array2vtk(z) # Test for memory leaks. self.assertEqual(vtk_arr.GetReferenceCount(), array_handler.BASE_REFERENCE_COUNT) self._check_arrays(z, vtk_arr) z1 = array_handler.vtk2array(vtk_arr) if len(z.shape) == 1: self.assertEqual(len(z1.shape), 1) if z.dtype.char != 'c': #print z1 self.assertEqual(sum(numpy.ravel(z) - numpy.ravel(z1)), 0) else: #print z1.astype('c') self.assertEqual(z, z1.astype('c')) # Check if type conversion works correctly. z = numpy.array([-128, 0, 127], numpy.int8) vtk_arr = vtk.vtkDoubleArray() ident = id(vtk_arr) vtk_arr = array_handler.array2vtk(z, vtk_arr) # Make sure this is the same array! self.assertEqual(ident, id(vtk_arr)) self._check_arrays(z, vtk_arr) # Check the vtkBitArray. vtk_arr = vtk.vtkBitArray() vtk_arr.InsertNextValue(0) vtk_arr.InsertNextValue(1) vtk_arr.InsertNextValue(0) vtk_arr.InsertNextValue(1) arr = array_handler.vtk2array(vtk_arr) self.assertEqual(numpy.sum(arr - [0, 1, 0, 1]), 0) vtk_arr = array_handler.array2vtk(arr, vtk_arr) self.assertEqual(vtk_arr.GetValue(0), 0) self.assertEqual(vtk_arr.GetValue(1), 1) self.assertEqual(vtk_arr.GetValue(2), 0) self.assertEqual(vtk_arr.GetValue(3), 1) # ---------------------------------------- # Test if the array is copied or not. a = numpy.array([[1, 2, 3], [4, 5, 6]], 'd') vtk_arr = array_handler.array2vtk(a) # Change the numpy array and see if the changes are # reflected in the VTK array. a[0] = [10.0, 20.0, 30.0] self.assertEqual(vtk_arr.GetTuple3(0), (10., 20., 30.)) # Make sure the cache is doing its job. key = vtk_arr.__this__ z = array_handler._array_cache.get(vtk_arr) self.assertEqual(numpy.sum(z - numpy.ravel(a)), 0.0) l1 = len(array_handler._array_cache) # del the Numeric array and see if this still works. del a self.assertEqual(vtk_arr.GetTuple3(0), (10., 20., 30.)) # Check the cache -- just making sure. self.assertEqual(len(array_handler._array_cache), l1) # Delete the VTK array and see if the cache is cleared. del vtk_arr self.assertEqual(len(array_handler._array_cache), l1 - 1) self.assertEqual(key in array_handler._array_cache._cache, False) # Make sure bit arrays are copied. vtk_arr = vtk.vtkBitArray() a = numpy.array([0, 1, 0, 1], numpy.int32) vtk_arr = array_handler.array2vtk(a, vtk_arr) del a self.assertEqual(vtk_arr.GetValue(0), 0) self.assertEqual(vtk_arr.GetValue(1), 1) self.assertEqual(vtk_arr.GetValue(2), 0) self.assertEqual(vtk_arr.GetValue(3), 1) # Make sure the code at least runs for all the non-complex # numerical dtypes in numpy. for dtype in (numpy.sctypes['int'] + numpy.sctypes['uint'] + numpy.sctypes['float']): array_handler.array2vtk(numpy.zeros((1, ), dtype=dtype))