fix cell_data writing for vtk
@@ -218,10 +218,43 @@ def write(filetype,
# For VT{K,U} files, no underscore is ever added.
pd.AddArray(_create_vtkarray(X, name))
- # add cell data
+ # Add cell data.
+ # The cell_data is structured like
+ #
+ # cell_type ->
+ # key -> array
+ # key -> array
+ # [...]
+ # cell_type ->
+ # key -> array
+ # key -> array
+ # [...]
+ # [...]
+ #
+ # VTK expects one array for each `key`, so assemble the keys across all
+ # mesh_types. This requires each key to be present for each mesh_type, of
+ # course.
+ all_keys = []
+ for cell_type in cell_data:
+ all_keys += cell_data[cell_type].keys()
+ # count all cells
+ num_cells = 0
+ for cell_type in cells:
+ num_cells += len(cells[cell_type])
+ # create unified cell data
+ unified_cell_data = {}
+ for key in all_keys:
+ unified_cell_data[key] = numpy.empty(num_cells, dtype=int)
+ index = 0
+ for cell_type in cell_data:
+ assert key in cell_data[cell_type]
+ n = len(cell_data[cell_type][key])
+ unified_cell_data[key][index:index+n] = cell_data[cell_type][key]
+ index += n
+ # add the array data to the mesh
cd = vtk_mesh.GetCellData()
- for key, value in cell_data.iteritems():
- cd.AddArray(_create_vtkarray(value, key))
+ for name, array in unified_cell_data.iteritems():
+ cd.AddArray(_create_vtkarray(array, name))
# add field data
fd = vtk_mesh.GetFieldData()
