Merge pull request #16 from janelia-cellmap/multiscale

extension of #15 
This PR consolidates several feature branches that support conversion of light imagery:

    Builds on top of the tif_volume branch
    Move to Pixi build
    Allow LSF directives on the command line
    Create .zattrs before writing data
    Use Z axis for slabs
    Optionally turn off read optimizations because they don't work in all cases
    Don't write empty units (breaks Neuroglancer)
Added:

    Integrated multiscale pyramid generation
    Added config as click option to specify all the necessary parameters in a yaml file
    Added pydantic model to validate input parameters
    Rewrite optimize reads for tiff file conversion. Read the smallest tiff segment but big enough to write into chunks only once (align with the chunk grid)

Author: yuriyzubov

.github/workflows/ci.yaml

@@ -12,7 +12,8 @@ jobs:
 
     strategy:
       matrix:
-        python-version: ['3.12'] #['3.10', '3.11', '3.12']
+        python-version: ['3.12']
+
 
     steps:
       - uses: actions/checkout@v4

.gitignore

@@ -7,4 +7,6 @@ dist/
 .python-version
 *.egg-info/
 build/
-.vscode/
+.vscode/
+.claude
+CLAUDE.md

README.md

@@ -1,24 +1,72 @@
-# zarrify
+# Zarrify
 
-Convert TIFF/TIFF Stacks, MRC, and N5 files to Zarr (v2) format with ome-ngff (0.4) metadata.
+Convert TIFF/TIFF Stacks, MRC, and N5 files to Zarr (v2) format with OME-NGFF 0.4 metadata.
 
-## Install
+## Prerequisites
+
+You must [install Pixi](https://pixi.sh/latest/installation/) to run this tool.
+
+## Install into local env
+
+After you check out the repo, run Pixi to download dependencies and create the local environment:
 
 ```bash
-pip install zarrify
+pixi run dev-install
 ```
 
 ## Usage
 
-### Local processing
-```bash
-zarrify --src input.tiff --dest output.zarr --cluster local --workers 20
+### Step 1: Convert image to OME-Zarr format
+
+To convert an OME-TIFF image with axes ZCYX:
+
+```
+pixi run zarrify --src=/path/to/input --dest=/path/to/output.zarr   --axes z c y x --units nanometer '' nanometer nanometer --zarr_chunks 64 3 128 128
 ```
 
-### LSF cluster processing
-```bash
-bsub -n 1 -J to_zarr 'zarrify --src input.tiff --dest output.zarr --cluster lsf --workers 20'
+### Step 2: Add multiscale pyramid
+
+To add multiscale pyramids to the above OME-Zarr:
+
 ```
+pixi run zarr-multiscale --config=/path/to/config.yaml
+```
+
+Where `config.yaml` contains the parameters, e.g.:
+
+```yaml
+src: "/path/to/output.zarr"
+workers: 50
+data_origin: "raw"
+cluster: "lsf"
+log_dir: "/path/to/logs/workers"
+antialiasing: false
+high_aspect_ratio: false
+custom_scale_factors: 
+  - [1, 1, 1, 1]
+  - [1, 1, 2, 2]
+  - [1, 1, 4, 4]
+  - [1, 1, 8, 8]
+  - [1, 2, 16, 16]
+  - [1, 4, 32, 32]
+  - [1, 8, 64, 64]
+  - [1, 16, 128, 128]
+  - [1, 32, 256, 256]
+  - [1, 64, 512, 512]
+  - [1, 128, 1024, 1024]
+```
+
+### IBM Platform LSF 
+
+You can run the pipeline on an LSF cluster by specifying the `--cluster=lsf` parameter. For example:
+
+```
+export INPUT=/path/to/input
+export OUTPUT=/path/to/output.zarr
+export LOGDIR=`pwd`
+bsub -n 4 -P myproject -o $LOGDIR/jobout.log -e $LOGDIR/joberr.log "/misc/sc/pixi run zarrify --src=$INPUT --dest=$OUTPUT --cluster $CLUSTER --log_dir $LOGDIR/workers --workers 500 --zarr_chunks 3 128 128 128 --extra_directives='-P myproject' --scale 1.0 1.0 0.116 0.116 --optimize_reads"
+```
+
 
 ## Python API
 

config_example.yaml

@@ -0,0 +1,25 @@
+src : "/path_to_file/input_file.tiff" # image source path
+dest : "/path_to_output/output_file.zarr" # zarr container path
+
+# Dask configs
+cluster : local # "local"/"lsf" cluster
+workers : 20 # number of dask workers
+log_dir : # specify log directory, otherwise get an email from each worker
+extra_directives : null # pass extra dask options
+
+# Output zarr array configs
+zarr_chunks : [128, 128, 128] # output chunksize 
+axes : [z, y, x] # axis order
+translation : [0.0, 0.0, 0.0] # array offset
+scale : [8.0, 8.0, 8.0] # voxel size
+units : [nanometer, nanometer, nanometer] # physical units 
+optimize_reads : True 
+
+# Multiscale configs
+multiscale : True # whether to create a multiscale pyramid from converted zarr array or not 
+ms_workers : 20 # number of dask workers
+data_origin : "raw" # "raw"/"labels" 
+antialiasing : False # True/False blur dataset with Gaussian filter before downsampling to reduce aliasing       
+normalize_voxel_size : False # True/False
+custom_scale_factors : null # specify custom multipliers for each level [[1,1,1], [2,4,8], [4, 8, 16],...]
+

pixi.lock

@@ -1,12 +1,12 @@
 [project]
 name = "zarrify"
-version = "0.0.8"
+version = "0.1.0"
 description = "Convert scientific image formats (TIFF, MRC, N5) to OME-Zarr"
 authors = [
     { name = "Yurii Zubov", email = "[email protected]" }
 ]
 readme = "README.md"
-requires-python = ">=3.10"
+requires-python = ">=3.12"
 license = { text = "MIT" }
 keywords = ["zarr", "ome-zarr", "tiff", "mrc", "n5", "scientific-imaging"]
 classifiers = [
@@ -17,32 +17,56 @@ classifiers = [
     "Programming Language :: Python :: 3.11",
     "Topic :: Scientific/Engineering :: Image Processing",
 ]
-
 dependencies = [
+    "bokeh>=3.3.2",
     "click>=8.1.8,<9.0.0",
-    "colorama>=0.4.6,<0.5.0",
+    "colorama>=0.4.6,<0.5.0", 
     "dask>=2024.12.1,<2025.0.0",
-    "dask-jobqueue==0.8.2",
+    "dask-jobqueue==0.9.0",
     "imagecodecs>=2024.12.30,<2025.0.0",
     "mrcfile>=1.5.3,<2.0.0",
     "natsort>=8.4.0,<9.0.0",
-    "numcodecs>=0.13.0,<0.16",
-    "pydantic-zarr>=0.4.0,<0.5.0",
+    "numcodecs>=0.15.1,<0.16.1",
+    "numpy>=1.26.2",
+    "pydantic>=2.11.7,<3",
+    "pydantic-zarr>=0.4.0,<0.9.0",
     "pint>=0.20.0,<1.0.0",
+    "scipy>=1.16.1",
     "tifffile>=2024.1.0,<2025.5.10",
-    "zarr==2.18.3",
-    "bokeh>=3.1.0"
+    "xarray-multiscale>=2.1.0",
+    "zarr>=2.18.3,<3.0"
 ]
 
 [project.scripts]
 zarrify = "zarrify.to_zarr:cli"
+zarr-multiscale = "zarrify.multiscale.to_multiscale:cli"
 
 [project.optional-dependencies]
 test = [
     "pytest"
 ]
 
+[tool.pixi.project]
+channels = ["conda-forge"]
+platforms = ["osx-64", "osx-arm64", "linux-64"]
+[tool.pixi.tasks]
+dev-install = "pip install -e ."
+zarrify = { cmd = "zarrify" }
+zarr-multiscale = { cmd = "zarr-multiscale" }
+test = "pytest"
+clean = "rm -rf build/ dist/ *.egg-info/ .pytest_cache/ __pycache__/ dask_dashboard_link.txt"
+
+[tool.pixi.feature.test.dependencies]
+pytest = "*"
+
+[tool.pixi.feature.test.tasks]
+test-install = "pip install -e ."
+test = { cmd = "pytest", depends-on = ["test-install"] }
+
+[tool.pixi.environments]
+default = { solve-group = "default" }
+test = { features = ["test"], solve-group = "default" }
 
-[build-system]
-requires = ["setuptools>=61.0"]
-build-backend = "setuptools.build_meta"
+[tool.pixi.dependencies]
+python = "3.12.*"
+pip = ">=25.2,<26"

pyproject.toml

@@ -38,10 +38,8 @@ def __init__(
 
     def write_to_zarr(
         self,
-        dest: str,
+        zarr_array: zarr.Array,
         client: Client,
-        zarr_chunks : list[int],
-        comp : ABCMeta = Zstd(level=6),
     ):
         """Use mrcfile memmap to access small parts of the mrc file and write them into zarr chunks.
 
@@ -55,10 +53,7 @@ def write_to_zarr(
 
         src_path = copy.copy(self.src_path)
 
-        if len(zarr_chunks) != len(self.shape):
-           zarr_chunks = self.reshape_to_arr_shape(zarr_chunks, self.shape)
-        
-        z_arr = self.get_output_array(dest, zarr_chunks, comp)
+        z_arr = zarr_array
 
         out_slices = slices_from_chunks(
             normalize_chunks(z_arr.chunks, shape=z_arr.shape)

src/zarrify/formats/mrc.py

@@ -4,15 +4,21 @@
 import os
 from dask.distributed import Client, wait
 import time
-import dask.array as da
 import copy
 from zarrify.utils.volume import Volume
 from abc import ABCMeta
 from numcodecs import Zstd
+from dask.array.core import slices_from_chunks, normalize_chunks
 import logging
 
+logging.basicConfig(
+    level=logging.INFO,
+    format="%(asctime)s %(levelname)s %(name)s: %(message)s"
+)
+logger = logging.getLogger(__name__)
 
-class Tiff3D(Volume):
+
+class Tiff(Volume):
 
     def __init__(
         self,
@@ -21,6 +27,7 @@ def __init__(
         scale: list[float],
         translation: list[float],
         units: list[str],
+        optimize_reads: bool = False,
     ):
         """Construct all the necessary attributes for the proper conversion of tiff to OME-NGFF Zarr.
 
@@ -35,57 +42,78 @@ def __init__(
         self.shape = self.zarr_arr.shape
         self.dtype = self.zarr_arr.dtype
         self.ndim = self.zarr_arr.ndim
-        
+        self.optimize_reads = optimize_reads
+
         # Scale metadata parameters to match data dimensionality
-        self.metadata["axes"] = self.metadata["axes"][-self.ndim:]
+        self.metadata["axes"] = list(self.metadata["axes"])[-self.ndim:]
         self.metadata["scale"] = self.metadata["scale"][-self.ndim:]
         self.metadata["translation"] = self.metadata["translation"][-self.ndim:]
         self.metadata["units"] = self.metadata["units"][-self.ndim:]
 
     def write_to_zarr(self,
-        dest: str,
+        zarr_array: zarr.Array,
         client: Client,
-        zarr_chunks : list[int],
-        comp : ABCMeta = Zstd(level=6),
         ):
 
-        # reshape chunk shape to align with arr shape
-        if len(zarr_chunks) != self.shape:
-           zarr_chunks = self.reshape_to_arr_shape(zarr_chunks, self.shape)
-             
-        z_arr = self.get_output_array(dest, zarr_chunks, comp)
-        chunks_list = np.arange(0, z_arr.shape[0], z_arr.chunks[0])
+        # Find slab axis based on metadata axes - use z axis for slabbing
+        axes = self.metadata["axes"]
+        slab_axis = axes.index('z') if 'z' in axes else 0
+            
+        z_arr = zarr_array
+        
+        slice_chunks = z_arr.chunks
+        if self.optimize_reads:
+            logger.info("Optimizing read chunking...")
+            # TODO: this works for some cases, doesn't work in others, need to understand why
+            #slicing
+            #(c, z, y, x) or (z, y, x) - combine (c, z) from zarr_chunks and (y, x) from tiff chunking
+            logger.info(f"Output Zarr array chunks: {z_arr.chunks}")
+            logger.info(f"Input Tiff array chunks: {self.zarr_arr.chunks}")
+            slice_chunks = list(z_arr.chunks[:slab_axis+1]).copy()
+            logger.info(f"Slice chunks: {slice_chunks}")
+        
+            # cast slab size to write into zarr:
+            for zarr_chunkdim, tiff_chunkdim, tiff_dim in zip(z_arr.chunks[slab_axis+1:], self.zarr_arr.chunks[slab_axis+1:], self.zarr_arr.shape[slab_axis+1:]):
+                if tiff_chunkdim < zarr_chunkdim:
+                    slice_chunks.append(zarr_chunkdim)
+                elif tiff_chunkdim/tiff_dim < 0.5:
+                    slice_chunks.append(int(tiff_chunkdim/zarr_chunkdim)*zarr_chunkdim)
+                else:
+                    slice_chunks.append(tiff_dim)
+            
+            logger.info(f"Slice chunks extended: {slice_chunks}")
+        
+        # compute size of the slab 
+        slab_size_bytes = np.prod(slice_chunks) * np.dtype(self.dtype).itemsize
+        
+        # get dask worker allocated memery size
+        dask_worker_memory_bytes = next(iter(client.scheduler_info()["workers"].values()))["memory_limit"]
+        
+        logger.info(f"Slab size: {slab_size_bytes / 1e9} GB")
+        logger.info(f"Dask memory limit: {dask_worker_memory_bytes / 1e9} GB")
+        if slab_size_bytes > dask_worker_memory_bytes:
+            raise ValueError("Tiff segment size exceeds Dask worker memory limit. Please reduce the chunksize of the output array.")
+        
+        logger.info(f"Zarr array shape: {self.zarr_arr.shape}")
+        normalized_chunks = normalize_chunks(slice_chunks, shape=self.zarr_arr.shape)
+        slice_tuples = slices_from_chunks(normalized_chunks)
 
         src_path = copy.copy(self.src_path)
 
         start = time.time()
         fut = client.map(
-            lambda v: write_volume_slab_to_zarr(v, z_arr, src_path), chunks_list
-        )
-        logging.info(
-            f"Submitted {len(chunks_list)} tasks to the scheduler in {time.time()- start}s"
+            lambda v: write_volume_slab_to_zarr(v, z_arr, src_path), slice_tuples
         )
+        logger.info(f"Submitted {len(slice_tuples)} tasks to the scheduler in {round(time.time()- start, 4)}s")
 
         # wait for all the futures to complete
         result = wait(fut)
-        logging.info(f"Completed {len(chunks_list)} tasks in {time.time() - start}s")
+        logger.info(f"Completed {len(slice_tuples)} tasks in {round(time.time() - start, 2)}s")
 
         return 0
 
 
-def write_volume_slab_to_zarr(chunk_num: int, zarray: zarr.Array, src_path: str):
-
-    # check if the slab is at the array boundary or not
-    if chunk_num + zarray.chunks[0] > zarray.shape[0]:
-        slab_thickness = zarray.shape[0] - chunk_num
-    else:
-        slab_thickness = zarray.chunks[0]
-
-    slab_shape = [slab_thickness] + list(zarray.shape[-2:])
-    np_slab = np.empty(slab_shape, zarray.dtype)
-
-    tiff_slab = imread(src_path, key=range(chunk_num, chunk_num + slab_thickness, 1))
-    np_slab[0 : zarray.chunks[0], :, :] = tiff_slab
-
-    # write a tiff stack slab into zarr array
-    zarray[chunk_num : chunk_num + zarray.chunks[0], :, :] = np_slab
+def write_volume_slab_to_zarr(slice: slice, zarray: zarr.Array, src_path: str):
+    tiff_store = imread(src_path, aszarr=True)
+    src_tiff_arr = zarr.open(tiff_store, mode='r')
+    zarray[slice] = src_tiff_arr[slice]

src/zarrify/formats/tiff.py

@@ -67,17 +67,11 @@ def write_tile_slab_to_zarr(
 
     # parallel writing of tiff stack into zarr array
     def write_to_zarr(self,
-        dest: str,
+        zarr_array: zarr.Array,
         client: Client,
-        zarr_chunks : list[int],
-        comp : ABCMeta = Zstd(level=6),
         ):
 
-        # reshape chunk shape to align with arr shape
-        if len(zarr_chunks) != len(self.shape):
-           zarr_chunks = self.reshape_to_arr_shape(zarr_chunks, self.shape)
-        
-        z_arr = self.get_output_array(dest, zarr_chunks, comp)
+        z_arr = zarr_array
         chunks_list = np.arange(0, z_arr.shape[0], z_arr.chunks[0])
 
         start = time.time()