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