From 561135c597f2e911130f34cc703515251f5ed79c Mon Sep 17 00:00:00 2001
From: Tony Tung <ttung@chanzuckerberg.com>
Date: Thu, 18 Apr 2019 15:13:23 -0400
Subject: [PATCH] Avoid use of _tile_data.keys()

`_tile_data.keys()` returns tiles outside the scope of the imagestack if it has been subselected.  Instead, use `_iter_axes(..)`, which uses the actual xarray coordinates to determine which tiles exist.

Add a test that selects and exports an imagestack.

Depends on #1189
Fixes #1154
---
 starfish/imagestack/imagestack.py      |  6 +++++-
 starfish/imagestack/test/test_index.py | 28 ++++++++++++++++++++++++++
 2 files changed, 33 insertions(+), 1 deletion(-)

diff --git a/starfish/imagestack/imagestack.py b/starfish/imagestack/imagestack.py
index e951eadf8..c8adbf570 100644
--- a/starfish/imagestack/imagestack.py
+++ b/starfish/imagestack/imagestack.py
@@ -1072,7 +1072,11 @@ def export(self,
             default_tile_shape={Axes.Y: self.tile_shape[0], Axes.X: self.tile_shape[1]},
             extras=self._tile_data.extras,
         )
-        for tilekey in self._tile_data.keys():
+        for axis_val_map in self._iter_axes({Axes.ROUND, Axes.CH, Axes.ZPLANE}):
+            tilekey = TileKey(
+                round=axis_val_map[Axes.ROUND],
+                ch=axis_val_map[Axes.CH],
+                zplane=axis_val_map[Axes.ZPLANE])
             round_, ch, zplane = tilekey.round, tilekey.ch, tilekey.z
             extras: dict = self._tile_data[tilekey]
 
diff --git a/starfish/imagestack/test/test_index.py b/starfish/imagestack/test/test_index.py
index d57a0b5c9..2e4217289 100644
--- a/starfish/imagestack/test/test_index.py
+++ b/starfish/imagestack/test/test_index.py
@@ -1,5 +1,10 @@
+import os
+import tempfile
 from collections import OrderedDict
 
+import numpy as np
+
+from starfish import ImageStack
 from starfish.imagestack import indexing_utils as iu
 from starfish.types import Axes, Coordinates, PhysicalCoordinateTypes
 from .factories import imagestack_with_coords_factory, synthetic_stack
@@ -225,3 +230,26 @@ def test_nonindexed_dimensions_restored():
                 (Axes.ZPLANE, Coordinates.Z),
         ):
             assert len(sel_xarray[primary_axis.value]) == len(sel_xarray[dependent_axis.value])
+
+
+def test_select_and_export():
+    """Tests selecting on an Imagestack with a shape (5, 5, 15, 200, 200)
+        1.) stack.sel(indexers)
+        2.) export stack
+    """
+    stack = synthetic_stack(
+        num_round=5, num_ch=5, num_z=15, tile_height=200, tile_width=200)
+
+    # select on range of rounds and single ch and Z
+    selected = stack.sel({Axes.ROUND: (1, None), Axes.CH: (2, 3), Axes.ZPLANE: 0})
+
+    with tempfile.TemporaryDirectory() as tfd:
+        path = os.path.join(tfd, "stack.json")
+        selected.export(path)
+
+        loaded = ImageStack.from_path_or_url(path)
+
+        assert np.array_equal(selected.xarray, loaded.xarray)
+        for coords in (
+                Axes.ROUND, Axes.CH, Axes.ZPLANE, Coordinates.X, Coordinates.Y, Coordinates.Z):
+            assert np.allclose(selected.xarray[coords.value], loaded.xarray[coords.value])