diff --git a/CHANGELOG.rst b/CHANGELOG.rst index eba7e9bd..f8ad1b15 100644 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -11,6 +11,12 @@ adheres to `Semantic Versioning `_. `Unreleased`_ ------------- +Added +~~~~~~~ + +- Added a boolean flag ``plot_walls`` to toggle wall rendering in ``Room.plot``. +- Set the ``plot_walls`` to ``False`` for ``AnechoicRoom`` by default. + Changed ~~~~~~~ diff --git a/pyroomacoustics/room.py b/pyroomacoustics/room.py index e0076581..5b3c2215 100644 --- a/pyroomacoustics/room.py +++ b/pyroomacoustics/room.py @@ -1470,6 +1470,7 @@ def plot( mic_marker_size=10, plot_directivity=True, ax=None, + plot_walls=True, **kwargs, ): """Plots the room with its walls, microphones, sources and images""" @@ -1502,15 +1503,16 @@ def plot( ax = fig.add_subplot(111, aspect="equal", **kwargs) # draw room - corners = np.array([wall.corners[:, 0] for wall in self.walls]).T - polygons = [Polygon(xy=corners.T, closed=True)] - p = PatchCollection( - polygons, - cmap=matplotlib.cm.jet, - facecolor=np.array([1, 1, 1]), - edgecolor=np.array([0, 0, 0]), - ) - ax.add_collection(p) + if plot_walls: + corners = np.array([wall.corners[:, 0] for wall in self.walls]).T + polygons = [Polygon(xy=corners.T, closed=True)] + p = PatchCollection( + polygons, + cmap=matplotlib.cm.jet, + facecolor=np.array([1, 1, 1]), + edgecolor=np.array([0, 0, 0]), + ) + ax.add_collection(p) if self.mic_array is not None: for i in range(self.mic_array.nmic): @@ -1654,12 +1656,13 @@ def plot( fig.add_axes(ax) # plot the walls - for w in self.walls: - tri = a3.art3d.Poly3DCollection([w.corners.T], alpha=0.5) - rng = random.get_rng() - tri.set_color(colors.rgb2hex(rng.uniform(size=3))) - tri.set_edgecolor("k") - ax.add_collection3d(tri) + if plot_walls: + for w in self.walls: + tri = a3.art3d.Poly3DCollection([w.corners.T], alpha=0.5) + rng = random.get_rng() + tri.set_color(colors.rgb2hex(rng.uniform(size=3))) + tri.set_edgecolor("k") + ax.add_collection3d(tri) # define some markers for different sources and colormap for damping markers = ["o", "s", "v", "."] @@ -3061,4 +3064,4 @@ def plot(self, **kwargs): """Overloaded function to issue warning when img_order is given.""" if "img_order" in kwargs.keys(): warnings.warn("Ignoring img_order argument for AnechoicRoom.", UserWarning) - ShoeBox.plot(self, **kwargs) + return ShoeBox.plot(self, plot_walls=False, **kwargs) diff --git a/tests/test_room_plot.py b/tests/test_room_plot.py index 1456601e..6f0f18fa 100644 --- a/tests/test_room_plot.py +++ b/tests/test_room_plot.py @@ -1,19 +1,69 @@ +import unittest + import matplotlib +matplotlib.use("Agg") + +from matplotlib.collections import PatchCollection +from mpl_toolkits.mplot3d import art3d + import pyroomacoustics as pra -matplotlib.use("Agg") +def _has_wall_collection(ax): + return any( + isinstance(c, (PatchCollection, art3d.Poly3DCollection)) + for c in ax.collections + ) + + +class TestRoomPlot(unittest.TestCase): + def test_room_2d(self): + room = pra.ShoeBox([3, 4], max_order=2) + room.add_source([1.4, 2.2]) + room.add_microphone([2.3, 3.5]) + room.plot() + + def test_room_3d(self): + room = pra.ShoeBox([3, 4, 5], max_order=2) + room.add_source([1.4, 2.2, 4.3]) + room.add_microphone([2.3, 3.5, 2.7]) + room.plot() + + def test_plot_walls_2d(self): + room = pra.ShoeBox([3, 4], max_order=2) + room.add_source([1.4, 2.2]) + room.add_microphone([2.3, 3.5]) + + fig, ax = room.plot() + self.assertTrue(_has_wall_collection(ax)) + + fig, ax = room.plot(plot_walls=False) + self.assertFalse(_has_wall_collection(ax)) + + def test_plot_walls_3d(self): + room = pra.ShoeBox([3, 4, 5], max_order=2) + room.add_source([1.4, 2.2, 4.3]) + room.add_microphone([2.3, 3.5, 2.7]) + + fig, ax = room.plot() + self.assertTrue(_has_wall_collection(ax)) + + fig, ax = room.plot(plot_walls=False) + self.assertFalse(_has_wall_collection(ax)) + + def test_anechoic_room_no_walls_2d(self): + room = pra.AnechoicRoom(dim=2) + room.add_source([1.4, 2.2]) + room.add_microphone([2.3, 3.5]) -def test_room_2d(): - room = pra.ShoeBox([3, 4], max_order=2) - room.add_source([1.4, 2.2]) - room.add_microphone([2.3, 3.5]) - room.plot() + fig, ax = room.plot() + self.assertFalse(_has_wall_collection(ax)) + def test_anechoic_room_no_walls_3d(self): + room = pra.AnechoicRoom(dim=3) + room.add_source([1.4, 2.2, 4.3]) + room.add_microphone([2.3, 3.5, 2.7]) -def test_room_3d(): - room = pra.ShoeBox([3, 4, 5], max_order=2) - room.add_source([1.4, 2.2, 4.3]) - room.add_microphone([2.3, 3.5, 2.7]) - room.plot() + fig, ax = room.plot() + self.assertFalse(_has_wall_collection(ax))