Making vtk python object from solution object in the same script

Question

I would like to generate a vtk file with more than one scalar field from fenics. I haven't found a way of doing this within fenics, what I can do now is generate vtk files and then open them again in the same script and use the following code:

from vtk import * 

reader = vtkPolyDataReader()
append = vtkAppendPolyData()

filenames = ['file1.vtk', 'file2.vtk']
for file in filenames:
    reader.SetFileName(file)
    reader.Update()
    polydata = vtkPolyData()
    polydata.ShallowCopy(reader.GetOutput())
    append.AddInputData(polydata)

append.Update()    
writer = vtkPolyDataWriter()
writer.SetFileName('output.vtk')
writer.SetInput(append.GetOutput())
writer.Write()

Which is not a bad solution but I would like to do this without having to make 2 vtk files for every iteration of a solution in fenics. Is there a way of changing a solution in fenics directly into a vtk object?

Answer 1

You can convert a dolfin.Function object to a vtkUnstructuredGrid in serial as follows:

import vtk
import dolfin as dlf
import numpy as np

from vtk.util import numpy_support as ns

mesh = dlf.UnitCubeMesh(10,10,10)
V = dlf.FunctionSpace(mesh, "CG", 2)
u = dlf.interpolate(dlf.Expression("1.0 + exp(x[0]*x[1])", element=V.ufl_element()), V)

dofmap = V.dofmap()
n_cells = mesh.num_cells()
nodes_per_element = dofmap.num_element_dofs(0)

# Create a numpy array to store connectivity matrix
np_connectivity = np.zeros([n_cells, nodes_per_element], dtype=np.int)

# Get the connectivity of each cell
for i in range(n_cells):
    np_connectivity[i,:] = dofmap.cell_dofs(i)

# Node numbering scheme in FEniCS is different from VTK for 2nd order
PERM_DOLFIN_TO_VTK = [0, 1, 2, 3, 9, 6, 8, 7, 5, 4]
np_connectivity = np_connectivity[:, PERM_DOLFIN_TO_VTK]

# Add left column specifying number of nodes per cell and flatten array
np_connectivity = np.hstack((np.ones([n_cells, 1], dtype=np.int)*nodes_per_element,
                             np_connectivity)).flatten()

# Convert connectivity to VTK array
vtk_connectivity = ns.numpy_to_vtkIdTypeArray(np_connectivity)

# Create cell array
vtk_cells = vtk.vtkCellArray()
vtk_cells.SetCells(n_cells, vtk_connectivity)

# Extract point coordinates
geo_dim = mesh.geometry().dim()
np_coordinates = V.tabulate_dof_coordinates().reshape([-1, geo_dim])

# Convert coordinates to VTK
vtk_coordinates = ns.numpy_to_vtk(np_coordinates)
vtk_points = vtk.vtkPoints()
vtk_points.SetData(vtk_coordinates)

# Create unstructured grid and set points and connectivity
u_grid = vtk.vtkUnstructuredGrid()
u_grid.SetCells(vtk.VTK_QUADRATIC_TETRA, vtk_cells) # Use vtk.VTK_TETRA for linear
u_grid.SetPoints(vtk_points)

# Convert function values and add as scalar data
vec = u.vector().array()
vtk_u = ns.numpy_to_vtk(vec)
vtk_u.SetName(u.name())

# There are multiple ways of adding this
u_grid.GetPointData().SetScalars(vtk_u)

You should be able to add other data arrays to the u_grid object. I have yet to figure out how to do this in parallel, and have asked a separate question here.