Distortion by z

Perform distortion correction by z in order to save memory

stpt_pipeline/geometric_distortion.py

+import shutil
 import traceback
+from contextlib import suppress
 from functools import partial
 from pathlib import Path
 from typing import Dict, List
@@ -10,10 +12,8 @@ from dask import delayed
 from distributed import Client, as_completed
 from skimage.transform import warp
 
-import zarr
 from owl_dev.logging import logger
 
-from .retry import retry
 from .settings import Settings
 from .stpt_displacement import mask_image
 from .utils import get_coords
@@ -101,53 +101,57 @@ def apply_geometric_transform(
     return res
 
 
-@retry(Exception)
 def _write_dataset(arr, *, dark, flat, output, cof_dist):
     logger.debug("Applying optical distortion and normalization %s", arr.name)
-    g = xr.Dataset()
-    tiles = []
-    for i in arr.tile:
-        channels = []
-        for j in arr.channel:
-            im = arr.sel(tile=i, channel=j) / Settings.norm_val
-
-            n = int(j.values)
-
-            d = dark[n] / Settings.norm_val
-            f = flat[n]
-            img = da.stack(
-                [
-                    apply_geometric_transform(im.sel(z=k).data, d, f, cof_dist,)
-                    for k in arr.z
-                ]
-            )
-            channels.append(img)
-
-        channels = da.stack(channels)
-        tiles.append(channels)
-    tiles = da.stack(tiles)
-    ntiles, nchannels, nz, ny, nx = tiles.shape
-    this = xr.DataArray(
-        tiles,
-        dims=("tile", "channel", "z", "y", "x"),
-        name=arr.name,
-        coords={
-            "tile": range(ntiles),
-            "channel": range(nchannels),
-            "z": range(nz),
-            "y": range(ny),
-            "x": range(nx),
-        },
-    )
-    this.attrs = arr.attrs
-    g[arr.name] = this
-
-    try:
-        g.to_zarr(output, mode="a")
-    except Exception:
-        zz = zarr.open(output)
-        del zz[arr.name]
-        raise Exception(f"Section {arr.name} already exists")
+    for k in list(arr.z.values):
+        g = xr.Dataset()
+        tiles = []
+        for i in arr.tile:
+            channels = []
+            for j in arr.channel:
+                im = arr.sel(tile=i, channel=j) / Settings.norm_val
+
+                n = int(j.values)
+
+                d = dark[n] / Settings.norm_val
+                f = flat[n]
+                img = apply_geometric_transform(im.sel(z=k).data, d, f, cof_dist)
+                channels.append([img])
+
+            channels = da.stack(channels)
+            tiles.append(channels)
+        tiles = da.stack(tiles)
+        ntiles, nchannels, nz, ny, nx = tiles.shape
+        this = xr.DataArray(
+            tiles,
+            dims=("tile", "channel", "z", "y", "x"),
+            name=arr.name,
+            coords={
+                "tile": range(ntiles),
+                "channel": range(nchannels),
+                "z": [k],
+                "y": range(ny),
+                "x": range(nx),
+            },
+        )
+        this.attrs = arr.attrs
+        g[arr.name] = this
+
+        try:
+            g.to_zarr(output, mode="a", group=f"G{arr.name}", append_dim="z")
+        except ValueError:
+            g.to_zarr(output, mode="a", group=f"G{arr.name}")
+
+        del this
+        del g
+
+    output = Path(output)
+    for d in ["x", "y", "z", "tile", "channel", arr.name]:
+        with suppress(FileExistsError):
+            (output / f"G{arr.name}" / d).replace(output / d)
+
+    shutil.rmtree(f"{output}/G{arr.name}", ignore_errors=True)
+
     return f"{output}[{arr.name}]"