Copyright (C) 2021-2022 Jørgen S. Dokken and Garth N. Wells

This file is part of DOLFINx (https://www.fenicsproject.org)

SPDX-License-Identifier: LGPL-3.0-or-later

Download sources

- `Python script`

- `Jupyter notebook`

# 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 pathlib import Path

    from mpi4py import MPI

    import numpy as np
    import pyvista

    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

    # Set some global options for all plots
    transparent = False
    figsize = 800
    out_folder = Path("out_pyvista")
    out_folder.mkdir(parents=True, exist_ok=True)

## 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(
                out_folder / "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(
                out_folder / "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(
                out_folder / 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(
                out_folder / "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(
                out_folder / 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()