Copyright (C) 2021 Jørgen S. Dokken

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

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

===============================
Using pyvista for visualization
===============================


::

  import dolfinx
  import dolfinx.io
  import dolfinx.plot
  import numpy as np
  import ufl
  from mpi4py import MPI
  
  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
  pyvista.rcParams["background"] = [0.5, 0.5, 0.5]
  
  
Plotting a 3D dolfinx.Function with pyvista
===========================================

::

  # Interpolate a simple scalar function in 3D
  def int_u(x):
      return x[0] + 3 * x[1] + 5 * x[2]
  
  
  mesh = dolfinx.UnitCubeMesh(MPI.COMM_WORLD, 4, 3, 5, cell_type=dolfinx.cpp.mesh.CellType.tetrahedron)
  V = dolfinx.FunctionSpace(mesh, ("CG", 1))
  u = dolfinx.Function(V)
  u.interpolate(int_u)
  
  # Extract mesh data from DOLFINx (only plot cells owned by the
  # processor) and create a pyvista UnstructuredGrid
  num_cells = mesh.topology.index_map(mesh.topology.dim).size_local
  cell_entities = np.arange(num_cells, dtype=np.int32)
  pyvista_cells, cell_types = dolfinx.plot.create_vtk_topology(mesh, mesh.topology.dim, cell_entities)
  grid = pyvista.UnstructuredGrid(pyvista_cells, cell_types, mesh.geometry.x)
  
  # Compute the function values at the vertices, this is equivalent to a
  # P1 Lagrange interpolation, and can be directly attached to the Pyvista
  # mesh. Discard complex value if running DOLFINx with complex PETSc as
  # backend
  vertex_values = u.compute_point_values()
  if np.iscomplexobj(vertex_values):
      vertex_values = vertex_values.real
  
  # Create point cloud of vertices, and add the vertex values to the cloud
  grid.point_arrays["u"] = vertex_values
  grid.set_active_scalars("u")
  
  # Create a pyvista plotter which is used to visualize the output
  plotter = pyvista.Plotter()
  plotter.add_text("Mesh and corresponding dof values",
                   position="upper_edge", font_size=14, color="black")
  
  # Some styling arguments for the colorbar
  sargs = dict(height=0.6, width=0.1, vertical=True, position_x=0.825, position_y=0.2, fmt="%1.2e",
               title_font_size=40, color="black", label_font_size=25)
  
  # Plot the mesh (as a wireframe) with the finite element function
  # visualized as the point cloud
  plotter.add_mesh(grid, style="wireframe", line_width=2, color="black")
  
  # To be able to visualize the mesh and nodes at the same time, we have
  # to copy the grid
  plotter.add_mesh(grid.copy(), style="points", render_points_as_spheres=True,
                   scalars=vertex_values, point_size=10)
  plotter.set_position([1.5, 0.5, 4])
  
  # Save as png if we are using a container with no rendering
  if pyvista.OFF_SCREEN:
      plotter.screenshot("3D_wireframe_with_nodes.png", transparent_background=transparent,
                         window_size=[figsize, figsize])
  else:
      plotter.show()
  
  # Create a new plotter, and plot the values as a surface over the mesh
  plotter = pyvista.Plotter()
  plotter.add_text("Function values over the surface of a mesh",
                   position="upper_edge", font_size=14, color="black")
  
  # Define some styling arguments for a colorbar
  sargs = dict(height=0.1, width=0.8, vertical=False, position_x=0.1,
               position_y=0.05, fmt="%1.2e",
               title_font_size=40, color="black", label_font_size=25)
  
  # Adjust camera to show the entire mesh
  plotter.set_position([-2, -2, 2.1])
  plotter.set_focus([1, 1, -0.01])
  plotter.set_viewup([0, 0, 1])
  
  # Add mesh with edges
  plotter.add_mesh(grid, show_edges=True, scalars="u", scalar_bar_args=sargs)
  if pyvista.OFF_SCREEN:
      plotter.screenshot("3D_function.png", transparent_background=transparent, window_size=[figsize, figsize])
  else:
      plotter.show()
  
Plotting a 2D dolfinx.Function with pyvista using warp by scalar
================================================================

As in the previous section, we interpolate a function into a CG
function space


::

  def int_2D(x):
      return np.sin(np.pi * x[0]) * np.sin(2 * x[1] * np.pi)
  
  
  mesh = dolfinx.UnitSquareMesh(MPI.COMM_WORLD, 12, 12, cell_type=dolfinx.cpp.mesh.CellType.quadrilateral)
  V = dolfinx.FunctionSpace(mesh, ("Lagrange", 1))
  u = dolfinx.Function(V)
  u.interpolate(int_2D)
  
  # As in the previous section, we extract the geometry and topology of
  # the mesh, and attach values to the vertices
  num_cells = mesh.topology.index_map(mesh.topology.dim).size_local
  cells = np.arange(num_cells, dtype=np.int32)
  pyvista_cells, cell_types = dolfinx.plot.create_vtk_topology(mesh, mesh.topology.dim, cells)
  grid = pyvista.UnstructuredGrid(pyvista_cells, cell_types, mesh.geometry.x)
  point_values = u.compute_point_values()
  if np.iscomplexobj(point_values):
      point_values = point_values.real
  grid.point_arrays["u"] = point_values
  
  # We set the function "u" as the active scalar for the mesh, and warp
  # the mesh in z-direction by its values
  grid.set_active_scalars("u")
  warped = grid.warp_by_scalar()
  
  # Plot mesh with scalar bar
  plotter = pyvista.Plotter()
  plotter.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)
  plotter.set_position([-3, 2.6, 0.3])
  plotter.set_focus([3, -1, -0.15])
  plotter.set_viewup([0, 0, 1])
  plotter.add_mesh(warped, show_edges=True, scalar_bar_args=sargs)
  if pyvista.OFF_SCREEN:
      plotter.screenshot("2D_function_warp.png", transparent_background=transparent, window_size=[figsize, figsize])
  else:
      plotter.show()
  
Plotting a 2D MeshTags and using subplots
=========================================

We continue using the mesh from the previous section, and find all
cells satisfying the condition below


::

  def in_circle(x):
      # Mark sphere with radius < sqrt(2)
      return np.array((x.T[0] - 0.5)**2 + (x.T[1] - 0.5)**2 < 0.2**2, dtype=np.int32)
  
  
  # Create a dolfinx.MeshTag for all cells. If midpoint is inside the
  # circle, it gets value 1, otherwise 0.
  num_cells = mesh.topology.index_map(mesh.topology.dim).size_local
  midpoints = dolfinx.cpp.mesh.midpoints(mesh, mesh.topology.dim, list(np.arange(num_cells, dtype=np.int32)))
  cell_tags = dolfinx.MeshTags(mesh, mesh.topology.dim, np.arange(num_cells), in_circle(midpoints))
  
  # As the dolfinx.MeshTag contains a value for every cell in the
  # geometry, we can attach it directly to the grid
  grid.cell_arrays["Marker"] = cell_tags.values
  grid.set_active_scalars("Marker")
  
  # We create a plotter consisting of two windows, and add a plot of the
  # Meshtags to the first 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 also visualize subsets of data, by creating a smaller topology,
  # only consisting of thos entities that has value one in the
  # dolfinx.MeshTag
  pyvista_cells, cell_types = dolfinx.plot.create_vtk_topology(
      mesh, mesh.topology.dim, cell_tags.indices[cell_tags.values == 1])
  
  # We add this grid to the second plotter
  sub_grid = pyvista.UnstructuredGrid(pyvista_cells, cell_types, mesh.geometry.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()
  
Plotting a dolfinx.Function
===========================

In the previous sections we have considered CG-1 spaces, which have a
1-1 correspondence with the vertices of the geometry. To be able to
plot higher order function spaces, both CG and DG spaces, we have to
adjust our plotting technique.

::

  # We start by projecting a discontinuous function into a second order DG
  # space Note that we use the `cell_tags` from the previous section to
  # restrict the integration domain on the RHS.
  dx = ufl.Measure("dx", subdomain_data=cell_tags)
  V = dolfinx.FunctionSpace(mesh, ("DG", 2))
  uh = dolfinx.Function(V)
  u = ufl.TrialFunction(V)
  v = ufl.TestFunction(V)
  x = ufl.SpatialCoordinate(mesh)
  a = ufl.inner(u, v) * dx
  L = ufl.inner(x[0], v) * dx(1) + ufl.inner(0.01, v) * dx(0)
  problem = dolfinx.fem.LinearProblem(a, L, u=uh, petsc_options={"ksp_type": "preonly", "pc_type": "lu"})
  problem.solve()
  
  # To get a topology that has a 1-1 correspondence with the degrees of
  # freedom in the function space, we call
  # `dolfinx.plot.create_vtk_topology`. We obtain the geometry for
  # the dofs owned on this process by tabulation of the dof coordinates.
  topology, cell_types = dolfinx.plot.create_vtk_topology(V)
  num_dofs_local = uh.function_space.dofmap.index_map.size_local
  geometry = uh.function_space.tabulate_dof_coordinates()[:num_dofs_local]
  
  # We discard the complex values if using PETSc in complex mode
  values = uh.vector.array.real if np.iscomplexobj(uh.vector.array) else uh.vector.array
  
  # We create a pyvista mesh from the topology and geometry, and attach
  # the coefficients of the degrees of freedom
  grid = pyvista.UnstructuredGrid(topology, cell_types, geometry)
  grid.point_arrays["DG"] = values
  grid.set_active_scalars("DG")
  
  # We would also like to visualize the underlying mesh and obtain that as
  # we have done previously
  num_cells = mesh.topology.index_map(mesh.topology.dim).size_local
  cell_entities = np.arange(num_cells, dtype=np.int32)
  pyvista_cells, cell_types = dolfinx.plot.create_vtk_topology(mesh, mesh.topology.dim, cell_entities)
  org_grid = pyvista.UnstructuredGrid(pyvista_cells, cell_types, mesh.geometry.x)
  
  
  # We visualize the data
  plotter = pyvista.Plotter()
  plotter.add_text("Second order 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()
  
  
Plotting a dolfinx.Function with vector values
===============================================

::

  # In the previous sections, we have considered how to plot scalar valued
  # functions. This section will show you how to plot vector valued
  # functions. We start by interpolating an expression into a second order
  # CG space.
  def vel(x):
      vals = np.zeros((2, x.shape[1]))
      vals[0] = np.sin(x[1])
      vals[1] = 0.1 * x[0]
      return vals
  
  
  mesh = dolfinx.UnitSquareMesh(MPI.COMM_WORLD, 6, 6, dolfinx.cpp.mesh.CellType.triangle)
  V = dolfinx.VectorFunctionSpace(mesh, ("CG", 2))
  uh = dolfinx.Function(V)
  uh.interpolate(vel)
  
  # We use the `dolfinx.plot.create_vtk_topology`
  # function, as in the previous section. However, we input a set of cell
  # entities, which can restrict the plotting to subsets of our mesh
  num_cells = mesh.topology.index_map(mesh.topology.dim).size_local
  cell_entities = np.arange(num_cells, dtype=np.int32)
  topology, cell_types = dolfinx.plot.create_vtk_topology(V, cell_entities)
  
  # As we deal with a vector function space, we need to adjust the values
  # in the underlying one dimensional array in dolfinx.Function, by
  # reshaping the data, and add an extra column to make it a 3D vector
  num_dofs_local = uh.function_space.dofmap.index_map.size_local
  geometry = uh.function_space.tabulate_dof_coordinates()[:num_dofs_local]
  values = np.zeros((V.dofmap.index_map.size_local, 3), dtype=np.float64)
  values[:, :mesh.geometry.dim] = uh.vector.array.real.reshape(V.dofmap.index_map.size_local, V.dofmap.index_map_bs)
  
  # Create a point cloud of glyphs
  function_grid = pyvista.UnstructuredGrid(topology, cell_types, geometry)
  function_grid["vectors"] = values
  function_grid.set_active_vectors("vectors")
  glyphs = function_grid.glyph(orient="vectors", factor=0.1)
  
  # To add streamlines in 2D, you need to supply the two points creating a
  # sampling line
  # NOTE: This crashes if ran in parallel, seems to be a pyvista issue
  # with pointa, pointb
  if MPI.COMM_WORLD.size == 0:
      streamlines = function_grid.streamlines(vectors="vectors", return_source=False,
                                              pointa=(0.5, 0.0, 0), pointb=(0.5, 1, 0))
  
  # Create pyvista mesh from the mesh
  pyvista_cells, cell_types = dolfinx.plot.create_vtk_topology(mesh, mesh.topology.dim, cell_entities)
  grid = pyvista.UnstructuredGrid(pyvista_cells, cell_types, mesh.geometry.x)
  
  # Add mesh, glyphs and streamlines to plotter
  plotter = pyvista.Plotter()
  plotter.add_text("Second order vector function.",
                   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=True, color="black", style="wireframe")
  plotter.add_mesh(glyphs)
  if MPI.COMM_WORLD.size == 0:
      plotter.add_mesh(streamlines.tube(radius=0.001), show_scalar_bar=False)
  plotter.view_xy()
  if pyvista.OFF_SCREEN:
      plotter.screenshot(f"vectors_{MPI.COMM_WORLD.rank}.png",
                         transparent_background=transparent, window_size=[figsize, figsize])
  else:
      plotter.show()
  
Plotting a dolfinx.Function with vector values
===============================================

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


::

  def vel(x):
      vals = np.zeros((3, x.shape[1]))
      vals[0] = -(x[1] - 0.5)
      vals[1] = x[0] - 0.5
      vals[2] = 0
      return vals
  
  
  mesh = dolfinx.UnitCubeMesh(MPI.COMM_WORLD, 4, 4, 4, dolfinx.cpp.mesh.CellType.hexahedron)
  V = dolfinx.VectorFunctionSpace(mesh, ("DG", 2))
  uh = dolfinx.Function(V)
  uh.interpolate(vel)
  
  num_cells = mesh.topology.index_map(mesh.topology.dim).size_local
  cell_entities = np.arange(num_cells, dtype=np.int32)
  
  topology, cell_types = dolfinx.plot.create_vtk_topology(V, cell_entities)
  num_dofs_local = uh.function_space.dofmap.index_map.size_local
  geometry = uh.function_space.tabulate_dof_coordinates()[:num_dofs_local]
  values = np.zeros((V.dofmap.index_map.size_local, 3), dtype=np.float64)
  values[:, :mesh.geometry.dim] = uh.vector.array.real.reshape(V.dofmap.index_map.size_local, V.dofmap.index_map_bs)
  
  # Create a point cloud of glyphs
  grid = pyvista.UnstructuredGrid(topology, cell_types, geometry)
  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("Vector function as streamlines.",
                   position="upper_edge", font_size=20, 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, 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()