def test_filter_out_dead(grid: RectangularGrid): cells = CellList(grid=grid) cells.extend([CellData.create_cell(dead=bool(i % 2)) for i in range(10)]) assert_array_equal(cells.alive(), [0, 2, 4, 6, 8]) assert_array_equal(cells.alive([1, 2, 3, 6]), [2, 6]) mask = np.arange(10) < 5 assert_array_equal(cells.alive(mask), [0, 2, 4])
def convert_cells_to_vtk(cells: CellList) -> vtkPolyData: cell_data: CellData = cells.cell_data live_cells = cells.alive() cell_data = cell_data[live_cells] fields = dict(cell_data.dtype.fields) # type: ignore fields.pop('point') points = vtkPoints() poly = vtkPolyData() poly.SetPoints(points) if not len(cell_data): return poly # vtk uses coordinate ordering x, y, z while we use z, y, x. points.SetData(numpy_to_vtk(np.flip(cell_data['point'], axis=1))) point_data = poly.GetPointData() for field, (dtype, *_) in fields.items(): data = cell_data[field] # numpy_to_vtk doesn't handle bool for some reason if dtype == np.dtype('bool'): data = data.astype(np.dtype('uint8')) # noinspection PyBroadException try: scalar = numpy_to_vtk(data) except Exception: print(f'Unhandled data type in field {field}') continue scalar.SetName(field) point_data.AddArray(scalar) return poly