fmi-faim/faim-ipa

Migrate `write_zarr` to here.

imagejan opened this issue · 0 comments

This code (from a currently private repository) would be useful to have in faim-hcs:

def write_zarr(client, zarr_path, stitched_da, worm_acquisition):
    """
    Write dask-array to ome-zarr multiscale.

    Parameters
    ----------
    client :
        Dask client used for the computation.
    zarr_path :
        Path to the output zarr file.
    stitched_da :
        Dask array containing the stitched data.
    worm_acquisition :
        Acquisition object.
    """
    store = parse_url(zarr_path, mode="w").store
    group = zarr.group(store)
    wait(
        client.persist(
            da.to_zarr(
                arr=stitched_da,
                url=group.store,
                compute=False,
                component=str(Path(group.path, "0")),
                storage_options=dict(
                    dimension_separator="/",
                    chunks=(
                        1,
                        1,
                    )
                    + stitched_da.shape[2:],
                    write_empty_chunks=False,
                ),
                compressor=Blosc(cname="zstd", clevel=3, shuffle=Blosc.SHUFFLE),
                dimension_separator=group._store._dimension_separator,
            )
        )
    )
    shapes = [stitched_da.shape]
    datasets = [{"path": "0"}]
    coordinate_transformations = get_coordinate_transformations(
        acquisition=worm_acquisition,
        max_layer=0,
        mip=stitched_da.ndim == 4,
    )
    fmt = CurrentFormat()
    dims = len(shapes[0])
    fmt.validate_coordinate_transformations(
        dims, len(datasets), coordinate_transformations
    )
    for dataset, transform in zip(datasets, coordinate_transformations):
        dataset["coordinateTransformations"] = transform

    if stitched_da.ndim == 4:
        axes = ["t", "c", "y", "x"]
    else:
        axes = ["t", "c", "z", "y", "x"]

    axes = _get_valid_axes(dims, axes, fmt)
    write_multiscales_metadata(
        group,
        datasets,
        fmt,
        axes,
    )

... with the necessary adaptations (e.g. WellAcquisition for worm_acquisition) and tests.