Re-factoring Categorical ROI

glue/clients/histogram_client.py

@@ -5,7 +5,7 @@
 from ..core.client import Client
 from ..core import message as msg
 from ..core.data import Data, CategoricalComponent
-from ..core.subset import RangeSubsetState
+from ..core.subset import RangeSubsetState, CategoricalRoiSubsetState
 from ..core.exceptions import IncompatibleDataException, IncompatibleAttribute
 from ..core.edit_subset_mode import EditSubsetMode
 from .layer_artist import HistogramLayerArtist, LayerArtistContainer
@@ -379,12 +379,19 @@ def apply_roi(self, roi):
             lo = 10 ** lo
             hi = 10 ** hi
 
-        state = RangeSubsetState(lo, hi)
-        state.att = self.component
-        mode = EditSubsetMode()
-        visible = [d for d in self.data if self.is_layer_visible(d)]
-        focus = visible[0] if len(visible) > 0 else None
-        mode.update(self.data, state, focus_data=focus)
+        comp = list(self._get_data_components('x'))
+        if comp:
+            comp = comp[0]
+            if comp.categorical:
+                state = CategoricalRoiSubsetState.from_range(comp, self.component,
+                                                             lo, hi)
+            else:
+                state = RangeSubsetState(lo, hi)
+                state.att = self.component
+            mode = EditSubsetMode()
+            visible = [d for d in self.data if self.is_layer_visible(d)]
+            focus = visible[0] if len(visible) > 0 else None
+            mode.update(self.data, state, focus_data=focus)
 
     def register_to_hub(self, hub):
         dfilter = lambda x: x.sender.data in self._artists

glue/clients/scatter_client.py

@@ -6,8 +6,8 @@
 
 from ..core.client import Client
 from ..core.data import Data, IncompatibleAttribute, ComponentID, CategoricalComponent
-from ..core.subset import RoiSubsetState, RangeSubsetState
-from ..core.roi import PolygonalROI, RangeROI
+from ..core.subset import RoiSubsetState, RangeSubsetState, CategoricalRoiSubsetState, AndState
+from ..core.roi import PolygonalROI, RangeROI, CategoricalRoi, RectangularROI
 from ..core.util import relim
 from ..core.edit_subset_mode import EditSubsetMode
 from ..core.message import ComponentReplacedMessage
@@ -251,20 +251,52 @@ def _set_xydata(self, coord, attribute, snap=True):
         self._pull_properties()
         self._redraw()
 
+    def _process_categorical_roi(self, roi):
+        """ Returns a RoiSubsetState object.
+        """
+
+        if isinstance(roi, RectangularROI):
+            subsets = []
+            axes = [('x', roi.xmin, roi.xmax),
+                    ('y', roi.ymin, roi.ymax)]
+            for coord, lo, hi in axes:
+                comp = list(self._get_data_components(coord))
+                if comp:
+                    if comp[0].categorical:
+                        subset = CategoricalRoiSubsetState.from_range(comp[0], self.xatt, lo, hi)
+                    else:
+                        subset = RangeSubsetState(lo, hi, self.xatt)
+                else:
+                    subset = None
+                subsets.append(subset)
+        else:
+            raise AssertionError
+        return AndState(*subsets)
+
     def apply_roi(self, roi):
         # every editable subset is updated
         # using specified ROI
 
         if isinstance(roi, RangeROI):
             lo, hi = roi.range()
             att = self.xatt if roi.ori == 'x' else self.yatt
-            subset_state = RangeSubsetState(lo, hi, att)
+            if self._check_categorical(att):
+                comp = list(self._get_data_components(roi.ori))
+                if comp:
+                    subset_state = CategoricalRoiSubsetState.from_range(comp[0], att, lo, hi)
+                else:
+                    subset_state = None
+            else:
+                subset_state = RangeSubsetState(lo, hi, att)
         else:
-            subset_state = RoiSubsetState()
-            subset_state.xatt = self.xatt
-            subset_state.yatt = self.yatt
-            x, y = roi.to_polygon()
-            subset_state.roi = PolygonalROI(x, y)
+            if self._check_categorical(self.xatt) or self._check_categorical(self.yatt):
+                subset_state = self._process_categorical_roi(roi)
+            else:
+                subset_state = RoiSubsetState()
+                subset_state.xatt = self.xatt
+                subset_state.yatt = self.yatt
+                x, y = roi.to_polygon()
+                subset_state.roi = PolygonalROI(x, y)
 
         mode = EditSubsetMode()
         visible = [d for d in self._data if self.is_visible(d)]

glue/clients/tests/test_histogram_client.py

@@ -3,6 +3,7 @@
 from __future__ import absolute_import, division, print_function
 
 import pytest
+import numpy as np
 
 from mock import MagicMock
 
@@ -12,7 +13,7 @@
 from ...core.data_collection import DataCollection
 from ...core.exceptions import IncompatibleDataException
 from ...core.data import Data, CategoricalComponent, ComponentID
-from ...core.subset import RangeSubsetState
+from ...core.subset import RangeSubsetState, CategoricalRoiSubsetState
 
 from .util import renderless_figure
 
@@ -369,6 +370,23 @@ def test_tick_labels(self):
         xlabels = [formatter.format_data(pos) for pos in range(6)]
         assert correct_labels == xlabels
 
+    def test_apply_roi(self):
+        self.client.add_layer(self.data)
+        self.client.set_component(self.data.id['x'])
+        # bins are 1...4
+
+        self.data.edit_subset = [self.data.subsets[0]]
+
+        roi = MagicMock()
+        roi.to_polygon.return_value = [1.2, 2, 4], [2, 3, 4]
+
+        self.client.apply_roi(roi)
+        state = self.data.subsets[0].subset_state
+        assert isinstance(state, CategoricalRoiSubsetState)
+        np.testing.assert_equal(self.data.subsets[0].subset_state.roi.categories,
+                                np.array(['a', 'b', 'c', 'd', 'e']))
+
+
     # REMOVED TESTS
     def test_xlog_axes_labels(self):
         """ log-scale doesn't make sense for categorical data"""

glue/clients/tests/test_scatter_client.py

@@ -677,6 +677,36 @@ def test_high_cardinatility_timing(self):
         timer = timeit(timer_func, number=1)
         assert timer < 3  # this is set for Travis speed
 
+    def test_apply_roi(self):
+        data = self.add_data_and_attributes()
+        roi = core.roi.RectangularROI()
+        roi.update_limits(*self.roi_limits)
+        x, y = self.roi_points
+        self.client.apply_roi(roi)
+
+    def test_range_rois_preserved(self):
+        data = self.add_data_and_attributes()
+        assert self.client.xatt is not self.client.yatt
+
+        roi = core.roi.XRangeROI()
+        roi.set_range(1, 2)
+        self.client.apply_roi(roi)
+        assert isinstance(data.edit_subset.subset_state,
+                          core.subset.CategoricalRoiSubsetState)
+        assert data.edit_subset.subset_state.att == self.client.xatt
+
+        roi = core.roi.YRangeROI()
+        roi.set_range(1, 2)
+        self.client.apply_roi(roi)
+        assert isinstance(data.edit_subset.subset_state,
+                          core.subset.RangeSubsetState)
+        assert data.edit_subset.subset_state.att == self.client.yatt
+        roi = core.roi.RectangularROI(xmin=1, xmax=2, ymin=1, ymax=2)
+
+        self.client.apply_roi(roi)
+        assert isinstance(data.edit_subset.subset_state,
+                          core.subset.AndState)
+
     # REMOVED TESTS
     def test_invalid_plot(self):
         """ This fails because the axis ticks shouldn't reset after

glue/core/data.py

@@ -217,6 +217,13 @@ def numeric(self):
         """
         return np.can_cast(self.data[0], np.complex)
 
+    @property
+    def categorical(self):
+        """
+        Whether or not the datatype is categorical
+        """
+        return False
+
     def __str__(self):
         return "Component with shape %s" % shape_to_string(self.shape)
 
@@ -389,6 +396,10 @@ def __init__(self, categorical_data, categories=None, jitter=None, units=None):
         else:
             self._update_data()
 
+    @property
+    def categorical(self):
+        return True
+
     def _update_categories(self, categories=None):
         """
         :param categories: A sorted array of categories to find in the dataset.

glue/core/roi.py

@@ -1102,3 +1102,80 @@ def finalize_selection(self, event):
         if self._patch is not None:
             self._patch.set_visible(False)
         self._axes.figure.canvas.draw()
+
+
+class CategoricalRoi(Roi):
+
+    """
+    A ROI abstraction to represent selections of categorical data.
+    """
+
+    def __init__(self, categories=None):
+        if categories is None:
+            self.categories = None
+        else:
+            self.update_categories(categories)
+
+    def _categorical_helper(self, indata):
+        """
+        A helper function to do the rigamaroll of getting categorical data.
+
+        :param indata: Any type of input data
+        :return: The best guess at the categorical data associated with indata
+        """
+
+        try:
+            if indata.categorical:
+                return indata._categorical_data
+            else:
+                return indata[:]
+        except AttributeError:
+            return np.asarray(indata)
+
+    def contains(self, x, y):
+        """
+        Test whether a set categorical elements fall within
+        the region of interest
+
+        :param x: Any array-like object of categories
+                 (includes CategoricalComponenets)
+        :param y: Unused but required for compatibility
+
+        *Returns*
+
+           A list of True/False values, for whether each x value falls
+           within the ROI
+
+        """
+
+        check = self._categorical_helper(x)
+        index = np.minimum(np.searchsorted(self.categories, check),
+                           len(self.categories)-1)
+        return self.categories[index] == check
+
+    def update_categories(self, categories):
+
+        self.categories = np.unique(self._categorical_helper(categories))
+
+    def defined(self):
+        """ Returns True if the ROI is defined """
+        return self.categories is not None
+
+    def reset(self):
+        self.categories = None
+
+    @staticmethod
+    def from_range(cat_comp, lo, hi):
+        """
+        Utility function to help construct the Roi from a range.
+
+        :param cat_comp: Anything understood by ._categorical_helper ... array, list or component
+        :param lo: lower bound of the range
+        :param hi: upper bound of the range
+        :return: CategoricalRoi object
+        """
+
+        roi = CategoricalRoi()
+        cat_data = cat_comp._categories
+        roi.update_categories(cat_data[np.floor(lo):np.ceil(hi)])
+        return roi

glue/core/subset.py

@@ -14,6 +14,7 @@
 from ..utils import view_shape
 from ..external.six import PY3
 from .contracts import contract
+from .roi import CategoricalRoi
 
 __all__ = ['Subset', 'SubsetState', 'RoiSubsetState', 'CompositeSubsetState',
            'OrState', 'AndState', 'XorState', 'InvertState',
@@ -446,6 +447,40 @@ def copy(self):
         return result
 
 
+class CategoricalRoiSubsetState(SubsetState):
+
+    def __init__(self, att=None, roi=None):
+        super(CategoricalRoiSubsetState, self).__init__()
+        self.att = att
+        self.roi = roi
+
+    @property
+    def attributes(self):
+        return self.att,
+
+    @memoize
+    @contract(data='isinstance(Data)', view='array_view')
+    def to_mask(self, data, view=None):
+        x = data.get_component(self.att)._categorical_data[view]
+        result = self.roi.contains(x, None)
+        assert x.shape == result.shape
+        return result.ravel()
+
+    def copy(self):
+        result = CategoricalRoiSubsetState()
+        result.att = self.att
+        result.roi = self.roi
+        return result
+
+    @staticmethod
+    def from_range(component, att, lo, hi):
+
+        roi = CategoricalRoi.from_range(component, lo, hi)
+        subset = CategoricalRoiSubsetState(roi=roi,
+                                           att=att)
+        return subset
+
+
 class RangeSubsetState(SubsetState):
 
     def __init__(self, lo, hi, att=None):