Time to give up on xarray

pydata/xarray#2159
y3nr1ng · Dec 6, 2019 · ae6a8d0 · ae6a8d0
1 parent c40e697
commit ae6a8d0
Show file tree

Hide file tree

Showing 2 changed files with 41 additions and 18 deletions.
diff --git a/utoolbox/io/dataset/base/data/dense.py b/utoolbox/io/dataset/base/data/dense.py
@@ -2,6 +2,7 @@
 import logging
 from itertools import product
 
+import numpy as np
 import xarray as xr
 
 from ..generic import GenericDataset
@@ -38,21 +39,32 @@ def preload(self):
         )
 
         logger.debug("preloading...")
+        arrays, coords = [], []
         for data_var in data_vars:
             logger.debug(f"> data_var: {data_var}")
-            for coord in product(*coord_values):
-                coord = {k: v for k, v in zip(coord_keys, coord)}
-                logger.debug(f">> coord: {coord}")
-                file_list = self._retrieve_file_list(data_var, coord)
+            for raw_coord in product(*coord_values):
+                coord_dict = {k: v for k, v in zip(coord_keys, raw_coord)}
+                logger.debug(f">> coord: {coord_dict}")
+                file_list = self._retrieve_file_list(data_var, coord_dict)
                 if file_list:
                     array = self.read_func(file_list, shape, dtype)
-                    array = array.assign_coords(coord)
-                    self.dataset[data_var] = array
+                    array = array.expand_dims(coord_keys)
+                    print(array)
+                    coord_dict = {
+                        k: np.array(v) for k, v in zip(coord_keys, raw_coord)
+                    }
+                    array = xr.Dataset({data_var: array}, coords=coord_dict)
+                    print(array)
+                    arrays.append(array)
+                    coords.append(coord_dict)
                 else:
                     logger.warning(
-                        f'missing data, DATA_VAR "{data_var}", COORD "{coord}"'
+                        f'missing data, DATA_VAR "{data_var}", COORD "{coord_dict}"'
                     )
-
+
+        logger.debug("merging...")
+        self._dataset = xr.combine_by_coords(arrays, coords=coord_keys, join="all")
+        # self.dataset.update({data_var: arrays})
 
     ##
 

diff --git a/workspace/test_xarray_dataset.py b/workspace/test_xarray_dataset.py
@@ -15,7 +15,7 @@
     src_dir = "C:/Users/Andy/Downloads/20191119_ExM_kidney_10XolympusNA06_zp5_7x8_DKO_4-2_Nkcc2_488_slice_2_1"
     ds = MicroManagerV1Dataset(src_dir)
 
-    print("dataset info")
+    print("== dataset info")
     print(ds.dataset)
     print()
 
@@ -28,16 +28,27 @@
 
     # OK
     # tile = ds.dataset.sel(tile_x=3370.7395, tile_y=-488.026)["488"][:, ::4, ::4]
-    #tile = ds.dataset.sel(tile_x=1919.65, tile_y=-488.026)["488"][:, ::4, ::4]
+    # tile = ds.dataset.sel(tile_x=1919.65, tile_y=-488.026)["488"][:, ::4, ::4]
     # ERROR
-    tile = ds.dataset.sel(tile_x=-2433.619, tile_y=-488.026)["488"][:, ::4, ::4]
-    print(tile)
-    print()
+    # tile = ds.dataset.sel(tile_x=-2433.619, tile_y=-488.026)["488"][:, ::4, ::4]
+    # print(".. selected")
+    # print(tile)
+    # print()
+
+    dst_dir = "_debug"
+    try:
+        os.makedirs(dst_dir)
+    except FileExistsError:
+        pass
 
     import imageio
 
-    with ProgressBar():
-        data = tile.compute()
-        print(f"{data.shape}, {data.dtype}")
-        imageio.volwrite("_debug.tif", data.values)
-    # print(f"tile (0, 0) mean: {result}")
+    for j, (_, ds_x) in enumerate(ds.dataset.groupby("tile_y")):
+        for i, (_, tile) in enumerate(ds_x.groupby("tile_x")):
+            print(f".. iter (i: {i}, j: {j})")
+
+            with ProgressBar():
+                data = tile["488"][:, ::4, ::4].max("z").compute()
+                print(f"{data.shape}, {data.dtype}")
+                dst_path = os.path.join(dst_dir, f"tile-{i:03d}-{j:03d}_mip.tif")
+                imageio.imwrite(dst_path, data.values)