Add MeshContainer.plot(), img = MeshContainer.screenshot() and ax = MeshContainer.imshow()

CHANGELOG.md

@@ -3,6 +3,9 @@ All notable changes to this project will be documented in this file. The format
 
 ## [Unreleased]
 
+### Added
+- Add `MeshContainer.plot()`, `img = MeshContainer.screenshot()` and `ax = MeshContainer.imshow()`. The default list of colors contains PyVista's default color as first item and then the list of matplotlib's named colors *C1*, *C2*, etc (excluding *C0*).
+
 ## [7.15.0] - 2024-02-11
 
 ### Added

src/felupe/mesh/_container.py

@@ -154,6 +154,93 @@ def copy(self):
         "Return a deepcopy of the mesh container."
         return deepcopy(self)
 
+    def plot(self, *args, colors=None, **kwargs):
+        """Plot the meshes of the mesh container.
+
+        See Also
+        --------
+        felupe.Scene.plot: Plot method of a scene.
+        """
+
+        if colors is None:
+            import pyvista
+            import matplotlib.colors as mcolors
+
+            colors = [
+                pyvista.global_theme.color,
+                *list(mcolors.TABLEAU_COLORS.values())[1:],
+            ]
+
+        plotter = None
+        for mesh, color in zip(self.meshes, colors):
+            plotter = mesh.view().plot(
+                *args,
+                show_undeformed=False,
+                color=color,
+                plotter=plotter,
+                opacity=0.99,
+                **kwargs,
+            )
+
+        return plotter
+
+    def screenshot(
+        self,
+        *args,
+        filename="mesh.png",
+        transparent_background=None,
+        scale=None,
+        colors=None,
+        **kwargs,
+    ):
+        """Take a screenshot of the meshes of the mesh container.
+
+        See Also
+        --------
+        pyvista.Plotter.screenshot: Take a screenshot of a PyVista plotter.
+        """
+
+        if colors is None:
+            import pyvista
+            import matplotlib.colors as mcolors
+
+            colors = [
+                pyvista.global_theme.color,
+                *list(mcolors.TABLEAU_COLORS.values())[1:],
+            ]
+
+        plotter = None
+        for mesh, color in zip(self.meshes, colors):
+            plotter = mesh.plot(
+                *args,
+                off_screen=True,
+                color=color,
+                plotter=plotter,
+                opacity=0.99,
+                **kwargs,
+            )
+
+        return plotter.screenshot(
+            filename=filename,
+            transparent_background=transparent_background,
+            scale=scale,
+        )
+
+    def imshow(self, ax=None, *args, **kwargs):
+        """Take a screenshot of the meshes of the mesh container, show the image data in
+        a figure and return the ax.
+        """
+
+        if ax is None:
+            import matplotlib.pyplot as plt
+
+            fig, ax = plt.subplots()
+
+        ax.imshow(self.screenshot(*args, filename=None, **kwargs))
+        ax.set_axis_off()
+
+        return ax
+
     def __iadd__(self, mesh):
         self.append(mesh)
         return self

src/felupe/mesh/_mesh.py

@@ -559,6 +559,9 @@ def sweep(self, decimals=None):
           Number of cells:
             quad: 100
             quad: 100
+
+        ..  image:: images/mesh_container.png
+            :width: 400px
 
         The meshes of the mesh container are :func:`stacked <felupe.mesh.stack>`.
 
@@ -573,10 +576,16 @@ def sweep(self, decimals=None):
         the number of cells is unchanged.
 
         >>> mesh = fem.mesh.sweep(stack)
+        >>> mesh
         <felupe Mesh object>
           Number of points: 220
           Number of cells:
             quad: 200
+
+        >>> ax = mesh.imshow(opacity=0.6)
+
+        ..  image:: images/mesh_sweep.png
+            :width: 400px
 
         ..  note::
             The :class:`~felupe.MeshContainer` may be directly created with

src/felupe/mesh/_tools.py

@@ -373,10 +373,16 @@ def sweep(points, cells, cell_type, decimals=None):
     the number of cells is unchanged.
 
     >>> mesh = fem.mesh.sweep(stack)
+    >>> mesh
     <felupe Mesh object>
       Number of points: 220
       Number of cells:
         quad: 200
+
+    >>> ax = mesh.imshow(opacity=0.6)
+
+    ..  image:: images/mesh_sweep.png
+        :width: 400px
 
     ..  note::
         The :class:`~felupe.MeshContainer` may be directly created with ``merge=True``.

tests/test_mesh.py

@@ -490,6 +490,12 @@ def test_view():
     # img = mesh.screenshot(transparent_background=True)
     # ax = mesh.imshow()
 
+    mesh2 = mesh.translate(move=0.8, axis=0)
+    container = fem.MeshContainer([mesh, mesh2])
+    plotter = container.plot(off_screen=True)
+    # img = container.screenshot(transparent_background=True)
+    # ax = container.imshow()
+
 
 def test_mesh_update():
     mesh = fem.Cube(n=11)