Merge pull request #27 from janelia-cellmap/expand_dims

# Changes
**New feature: `--expand_dims` flag**
Prepends a size-1 channel dimension to the output array, converting `(z,y,x)` → `(1,z,y,x)`. Supported for all input formats: TIFF, TIFF stack, MRC, zarr v2, and N5.

**Bug fixes bundled in**
- `TiffStack.__init__`: replaced `da.from_zarr(ZarrTiffStore)` with direct `.zarray` JSON read — `ZarrTiffStore` is incompatible with zarr v3
- `N5Group.apply_ome_template`: `coordinateTransformations` scale/translation vectors had wrong length for non-3D arrays (ndim was computed before the channel axis was prepended)

# Implementation
Each format's `write_to_zarr` and its worker function (`save_chunk` / `write_tile_slab` / `write_volume_slab`) receives `expand_dims`. The worker reads from the source using `chunk_slice[1:]` (skipping the channel dim absent in the source) and writes `data[np.newaxis]` into the destination.

OME multiscales metadata is updated per format:
- **Tiff/TiffStack/MRC**: channel axis prepended to `axes/scale/translation/units` in `to_zarr` before `add_ome_metadata` is called
- **zarr v2**: post-processing walk (`_patch_group_multiscales`) rewrites the copied multiscales attrs
- **N5**: `expand_dims` threaded through `normalize_to_omengff` → `apply_ome_template` / `ome_dataset_metadata`

# Tests
`tests/test_expand_dims_and_4d.py` - 9 tests
- `test_expand_dims_{tiff,tiff_stack,mrc,zarr2,n5}`: verify output shape is `(1, *src_shape)` and OME axes metadata has a leading channel axis
- `test_4d_{tiff,mrc,zarr2,n5}`: verify native 4D inputs pass through with correct shape and axes

Author: yuriyzubov

pixi.lock

@@ -45,7 +45,7 @@ def __init__(
         self.shape = np.squeeze(self.memmap.data.shape)
         self.dtype = self.memmap.data.dtype
 
-    def write_to_zarr(self, dst_spec: dict, client: Client) -> None:
+    def write_to_zarr(self, dst_spec: dict, client: Client, expand_dims: bool = False) -> None:
         """Read the MRC file in chunks and write each chunk to a zarr3 array via TensorStore.
 
         Chunks that are entirely zero are skipped to avoid unnecessary writes.
@@ -57,6 +57,9 @@ def write_to_zarr(self, dst_spec: dict, client: Client) -> None:
             by :func:`~zarrify.utils.ts_utils.zarr3_spec`.
         client:
             Dask distributed client used to parallelise chunk writes.
+        expand_dims:
+            When True, the destination array has a leading size-1 channel
+            dimension; the leading slice is stripped when reading the source.
         """
         logging.basicConfig(
             level=logging.INFO,
@@ -78,15 +81,16 @@ def write_to_zarr(self, dst_spec: dict, client: Client) -> None:
         for idx, part in enumerate(out_slices_partitioned):
             logging.info(f"{idx + 1} / {len(out_slices_partitioned)}")
             start = time.time()
-            fut = client.map(lambda v: save_chunk(src_path, dst_spec, v), part)
+            fut = client.map(lambda v: save_chunk(src_path, dst_spec, v, expand_dims), part)
             logging.info(
                 f"Submitted {len(part)} tasks to the scheduler in {time.time() - start:.2f}s"
             )
             wait(fut)
             logging.info(f"Completed {len(part)} tasks in {time.time() - start:.2f}s")
 
 
-def save_chunk(src_path: str, dst_spec: dict, chunk_slice: Tuple[slice, ...]) -> None:
+def save_chunk(src_path: str, dst_spec: dict, chunk_slice: Tuple[slice, ...],
+               expand_dims: bool = False) -> None:
     """Copy one chunk from an MRC file into a zarr3 TensorStore array.
 
     Chunks that are entirely zero are skipped to avoid unnecessary I/O.
@@ -98,10 +102,14 @@ def save_chunk(src_path: str, dst_spec: dict, chunk_slice: Tuple[slice, ...]) ->
     dst_spec:
         TensorStore zarr3 spec dict for the destination array.
     chunk_slice:
-        The slice tuple identifying the chunk region to copy.
+        The slice tuple identifying the chunk region in the destination array.
+    expand_dims:
+        When True, strip the leading slice when reading the source (which has
+        no channel dimension) and prepend np.newaxis before writing.
     """
     mrc_file = mrcfile.mmap(src_path, mode="r")
-    data = mrc_file.data[chunk_slice]
+    src_slice = chunk_slice[1:] if expand_dims else chunk_slice
+    data = mrc_file.data[src_slice]
     if not (data == 0).all():
         dest_arr = open_ts(dst_spec)
-        dest_arr[chunk_slice].write(data).result()
+        dest_arr[chunk_slice].write(data[np.newaxis] if expand_dims else data).result()

src/zarrify/formats/mrc.py

@@ -177,13 +177,16 @@ def _copy_n5_array_attrs(self, n5_root: str, dest: str, array_paths: list[str])
             for k, v in attrs.items():
                 z_arr.attrs[k] = v
 
-    def apply_ome_template(self, zgroup: zarr.Group) -> dict:
+    def apply_ome_template(self, zgroup: zarr.Group, expand_dims: bool = False) -> dict:
         """Build an OME-NGFF v0.4 multiscales attribute dict from N5 group attributes.
 
         Parameters
         ----------
         zgroup:
             Zarr group with N5-style "axes", "units", and "scales" attributes.
+        expand_dims:
+            When True, prepend a channel axis to the axes list and an extra
+            element to the top-level coordinateTransformations.
 
         Returns
         -------
@@ -196,40 +199,44 @@ def apply_ome_template(self, zgroup: zarr.Group) -> dict:
         ureg = pint.UnitRegistry()
         units_list = [str(ureg.Unit(unit)) for unit in zgroup.attrs['units']]
 
+        axes = [{"name": axis, "type": "space", "unit": unit}
+                for (axis, unit) in zip(zgroup.attrs['axes'], units_list)]
+        if expand_dims:
+            axes = [{"name": "c", "type": "channel"}] + axes
+
         #populate .zattrs
-        z_attrs['multiscales'][0]['axes'] = [{"name": axis,
-                                            "type": "space",
-                                            "unit": unit} for (axis, unit) in zip(zgroup.attrs['axes'],
-                                                                                    units_list)]
+        z_attrs['multiscales'][0]['axes'] = axes
         z_attrs['multiscales'][0]['version'] = '0.4'
         z_attrs['multiscales'][0]['name'] = zgroup.name
-        ndim = len(zgroup.attrs['axes'])
+        ndim = len(axes)
         z_attrs['multiscales'][0]['coordinateTransformations'] = [
             {"type": "scale", "scale": [1.0] * ndim},
             {"type": "translation", "translation": [0.0] * ndim},
         ]
 
         return z_attrs
 
-    def normalize_to_omengff(self, zgroup: zarr.Group) -> None:
+    def normalize_to_omengff(self, zgroup: zarr.Group, expand_dims: bool = False) -> None:
         """Recursively convert N5 metadata to OME-NGFF multiscales attributes.
 
         Parameters
         ----------
         zgroup:
             Root zarr group of the output zarr store.
+        expand_dims:
+            When True, prepend a channel axis to axes and coordinate transforms.
         """
         group_keys = zgroup.keys()
 
         for key in chain(group_keys, '/'):
             if isinstance(zgroup[key], zarr.Group):
                 if key!='/':
-                    self.normalize_to_omengff(zgroup[key])
+                    self.normalize_to_omengff(zgroup[key], expand_dims)
                 if 'scales' in zgroup[key].attrs.asdict():
-                    zattrs = self.apply_ome_template(zgroup[key])
+                    zattrs = self.apply_ome_template(zgroup[key], expand_dims)
                     unsorted_datasets = []
                     for arr in self._iter_arrays(zgroup[key]):
-                        unsorted_datasets.append(self.ome_dataset_metadata(arr[1], zgroup[key]))
+                        unsorted_datasets.append(self.ome_dataset_metadata(arr[1], zgroup[key], expand_dims))
 
                     #1.apply natural sort to organize datasets metadata array for different resolution degrees (s0 -> s10)
                     #2.add datasets metadata to the omengff template
@@ -246,7 +253,8 @@ def _iter_arrays(group: zarr.Group):
                 yield from N5Group._iter_arrays(node)
 
     @staticmethod
-    def ome_dataset_metadata(n5arr: zarr.Array, group: zarr.Group) -> dict:
+    def ome_dataset_metadata(n5arr: zarr.Array, group: zarr.Group,
+                             expand_dims: bool = False) -> dict:
         """Build one OME-NGFF dataset metadata entry from an N5 array.
 
         Parameters
@@ -255,23 +263,27 @@ def ome_dataset_metadata(n5arr: zarr.Array, group: zarr.Group) -> dict:
             Source N5 array with a "transform" attribute.
         group:
             Parent group used to compute the relative path.
+        expand_dims:
+            When True, prepend 1.0/0.0 to the scale/translation vectors.
 
         Returns
         -------
         dict
             A single entry suitable for the "datasets" list in multiscales.
         """
-
         arr_attrs_n5 = n5arr.attrs['transform']
-        dataset_meta =  {
-                        "path": os.path.relpath(n5arr.path, group.path),
-                        "coordinateTransformations": [{
-                            'type': 'scale',
-                            'scale': arr_attrs_n5['scale']},{
-                            'type': 'translation',
-                            'translation' : arr_attrs_n5['translate']
-                        }]}
-
+        scale = arr_attrs_n5['scale']
+        translate = arr_attrs_n5['translate']
+        if expand_dims:
+            scale = [1.0] + list(scale)
+            translate = [0.0] + list(translate)
+        dataset_meta = {
+            "path": os.path.relpath(n5arr.path, group.path),
+            "coordinateTransformations": [
+                {'type': 'scale', 'scale': scale},
+                {'type': 'translation', 'translation': translate},
+            ],
+        }
         return dataset_meta
 
     def write_to_zarr(
@@ -281,6 +293,7 @@ def write_to_zarr(
         chunk_shape: list[int],
         shard_shape: list[int] | None = None,
         codec: dict = zstd_codec(level=6),
+        expand_dims: bool = False,
     ) -> None:
         """Copy all N5 arrays into a zarr3 store via TensorStore.
 
@@ -325,12 +338,14 @@ def write_to_zarr(
             src_arr = open_ts(src_spec)
             shape = src_arr.shape
             dtype = np.dtype(src_arr.dtype.numpy_dtype)
+            out_shape = (1, *shape) if expand_dims else shape
 
             # trim chunk/shard shapes to array ndim; N5 trees can hold mixed-dimensionality arrays
-            arr_chunk_shape = [min(c, s) for c, s in zip(list(chunk_shape)[-len(shape):], shape)]
+            arr_chunk_shape_base = [min(c, s) for c, s in zip(list(chunk_shape)[-len(shape):], shape)]
+            arr_chunk_shape = ([1] + arr_chunk_shape_base) if expand_dims else arr_chunk_shape_base
             arr_shard_shape = (
                 align_shard_to_chunks(
-                    [min(s, dim) for s, dim in zip(list(shard_shape)[-len(shape):], shape)],
+                    [min(s, dim) for s, dim in zip(list(shard_shape)[-len(out_shape):], out_shape)],
                     arr_chunk_shape,
                 )
                 if shard_shape is not None else None
@@ -341,7 +356,7 @@ def write_to_zarr(
             dst_spec = zarr3_spec(
                 store_path=dest,
                 array_path=rel_path,
-                shape=shape,
+                shape=out_shape,
                 dtype=dtype,
                 chunk_shape=arr_chunk_shape,
                 shard_shape=arr_shard_shape,
@@ -353,7 +368,7 @@ def write_to_zarr(
             dest_chunks = dest_arr.chunk_layout.write_chunk.shape
 
             out_slices = slices_from_chunks(
-                normalize_chunks(dest_chunks, shape=shape)
+                normalize_chunks(dest_chunks, shape=out_shape)
             )
             # break the slices up into batches, to make things easier for the dask scheduler
             out_slices_partitioned = tuple(partition_all(100000, out_slices))
@@ -362,7 +377,7 @@ def write_to_zarr(
                 logging.info(f"{idx + 1} / {len(out_slices_partitioned)}")
                 start = time.time()
                 fut = client.map(
-                    lambda v: save_chunk(src_spec, dst_spec, v, invert=False), part
+                    lambda v: save_chunk(src_spec, dst_spec, v, invert=False, expand_dims=expand_dims), part
                 )
                 logging.info(
                     f"Submitted {len(part)} tasks to the scheduler in {time.time() - start:.2f}s"
@@ -376,14 +391,15 @@ def write_to_zarr(
         # copy array-level N5 attributes (e.g. transform) then build OME metadata
         self._copy_n5_array_attrs(n5_root_path, dest, n5_array_paths)
         z_root = zarr.open_group(store=z_store, mode='a')
-        self.normalize_to_omengff(z_root)
+        self.normalize_to_omengff(z_root, expand_dims)
 
 
 def save_chunk(
     src_spec: dict,
     dst_spec: dict,
     chunk_slice: Tuple[slice, ...],
     invert: bool = False,
+    expand_dims: bool = False,
 ) -> None:
     """Copy one chunk from an N5 array into a zarr3 TensorStore array.
 
@@ -396,16 +412,22 @@ def save_chunk(
     dst_spec:
         TensorStore zarr3 driver spec for the destination array.
     chunk_slice:
-        Slice tuple identifying the chunk region to copy.
+        Slice tuple identifying the chunk region in the destination array.
     invert:
         When True, apply bitwise inversion to the data before writing.
+    expand_dims:
+        When True, strip the leading slice when reading the source (which has
+        no channel dimension) and prepend np.newaxis before writing.
     """
     src = open_ts(src_spec)
-    data = src[chunk_slice].read().result()
+    src_slice = chunk_slice[1:] if expand_dims else chunk_slice
+    data = src[src_slice].read().result()
     # only store data if it is not all 0s
     if (data == 0).all():
         return
     if invert:
         data = np.invert(data)
+    if expand_dims:
+        data = data[np.newaxis]
     dest = open_ts(dst_spec)
     dest[chunk_slice].write(data).result()

src/zarrify/formats/n5.py

@@ -69,7 +69,7 @@ def __init__(
         self.metadata["translation"] = self.metadata["translation"][-self.ndim:]
         self.metadata["units"] = self.metadata["units"][-self.ndim:]
 
-    def write_to_zarr(self, dst_spec: dict, client: Client) -> None:
+    def write_to_zarr(self, dst_spec: dict, client: Client, expand_dims: bool = False) -> None:
         """Read the TIFF file in slabs and write each slab to a zarr3 array via TensorStore.
 
         Parameters
@@ -79,6 +79,9 @@ def write_to_zarr(self, dst_spec: dict, client: Client) -> None:
             by :func:`~zarrify.utils.ts_utils.zarr3_spec`.
         client:
             Dask distributed client used to parallelise slab writes.
+        expand_dims:
+            When True, the destination array has a leading size-1 channel
+            dimension; each slab is written with an extra np.newaxis prepended.
         """
         # TODO: With large shard shapes (e.g. 1024^3) the slab thickness along
         # the slab axis becomes 1024 voxels, potentially exhausting worker memory.
@@ -87,21 +90,25 @@ def write_to_zarr(self, dst_spec: dict, client: Client) -> None:
         # can process multiple passes concurrently without holding the full shard
         # in memory at once.
 
-        axes = self.metadata["axes"]
-        slab_axis = axes.index("z") if "z" in axes else 0
+        # axes in metadata are already trimmed to source ndim; use them to find z.
+        src_axes = self.metadata["axes"][1:] if expand_dims else self.metadata["axes"]
+        slab_axis = src_axes.index("z") if "z" in src_axes else 0
 
         dest_arr = open_ts(dst_spec)
+        # dest_chunks includes the leading 1 when expand_dims; strip it for
+        # slice computation which is done against the source shape.
         dest_chunks = list(dest_arr.chunk_layout.write_chunk.shape)
+        src_chunks = dest_chunks[1:] if expand_dims else dest_chunks
 
-        slice_chunks = dest_chunks
+        slice_chunks = src_chunks
         if self.optimize_reads:
             logger.info("Optimizing read chunking...")
             logger.info(f"Output zarr3 write-chunk shape: {dest_chunks}")
             logger.info(f"Input TIFF chunk shape: {self._tiff_chunks}")
-            slice_chunks = dest_chunks[: slab_axis + 1].copy()
+            slice_chunks = src_chunks[: slab_axis + 1].copy()
 
             for dest_chunkdim, tiff_chunkdim, tiff_dim in zip(
-                dest_chunks[slab_axis + 1:],
+                src_chunks[slab_axis + 1:],
                 self._tiff_chunks[slab_axis + 1:],
                 self.shape[slab_axis + 1:],
             ):
@@ -123,7 +130,7 @@ def write_to_zarr(self, dst_spec: dict, client: Client) -> None:
 
         start = time.time()
         fut = client.map(
-            lambda v: write_volume_slab(v, dst_spec, src_path), slice_tuples
+            lambda v: write_volume_slab(v, dst_spec, src_path, expand_dims), slice_tuples
         )
         logger.info(
             f"Submitted {len(slice_tuples)} tasks to the scheduler in "
@@ -134,7 +141,8 @@ def write_to_zarr(self, dst_spec: dict, client: Client) -> None:
         logger.info(f"Completed {len(slice_tuples)} tasks in {time.time() - start:.2f}s")
 
 
-def write_volume_slab(slice_tuple: tuple, dst_spec: dict, src_path: str) -> None:
+def write_volume_slab(slice_tuple: tuple, dst_spec: dict, src_path: str,
+                      expand_dims: bool = False) -> None:
     """Copy one slab from a TIFF file into a zarr3 TensorStore array.
 
     The TIFF is read directly via tifffile into NumPy (no TensorStore TIFF
@@ -143,12 +151,18 @@ def write_volume_slab(slice_tuple: tuple, dst_spec: dict, src_path: str) -> None
     Parameters
     ----------
     slice_tuple:
-        The slice tuple identifying the slab region to copy.
+        The slice tuple identifying the slab region in the source array.
     dst_spec:
         TensorStore zarr3 spec dict for the destination array.
     src_path:
         Path to the source TIFF file.
+    expand_dims:
+        When True, prepend a size-1 channel dimension to the data before
+        writing and offset the destination slice accordingly.
     """
-    data = imread(src_path)[slice_tuple]
+    data = np.asarray(imread(src_path)[slice_tuple])
     dest_arr = open_ts(dst_spec)
-    dest_arr[slice_tuple].write(np.asarray(data)).result()
+    if expand_dims:
+        dest_arr[(slice(0, 1), *slice_tuple)].write(data[np.newaxis]).result()
+    else:
+        dest_arr[slice_tuple].write(data).result()