Fix imshow set ax

TESTS/unitTests.py

@@ -1467,6 +1467,18 @@ def test_imshow25(self):
             fig = pt.imshow(im, vrange="auto1")
         _check_ranges(fig, ranges)
 
+    def test_imshow26(self):
+        im = np.random.randn(32, 32)
+        fig, ax = plt.subplots(1, 1)
+        fig = pt.imshow(im, ax=ax)
+        plt.close(fig)
+
+    def test_imshow27(self):
+        im = np.random.randn(32, 32)
+        fig, ax = plt.subplots(1, 1)
+        with self.assertRaises(ValueError, msg="Can only display one image"):
+            pt.imshow([im, im], ax=ax)
+
 class TestAnimshow(unittest.TestCase):
 
     def test_animshow0(self):
@@ -1624,6 +1636,18 @@ def test_animshow19(self):
         fig = pt.animshow([vid, 2*vid, 3*vid], vrange="auto1col", col_wrap=2, as_html5=False)._fig
         _check_ranges(fig, ranges)
 
+    def test_animshow20(self):
+        vid = np.random.randn(3, 32, 32)
+        fig, ax = plt.subplots(1, 1)
+        fig = pt.animshow(vid, ax=ax, as_html5=False)._fig
+        plt.close(fig)
+
+    def test_animshow21(self):
+        vid = np.random.randn(3, 32, 32)
+        fig, ax = plt.subplots(1, 1)
+        with self.assertRaises(ValueError, msg="Can only display one image"):
+            pt.animshow([vid, vid], ax=ax, as_html5=False)._fig
+
 
 class TestPyrshow(unittest.TestCase):
 

src/pyrtools/tools/display.py

@@ -705,7 +705,10 @@ def _setup_figure(ax, col_wrap, image, zoom, max_shape, vert_pct):
         axes = fig.axes
     else:
         fig = ax.figure
+        if len(image) > 1:
+            raise ValueError("Can only display one image if passing an axis!")
         axes = [reshape_axis(ax,  zoom * max_shape)]
+        n_cols = 1
     return fig, axes, n_cols
 
 
@@ -1039,6 +1042,9 @@ def animate_video(t):
 def pyrshow(pyr_coeffs, vrange='auto1pyr', col_wrap=None, zoom=1, show_residuals=True, **kwargs):
     """Display the coefficients of the pyramid in an orderly fashion
 
+    .. versionchanged:: 1.1
+       Removed ``"is_complex"`` argument -- complexity of input is now determined automatically.
+
     Arguments
     ---------
     pyr_coeffs : `dict`