fix: fix the chunking along Z (can handle prime numbers now) (#376)

* fix: fix the chunking along Z (can handle prime numbers now)

Signed-off-by: Sricharan Reddy Varra <sricharan.varra@biohub.org>

* fix: simplified chunk sizes for conversion

Signed-off-by: Sricharan Reddy Varra <sricharan.varra@biohub.org>

---------

Signed-off-by: Sricharan Reddy Varra <sricharan.varra@biohub.org>

Author: srivarra

src/iohub/convert.py

@@ -4,7 +4,6 @@
 from pathlib import Path
 from typing import Literal
 
-import dask
 import numpy as np
 from tqdm import tqdm
 from tqdm.contrib.itertools import product
@@ -13,7 +12,7 @@
 from iohub.ngff.models import TransformationMeta
 from iohub.ngff.nodes import Position, open_ome_zarr
 from iohub.ngff.utils import (
-    _adjust_chunks_for_divisibility,
+    _clamp_chunks_to_shape,
     _limit_zyx_chunk_size,
 )
 from iohub.reader import MMStack, NDTiffDataset, read_images
@@ -265,11 +264,11 @@ def _gen_chunks(self, input_chunks):
         chunk_zyx_shape = _limit_zyx_chunk_size(shape, bytes_per_pixel, MAX_CHUNK_SIZE, chunks=chunks)
         chunks[-3:] = list(chunk_zyx_shape)
 
-        # Adjust chunks to divide evenly into dimensions
-        chunks = _adjust_chunks_for_divisibility(shape, chunks)
+        # Clamp chunks so they don't exceed dimension sizes
+        chunks = _clamp_chunks_to_shape(shape, chunks)
         for i, (orig, adj, dim) in enumerate(zip(original_chunks, chunks, shape, strict=False)):
-            if orig != adj:
-                _logger.warning(f"Chunk size {orig} on axis {i} adjusted to {adj} (dimension {dim}).")
+            if adj < orig:
+                _logger.warning(f"Chunk size {orig} on axis {i} clamped to {adj} (dimension size {dim}).")
 
         _logger.debug(f"Zarr store chunk size will be set to {chunks}.")
 
@@ -384,13 +383,9 @@ def __call__(self) -> None:
         _logger.debug("Setting up Zarr store.")
 
         self._init_zarr_arrays()
-        # Calculate chunk size in bytes for dask config
-        # This prevents data loss when rechunking to zarr chunks
-        # See: https://github.com/czbiohub-sf/iohub/issues/367
-        chunk_size_bytes = int(np.prod(self.chunks) * np.dtype(self.reader.dtype).itemsize)
 
         _logger.debug("Converting images.")
-        with logging_redirect_tqdm(), dask.config.set({"array.chunk-size": chunk_size_bytes}):
+        with logging_redirect_tqdm():
             for zarr_pos_name, (_, fov) in tqdm(
                 zip(self.zarr_position_names, self.reader, strict=True),
                 total=len(self.zarr_position_names),

src/iohub/core/implementations/zarr_python.py

@@ -139,9 +139,10 @@ def to_dask(self, handle: zarr.Array) -> Any:
         return da.from_zarr(handle)
 
     def write_from_dask(self, handle: zarr.Array, dask_array: Any) -> None:
-        from dask.array import to_zarr
+        from iohub.core.downsample import iter_work_regions
 
-        to_zarr(dask_array, handle)
+        for region in iter_work_regions(handle.shape, handle.chunks):
+            handle[region] = dask_array[region].compute()
 
     # -- High-performance operations ---------------------------------------
 

src/iohub/ngff/utils.py

@@ -447,8 +447,8 @@ def _limit_zyx_chunk_size(
     return chunk_zyx_shape
 
 
-def _adjust_chunks_for_divisibility(shape: tuple[int, ...], chunks: tuple[int, ...]) -> tuple[int, ...]:
-    """Adjust chunks to divide evenly into dimensions for Dask.
+def _clamp_chunks_to_shape(shape: tuple[int, ...], chunks: tuple[int, ...]) -> tuple[int, ...]:
+    """Clamp chunk sizes so they don't exceed dimension sizes.
 
     Parameters
     ----------
@@ -460,18 +460,6 @@ def _adjust_chunks_for_divisibility(shape: tuple[int, ...], chunks: tuple[int, .
     Returns
     -------
     tuple[int, ...]
-        Adjusted chunk sizes that divide evenly into dimensions.
+        Chunk sizes clamped to at most the dimension size.
     """
-    shape = list(shape)
-    chunks = list(chunks)
-    adjusted = []
-    for chunk, dim in zip(chunks, shape, strict=False):
-        if chunk > dim:
-            adjusted.append(dim)
-        elif dim % chunk != 0:
-            while chunk > 1 and dim % chunk != 0:
-                chunk -= 1
-            adjusted.append(chunk)
-        else:
-            adjusted.append(chunk)
-    return tuple(adjusted)
+    return tuple(min(c, d) for c, d in zip(chunks, shape, strict=False))

tests/test_converter.py

@@ -8,7 +8,7 @@
 from ndtiff import Dataset
 from tifffile import TiffFile
 
-from iohub.convert import TIFFConverter, _adjust_chunks_for_divisibility
+from iohub.convert import TIFFConverter, _clamp_chunks_to_shape
 from iohub.ngff import Position, open_ome_zarr
 from iohub.reader import MMStack, NDTiffDataset
 from tests.conftest import (
@@ -54,7 +54,7 @@ def _check_chunks(position: Position, chunks: Literal["XY", "XYZ"] | tuple[int]
         case "XYZ" | None:
             assert img.chunks == (1,) * 2 + img.shape[-3:]
         case tuple():
-            expected = _adjust_chunks_for_divisibility(img.shape, chunks)
+            expected = _clamp_chunks_to_shape(img.shape, chunks)
             assert img.chunks == expected
         case _:
             raise AssertionError()
@@ -264,34 +264,41 @@ class MockReader:
     @pytest.mark.parametrize(
         ("z_dim", "input_chunk", "expected_z"),
         [
-            (15, 10, 5),  # 10 doesn't divide 15, adjust to 5
-            (7, 5, 1),  # prime - falls to 1
-            (16, 8, 8),  # divides evenly, no change
-            (100, 50, 50),  # divides evenly
-            (9, 4, 3),  # odd non-prime, adjust to 3
+            (15, 10, 10),  # 10 < 15, no clamping needed
+            (7, 5, 5),  # 5 < 7, no clamping needed
+            (16, 8, 8),  # 8 < 16, no clamping needed
+            (100, 50, 50),  # 50 < 100, no clamping needed
+            (9, 4, 4),  # 4 < 9, no clamping needed
+            (5, 10, 5),  # 10 > 5, clamped to 5
         ],
     )
-    def test_gen_chunks_divisibility(self, make_mock_converter, z_dim, input_chunk, expected_z):
-        """Test chunks are adjusted to divide evenly into dimensions."""
+    def test_gen_chunks_clamping(self, make_mock_converter, z_dim, input_chunk, expected_z):
+        """Test chunks are clamped to dimension size but not reduced for divisibility."""
         converter = make_mock_converter(z=z_dim)
         chunks = converter._gen_chunks((1, 1, input_chunk, 256, 256))
         assert chunks[2] == expected_z
-        assert z_dim % chunks[2] == 0
+        assert chunks[2] <= z_dim
 
-    def test_gen_chunks_max_chunk_size_then_divisibility(self, make_mock_converter):
-        """Test divisibility check happens AFTER MAX_CHUNK_SIZE adjustment."""
-        # Large chunks that trigger MAX_CHUNK_SIZE, then need divisibility fix
-        # Z=15, large XY to trigger size limit, chunk should end up dividing 15
+    def test_gen_chunks_max_chunk_size_then_clamp(self, make_mock_converter):
+        """Test clamping happens AFTER MAX_CHUNK_SIZE adjustment."""
         converter = make_mock_converter(z=15, y=2048, x=2048)
         chunks = converter._gen_chunks((1, 1, 15, 2048, 2048))
-        assert 15 % chunks[2] == 0
+        assert chunks[2] <= 15
 
     @pytest.mark.parametrize("z_dim", [7, 11, 13, 17, 19, 23])  # prime numbers
     def test_gen_chunks_prime_dimensions(self, make_mock_converter, z_dim):
-        """Test handling of prime dimension sizes - should fall to 1."""
+        """Test prime dimensions do NOT reduce chunk to 1."""
         converter = make_mock_converter(z=z_dim)
         chunks = converter._gen_chunks((1, 1, z_dim - 1, 256, 256))
-        assert z_dim % chunks[2] == 0
+        assert chunks[2] == z_dim - 1  # should stay as-is (< dim)
+        assert chunks[2] > 1
+
+    def test_gen_chunks_prime_z_1187(self, make_mock_converter):
+        """Regression: Z=1187 (prime) should not reduce chunk to 1."""
+        converter = make_mock_converter(z=1187, y=2048, x=2048)
+        chunks = converter._gen_chunks("XYZ")
+        assert chunks[2] > 1
+        assert chunks[2] <= 1187
 
     @pytest.mark.parametrize("input_chunks", ["XY", "XYZ", None])
     def test_gen_chunks_string_inputs(self, make_mock_converter, input_chunks):
@@ -342,6 +349,43 @@ def test_rechunk_xy_to_xyz_preserves_data(shape, target_chunks, tmpdir):
     np.testing.assert_array_equal(data, written, err_msg="Data lost during XY to XYZ rechunking")
 
 
+@pytest.mark.parametrize(
+    ("shape", "target_chunks"),
+    [
+        # Non-divisible Z: 1187 is prime, chunk 512 doesn't divide evenly
+        ((1, 1, 1187, 256, 256), (1, 1, 512, 256, 256)),
+        # Another non-divisible case
+        ((1, 1, 50, 512, 512), (1, 1, 32, 512, 512)),
+    ],
+)
+def test_rechunk_non_divisible_z_preserves_data(shape, target_chunks, tmpdir):
+    """Test that rechunking with non-divisible Z chunks preserves all data.
+
+    Regression test for https://github.com/czbiohub-sf/iohub/issues/374
+    """
+    import dask
+    import dask.array as da
+
+    data = np.arange(np.prod(shape), dtype=np.uint16).reshape(shape)
+
+    source_chunks = (1, 1, 1, shape[3], shape[4])
+    dask_data = da.from_array(data, chunks=source_chunks)
+
+    output = tmpdir / "test_rechunk_nondiv.zarr"
+    with open_ome_zarr(output, layout="hcs", mode="w-", channel_names=["test"], version="0.5") as writer:
+        pos = writer.create_position("0", "0", "0")
+        zarr_img = pos.create_zeros(name="0", shape=shape, dtype=np.uint16, chunks=target_chunks)
+
+        chunk_size_bytes = int(np.prod(target_chunks) * 2)
+        with dask.config.set({"array.chunk-size": chunk_size_bytes}):
+            zarr_img.write_from_dask(dask_data.rechunk(target_chunks))
+
+    with open_ome_zarr(output, layout="hcs", mode="r") as reader:
+        written = reader["0/0/0"]["0"][:]
+
+    np.testing.assert_array_equal(data, written, err_msg="Data lost during non-divisible Z rechunking")
+
+
 class TestVersionParameter:
     """Tests for the OME-NGFF version parameter on TIFFConverter."""