ImageStack select on Physical Coordinates

starfish/imagestack/imagestack.py

@@ -356,7 +356,7 @@ def sel(self, indexers: Mapping[Axes, Union[int, tuple]]):
 
         Parameters
         ----------
-        indexers : Dict[Axes, (int/tuple)]
+        indexers : Mapping[Axes, Union[int, tuple]]
             A dictionary of dim:index where index is the value or range to index the dimension
 
         Examples
@@ -382,14 +382,67 @@ def sel(self, indexers: Mapping[Axes, Union[int, tuple]]):
         ImageStack :
             a new image stack indexed by given value or range.
         """
-
-        # convert indexers to Dict[str, (int/slice)] format
-        # TODO shanaxel42 check if this can be changed to xarray.copy(deep=false)
         stack = deepcopy(self)
         selector = indexing_utils.convert_to_selector(indexers)
         stack._data._data = indexing_utils.index_keep_dimensions(self.xarray, selector)
         return stack
 
+    def isel(self, indexers: Mapping[Axes, Union[int, tuple]]):
+        """Given a dictionary mapping the index name to either a value or a range represented as a
+        tuple, return an Imagestack with each dimension indexed by position accordingly
+
+        Parameters
+        ----------
+        indexers : Dict[Axes, (int/tuple)]
+            A dictionary of dim:index where index is the value or range to index the dimension
+
+        Examples
+        --------
+
+        Create an Imagestack using the ``synthetic_stack`` method
+            >>> from starfish import ImageStack
+            >>> from starfish.types import Axes
+            >>> stack = ImageStack.synthetic_stack(5, 5, 15, 200, 200)
+            >>> stack
+            <starfish.ImageStack (r: 5, c: 5, z: 15, y: 200, x: 200)>
+            >>> stack.isel({Axes.ROUND: (1, None), Axes.CH: 0, Axes.ZPLANE: 0})
+            <starfish.ImageStack (r: 4, c: 1, z: 1, y: 200, x: 200)>
+            >>> stack.isel({Axes.ROUND: 0, Axes.CH: 0, Axes.ZPLANE: 1,
+            ...Axes.Y: 100, Axes.X: (None, 100)})
+            <starfish.ImageStack (r: 1, c: 1, z: 1, y: 1, x: 100)>
+            and the imagestack's physical coordinates
+            xarray also indexed and recalculated according to the x,y slicing.
+
+        Returns
+        -------
+        ImageStack :
+            a new image stack indexed by given value or range.
+        """
+        stack = deepcopy(self)
+        selector = indexing_utils.convert_to_selector(indexers)
+        stack._data._data = indexing_utils.index_keep_dimensions(self.xarray, selector, by_pos=True)
+        return stack
+
+    def sel_by_physical_coords(
+            self, indexers: Mapping[Coordinates, Union[Number, Tuple[Number, Number]]]):
+        """
+        Given a dictionary mapping the coordinate name to either a value or a range represented as a
+        tuple, return an Imagestack with each the Coordinate dimension indexed accordingly.
+
+        Parameters
+        ----------
+        indexers : Mapping[Coordinates, Union[Number, Tuple[Number, Number]]]:
+            A dictionary of coord:index where index is the value or range to index the coordinate
+            dimension.
+
+        Returns
+        -------
+        ImageStack :
+            a new image stack indexed by given value or range.
+        """
+        new_indexers = indexing_utils.convert_coords_to_indices(self.xarray, indexers)
+        return self.isel(new_indexers)
+
     def get_slice(
             self,
             selector: Mapping[Axes, Union[int, slice]]

starfish/imagestack/indexing_utils.py

@@ -1,8 +1,9 @@
-from typing import Mapping, MutableMapping, Union
+from typing import Dict, Mapping, MutableMapping, Tuple, Union
 
+import numpy as np
 import xarray as xr
 
-from starfish.types import Axes
+from starfish.types import Axes, Coordinates, Number
 
 
 def convert_to_selector(
@@ -26,16 +27,83 @@ def convert_to_selector(
     return return_dict
 
 
-def index_keep_dimensions(
-        data: xr.DataArray, indexers: Mapping[str, Union[int, slice]]) -> xr.DataArray:
+def convert_coords_to_indices(array: xr.DataArray,
+                              indexers: Mapping[Coordinates, Union[Number, Tuple[Number, Number]]]
+                              ) -> Dict[Axes, Union[int, Tuple[Number, Number]]]:
+    """
+    Convert mapping of physical coordinates to value or range to mapping of corresponding Axes and
+    positional coordinates.
+
+    Parameters
+    ----------
+    array : xr.DataArray
+        The xarray with both physical and positional coordinates.
+    indexers: Mapping[Coordinates, Union[Number, Tuple[Number, Number]]]
+        Mapping of physical coordinates to value or range
+
+    Returns
+    -------
+    Mapping[Axes, Union[int, Tuple[Number, Number]]]:
+        Mapping of Axes and positional indices that correspond to the given physical indices.
+
+    """
+    axes_indexers: Dict[Axes, Union[int, Tuple[Number, Number]]] = dict()
+    if Coordinates.X in indexers:
+        idx_x = find_nearest(array[Coordinates.X.value], indexers[Coordinates.X])
+        axes_indexers[Axes.X] = idx_x
+    if Coordinates.Y in indexers:
+        idx_y = find_nearest(array[Coordinates.Y.value], indexers[Coordinates.Y])
+        axes_indexers[Axes.Y] = idx_y
+    if Coordinates.Z in indexers:
+        idx_z = find_nearest(array[Coordinates.Z.value], indexers[Coordinates.Z])
+        axes_indexers[Axes.ZPLANE] = idx_z
+    return axes_indexers
+
+
+def index_keep_dimensions(data: xr.DataArray,
+                          indexers: Mapping[str, Union[int, slice]],
+                          by_pos: bool=False
+                          ) -> xr.DataArray:
     """Takes an xarray and key to index it. Indexes then adds back in lost dimensions"""
     # store original dims
     original_dims = data.dims
     # index
-    data = data.sel(indexers)
+    if by_pos:
+        data = data.isel(indexers)
+    else:
+        data = data.sel(indexers)
     # find missing dims
     missing_dims = set(original_dims) - set(data.dims)
     # Add back in missing dims
     data = data.expand_dims(tuple(missing_dims))
     # Reorder to correct format
     return data.transpose(*original_dims)
+
+
+def find_nearest(array: xr.DataArray,
+                 value: Union[Number, Tuple[Number, Number]]
+                 ) -> Union[int, Tuple[int, int]]:
+    """
+    Given an xarray and value or tuple range return the indices of the closest corresponding
+    value/values in the array.
+
+    Parameters
+    ----------
+    array: xr.DataArray
+        The array to do lookups in.
+
+    value : Union[Number, Tuple[Number, Number]]
+        The value or values to lookup.
+
+    Returns
+    -------
+    Union[int, Tuple[int, int]]:
+        The index or indicies of the entries closest to the given values in the array.
+    """
+    array = np.asarray(array)
+    if isinstance(value, tuple):
+        idx1 = (np.abs(array - value[0])).argmin()
+        idx2 = (np.abs(array - value[1])).argmin()
+        return idx1, idx2
+    idx = (np.abs(array - value)).argmin()
+    return idx

starfish/imagestack/test/test_index.py

@@ -1,7 +1,8 @@
 from collections import OrderedDict
 
-from starfish.types import Axes
-from .factories import synthetic_stack
+from starfish.imagestack import indexing_utils as iu
+from starfish.types import Axes, Coordinates, PhysicalCoordinateTypes
+from .factories import imagestack_with_coords_factory, synthetic_stack
 
 
 def test_imagestack_indexing():
@@ -65,3 +66,138 @@ def test_imagestack_indexing():
     expected_shape = OrderedDict([(Axes.ROUND, 1), (Axes.CH, 1),
                                   (Axes.ZPLANE, 1), (Axes.Y, 190), (Axes.X, 190)])
     assert indexed_stack.shape == expected_shape
+
+
+X_COORDS = 1, 2
+Y_COORDS = 4, 6
+Z_COORDS = 1, 3
+
+
+def test_find_nearest():
+    """
+    Set up ImageStack with physical coordinates:
+        x_coords = [1. 1.11111111 1.22222222 1.33333333 1.44444444 1.55555556
+        1.66666667 1.77777778 1.88888889 2.]
+
+        y_coords = [4. 4.22222222 4.44444444 4.66666667 4.88888889 5.11111111
+        5.33333333 5.55555556 5.77777778 6. ]
+
+     Test that find_nearest() finds the correct corresponding positional index values
+    """
+    stack_shape = OrderedDict([(Axes.ROUND, 3), (Axes.CH, 2),
+                               (Axes.ZPLANE, 1), (Axes.Y, 10), (Axes.X, 10)])
+
+    physical_coords = OrderedDict([(PhysicalCoordinateTypes.X_MIN, X_COORDS[0]),
+                                   (PhysicalCoordinateTypes.X_MAX, X_COORDS[1]),
+                                   (PhysicalCoordinateTypes.Y_MIN, Y_COORDS[0]),
+                                   (PhysicalCoordinateTypes.Y_MAX, Y_COORDS[1]),
+                                   (PhysicalCoordinateTypes.Z_MIN, Z_COORDS[0]),
+                                   (PhysicalCoordinateTypes.Z_MAX, Z_COORDS[1])])
+
+    stack = imagestack_with_coords_factory(stack_shape, physical_coords)
+    assert iu.find_nearest(stack.xarray[Coordinates.X.value], 1.2) == 2
+    assert iu.find_nearest(stack.xarray[Coordinates.X.value], 1.5) == 4
+    assert iu.find_nearest(stack.xarray[Coordinates.X.value], (1.2, 1.5)) == (2, 4)
+
+    assert iu.find_nearest(stack.xarray[Coordinates.Y.value], 4) == 0
+    assert iu.find_nearest(stack.xarray[Coordinates.Y.value], 5.1) == 5
+    assert iu.find_nearest(stack.xarray[Coordinates.Y.value], (4, 5.1)) == (0, 5)
+
+    # assert values outside the range are given min/max of array
+    assert iu.find_nearest(stack.xarray[Coordinates.X.value], 5) == 9
+    assert iu.find_nearest(stack.xarray[Coordinates.X.value], -5) == 0
+
+
+def test_convert_coords_to_indices():
+    """
+    Set up ImageStack with physical coordinates:
+        x_coords = [1. 1.11111111 1.22222222 1.33333333 1.44444444 1.55555556
+        1.66666667 1.77777778 1.88888889 2.]
+
+        y_coords = [4. 4.22222222 4.44444444 4.66666667 4.88888889 5.11111111
+        5.33333333 5.55555556 5.77777778 6. ]
+
+    Test that convert_coords_to_indices() correctly converts Coordinate indices
+    to their corresponding positional indices
+    """
+    stack_shape = OrderedDict([(Axes.ROUND, 3), (Axes.CH, 2),
+                               (Axes.ZPLANE, 1), (Axes.Y, 10), (Axes.X, 10)])
+
+    physical_coords = OrderedDict([(PhysicalCoordinateTypes.X_MIN, X_COORDS[0]),
+                                   (PhysicalCoordinateTypes.X_MAX, X_COORDS[1]),
+                                   (PhysicalCoordinateTypes.Y_MIN, Y_COORDS[0]),
+                                   (PhysicalCoordinateTypes.Y_MAX, Y_COORDS[1]),
+                                   (PhysicalCoordinateTypes.Z_MIN, Z_COORDS[0]),
+                                   (PhysicalCoordinateTypes.Z_MAX, Z_COORDS[1])])
+
+    stack = imagestack_with_coords_factory(stack_shape, physical_coords)
+    coordinate_indices = {Coordinates.X: (1.2, 1.5), Coordinates.Y: (4, 5.1)}
+    positional_indices = iu.convert_coords_to_indices(stack.xarray, coordinate_indices)
+
+    assert positional_indices[Axes.X] == iu.find_nearest(
+        stack.xarray[Coordinates.X.value], (1.2, 1.5))
+    assert positional_indices[Axes.Y] == iu.find_nearest(
+        stack.xarray[Coordinates.Y.value], (4, 5.1))
+
+
+def test_sel_by_physical_coords():
+    """
+    Set up ImageStack with physical coordinates:
+        x_coords = [1. 1.11111111 1.22222222 1.33333333 1.44444444 1.55555556
+        1.66666667 1.77777778 1.88888889 2.]
+
+        y_coords = [4. 4.22222222 4.44444444 4.66666667 4.88888889 5.11111111
+        5.33333333 5.55555556 5.77777778 6. ]
+
+    Test that sel_by_physical_coords() correctly indexes the imagestack by the
+    corresponding positional indexers
+    """
+    stack_shape = OrderedDict([(Axes.ROUND, 3), (Axes.CH, 2),
+                               (Axes.ZPLANE, 1), (Axes.Y, 10), (Axes.X, 10)])
+
+    physical_coords = OrderedDict([(PhysicalCoordinateTypes.X_MIN, X_COORDS[0]),
+                                   (PhysicalCoordinateTypes.X_MAX, X_COORDS[1]),
+                                   (PhysicalCoordinateTypes.Y_MIN, Y_COORDS[0]),
+                                   (PhysicalCoordinateTypes.Y_MAX, Y_COORDS[1]),
+                                   (PhysicalCoordinateTypes.Z_MIN, Z_COORDS[0]),
+                                   (PhysicalCoordinateTypes.Z_MAX, Z_COORDS[1])])
+
+    stack = imagestack_with_coords_factory(stack_shape, physical_coords)
+
+    indexed_stack_by_coords = stack.sel_by_physical_coords({Coordinates.X: (1.2, 1.5),
+                                                            Coordinates.Y: (4, 5.1)})
+    indexed_stack_by_pos = stack.sel({Axes.X: (2, 4), Axes.Y: (0, 5)})
+
+    # assert that the resulting xarrays are the same
+    assert indexed_stack_by_coords.xarray.equals(indexed_stack_by_pos.xarray)
+
+
+def test_sel_by_physical_and_axes():
+    """
+       Set up ImageStack with physical coordinates:
+           x_coords = [1. 1.11111111 1.22222222 1.33333333 1.44444444 1.55555556
+           1.66666667 1.77777778 1.88888889 2.]
+
+           y_coords = [4. 4.22222222 4.44444444 4.66666667 4.88888889 5.11111111
+           5.33333333 5.55555556 5.77777778 6. ]
+
+       Test that sel_by_physical_coords() correctly indexes the imagestack by the
+       corresponding positional indexers
+       """
+    stack_shape = OrderedDict([(Axes.ROUND, 3), (Axes.CH, 2),
+                               (Axes.ZPLANE, 1), (Axes.Y, 10), (Axes.X, 10)])
+
+    physical_coords = OrderedDict([(PhysicalCoordinateTypes.X_MIN, X_COORDS[0]),
+                                   (PhysicalCoordinateTypes.X_MAX, X_COORDS[1]),
+                                   (PhysicalCoordinateTypes.Y_MIN, Y_COORDS[0]),
+                                   (PhysicalCoordinateTypes.Y_MAX, Y_COORDS[1]),
+                                   (PhysicalCoordinateTypes.Z_MIN, Z_COORDS[0]),
+                                   (PhysicalCoordinateTypes.Z_MAX, Z_COORDS[1])])
+
+    stack = imagestack_with_coords_factory(stack_shape, physical_coords)
+    indexed_stack_by_coords = stack.sel_by_physical_coords({Coordinates.X: (1.2, 1.5),
+                                                            Coordinates.Y: (4, 5.1)})
+    indexed_stack = indexed_stack_by_coords.sel({Axes.ROUND: 2, Axes.CH: 1, Axes.ZPLANE: 0})
+    expected_shape = OrderedDict([(Axes.ROUND, 1), (Axes.CH, 1),
+                                  (Axes.ZPLANE, 1), (Axes.Y, 5), (Axes.X, 2)])
+    assert indexed_stack.shape == expected_shape