Support for 3d subset and a lasso 3d selection

glue/core/roi.py

@@ -74,6 +74,31 @@ def contains(self, x, y):
         """
         raise NotImplementedError()
 
+    def contains3d(self, x, y, z):
+        """Return true/false for each x/y/z pair.
+
+        Parameters
+        ----------
+        x : :class:`numpy.ndarray`
+            Array of x locations
+        y : :class:`numpy.ndarray`
+            Array of y locations
+        z : :class:`numpy.ndarray`
+            Array of z locations
+
+        Returns
+        -------
+        :class:`numpy.ndarray`
+            A boolean array, where each element is `True` if the corresponding
+            (x,y,z) tuple is inside the Roi.
+
+        Raises
+        ------
+        UndefinedROI
+            if not defined
+        """
+        raise NotImplementedError()
+
     def center(self):
         """Return the (x,y) coordinates of the ROI center"""
         raise NotImplementedError()
@@ -561,6 +586,67 @@ def move_to(self, xdelta, ydelta):
         self.vy = list(map(lambda y: y + ydelta, self.vy))
 
 
+def _project(projection_matrix, x, y, z):
+    """Projects 3d coordinates to 2d coordinates using a 4x4 matrix"""
+    x = np.asarray(x)
+    y = np.asarray(y)
+    z = np.asarray(z)
+    # work in homogeneous coordinates so we can support perspective
+    # projections as well
+    vertices = np.array([x, y, z, np.ones(x.shape)])
+    # homogeneous screen coordinates
+    screen_h = np.tensordot(projection_matrix, vertices, axes=(1, 0))
+    # convert to screen coordinates, and we don't care about z
+    x, y = screen_h[:2] / screen_h[3]
+    return x, y
+
+class Projected3dROI(Roi):
+    """"A region of interest defined in screen coordinates.
+
+    The screen coordinates are defined by the projection matrix.
+    The projection matrix converts homogeneous coordinates (x, y, z, w), where
+    w is implicitly 1, to homogeneous screen coordinates (usually the product
+    of the world and projection matrix).
+    """
+    def __init__(self, roi_2d=None, projection_matrix=None):
+        super(Projected3dROI, self).__init__()
+        self.roi_2d = roi_2d
+        self.projection_matrix = np.asarray(projection_matrix)
+
+    def contains3d(self, x, y, z):
+        """"Test if the projected coordinates are contained in the 2d roi."""
+        if not self.defined():
+            raise UndefinedROI
+        x, y = _project(self.projection_matrix, x, y, z)
+        return self.roi_2d.contains(x, y)
+
+    def __gluestate__(self, context):
+        return dict(roi_2d=context.id(self.roi_2d), projection_matrix=self.projection_matrix.tolist())
+
+    @classmethod
+    def __setgluestate__(cls, rec, context):
+        return cls(roi_2d=context.object(rec['roi_2d']), projection_matrix=np.asarray(rec['projection_matrix']))
+
+    # TODO: these methods forward directly to roi_2d, not sure if this makes sense for all
+
+    def contains(self, x, y):
+        return self.roi_2d.contains(x, y)
+
+    def center(self):
+        return self.roi_2d.center()
+
+    def move_to(self, x, y):
+        return self.roi_2d.move_to(x, y)
+
+    def defined(self):
+        return self.roi_2d.defined()
+
+    def to_polygon(self):
+        return self.roi_2d.to_polygon()
+
+    def transformed(self, xfunc=None, yfunc=None):
+        return self.roi_2d.transformed(xfunc, yfunc)
+
 class Path(VertexROIBase):
 
     def __str__(self):

glue/core/subset.py

@@ -1091,6 +1091,49 @@ def __str__(self):
     def __repr__(self):
         return '<%s: %s>' % (self.__class__.__name__, self)
 
+class RoiSubsetState3d(SubsetState):
+    """Subset state for a roi that implements .contains3d
+    """
+
+    @contract(xatt='isinstance(ComponentID)', yatt='isinstance(ComponentID)', zatt='isinstance(ComponentID)')
+    def __init__(self, xatt=None, yatt=None, zatt=None, roi=None):
+        super(RoiSubsetState3d, self).__init__()
+        self.xatt = xatt
+        self.yatt = yatt
+        self.zatt = zatt
+        self.roi = roi
+
+    @property
+    def attributes(self):
+        return (self.xatt, self.yatt, self.zatt)
+
+    @contract(data='isinstance(Data)', view='array_view')
+    def to_mask(self, data, view=None):
+
+        # TODO: make sure that pixel components don't actually take up much
+        #       memory and are just views
+        x = data[self.xatt, view]
+        y = data[self.yatt, view]
+        z = data[self.zatt, view]
+
+        if self.roi.defined():
+            result = self.roi.contains3d(x, y, z)
+        else:
+            result = np.zeros(x.shape, dtype=bool)
+
+        if result.shape != x.shape:
+            raise ValueError("Unexpected error: boolean mask has incorrect dimensions")
+
+        return result
+
+    def copy(self):
+        result = RoiSubsetState3d()
+        result.xatt = self.xatt
+        result.yatt = self.yatt
+        result.zatt = self.zatt
+        result.roi = self.roi
+        return result
+
 
 @contract(subsets='list(isinstance(Subset))', returns=Subset)
 def _combine(subsets, operator):

glue/core/tests/test_roi.py

@@ -12,7 +12,7 @@
 from ..component import CategoricalComponent
 from ..roi import (RectangularROI, UndefinedROI, CircularROI, PolygonalROI, CategoricalROI,
                    MplCircularROI, MplRectangularROI, MplPolygonalROI, MplPickROI, PointROI,
-                   XRangeROI, MplXRangeROI, YRangeROI, MplYRangeROI, RangeROI)
+                   XRangeROI, MplXRangeROI, YRangeROI, MplYRangeROI, RangeROI, Projected3dROI)
 
 
 FIG = Figure()
@@ -373,6 +373,44 @@ def test_str(self):
         assert type(str(self.roi)) == str
 
 
+class TestProjected3dROI(object):
+    # matrix that converts xyzw to yxzw
+    xyzw2yxzw = np.array([
+             [0, 1, 0, 0],
+             [0, 0, 1, 0],
+             [1, 0, 0, 0],
+             [0, 0, 0, 1]
+            ])
+    x = [1, 2, 3]
+    y = [2, 3, 4]
+    z = [5, 6, 7]
+    # repeat the arrays, 'rolled over' by 1
+    x_nd = [[1, 3], [2, 1], [3, 2]]
+    y_nd = [[2, 3], [3, 2], [4, 3]]
+    z_nd = [[5, 7], [6, 5], [7, 6]]
+
+
+    def test_contains2d(self):
+        roi_2d = PolygonalROI(vx=[0.5, 2.5, 2.5, 0.5], vy=[1, 1, 3.5, 3.5])
+        roi = Projected3dROI(roi_2d=roi_2d, projection_matrix=np.eye(4))
+        assert roi.contains(self.x, self.y).tolist() == [True, True, False]
+
+    def test_contains3d(self):
+        roi_2d = PolygonalROI(vx=[1.5, 3.5, 3.5, 1.5], vy=[4, 4, 6.5, 6.5])
+        roi = Projected3dROI(roi_2d=roi_2d, projection_matrix=self.xyzw2yxzw)
+        assert roi.contains3d(self.x, self.y, self.z).tolist() == [True, True, False]
+        assert roi.contains3d(self.x_nd, self.y_nd, self.z_nd).tolist() == [[True, False], [True, True], [False, True]]
+
+
+    def test_forward(self):
+        # testing the calls that should be forwarded to roi_2d
+        roi_2d = PolygonalROI(vx=[0.5, 2.5, 2.5, 0.5], vy=[1, 1, 3.5, 3.5])
+        roi = Projected3dROI(roi_2d=roi_2d, projection_matrix=np.eye(4))
+
+        assert roi.contains(self.x, self.y).tolist() == roi_2d.contains(self.x, self.y).tolist()
+        assert roi.to_polygon() == roi_2d.to_polygon()
+        assert roi.defined() == roi_2d.defined()
+
 class TestCategorical(object):
 
     def test_empty(self):

glue/core/tests/test_state.py

@@ -229,6 +229,15 @@ def test_polygonal_roi():
     assert_equal(r2.vx, [0, 0, 1])
     assert_equal(r2.vy, [0, 1, 0])
 
+def test_projected3d_roi():
+    roi_2d = core.roi.PolygonalROI(vx=[0.5, 2.5, 2.5, 0.5], vy=[1, 1, 3.5, 3.5])
+    roi = core.roi.Projected3dROI(roi_2d=roi_2d, projection_matrix=np.eye(4))
+    roi_clone = clone(roi)
+    x = [1, 2, 3]
+    y = [2, 3, 4]
+    z = [5, 6, 7]
+    assert roi.contains(x, y).tolist() == roi_clone.contains(x, y).tolist()
+
 
 def test_matplotlib_cmap():
     from matplotlib import cm