Add BoundaryDict (#915)

* Add `BoundaryDict`

 with methods to plot, screenshot and imshow

* Update test_dof.py

* Update test_dof.py

* Ensure 3d-data in `Boundary.plot()` for PyVista

* Update test_dof.py

* Add `Boundary.plot(..., show_points=True, show_lines=True)`

CHANGELOG.md

@@ -8,6 +8,7 @@ All notable changes to this project will be documented in this file. The format
 - Add `SolidBody.evaluate.stress(field)` to evaluate the (first Piola-Kirchhoff) stress tensor (engineering stress in linear elasticity).
 - Add a free-vibration modal analysis Step/Job `FreeVibration(items, boundaries)` with methods to evaluate `FreeVibration.evaluate()` and to extract `field, frequency = FreeVibration.extract(n)` its n-th result.
 - Add `Boundary.plot()` to plot the points and prescribed directions of a boundary.
+- Add `BoundaryDict` as a subclassed dict with methods to `plot()`, `screenshot()` and `imshow()`.
 
 ### Changed
 - The first Piola-Kirchhoff stress tensor is evaluated if `ViewSolid(stress_type=None)`.

docs/felupe/dof.rst

@@ -10,6 +10,7 @@ This module contains the definition of a boundary condition, tools related to th
 .. autosummary::
 
    Boundary
+   BoundaryDict
 
 
 **Tools**
@@ -37,6 +38,11 @@ This module contains the definition of a boundary condition, tools related to th
    :undoc-members:
    :inherited-members:
 
+.. autoclass:: felupe.BoundaryDict
+   :members:
+   :undoc-members:
+   :inherited-members:
+
 .. autofunction:: felupe.dof.partition
 
 .. autofunction:: felupe.dof.apply

docs/tutorial/examples/extut03_building_blocks.py

@@ -171,14 +171,13 @@ def W(C, mu, bulk):
 f0 = lambda x: np.isclose(x, 0)
 f1 = lambda x: np.isclose(x, 1)
 
-boundaries = {}
+boundaries = fem.BoundaryDict()
+
 boundaries["left"] = fem.Boundary(displacement, fx=f0)
 boundaries["right"] = fem.Boundary(displacement, fx=f1, skip=(1, 0, 0))
 boundaries["move"] = fem.Boundary(displacement, fx=f1, skip=(0, 1, 1), value=0.5)
 
-plotter = boundaries["left"].plot(color="green")
-plotter = boundaries["right"].plot(plotter=plotter, color="red")
-boundaries["move"].plot(plotter=plotter, color="blue", point_size=5).show()
+boundaries.plot().show()
 
 # %%
 # Partition of deegrees of freedom

src/felupe/__init__.py

@@ -40,7 +40,7 @@
     linear_elastic,
     linear_elastic_plastic_isotropic_hardening,
 )
-from .dof import Boundary
+from .dof import Boundary, BoundaryDict
 from .element import ArbitraryOrderLagrange as ArbitraryOrderLagrangeElement
 from .element import (
     BiQuadraticQuad,
@@ -168,6 +168,7 @@
     "linear_elastic",
     "linear_elastic_plastic_isotropic_hardening",
     "Boundary",
+    "BoundaryDict",
     "ArbitraryOrderLagrangeElement",
     "BiQuadraticQuad",
     "ConstantHexahedron",

src/felupe/dof/__init__.py

@@ -1,9 +1,11 @@
 from ._boundary import Boundary
+from ._dict import BoundaryDict
 from ._loadcase import biaxial, shear, symmetry, uniaxial
 from ._tools import apply, get_dof0, get_dof1, partition
 
 __all__ = [
     "Boundary",
+    "BoundaryDict",
     "biaxial",
     "shear",
     "symmetry",

src/felupe/dof/_boundary.py

@@ -263,7 +263,16 @@ def update(self, value):
         self.value = value  #
 
     def plot(
-        self, plotter=None, color="black", scale=0.125, point_size=10, width=3, **kwargs
+        self,
+        plotter=None,
+        color="black",
+        scale=0.125,
+        point_size=10,
+        width=3,
+        label=None,
+        show_points=True,
+        show_lines=True,
+        **kwargs,
     ):
         "Plot the points and their prescribed directions of a boundary condition."
 
@@ -283,24 +292,32 @@ def plot(
             )
 
         if len(self.points) > 0:
-            magnitude = min(mesh.points.max(axis=0) - mesh.points.min(axis=0)) * scale
+            if label is None:
+                label = self.name
+
+            if show_points or show_lines:
+                points = np.pad(mesh.points, ((0, 0), (0, 3 - mesh.dim)))
+
+            if show_points:
+                _ = plotter.add_points(
+                    points[self.points],
+                    color=color,
+                    point_size=point_size,
+                    label=label,
+                )
 
-            _ = plotter.add_points(
-                mesh.points[self.points],
-                color=color,
-                point_size=point_size,
-                label=self.name,
-            )
+            if show_lines:
+                magnitude = (
+                    min(mesh.points.max(axis=0) - mesh.points.min(axis=0)) * scale
+                )
 
-            for skip, direction in zip(self.skip, np.eye(mesh.dim)):
-                if not skip:
-                    end = mesh.points[self.points] + direction * magnitude
-                    _ = plotter.add_lines(
-                        np.hstack([mesh.points[self.points], end]).reshape(
-                            -1, mesh.dim
-                        ),
-                        color=color,
-                        width=width,
-                    )
+                for skip, direction in zip(self.skip, np.eye(3)):
+                    if not skip:
+                        end = points[self.points] + direction * magnitude
+                        _ = plotter.add_lines(
+                            np.hstack([points[self.points], end]).reshape(-1, 3),
+                            color=color,
+                            width=width,
+                        )
 
         return plotter

src/felupe/dof/_dict.py

@@ -0,0 +1,79 @@
+class BoundaryDict(dict):
+    "A dict of boundary conditions."
+
+    def plot(
+        self,
+        plotter=None,
+        colors=None,
+        size=(0.1, 0.1),
+        show_points=True,
+        show_lines=True,
+        **kwargs,
+    ):
+        "Plot the boundary conditions."
+
+        if colors is None:
+            import matplotlib.colors as mcolors
+
+            colors = list(mcolors.TABLEAU_COLORS.values())
+
+        colors_list = []
+        while len(colors_list) < len(self.keys()):
+            colors_list = [*colors_list, *colors]
+
+        for (key, boundary), color in zip(self.items(), colors_list):
+            label = key
+            if boundary.name != "default":
+                label = boundary.name
+
+            plotter = boundary.plot(
+                label=label,
+                color=color,
+                plotter=plotter,
+                show_points=show_points,
+                show_lines=show_lines,
+                **kwargs,
+            )
+
+        plotter.add_legend(size=size, bcolor="white")
+
+        return plotter
+
+    def screenshot(
+        self,
+        *args,
+        filename="boundaries.png",
+        transparent_background=None,
+        scale=None,
+        colors=None,
+        plotter=None,
+        **kwargs,
+    ):
+        "Take a screenshot of the boundary conditions."
+
+        if plotter is None:
+            mesh = self[list(self.keys())[0]].field.region.mesh
+            plotter = mesh.plot(off_screen=True)
+
+        plotter = self.plot(plotter=plotter, colors=colors, **kwargs)
+
+        return plotter.screenshot(
+            filename=filename,
+            transparent_background=transparent_background,
+            scale=scale,
+        )
+
+    def imshow(self, *args, ax=None, dpi=None, **kwargs):
+        """Take a screenshot of the boundary conditions, show the image data in
+        a figure and return the ax.
+        """
+
+        if ax is None:
+            import matplotlib.pyplot as plt
+
+            fig, ax = plt.subplots(dpi=dpi)
+
+        ax.imshow(self.screenshot(*args, filename=None, **kwargs))
+        ax.set_axis_off()
+
+        return ax

tests/test_dof.py

@@ -180,8 +180,22 @@ def test_boundary_plot():
         _ = boundaries["left"].plot()
 
 
+def test_boundary_dict():
+    field = fem.FieldPlaneStrain(
+        fem.RegionQuad(fem.Rectangle(b=(3, 1), n=5)), dim=2
+    ).as_container()
+    boundaries = fem.BoundaryDict(
+        left=fem.Boundary(field[0], fx=0, skip=(0, 0)),
+        right=fem.Boundary(field[0], name="my_label", fx=3, skip=(0, 1)),
+    )
+    plotter = boundaries.plot()
+    # img = boundaries.screenshot()
+    # ax = boundaries.imshow()
+
+
 if __name__ == "__main__":
     test_boundary()
+    test_boundary_dict()
     test_boundary_multiaxial()
     test_boundary_plot()
     test_loadcase()