Repository: cta-observatory/ctapipe
Pull request #2369 (2023-06-26 14:38:39)

Allow passing of ax to Subarray.peek, do not always call tight_layout

kosack: What is the use case here?  Really, you should just use an ArrayDisplay directly to do any more fancy plotting, the peek() function is for debugging or a quick look at what is there, not to make useful figures.  One of the design considerations early on in ctapipe was to avoid having plotting functions inside classes (after lessons learned about deep circular dependencies), which I agree is sort of violated already with this and the atmosphere module... 
maxnoe: The main thing I wanted to address was that I got a large warning from matplotlib from the `.peek()` methd in a notebook where I had constrained layout enabled by default for all figures.

Then, I am of the strong opinion that any matplotlib related function should allow overriding in which axes the thing is plotted, so I added that.

Patch:
diff --git a/ctapipe/instrument/subarray.py b/ctapipe/instrument/subarray.py
index 3382245a..78aad41a 100644
--- a/ctapipe/instrument/subarray.py
+++ b/ctapipe/instrument/subarray.py
@@ -392,20 +392,39 @@ class SubarrayDescription:
         )
         return newsub
 
-    def peek(self):
+    def peek(self, ax=None):
         """
         Draw a quick matplotlib plot of the array
+
+        Parameters
+        ----------
+        ax : matplotlib.axes.Axes or None
+            If given, the subarray will be plotted into this ax.
         """
         from matplotlib import pyplot as plt
 
         from ctapipe.coordinates.ground_frames import EastingNorthingFrame
         from ctapipe.visualization import ArrayDisplay
 
-        plt.figure(figsize=(8, 8))
-        ad = ArrayDisplay(subarray=self, frame=EastingNorthingFrame(), tel_scale=0.75)
+        if ax is None:
+            fig = plt.figure(figsize=(8, 8))
+            ax = fig.add_subplot(1, 1, 1)
+        else:
+            fig = ax.figure
+
+        ad = ArrayDisplay(
+            subarray=self, frame=EastingNorthingFrame(), tel_scale=0.75, axes=ax
+        )
         ad.add_labels()
-        plt.title(self.name)
-        plt.tight_layout()
+        ax.set_title(self.name)
+
+        # This avoids a warning if users have e.g. "figure.constrained_layout"
+        # enabled through their matplotlib config. Only call tight layout here
+        # if the layout engine has not been set already
+        if fig.get_layout_engine() is None:
+            fig.tight_layout()
+
+        return ad
 
     @lazyproperty
     def telescope_types(self) -> Tuple[TelescopeDescription]:
diff --git a/docs/changes/2369.feature.rst b/docs/changes/2369.feature.rst
new file mode 100644
index 00000000..620120de
--- /dev/null
+++ b/docs/changes/2369.feature.rst
@@ -0,0 +1,2 @@
+Allow passing the matplotlib axes to the ``SubarrayDescription.peek`` function,
+fix warnings in case of layout engine being already defined.


Test patch:
diff --git a/ctapipe/visualization/tests/test_mpl.py b/ctapipe/visualization/tests/test_mpl.py
index 4cc09835..70776678 100644
--- a/ctapipe/visualization/tests/test_mpl.py
+++ b/ctapipe/visualization/tests/test_mpl.py
@@ -377,3 +377,14 @@ def test_overlay_coord(tmp_path, subarray_and_event_gamma_off_axis_500_gev):
     )
 
     fig.savefig(tmp_path / "coord_overlay.png", dpi=300)
+
+
+@pytest.mark.parametrize("layout", (None, "constrained"))
+def test_array_display_axes(tmp_path, subarray_prod5_paranal, layout):
+    """Test passing axes to peek"""
+    fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(16, 8), layout=layout)
+
+    subarray_prod5_paranal.peek(ax=ax1)
+    subarray_prod5_paranal.peek(ax=ax2)
+
+    fig.savefig(tmp_path / "double_peek.png")