Merge pull request #2343 from astrofrog/speedup-composite-array

Performance improvements for CompositeArray

.github/workflows/ci_workflows.yml

@@ -58,7 +58,6 @@ jobs:
         # Test a few configurations on macOS
         - macos: py38-test-pyqt514-all
         - macos: py310-test-pyqt515
-        - macos: py310-test-pyside63
         - macos: py310-test-pyqt64
         - macos: py311-test-pyqt515
 
@@ -92,9 +91,10 @@ jobs:
         # Python 3.11.0 failing on Windows in test_image.py on
         # > assert df.find_factory(fname) is df.img_data
         - linux: py310-test-pyside64
-        - windows: py310-test-pyside64
         - linux: py311-test-pyside64
+        - macos: py310-test-pyside63
         - macos: py311-test-pyside64
+        - windows: py310-test-pyside64
         - windows: py311-test-pyqt515
 
         # Windows docs build

glue/plugins/tools/pv_slicer/qt/tests/test_pv_slicer.py

@@ -77,7 +77,7 @@ def test_set_cmap(self):
         cm_mode = self.w.toolbar.tools['image:colormap']
         act = cm_mode.menu_actions()[1]
         act.trigger()
-        assert self.w._composite.layers['image']['color'] is act.cmap
+        assert self.w._composite.layers['image']['cmap'] is act.cmap
 
     def test_double_set_image(self):
         assert len(self.w._axes.images) == 1

glue/viewers/image/composite_array.py

@@ -1,8 +1,12 @@
 # This artist can be used to deal with the sampling of the data as well as any
 # RGB blending.
 
+import warnings
+
 import numpy as np
 
+from glue.config import colormaps
+
 from matplotlib.colors import ColorConverter, Colormap
 from astropy.visualization import (LinearStretch, SqrtStretch, AsinhStretch,
                                    LogStretch, ManualInterval, ContrastBiasStretch)
@@ -12,6 +16,8 @@
 
 COLOR_CONVERTER = ColorConverter()
 
+CMAP_SAMPLING = np.linspace(0, 1, 256)
+
 STRETCHES = {
     'linear': LinearStretch,
     'sqrt': SqrtStretch,
@@ -30,13 +36,25 @@ def __init__(self, **kwargs):
         self.layers = {}
 
         self._first = True
+        self._mode = 'color'
+
+    @property
+    def mode(self):
+        return self._mode
+
+    @mode.setter
+    def mode(self, value):
+        if value not in ['color', 'colormap']:
+            raise ValueError("mode should be one of 'color' or 'colormap'")
+        self._mode = value
 
     def allocate(self, uuid):
         self.layers[uuid] = {'zorder': 0,
                              'visible': True,
                              'array': None,
                              'shape': None,
                              'color': '0.5',
+                             'cmap': colormaps.members[0][1],
                              'alpha': 1,
                              'clim': (0, 1),
                              'contrast': 1,
@@ -50,6 +68,9 @@ def set(self, uuid, **kwargs):
         for key, value in kwargs.items():
             if key not in self.layers[uuid]:
                 raise KeyError("Unknown key: {0}".format(key))
+            elif key == 'color' and isinstance(value, Colormap):
+                warnings.warn('Setting colormap using "color" key is deprecated, use "cmap" instead.', UserWarning)
+                self.layers[uuid]['cmap'] = value
             else:
                 self.layers[uuid][key] = value
 
@@ -80,7 +101,23 @@ def __call__(self, bounds=None):
         img = None
         visible_layers = 0
 
-        for uuid in sorted(self.layers, key=lambda x: self.layers[x]['zorder']):
+        # Get a sorted list of UUIDs with the top layers last
+        sorted_uuids = sorted(self.layers, key=lambda x: self.layers[x]['zorder'])
+
+        # We first check that layers are either all colormaps or all single colors.
+        # In the case where we are dealing with colormaps, we can start from
+        # the last layer that has an opacity of 1 because layers below will not
+        # affect the output, assuming also that the colormaps do not change the
+        # alpha
+        if self.mode == 'colormap':
+            for i in range(len(sorted_uuids) - 1, -1, -1):
+                layer = self.layers[sorted_uuids[i]]
+                if layer['visible']:
+                    if layer['alpha'] == 1 and layer['cmap'](CMAP_SAMPLING)[:, 3].min() == 1:
+                        sorted_uuids = sorted_uuids[i:]
+                        break
+
+        for uuid in sorted_uuids:
 
             layer = self.layers[uuid]
 
@@ -89,6 +126,7 @@ def __call__(self, bounds=None):
 
             interval = ManualInterval(*layer['clim'])
             contrast_bias = ContrastBiasStretch(layer['contrast'], layer['bias'])
+            stretch = STRETCHES[layer['stretch']]()
 
             if callable(layer['array']):
                 array = layer['array'](bounds=bounds)
@@ -104,27 +142,45 @@ def __call__(self, bounds=None):
             else:
                 scalar = False
 
-            data = STRETCHES[layer['stretch']]()(contrast_bias(interval(array)))
+            data = interval(array)
+            data = contrast_bias(data, out=data)
+            data = stretch(data, out=data)
             data[np.isnan(data)] = 0
 
-            if isinstance(layer['color'], Colormap):
+            if self.mode == 'colormap':
 
                 if img is None:
                     img = np.ones(data.shape + (4,))
 
                 # Compute colormapped image
-                plane = layer['color'](data)
+                plane = layer['cmap'](data)
+
+                # Check what the smallest colormap alpha value for this layer is
+                # - if it is 1 then this colormap does not change transparency,
+                # and this allows us to speed things up a little.
+
+                if layer['cmap'](CMAP_SAMPLING)[:, 3].min() == 1:
+
+                    if layer['alpha'] == 1:
+                        img[...] = 0
+                    else:
+                        plane *= layer['alpha']
+                        img *= (1 - layer['alpha'])
+
+                else:
+
+                    # Use traditional alpha compositing
+
+                    alpha_plane = layer['alpha'] * plane[:, :, 3]
 
-                alpha_plane = layer['alpha'] * plane[:, :, 3]
+                    plane[:, :, 0] = plane[:, :, 0] * alpha_plane
+                    plane[:, :, 1] = plane[:, :, 1] * alpha_plane
+                    plane[:, :, 2] = plane[:, :, 2] * alpha_plane
 
-                # Use traditional alpha compositing
-                plane[:, :, 0] = plane[:, :, 0] * alpha_plane
-                plane[:, :, 1] = plane[:, :, 1] * alpha_plane
-                plane[:, :, 2] = plane[:, :, 2] * alpha_plane
+                    img[:, :, 0] *= (1 - alpha_plane)
+                    img[:, :, 1] *= (1 - alpha_plane)
+                    img[:, :, 2] *= (1 - alpha_plane)
 
-                img[:, :, 0] *= (1 - alpha_plane)
-                img[:, :, 1] *= (1 - alpha_plane)
-                img[:, :, 2] *= (1 - alpha_plane)
                 img[:, :, 3] = 1
 
             else:
@@ -155,7 +211,7 @@ def __call__(self, bounds=None):
         if img is None:
             return None
         else:
-            img = np.clip(img, 0, 1)
+            img = np.clip(img, 0, 1, out=img)
 
         return img
 

glue/viewers/image/layer_artist.py

@@ -156,15 +156,16 @@ def _update_visual_attributes(self):
             return
 
         if self._viewer_state.color_mode == 'Colormaps':
-            color = self.state.cmap
+            self.composite.mode = 'colormap'
         else:
-            color = self.state.color
+            self.composite.mode = 'color'
 
         self.composite.set(self.uuid,
                            clim=(self.state.v_min, self.state.v_max),
                            visible=self.state.visible,
                            zorder=self.state.zorder,
-                           color=color,
+                           color=self.state.color,
+                           cmap=self.state.cmap,
                            contrast=self.state.contrast,
                            bias=self.state.bias,
                            alpha=self.state.alpha,

glue/viewers/image/python_export.py

@@ -30,15 +30,16 @@ def python_export_image_layer(layer, *args):
     script += "composite.allocate('{0}')\n".format(layer.uuid)
 
     if layer._viewer_state.color_mode == 'Colormaps':
-        color = code('plt.cm.' + layer.state.cmap.name)
+        script += "composite.mode = 'colormap'\n"
     else:
-        color = layer.state.color
+        script += "composite.mode = 'color'\n"
 
     options = dict(array=code('array_maker'),
                    clim=(layer.state.v_min, layer.state.v_max),
                    visible=layer.state.visible,
                    zorder=layer.state.zorder,
-                   color=color,
+                   color=layer.state.color,
+                   cmap=code('plt.cm.' + layer.state.cmap.name),
                    contrast=layer.state.contrast,
                    bias=layer.state.bias,
                    alpha=layer.state.alpha,

glue/viewers/image/tests/test_composite_array.py

@@ -49,14 +49,16 @@ def test_shape_function(self):
 
     def test_cmap_blending(self):
 
+        self.composite.mode = 'colormap'
+
         self.composite.allocate('a')
         self.composite.allocate('b')
 
         self.composite.set('a', zorder=0, visible=True, array=self.array1,
-                           color=cm.Blues, clim=(0, 2))
+                           cmap=cm.Blues, clim=(0, 2))
 
         self.composite.set('b', zorder=1, visible=True, array=self.array2,
-                           color=cm.Reds, clim=(0, 1))
+                           cmap=cm.Reds, clim=(0, 1))
 
         # Determine expected result for each layer individually in the absence
         # of transparency
@@ -83,6 +85,66 @@ def test_cmap_blending(self):
 
         assert_allclose(self.composite(bounds=self.default_bounds), 0.5 * (expected_b + expected_a))
 
+    def test_cmap_alphas(self):
+
+        self.composite.mode = 'colormap'
+
+        self.composite.allocate('a')
+        self.composite.allocate('b')
+
+        self.composite.set('a', zorder=0, visible=True, array=self.array1,
+                           cmap=cm.Blues, clim=(0, 2))
+
+        self.composite.set('b', zorder=1, visible=True, array=self.array2,
+                           cmap=lambda x: cm.Reds(x, alpha=abs(np.nan_to_num(x))), clim=(0, 1))
+
+        # Determine expected result for each layer individually in the absence
+        # of transparency
+
+        expected_a = np.array([[cm.Blues(1.), cm.Blues(0.5)],
+                               [cm.Blues(0.), cm.Blues(0.)]])
+
+        expected_b = np.array([[cm.Reds(0.), cm.Reds(1.)],
+                               [cm.Reds(0.), cm.Reds(0.)]])
+
+        # If the top layer has alpha=1 with a colormap alpha fading proportional to absval,
+        # it should be visible only at the nonzero value [0, 1]
+
+        assert_allclose(self.composite(bounds=self.default_bounds),
+                        [[expected_a[0, 0], expected_b[0, 1]], expected_a[1]])
+
+        # For the same case with the top layer alpha=0.5 that value should become an equal
+        # blend of both layers again
+
+        self.composite.set('b', alpha=0.5)
+
+        assert_allclose(self.composite(bounds=self.default_bounds),
+                        [[expected_a[0, 0], 0.5 * (expected_a[0, 1] + expected_b[0, 1])],
+                         expected_a[1]])
+
+        # A third layer added at the bottom should not be visible in the output
+
+        self.composite.allocate('c')
+        self.composite.set('c', zorder=-1, visible=True, array=self.array3,
+                           cmap=cm.Greens, clim=(0, 2))
+
+        assert_allclose(self.composite(bounds=self.default_bounds),
+                        [[expected_a[0, 0], 0.5 * (expected_a[0, 1] + expected_b[0, 1])],
+                         expected_a[1]])
+
+        # For only the bottom layer having such colormap, the top layer should appear just the same
+
+        self.composite.set('a', alpha=1., cmap=lambda x: cm.Blues(x, alpha=abs(np.nan_to_num(x))))
+        self.composite.set('b', alpha=1., cmap=cm.Reds)
+
+        assert_allclose(self.composite(bounds=self.default_bounds), expected_b)
+
+        # Settin the third layer on top with alpha=0 should not affect the appearance
+
+        self.composite.set('c', zorder=2, alpha=0.)
+
+        assert_allclose(self.composite(bounds=self.default_bounds), expected_b)
+
     def test_color_blending(self):
 
         self.composite.allocate('a')