Visualization with PyVista

PyVista can be used with DOLFINx for interactive visualisation.

To start, the required modules are imported and some PyVista parameters set.

from mpi4py import MPI

import numpy as np

import dolfinx.plot as plot
from dolfinx.fem import Function, functionspace
from dolfinx.mesh import CellType, compute_midpoints, create_unit_cube, create_unit_square, meshtags

try:
    import pyvista
except ModuleNotFoundError:
    print("pyvista is required for this demo")
    exit(0)

# If environment variable PYVISTA_OFF_SCREEN is set to true save a png
# otherwise create interactive plot
if pyvista.OFF_SCREEN:
    pyvista.start_xvfb(wait=0.1)

# Set some global options for all plots
transparent = False
figsize = 800

Plotting a finite element Function using warp by scalar

def plot_scalar():
    # We start by creating a unit square mesh and interpolating a
    # function into a degree 1 Lagrange space
    msh = create_unit_square(
        MPI.COMM_WORLD, 12, 12, cell_type=CellType.quadrilateral, dtype=np.float64
    )
    V = functionspace(msh, ("Lagrange", 1))
    u = Function(V, dtype=np.float64)
    u.interpolate(lambda x: np.sin(np.pi * x[0]) * np.sin(2 * x[1] * np.pi))

    # To visualize the function u, we create a VTK-compatible grid to
    # values of u to
    cells, types, x = plot.vtk_mesh(V)
    grid = pyvista.UnstructuredGrid(cells, types, x)
    grid.point_data["u"] = u.x.array

    # The function "u" is set as the active scalar for the mesh, and
    # warp in z-direction is set
    grid.set_active_scalars("u")
    warped = grid.warp_by_scalar()

    # A plotting window is created with two sub-plots, one of the scalar
    # values and the other of the mesh is warped by the scalar values in
    # z-direction
    subplotter = pyvista.Plotter(shape=(1, 2))
    subplotter.subplot(0, 0)
    subplotter.add_text("Scalar contour field", font_size=14, color="black", position="upper_edge")
    subplotter.add_mesh(grid, show_edges=True, show_scalar_bar=True)
    subplotter.view_xy()

    subplotter.subplot(0, 1)
    subplotter.add_text("Warped function", position="upper_edge", font_size=14, color="black")
    sargs = dict(
        height=0.8,
        width=0.1,
        vertical=True,
        position_x=0.05,
        position_y=0.05,
        fmt="%1.2e",
        title_font_size=40,
        color="black",
        label_font_size=25,
    )
    subplotter.set_position([-3, 2.6, 0.3])
    subplotter.set_focus([3, -1, -0.15])
    subplotter.set_viewup([0, 0, 1])
    subplotter.add_mesh(warped, show_edges=True, scalar_bar_args=sargs)
    if pyvista.OFF_SCREEN:
        subplotter.screenshot(
            "2D_function_warp.png",
            transparent_background=transparent,
            window_size=[figsize, figsize],
        )
    else:
        subplotter.show()

Mesh tags and using subplots

def plot_meshtags():
    # Create a mesh
    msh = create_unit_square(
        MPI.COMM_WORLD, 25, 25, cell_type=CellType.quadrilateral, dtype=np.float64
    )

    # Create a geometric indicator function
    def in_circle(x):
        return np.array((x.T[0] - 0.5) ** 2 + (x.T[1] - 0.5) ** 2 < 0.2**2, dtype=np.int32)

    # Create cell tags - if midpoint is inside circle, it gets value 1,
    # otherwise 0
    num_cells = msh.topology.index_map(msh.topology.dim).size_local
    midpoints = compute_midpoints(msh, msh.topology.dim, np.arange(num_cells, dtype=np.int32))
    cell_tags = meshtags(msh, msh.topology.dim, np.arange(num_cells), in_circle(midpoints))

    # Create VTK mesh
    cells, types, x = plot.vtk_mesh(msh)
    grid = pyvista.UnstructuredGrid(cells, types, x)

    # Attach the cells tag data to the pyvita grid
    grid.cell_data["Marker"] = cell_tags.values
    grid.set_active_scalars("Marker")

    # Create a plotter with two subplots, and add mesh tag plot to the
    # first sub-window
    subplotter = pyvista.Plotter(shape=(1, 2))
    subplotter.subplot(0, 0)
    subplotter.add_text("Mesh with markers", font_size=14, color="black", position="upper_edge")
    subplotter.add_mesh(grid, show_edges=True, show_scalar_bar=False)
    subplotter.view_xy()

    # We can visualize subsets of data, by creating a smaller topology
    # (set of cells). Here we create VTK mesh data for only cells with
    # that tag '1'.
    cells, types, x = plot.vtk_mesh(msh, entities=cell_tags.find(1))

    # Add this grid to the second plotter window
    sub_grid = pyvista.UnstructuredGrid(cells, types, x)
    subplotter.subplot(0, 1)
    subplotter.add_text("Subset of mesh", font_size=14, color="black", position="upper_edge")
    subplotter.add_mesh(sub_grid, show_edges=True, edge_color="black")

    if pyvista.OFF_SCREEN:
        subplotter.screenshot(
            "2D_markers.png", transparent_background=transparent, window_size=[2 * figsize, figsize]
        )
    else:
        subplotter.show()

Higher-order Functions

Higher-order finite element function can also be plotted.

def plot_higher_order():
    # Create a mesh
    msh = create_unit_square(
        MPI.COMM_WORLD, 12, 12, cell_type=CellType.quadrilateral, dtype=np.float64
    )

    # Define a geometric indicator function
    def in_circle(x):
        return np.array((x.T[0] - 0.5) ** 2 + (x.T[1] - 0.5) ** 2 < 0.2**2, dtype=np.int32)

    # Create mesh tags for all cells. If midpoint is inside the circle,
    # it gets value 1, otherwise 0.
    num_cells = msh.topology.index_map(msh.topology.dim).size_local
    midpoints = compute_midpoints(msh, msh.topology.dim, np.arange(num_cells, dtype=np.int32))
    cell_tags = meshtags(msh, msh.topology.dim, np.arange(num_cells), in_circle(midpoints))

    # We start by interpolating a discontinuous function (discontinuous
    # between cells with different mesh tag values) into a degree 2
    # discontinuous Lagrange space.
    V = functionspace(msh, ("Discontinuous Lagrange", 2))
    u = Function(V, dtype=msh.geometry.x.dtype)
    u.interpolate(lambda x: x[0], cell_tags.find(0))
    u.interpolate(lambda x: x[1] + 1, cell_tags.find(1))
    u.x.scatter_forward()

    # Create a topology that has a 1-1 correspondence with the
    # degrees-of-freedom in the function space V
    cells, types, x = plot.vtk_mesh(V)

    # Create a pyvista mesh and attach the values of u
    grid = pyvista.UnstructuredGrid(cells, types, x)
    grid.point_data["u"] = u.x.array
    grid.set_active_scalars("u")

    # We would also like to visualize the underlying mesh and obtain
    # that as we have done previously
    num_cells = msh.topology.index_map(msh.topology.dim).size_local
    cell_entities = np.arange(num_cells, dtype=np.int32)
    cells, types, x = plot.vtk_mesh(msh, entities=cell_entities)
    org_grid = pyvista.UnstructuredGrid(cells, types, x)

    # We visualize the data
    plotter = pyvista.Plotter()
    plotter.add_text(
        "Second-order (P2) discontinuous elements",
        position="upper_edge",
        font_size=14,
        color="black",
    )
    sargs = dict(height=0.1, width=0.8, vertical=False, position_x=0.1, position_y=0, color="black")
    plotter.add_mesh(grid, show_edges=False, scalar_bar_args=sargs, line_width=0)
    plotter.add_mesh(org_grid, color="white", style="wireframe", line_width=5)
    plotter.add_mesh(
        grid.copy(), style="points", point_size=15, render_points_as_spheres=True, line_width=0
    )
    plotter.view_xy()
    if pyvista.OFF_SCREEN:
        plotter.screenshot(
            f"DG_{MPI.COMM_WORLD.rank}.png",
            transparent_background=transparent,
            window_size=[figsize, figsize],
        )
    else:
        plotter.show()

Vector-element functions

In this section we will consider how to plot vector-element functions, e.g. Raviart-Thomas or Nédélec elements.

def plot_nedelec():
    msh = create_unit_cube(
        MPI.COMM_WORLD, 4, 3, 5, cell_type=CellType.tetrahedron, dtype=np.float64
    )

    # We create a pyvista plotter
    plotter = pyvista.Plotter()
    plotter.add_text(
        "Mesh and corresponding vectors", position="upper_edge", font_size=14, color="black"
    )

    # Next, we create a pyvista.UnstructuredGrid based on the mesh
    pyvista_cells, cell_types, x = plot.vtk_mesh(msh)
    grid = pyvista.UnstructuredGrid(pyvista_cells, cell_types, x)

    # Add this grid (as a wireframe) to the plotter
    plotter.add_mesh(grid, style="wireframe", line_width=2, color="black")

    # Create a function space consisting of first order Nédélec (first kind)
    # elements and interpolate a vector-valued expression
    V = functionspace(msh, ("N1curl", 2))
    u = Function(V, dtype=np.float64)
    u.interpolate(lambda x: (x[2] ** 2, np.zeros(x.shape[1]), -x[0] * x[2]))

    # Exact visualisation of the Nédélec spaces requires a Lagrange or
    # discontinuous Lagrange finite element functions. Therefore, we
    # interpolate the Nédélec function into a first-order discontinuous
    # Lagrange space.
    gdim = msh.geometry.dim
    V0 = functionspace(msh, ("Discontinuous Lagrange", 2, (gdim,)))
    u0 = Function(V0, dtype=np.float64)
    u0.interpolate(u)

    # Create a second grid, whose geometry and topology is based on the
    # output function space
    cells, cell_types, x = plot.vtk_mesh(V0)
    grid = pyvista.UnstructuredGrid(cells, cell_types, x)

    # Create point cloud of vertices, and add the vertex values to the cloud
    grid.point_data["u"] = u0.x.array.reshape(x.shape[0], V0.dofmap.index_map_bs)
    glyphs = grid.glyph(orient="u", factor=0.1)

    # We add in the glyphs corresponding to the plotter
    plotter.add_mesh(glyphs)

    # Save as png if we are using a container with no rendering
    if pyvista.OFF_SCREEN:
        plotter.screenshot(
            "3D_wireframe_with_vectors.png",
            transparent_background=transparent,
            window_size=[figsize, figsize],
        )
    else:
        plotter.show()

Plotting streamlines

In this section we illustrate how to visualize streamlines in 3D

def plot_streamlines():
    msh = create_unit_cube(MPI.COMM_WORLD, 4, 4, 4, CellType.hexahedron, dtype=np.float64)
    gdim = msh.geometry.dim
    V = functionspace(msh, ("Discontinuous Lagrange", 2, (gdim,)))
    u = Function(V, dtype=np.float64)
    u.interpolate(lambda x: np.vstack((-(x[1] - 0.5), x[0] - 0.5, np.zeros(x.shape[1]))))

    cells, types, x = plot.vtk_mesh(V)
    num_dofs = x.shape[0]
    values = np.zeros((num_dofs, 3), dtype=np.float64)
    values[:, : msh.geometry.dim] = u.x.array.reshape(num_dofs, V.dofmap.index_map_bs)

    # Create a point cloud of glyphs
    grid = pyvista.UnstructuredGrid(cells, types, x)
    grid["vectors"] = values
    grid.set_active_vectors("vectors")
    glyphs = grid.glyph(orient="vectors", factor=0.1)
    streamlines = grid.streamlines(
        vectors="vectors", return_source=False, source_radius=1, n_points=150
    )

    # Create Create plotter
    plotter = pyvista.Plotter()
    plotter.add_text("Streamlines.", position="upper_edge", font_size=20, color="black")
    plotter.add_mesh(grid, style="wireframe")
    plotter.add_mesh(glyphs)
    plotter.add_mesh(streamlines.tube(radius=0.001))
    plotter.view_xy()
    if pyvista.OFF_SCREEN:
        plotter.screenshot(
            f"streamlines_{MPI.COMM_WORLD.rank}.png",
            transparent_background=transparent,
            window_size=[figsize, figsize],
        )
    else:
        plotter.show()

plot_scalar()
plot_meshtags()
plot_higher_order()
plot_nedelec()
plot_streamlines()