Add parallel plotting in single window

.gitignore

@@ -9,4 +9,5 @@ _build
 *.png
 *.pvtu
 *.msh
-*.bp
+*.bp
+*.pyc

chapter1/complex_mode.py

@@ -180,6 +180,10 @@
 # Finally, we plot the real and imaginary solutions.
 #
 
+# Each process builds the grid of the cells it owns, the pieces are gathered on one process and
+# merged into a single grid, which is then drawn, see
+# [Plotting in parallel](./fundamentals_code).
+
 # +
 import pyvista
 
@@ -188,19 +192,34 @@
 grid = pyvista.UnstructuredGrid(pyvista_cells, cell_types, geometry)
 grid.point_data["u_real"] = uh.x.array.real
 grid.point_data["u_imag"] = uh.x.array.imag
-_ = grid.set_active_scalars("u_real")
-
-p_real = pyvista.Plotter()
-p_real.add_text("uh real", position="upper_edge", font_size=14, color="black")
-p_real.add_mesh(grid, show_edges=True)
-p_real.view_xy()
-if not pyvista.OFF_SCREEN:
-    p_real.show()
-
-grid.set_active_scalars("u_imag")
-p_imag = pyvista.Plotter()
-p_imag.add_text("uh imag", position="upper_edge", font_size=14, color="black")
-p_imag.add_mesh(grid, show_edges=True)
-p_imag.view_xy()
-if not pyvista.OFF_SCREEN:
-    p_imag.show()
+
+root = 0
+assert root < mesh.comm.size, f"Cannot gather on process {root} of {mesh.comm.size}"
+
+clim_real = [
+    mesh.comm.reduce(uh.x.array.real.min(), op=MPI.MIN, root=root),
+    mesh.comm.reduce(uh.x.array.real.max(), op=MPI.MAX, root=root),
+]
+clim_imag = [
+    mesh.comm.reduce(uh.x.array.imag.min(), op=MPI.MIN, root=root),
+    mesh.comm.reduce(uh.x.array.imag.max(), op=MPI.MAX, root=root),
+]
+pieces = mesh.comm.gather(grid, root=root)
+# -
+
+if pieces is not None:
+    merged_grid = pyvista.merge(pieces)
+
+    p_real = pyvista.Plotter()
+    p_real.add_text("uh real", position="upper_edge", font_size=14, color="black")
+    p_real.add_mesh(merged_grid, scalars="u_real", show_edges=True, clim=clim_real)
+    p_real.view_xy()
+    if not pyvista.OFF_SCREEN:
+        p_real.show()
+
+    p_imag = pyvista.Plotter()
+    p_imag.add_text("uh imag", position="upper_edge", font_size=14, color="black")
+    p_imag.add_mesh(merged_grid, scalars="u_imag", show_edges=True, clim=clim_imag)
+    p_imag.view_xy()
+    if not pyvista.OFF_SCREEN:
+        p_imag.show()

chapter1/fundamentals_code.py

@@ -317,13 +317,46 @@
 import pyvista
 
 print(pyvista.global_theme.jupyter_backend)
+# -
+
+# ### Plotting in parallel
+# When the program runs on several processes, each process holds one piece of the mesh, and knows
+# nothing about the rest of it, so no process can draw the whole figure on its own.
+# We therefore let each process build the grid of its own piece, send the pieces to a single
+# process with {py:meth}`gather<mpi4py.MPI.Comm.gather>`, and let that process join them into a
+# single grid with {py:func}`pyvista.merge`.
+# {py:func}`dolfinx.plot.vtk_mesh` uses only the cells that the process *owns*, so a cell that is
+# shared between two processes is not drawn twice.
+# `merge` welds the points that the pieces have in common back together, so we get a single mesh
+# that is drawn with a single color bar, and that filters can be applied to as a whole.
+#
+# ```{admonition} Points on a partition boundary
+# :class: tip
+# The processes on either side of a partition boundary compute the coordinates of the points they
+# share independently, so these can differ in the last bits, and `merge` then keeps both copies.
+# The picture is the same either way, but a filter that follows the field from one cell into the
+# next, such as {py:meth}`streamlines<pyvista.DataSetFilters.streamlines>`, would stop at such a
+# point. Merge with a tolerance to weld them as well:
+# `pieces[0].merge(pieces[1:], tolerance=1e-14)`.
+# ```
+#
+# The process that collects the pieces is chosen with `root`, which has to be one of the processes
+# we run on.
+# `gather` returns the list of pieces on that process, and `None` on every other process, so the
+# returned value tells each process whether it is the one that draws.
+# The same code runs on a single process, where the list holds a single piece.
 
 # +
 from dolfinx import plot
 
 domain.topology.create_connectivity(tdim, tdim)
 topology, cell_types, geometry = plot.vtk_mesh(domain, tdim)
 grid = pyvista.UnstructuredGrid(topology, cell_types, geometry)
+
+root = 0
+assert root < domain.comm.size, f"Cannot gather on process {root} of {domain.comm.size}"
+
+pieces = domain.comm.gather(grid, root=root)
 # -
 
 # There are several backends that can be used with pyvista, and they have different benefits and drawbacks.
@@ -334,13 +367,15 @@
 # In the jupyter notebook environment, we use the default setting of `pyvista.OFF_SCREEN=False`,
 # which will render plots directly in the notebook.
 
-plotter = pyvista.Plotter()
-plotter.add_mesh(grid, show_edges=True)
-plotter.view_xy()
-if not pyvista.OFF_SCREEN:
-    plotter.show()
-else:
-    figure = plotter.screenshot("fundamentals_mesh.png")
+if pieces is not None:
+    merged_grid = pyvista.merge(pieces)
+    plotter = pyvista.Plotter()
+    plotter.add_mesh(merged_grid, show_edges=True)
+    plotter.view_xy()
+    if not pyvista.OFF_SCREEN:
+        plotter.show()
+    else:
+        figure = plotter.screenshot("fundamentals_mesh.png")
 
 # ## Plotting a function using pyvista
 # We want to plot the solution `uh`.
@@ -352,23 +387,47 @@
 u_topology, u_cell_types, u_geometry = plot.vtk_mesh(V)
 
 # Next, we create the {py:class}`pyvista.UnstructuredGrid` and add the dof-values to the mesh.
+# The values are attached before the grids are gathered, so that each piece carries its own data
+# into the merged grid.
+# We also compute the range of `uh` over all processes with
+# {py:meth}`reduce<mpi4py.MPI.Comm.reduce>`, and pass it to the plotter as `clim`.
+# Only the root process draws, so the result is only needed there, and `reduce` leaves it as `None`
+# on the other processes.
 
+# +
 u_grid = pyvista.UnstructuredGrid(u_topology, u_cell_types, u_geometry)
 u_grid.point_data["u"] = uh.x.array.real
 u_grid.set_active_scalars("u")
-u_plotter = pyvista.Plotter()
-u_plotter.add_mesh(u_grid, show_edges=True)
-u_plotter.view_xy()
-if not pyvista.OFF_SCREEN:
-    u_plotter.show()
 
-# We can also warp the mesh by scalar to make use of the 3D plotting.
+u_clim = [
+    domain.comm.reduce(uh.x.array.real.min(), op=MPI.MIN, root=root),
+    domain.comm.reduce(uh.x.array.real.max(), op=MPI.MAX, root=root),
+]
+u_pieces = domain.comm.gather(u_grid, root=root)
+# -
+
+if u_pieces is not None:
+    merged_u_grid = pyvista.merge(u_pieces)
+    u_plotter = pyvista.Plotter()
+    u_plotter.add_mesh(merged_u_grid, show_edges=True, clim=u_clim)
+    u_plotter.view_xy()
+    if not pyvista.OFF_SCREEN:
+        u_plotter.show()
 
-warped = u_grid.warp_by_scalar()
-plotter2 = pyvista.Plotter()
-plotter2.add_mesh(warped, show_edges=True, show_scalar_bar=True)
-if not pyvista.OFF_SCREEN:
-    plotter2.show()
+# We can also warp the mesh by scalar to make use of the 3D plotting.
+# As the grid has already been merged, filters such as
+# {py:meth}`warp_by_scalar<pyvista.DataSetFilters.warp_by_scalar>` are applied to the whole mesh.
+
+if u_pieces is not None:
+    plotter2 = pyvista.Plotter()
+    plotter2.add_mesh(
+        merged_u_grid.warp_by_scalar(),
+        show_edges=True,
+        show_scalar_bar=True,
+        clim=u_clim,
+    )
+    if not pyvista.OFF_SCREEN:
+        plotter2.show()
 
 # ## External post-processing
 # For post-processing outside the python code, it is suggested to save the solution to file using either

chapter1/membrane_code.py

@@ -180,7 +180,10 @@ def on_boundary(x):
 from dolfinx.plot import vtk_mesh
 import pyvista
 
-# Extract topology from mesh and create {py:class}`pyvista.UnstructuredGrid`
+# Extract topology from mesh and create {py:class}`pyvista.UnstructuredGrid`.
+# Each process builds the grid of the cells it owns, the pieces are gathered on one process and
+# merged into a single grid, which is then drawn, see
+# [Plotting in parallel](./fundamentals_code).
 
 topology, cell_types, x = vtk_mesh(V)
 grid = pyvista.UnstructuredGrid(topology, cell_types, x)
@@ -189,29 +192,59 @@ def on_boundary(x):
 
 # +
 grid.point_data["u"] = uh.x.array
-warped = grid.warp_by_scalar("u", factor=25)
-
-plotter = pyvista.Plotter()
-plotter.add_mesh(warped, show_edges=True, show_scalar_bar=True, scalars="u")
-if not pyvista.OFF_SCREEN:
-    plotter.show()
-else:
-    plotter.screenshot("deflection.png")
+
+root = 0
+assert root < domain.comm.size, f"Cannot gather on process {root} of {domain.comm.size}"
+
+u_clim = [
+    domain.comm.reduce(uh.x.array.real.min(), op=MPI.MIN, root=root),
+    domain.comm.reduce(uh.x.array.real.max(), op=MPI.MAX, root=root),
+]
+u_pieces = domain.comm.gather(grid, root=root)
+
+if u_pieces is not None:
+    merged_grid = pyvista.merge(u_pieces)
+    plotter = pyvista.Plotter()
+    plotter.add_mesh(
+        merged_grid.warp_by_scalar("u", factor=25),
+        show_edges=True,
+        show_scalar_bar=True,
+        scalars="u",
+        clim=u_clim,
+    )
+    if not pyvista.OFF_SCREEN:
+        plotter.show()
+    else:
+        plotter.screenshot("deflection.png")
 # -
 
 # We next plot the load on the domain
 
-load_plotter = pyvista.Plotter()
+# +
 p_grid = pyvista.UnstructuredGrid(*vtk_mesh(Q))
 p_grid.point_data["p"] = pressure.x.array.real
-warped_p = p_grid.warp_by_scalar("p", factor=0.5)
-warped_p.set_active_scalars("p")
-load_plotter.add_mesh(warped_p, show_scalar_bar=True)
-load_plotter.view_xy()
-if not pyvista.OFF_SCREEN:
-    load_plotter.show()
-else:
-    load_plotter.screenshot("load.png")
+
+p_clim = [
+    domain.comm.reduce(pressure.x.array.real.min(), op=MPI.MIN, root=root),
+    domain.comm.reduce(pressure.x.array.real.max(), op=MPI.MAX, root=root),
+]
+p_pieces = domain.comm.gather(p_grid, root=root)
+
+if p_pieces is not None:
+    merged_p_grid = pyvista.merge(p_pieces)
+    load_plotter = pyvista.Plotter()
+    load_plotter.add_mesh(
+        merged_p_grid.warp_by_scalar("p", factor=0.5),
+        show_scalar_bar=True,
+        scalars="p",
+        clim=p_clim,
+    )
+    load_plotter.view_xy()
+    if not pyvista.OFF_SCREEN:
+        load_plotter.show()
+    else:
+        load_plotter.screenshot("load.png")
+# -
 
 # ## Making curve plots throughout the domain
 # Another way to compare the deflection and the load is to make a plot along the line $x=0$.

chapter1/nitsche.py

@@ -127,19 +127,35 @@
 # we no longer fullfill the equation to machine precision at the mesh vertices.
 # We also plot the solution using {py:mod}`pyvista`
 
+# Each process builds the grid of the cells it owns, the pieces are gathered on one process and
+# merged into a single grid, which is then drawn, see
+# [Plotting in parallel](./fundamentals_code).
+
 # +
 import pyvista
 
 grid = pyvista.UnstructuredGrid(*plot.vtk_mesh(V))
 grid.point_data["u"] = uh.x.array.real
 grid.set_active_scalars("u")
-plotter = pyvista.Plotter()
-plotter.add_mesh(grid, show_edges=True, show_scalar_bar=True)
-plotter.view_xy()
-if not pyvista.OFF_SCREEN:
-    plotter.show()
-else:
-    figure = plotter.screenshot("nitsche.png")
+
+root = 0
+assert root < domain.comm.size, f"Cannot gather on process {root} of {domain.comm.size}"
+
+clim = [
+    domain.comm.reduce(uh.x.array.real.min(), op=MPI.MIN, root=root),
+    domain.comm.reduce(uh.x.array.real.max(), op=MPI.MAX, root=root),
+]
+pieces = domain.comm.gather(grid, root=root)
+
+if pieces is not None:
+    grid = pyvista.merge(pieces)
+    plotter = pyvista.Plotter()
+    plotter.add_mesh(grid, show_edges=True, show_scalar_bar=True, clim=clim)
+    plotter.view_xy()
+    if not pyvista.OFF_SCREEN:
+        plotter.show()
+    else:
+        figure = plotter.screenshot("nitsche.png")
 # -
 
 # ```{bibliography}