Add colormaps and stretches to ipyvolume viewer

glue_jupyter/ipyvolume/tests/test_ipyvolume.py

@@ -2,7 +2,9 @@
 
 import nbformat
 import numpy as np
+from glue.config import stretches
 from glue.core.roi import PolygonalROI, Projected3dROI
+from matplotlib import colormaps
 from nbconvert.preprocessors import ExecutePreprocessor
 
 DATA = os.path.join(os.path.dirname(__file__), 'data')
@@ -177,6 +179,49 @@ def test_volshow_multiple_subsets(app, data_unlinked, data_volume):
     assert not viewer.layers[2].enabled
 
 
+def test_volshow_cmap_mode(app, data_volume):
+
+    assert data_volume in app.data_collection
+    v = app.volshow(data=data_volume)
+
+    layer = v.layers[0]
+    layer_widget = v.layer_options.layers[-1]['layer_panel']
+
+    assert layer.state.color_mode == 'Fixed'
+    assert layer.state.cmap.name == 'gray'
+
+    layer.state.color_mode = 'Linear'
+    assert layer_widget.widget_color.widget_cmap_mode.label == 'Linear'
+    assert layer_widget.widget_color.widget_cmap.label == 'Gray'
+    assert layer.state.cmap.name == 'gray'
+
+    layer.state.cmap = colormaps['viridis']
+    assert layer_widget.widget_color.widget_cmap.label == 'Viridis'
+
+    layer_widget.widget_color.widget_cmap.label = 'Hot'
+    assert layer.state.cmap == colormaps['hot']
+
+
+def test_volshow_stretch(app, data_volume):
+
+    assert data_volume in app.data_collection
+    v = app.volshow(data=data_volume)
+
+    layer = v.layers[0]
+    layer_widget = v.layer_options.layers[-1]['layer_panel']
+
+    assert layer.state.stretch == 'linear'
+    assert [item[1] for item in layer_widget.widget_stretch.options] == \
+           [item for item in stretches.members]
+    assert layer_widget.widget_stretch.value == 'linear'
+
+    layer.state.stretch = 'log'
+    assert layer_widget.widget_stretch.value == 'log'
+
+    layer_widget.widget_stretch.value = 'sqrt'
+    assert layer.state.stretch == 'sqrt'
+
+
 def test_notebook():
 
     # Run an actual notebook

glue_jupyter/ipyvolume/volume/layer_artist.py

@@ -36,17 +36,31 @@ def __init__(self, layer=None, **kwargs):
         self.clamp_max = False
 
 
-def _transfer_func_rgba(color, N=256, max_opacity=1):
+def _transfer_func_rgba(color, N=256, max_opacity=1, stretch=None):
     r, g, b = matplotlib.colors.to_rgb(color)
     data = np.zeros((N, 4), dtype=np.float32)
     ramp = np.linspace(0, 1, N)
+    if stretch is not None:
+        ramp = stretch(ramp)
     data[..., 0] = r
     data[..., 1] = g
     data[..., 2] = b
     data[..., 3] = ramp*max_opacity
     return data
 
 
+def _transfer_func_cmap(cmap, N=256, max_opacity=1, stretch=None):
+    data = np.zeros((N, 4), dtype=np.float32)
+    ramp = np.linspace(0, 1, N)
+    if stretch is not None:
+        ramp = stretch(ramp)
+    colors = cmap(ramp)
+    for i in range(3):
+        data[..., i] = [c[i] for c in colors]
+    data[..., 3] = ramp*max_opacity
+    return data
+
+
 data0 = [[[1, 2]] * 2] * 2
 
 
@@ -84,7 +98,9 @@ def __init__(self, ipyvolume_viewer=None, state=None, layer=None, layer_state=No
 
         link((self.state, 'opacity_scale'), (self.volume, 'opacity_scale'))
 
-        on_change([(self.state, 'color', 'alpha')])(self._update_transfer_function)
+        on_change([(self.state, 'color', 'alpha', 'color_mode',
+                    'cmap', 'stretch', 'stretch_parameters'
+                    )])(self._update_transfer_function)
 
     def clear(self):
         pass
@@ -136,5 +152,13 @@ def update(self):
             self.state.percentile = 100
 
     def _update_transfer_function(self):
-        self.transfer_function.rgba = _transfer_func_rgba(self.state.color,
-                                                          max_opacity=self.state.alpha)
+        def stretch(x):
+            return self.state.stretch_object(x, **self.state.stretch_parameters)
+        if self.state.color_mode == "Fixed":
+            self.transfer_function.rgba = _transfer_func_rgba(self.state.color,
+                                                              max_opacity=self.state.alpha,
+                                                              stretch=stretch)
+        else:
+            self.transfer_function.rgba = _transfer_func_cmap(self.state.cmap,
+                                                              max_opacity=self.state.alpha,
+                                                              stretch=stretch)

glue_jupyter/ipyvolume/volume/layer_style_widget.py

@@ -1,8 +1,12 @@
 from ipywidgets import (Checkbox, VBox, ColorPicker, Dropdown, FloatSlider,
                         FloatLogSlider)
 
+from glue.core.subset import Subset
 from glue.utils import color2hex
 
+from glue_jupyter.widgets import Color
+from glue_jupyter.widgets.linked_dropdown import LinkedDropdown
+
 from ...link import link, dlink
 
 __all__ = ['Volume3DLayerStateWidget']
@@ -53,8 +57,11 @@ def __init__(self, layer_state):
         self.widget_clamp_max = Checkbox(description='clamp maximum', value=self.state.clamp_max)
         link((self.state, 'clamp_max'), (self.widget_clamp_max, 'value'))
 
-        self.widget_color = ColorPicker(value=color2hex(self.state.color), description='color')
-        link((self.state, 'color'), (self.widget_color, 'value'), color2hex)
+        if isinstance(layer_state.layer, Subset):
+            self.widget_color = ColorPicker(value=color2hex(self.state.color), description='color')
+            link((self.state, 'color'), (self.widget_color, 'value'), color2hex)
+        else:
+            self.widget_color = Color(state=self.state, cmap_mode_attr='color_mode', cmap_att=None)
 
         if self.state.alpha is None:
             self.state.alpha = 1
@@ -68,6 +75,10 @@ def __init__(self, layer_state):
                                                    value=self.state.opacity_scale)
         link((self.state, 'opacity_scale'), (self.widget_opacity_scale, 'value'))
 
+        self.widget_stretch = LinkedDropdown(self.state, 'stretch',
+                                             ui_name='stretch',
+                                             label='stretch')
+
         # FIXME: this should be fixed
         # self.widget_reset_zoom = Button(description="Reset zoom")
         # self.widget_reset_zoom.on_click(self.state.viewer_state.reset_limits)
@@ -77,4 +88,4 @@ def __init__(self, layer_state):
                           self.widget_clamp_min, self.widget_clamp_max,
                           self.widget_max_resolution,  # self.widget_reset_zoom,
                           self.widget_color, self.widget_opacity,
-                          self.widget_opacity_scale])
+                          self.widget_opacity_scale, self.widget_stretch])

glue_jupyter/widgets/color.py

@@ -13,36 +13,46 @@ def __init__(self, state, **kwargs):
         super(Color, self).__init__(**kwargs)
         self.state = state
 
+        self.cmap_att = kwargs.get('cmap_att', 'cmap_att')
+        self.cmap_mode_attr = kwargs.get('cmap_mode_attr', 'cmap_mode')
+
         self.widget_color = widgets.ColorPicker(description='color')
         link((self.state, 'color'), (self.widget_color, 'value'), color2hex)
 
-        cmap_mode_options = type(self.state).cmap_mode.get_choice_labels(self.state)
+        cmap_mode_options = getattr(type(self.state),
+                                    self.cmap_mode_attr).get_choice_labels(self.state)
         self.widget_cmap_mode = widgets.RadioButtons(options=cmap_mode_options,
                                                      description='cmap mode')
-        link((self.state, 'cmap_mode'), (self.widget_cmap_mode, 'value'))
-
-        self.widget_cmap_att = LinkedDropdown(self.state, 'cmap_att',
-                                              ui_name='color attribute',
-                                              label='color attribute')
-
-        self.widget_cmap_vmin = widgets.FloatText(description='color min')
-        self.widget_cmap_vmax = widgets.FloatText(description='color max')
-        self.widget_cmap_v = widgets.VBox([self.widget_cmap_vmin, self.widget_cmap_vmax])
-        link((self.state, 'cmap_vmin'), (self.widget_cmap_vmin, 'value'), lambda value: value or 0)
-        link((self.state, 'cmap_vmax'), (self.widget_cmap_vmax, 'value'), lambda value: value or 1)
+        link((self.state, self.cmap_mode_attr), (self.widget_cmap_mode, 'value'))
+
+        children = [self.widget_cmap_mode, self.widget_color]
+        if self.cmap_att is not None:
+            self.widget_cmap_att = LinkedDropdown(self.state, 'cmap_att',
+                                                  ui_name='color attribute',
+                                                  label='color attribute')
+            self.widget_cmap_vmin = widgets.FloatText(description='color min')
+            self.widget_cmap_vmax = widgets.FloatText(description='color max')
+            self.widget_cmap_v = widgets.VBox([self.widget_cmap_vmin, self.widget_cmap_vmax])
+            link((self.state, 'cmap_vmin'),
+                 (self.widget_cmap_vmin, 'value'),
+                 lambda value: value or 0)
+            link((self.state, 'cmap_vmax'),
+                 (self.widget_cmap_vmax, 'value'),
+                 lambda value: value or 1)
+            children.extend((self.widget_cmap_att, self.widget_cmap_v))
 
         self.widget_cmap = widgets.Dropdown(options=colormaps, description='colormap')
+        children.append(self.widget_cmap)
         link((self.state, 'cmap'), (self.widget_cmap, 'label'),
              lambda cmap: colormaps.name_from_cmap(cmap), lambda name: colormaps[name])
 
         dlink((self.widget_cmap_mode, 'value'), (self.widget_color.layout, 'display'),
               lambda value: None if value == cmap_mode_options[0] else 'none')
         dlink((self.widget_cmap_mode, 'value'), (self.widget_cmap.layout, 'display'),
               lambda value: None if value == cmap_mode_options[1] else 'none')
-        dlink((self.widget_cmap_mode, 'value'), (self.widget_cmap_att.layout, 'display'),
-              lambda value: None if value == cmap_mode_options[1] else 'none')
-        dlink((self.widget_cmap_mode, 'value'), (self.widget_cmap_v.layout, 'display'),
-              lambda value: None if value == cmap_mode_options[1] else 'none')
-        self.children = (self.widget_cmap_mode, self.widget_color,
-                         self.widget_cmap_att, self.widget_cmap_v,
-                         self.widget_cmap)
+        if self.cmap_att is not None:
+            dlink((self.widget_cmap_mode, 'value'), (self.widget_cmap_att.layout, 'display'),
+                  lambda value: None if value == cmap_mode_options[1] else 'none')
+            dlink((self.widget_cmap_mode, 'value'), (self.widget_cmap_v.layout, 'display'),
+                  lambda value: None if value == cmap_mode_options[1] else 'none')
+        self.children = tuple(children)