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

adtzlr · Dec 2, 2024 · 11d5233 · 11d5233
1 parent 6893a50
commit 11d5233
Show file tree

Hide file tree

Showing 2 changed files with 42 additions and 21 deletions.
diff --git a/src/felupe/dof/_boundary.py b/src/felupe/dof/_boundary.py
@@ -270,6 +270,8 @@ def plot(
         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."
@@ -293,25 +295,29 @@ def plot(
             if label is None:
                 label = self.name
 
-            points = np.pad(mesh.points, ((0, 0), (0, 3 - mesh.dim)))
-            magnitude = min(mesh.points.max(axis=0) - mesh.points.min(axis=0)) * scale
+            if show_points or show_lines:
+                points = np.pad(mesh.points, ((0, 0), (0, 3 - mesh.dim)))
 
-            _ = plotter.add_points(
-                points[self.points],
-                color=color,
-                point_size=point_size,
-                label=label,
-            )
+            if show_points:
+                _ = plotter.add_points(
+                    points[self.points],
+                    color=color,
+                    point_size=point_size,
+                    label=label,
+                )
+
+            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(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,
-                    )
+                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
diff --git a/src/felupe/dof/_dict.py b/src/felupe/dof/_dict.py
@@ -1,7 +1,15 @@
 class BoundaryDict(dict):
     "A dict of boundary conditions."
 
-    def plot(self, plotter=None, colors=None, **kwargs):
+    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:
@@ -18,9 +26,16 @@ def plot(self, plotter=None, colors=None, **kwargs):
             if boundary.name != "default":
                 label = boundary.name
 
-            plotter = boundary.plot(label=label, color=color, plotter=plotter, **kwargs)
+            plotter = boundary.plot(
+                label=label,
+                color=color,
+                plotter=plotter,
+                show_points=show_points,
+                show_lines=show_lines,
+                **kwargs,
+            )
 
-        plotter.add_legend(size=(0.1, 0.1), bcolor="white", face="rectangle")
+        plotter.add_legend(size=size, bcolor="white")
 
         return plotter