Python: dump VTK unstructured grid, creating it ad-hoc from physics calculation data

python_vtk_dump_unstructured_grid.py

import vtk

def dump_model(model, name, path="."):
    # Data from physics calculation
    # All nodes
    nodes = model["nodes"]
    # Triangles of the first body
    bar_elements = model["bar_elements"]
    # Triangles of the second body
    surf_elements = model["surf_elements"]

    # Create VTK object to represent grid
    unstructuredGrid = vtk.vtkUnstructuredGrid()

    # Create VTK points array
    points = vtk.vtkPoints()

    # Create 4 VTK arrays to represent 4 scalar fields
    # (tied to nodes, not triangles)

    stress = vtk.vtkDoubleArray()
    stress.SetNumberOfValues(len(nodes))
    stress.SetName("stress")

    dx = vtk.vtkDoubleArray()
    dx.SetNumberOfValues(len(nodes))
    dx.SetName("dx")

    dy = vtk.vtkDoubleArray()
    dy.SetNumberOfValues(len(nodes))
    dy.SetName("dy")

    dz = vtk.vtkDoubleArray()
    dz.SetNumberOfValues(len(nodes))
    dz.SetName("dz")

    # Helpers to renumerate nodes according with VTK rules
    vtk_idx = 0
    remap = {}

    # Iterate over all nodes
    for node_id in nodes.keys():
        remap[node_id] = vtk_idx
        # Insert the position of the next point into VTK points
        points.InsertNextPoint(float(nodes[node_id]['x']) + float(nodes[node_id]['dx']),
                               float(nodes[node_id]['y']) + float(nodes[node_id]['dy']),
                               float(nodes[node_id]['z']) + float(nodes[node_id]['dz']))
        # Insert deformation and stress point data into VTK scalar fields
        stress.SetValue(vtk_idx, nodes[node_id]['max_stress'])
        dx.SetValue(vtk_idx, nodes[node_id]['dx'])
        dy.SetValue(vtk_idx, nodes[node_id]['dy'])
        dz.SetValue(vtk_idx, nodes[node_id]['dz'])
        vtk_idx += 1
    
    # Attach points and scalar fields to VTK grid
    unstructuredGrid.SetPoints(points)
    unstructuredGrid.GetPointData().AddArray(stress)
    unstructuredGrid.GetPointData().AddArray(dx)
    unstructuredGrid.GetPointData().AddArray(dy)
    unstructuredGrid.GetPointData().AddArray(dz)

    # Write indentor geometry cells
    # Iterate over triangles used in physics calculation
    for element_id in bar_elements:
        # Create VTK triangle
        triangle = vtk.vtkTriangle()
        # Set VTK triangle vertices numbers to *VTK* ids of correspoding points
        for i in range(0, 3):
            triangle.GetPointIds().SetId(i, remap[bar_elements[element_id][i]])
        # Insert the next triangle into VTK grid
        unstructuredGrid.InsertNextCell(triangle.GetCellType(), triangle.GetPointIds())

    # Write box geometry cells
    # (the same procedure as for indentor)
    for element_id in surf_elements:
        triangle = vtk.vtkTriangle()
        for i in range(0, 3):
            triangle.GetPointIds().SetId(i, remap[surf_elements[element_id][i]['id']])
        unstructuredGrid.InsertNextCell(triangle.GetCellType(), triangle.GetPointIds())

    # Create writer for unstructured grid
    writer = vtk.vtkXMLUnstructuredGridWriter()
    # Set input grid to write to the file
    writer.SetInputDataObject(unstructuredGrid)
    # Set file name
    writer.SetFileName(path + "/" + name + ".vtu")
    # Do file write
    writer.Write()