Allow specifying visual attributes and colormap when faceting subsets

glue/core/tests/test_util.py

@@ -3,8 +3,8 @@
 import numpy as np
 from matplotlib.cm import gray
 
-from .. import Data, DataCollection
-from ..util import facet_subsets, colorize_subsets
+from .. import Data, DataCollection, VisualAttributes
+from ..util import facet_subsets, colorize_subsets, sample_colormap
 
 
 class TestRelim(object):
@@ -81,6 +81,21 @@ def test_facet_reversed(self):
                                       [True, True, False, False, False,
                                        False, False])
 
+    def test_facet_styling(self):
+        visual_attrs = dict(alpha=0.7, markersize=8, marker='o',
+                            linewidth=2, linestyle='dashed')
+        style = VisualAttributes(**visual_attrs)
+
+        lo, hi, steps = 3, 6, 3
+        grps = facet_subsets(self.collect, self.data.id['x'],
+                             lo=lo, hi=hi, steps=steps,
+                             style=visual_attrs, cmap=gray)
+
+        colors = sample_colormap(steps, gray)
+        for sg, color in zip(grps, colors):
+            style.color = color
+            assert sg.style == style
+
 
 def test_colorize_subsets():
 

glue/core/util.py

@@ -194,7 +194,7 @@ def join_component_view(component, view):
 
 
 def facet_subsets(data_collection, cid, lo=None, hi=None, steps=5,
-                  prefix='', log=False):
+                  prefix='', log=False, style=None, cmap=None):
     """
     Create a series of subsets that partition the values of a particular
     attribute into several bins
@@ -218,6 +218,11 @@ def facet_subsets(data_collection, cid, lo=None, hi=None, steps=5,
         If present, the new subset labels will begin with `prefix`
     log : bool, optional
         If `True`, space divisions logarithmically. Default is `False`
+    style : dict, optional
+        Any visual attributes specified here will be used when creating subsets
+    cmap : :class:`matplotlib.colors.Colormap`, optional
+        Matplotlib colormap instance used to generate colors.
+        If specified, this will override any colors set in `style`
 
     Returns
     -------
@@ -247,6 +252,11 @@ def facet_subsets(data_collection, cid, lo=None, hi=None, steps=5,
     Note that the last range is inclusive on both sides. For example, if ``lo``
     is 0 and ``hi`` is 5, and ``steps`` is 5, then the intervals for the subsets
     are [0,1), [1,2), [2,3), [3,4), and [4,5].
+
+        facet_subset(data, data.id['mass'], lo=0, hi=10, steps=2, style=dict(alpha=0.5, markersize=5))
+
+    Performs the same faceting operation in the first example, but now each created subset
+    will have an alpha of 0.5 and a marker size of 5.
     """
     from glue.core.exceptions import IncompatibleAttribute
     if lo is None or hi is None:
@@ -294,27 +304,32 @@ def facet_subsets(data_collection, cid, lo=None, hi=None, steps=5,
                 labels.append(prefix + '{0}<={1}<={2}'.format(rng[i], cid, rng[i + 1]))
 
     result = []
+    style = style or {}
+    use_cmap = cmap is not None
+    if use_cmap:
+        colors = iter(sample_colormap(len(states), cmap))
+
     for lbl, s in zip(labels, states):
-        sg = data_collection.new_subset_group(label=lbl, subset_state=s)
+        kwargs = dict(label=lbl, subset_state=s, **style)
+        if use_cmap:
+            kwargs.update(color=next(colors))
+        sg = data_collection.new_subset_group(**kwargs)
         result.append(sg)
 
     return result
 
 
-def colorize_subsets(subsets, cmap, lo=0, hi=1):
+def sample_colormap(nsamples, cmap, lo=0, hi=1):
     """
-    Re-color a list of subsets according to a colormap.
+    Sample a colormap. The colormap will be sampled at `nsamples` even intervals
+    between `lo` and `hi`. Results are returned as hexadecimal strings.
 
-    The colormap will be sampled at `len(subsets)` even intervals
-    between `lo` and `hi`. The color at the `ith` interval will be
-    applied to `subsets[i]`
-
-    Parameters
+     Parameters
     ----------
-    subsets : list
-        List of subsets
+    nsamples: int
+        The number of samples
     cmap : :class:`matplotlib.colors.Colormap`
-        Matplotlib colormap instance
+        Matplotlib colormap instancecmap is not None
     lo : float, optional
         Start location in colormap. 0-1. Defaults to 0
     hi : float, optional
@@ -326,15 +341,43 @@ def colorize_subsets(subsets, cmap, lo=0, hi=1):
     sm.norm.vmin = 0
     sm.norm.vmax = 1
 
-    vals = np.linspace(lo, hi, len(subsets))
+    vals = np.linspace(lo, hi, nsamples)
     rgbas = sm.to_rgba(vals)
 
-    for color, subset in zip(rgbas, subsets):
+    samples = []
+    for color in rgbas:
         r, g, b, a = color
         r = int(255 * r)
         g = int(255 * g)
         b = int(255 * b)
-        subset.style.color = '#%2.2x%2.2x%2.2x' % (r, g, b)
+        samples.append('#%2.2x%2.2x%2.2x' % (r, g, b))
+
+    return samples
+
+
+def colorize_subsets(subsets, cmap, lo=0, hi=1):
+    """
+    Re-color a list of subsets according to a colormap.
+
+    The colormap will be sampled at `len(subsets)` even intervals
+    between `lo` and `hi`. The color at the `ith` interval will be
+    applied to `subsets[i]`.
+
+    Parameters
+    ----------
+    subsets : list
+        List of subsets
+    cmap : :class:`matplotlib.colors.Colormap`
+        Matplotlib colormap instance
+    lo : float, optional
+        Start location in colormap. 0-1. Defaults to 0
+    hi : float, optional
+        End location in colormap. 0-1. Defaults to 1
+    """
+
+    rgbas = sample_colormap(len(subsets), cmap, lo=lo, hi=hi)
+    for color, subset in zip(rgbas, subsets):
+        subset.style.color = color
 
 
 def disambiguate(label, taken):

glue/dialogs/subset_facet/qt/subset_facet.py

@@ -3,7 +3,7 @@
 from matplotlib import cm
 
 from qtpy import QtWidgets
-from glue.core.util import colorize_subsets, facet_subsets
+from glue.core.util import facet_subsets
 from glue.core.state_objects import State
 from echo import CallbackProperty, SelectionCallbackProperty
 from glue.utils.qt import load_ui
@@ -83,9 +83,9 @@ def _apply(self):
         if not np.isfinite(lo) or not np.isfinite(hi):
             return
 
-        subsets = facet_subsets(self._collect, self.state.att, lo=lo, hi=hi,
-                                steps=self.state.steps, log=self.state.log)
-        colorize_subsets(subsets, self.state.cmap)
+        facet_subsets(self._collect, self.state.att, lo=lo, hi=hi,
+                      steps=self.state.steps, log=self.state.log,
+                      cmap=self.state.cmap)
 
     @classmethod
     def facet(cls, collect, default=None, parent=None):

glue/dialogs/subset_facet/qt/tests/test_subset_facet.py

@@ -36,4 +36,5 @@ def test_apply(self):
             s._apply()
             p.assert_called_once_with(self.collect, s.state.att,
                                       lo=1, hi=3,
-                                      steps=5, log=False)
+                                      steps=5, log=False,
+                                      cmap=cm.RdYlBu)