Add column / row wise normalization to imshow and pyrshow

TESTS/unitTests.py

@@ -8,6 +8,8 @@
 import matplotlib.pyplot as plt
 
 import pyrtools as pt
+from pyrtools.pyramids.pyramid import Pyramid
+from pyrtools.tools.display import colormap_range
 
 import scipy.io
 import os
@@ -1599,7 +1601,91 @@ def test_pyrshow_2d_shape_err(self):
         with self.assertRaises(ValueError):
             pt.pyrshow(pyr.pyr_coeffs)
 
-
+class TestVrange(unittest.TestCase):
+
+    def _get_clims(self, fig):
+        """get vmin vmax for each image in fig as list of tuples (vmin, vmax)"""
+        return [ax.images[0].get_clim() for ax in fig.axes if ax.images]
+
+    def _get_title_clims(self, fig):
+        """get vmin vmax for each image in fig as list of tuples (vmin, vmax)"""
+        clims = []
+        for ax in fig.axes:
+            title = ax.get_title()
+            vmin, vmax  = title.split('[')[1].split(']')[0].split(',')
+            clims.append((float(vmin.strip()), float(vmax.strip())))
+        return clims
+
+    def _get_images(self):
+        # define test images such that each image has a different range of values,
+        # so we can test that the correct vrange is applied to each one
+        images = [np.arange(4 * i, 4 * i + 4, dtype=float).reshape(2, 2) for i in range(4)]
+        return images
+
+    def _imshow(self, vrange):
+        return pt.imshow(self._get_images(), vrange=vrange, zoom=1, col_wrap=2)
+
+    def _expected_clims(self, vrange):
+        clims, _ = colormap_range(image=self._get_images(), contains_rgb= [False]*len(self._get_images()), 
+                                  vrange=vrange, cmap=None, n_cols=2)
+        return clims
+
+    def test_vrange_max_vmin_matches_expected(self):
+        for mode in ['auto', 'indep', 'autoNrow', 'autoNcol','autoNcolcomplex']:
+            for mode_num in range(4):
+                for img_idx in range(4):
+                    with self.subTest(mode=mode, img_idx=img_idx):
+                        if 'N' in mode:
+                            mode_full = f"{mode.replace('N',str(mode_num))}"
+                        else:
+                            mode_full = f"{mode}{mode_num}"
+                        vmin, vmax = self._get_clims(self._imshow(mode_full))[img_idx]
+                        exp_vmin, exp_vmax = self._expected_clims(mode_full)[img_idx]
+                        self.assertTrue(np.isclose(vmax, exp_vmax, atol=1e-6))
+                        self.assertTrue(np.isclose(vmin, exp_vmin, atol=1e-6))
+
+    def test_vrange_title_matches_clim(self):
+        for mode in ['auto', 'indep', 'autoNrow', 'autoNcol','autoNcolcomplex']:
+            for mode_num in range(4):
+                for img_idx in range(4):
+                    with self.subTest(mode=mode, img_idx=img_idx):
+                        if 'N' in mode:
+                            mode_full = f"{mode.replace('N',str(mode_num))}"
+                        else:
+                            mode_full = f"{mode}{mode_num}"
+                        fig = self._imshow(mode_full)
+                        clim_vmin, clim_vmax = self._get_clims(fig)[img_idx]
+                        title_vmin, title_vmax = self._get_title_clims(fig)[img_idx]
+                        self.assertEqual("{:.1e}".format(clim_vmin), "{:.1e}".format(title_vmin))
+                        self.assertEqual("{:.1e}".format(clim_vmax), "{:.1e}".format(title_vmax))
+
+    def test_row_vrange_same_row_shares_clim(self):
+        for mode in range(4):
+            with self.subTest(mode=mode):
+                clims = self._get_clims(self._imshow(f"auto{mode}row"))
+                self.assertEqual(clims[0], clims[1], "row 0 images differ")
+                self.assertEqual(clims[2], clims[3], "row 1 images differ")
+
+    def test_col_vrange_same_col_shares_clim(self):
+        for mode in range(4):
+            with self.subTest(mode=mode):
+                clims = self._get_clims(self._imshow(f"auto{mode}col"))
+                self.assertEqual(clims[0], clims[2], "col 0 images differ")
+                self.assertEqual(clims[1], clims[3], "col 1 images differ")
+
+    def test_global_vrange_all_images_share_clim(self):
+        for mode in range(4):
+            with self.subTest(mode=mode):
+                clims = self._get_clims(self._imshow(f"auto{mode}"))
+                self.assertTrue(all(c == clims[0] for c in clims))
+
+    def test_indep_vrange_all_images_have_different_clims(self):
+        for mode in range(4):
+            with self.subTest(mode=mode):
+                clims = self._get_clims(self._imshow(f"indep{mode}"))
+                self.assertTrue(len(set(clims)) == len(clims),
+                                 "not all images have different clims")
+
 def main():
     unittest.main()