irregular chunking

[Lem-P]

After cropping the original spatialData object (sdata_full = spatialdata_io.xenium(dir)) with
sdata = sdata_full.query.bounding_box( axes=["y", "x"], min_coordinate=[25000, 25000], max_coordinate=[35000, 35000], target_coordinate_system="global" )

I can't save the new spatialData object. I get this error:

ValueError Traceback (most recent call last)
Cell In[7], line 3
1 import os
2 f = os.path.join(dir,"subset.zarr")
----> 3 sdata.write(f, overwrite=True)
4 sdata = spatialdata.read_zarr("subset.zarr") # read the data from the zarr directory

File ~/miniconda3/envs/sopa_env/lib/python3.11/site-packages/spatialdata/_core/spatialdata.py:1219, in SpatialData.write(self, file_path, overwrite, consolidate_metadata, format)
1216 store.close()
1218 for element_type, element_name, element in self.gen_elements():
-> 1219 self._write_element(
1220 element=element,
1221 zarr_container_path=file_path,
1222 element_type=element_type,
1223 element_name=element_name,
1224 overwrite=False,
1225 format=format,
1226 )
1228 if self.path != file_path:
1229 old_path = self.path

File ~/miniconda3/envs/sopa_env/lib/python3.11/site-packages/spatialdata/_core/spatialdata.py:1263, in SpatialData._write_element(self, element, zarr_container_path, element_type, element_name, overwrite, format)
1260 parsed = _parse_formats(formats=format)
1262 if element_type == "images":
-> 1263 write_image(image=element, group=element_type_group, name=element_name, format=parsed["raster"])
1264 elif element_type == "labels":
1265 write_labels(labels=element, group=root_group, name=element_name, format=parsed["raster"])

File ~/miniconda3/envs/sopa_env/lib/python3.11/site-packages/spatialdata/_io/io_raster.py:242, in write_image(image, group, name, format, storage_options, **metadata)
234 def write_image(
235 image: DataArray | DataTree,
236 group: zarr.Group,
(...) 240 **metadata: str | JSONDict | list[JSONDict],
241 ) -> None:
--> 242 _write_raster(
243 raster_type="image",
244 raster_data=image,
245 group=group,
246 name=name,
247 format=format,
248 storage_options=storage_options,
249 **metadata,
250 )

File ~/miniconda3/envs/sopa_env/lib/python3.11/site-packages/spatialdata/_io/io_raster.py:202, in _write_raster(raster_type, raster_data, group, name, format, storage_options, label_metadata, **metadata)
200 parsed_axes = _get_valid_axes(axes=list(input_axes), fmt=format)
201 storage_options = [{"chunks": chunk} for chunk in chunks]
--> 202 dask_delayed = write_multi_scale_ngff(
203 pyramid=data,
204 group=group_data,
205 fmt=format,
206 axes=parsed_axes,
207 coordinate_transformations=None,
208 storage_options=storage_options,
209 **metadata,
210 compute=False,
211 )
212 # Compute all pyramid levels at once to allow Dask to optimize the computational graph.
213 da.compute(*dask_delayed)

File ~/miniconda3/envs/sopa_env/lib/python3.11/site-packages/ome_zarr/writer.py:254, in write_multiscale(pyramid, group, chunks, fmt, axes, coordinate_transformations, storage_options, name, compute, **metadata)
252 data = da.array(data).rechunk(chunks=chunks_opt)
253 options["chunks"] = chunks_opt
--> 254 da_delayed = da.to_zarr(
255 arr=data,
256 url=group.store,
257 component=str(Path(group.path, str(path))),
258 storage_options=options,
259 compressor=options.get("compressor", zarr.storage.default_compressor),
260 dimension_separator=group._store._dimension_separator,
261 compute=compute,
262 )
264 if not compute:
265 dask_delayed.append(da_delayed)

File ~/miniconda3/envs/sopa_env/lib/python3.11/site-packages/dask/array/core.py:3875, in to_zarr(arr, url, component, storage_options, overwrite, region, compute, return_stored, **kwargs)
3872 raise ValueError("Cannot use region keyword when url is not a zarr.Array.")
3874 if not _check_regular_chunks(arr.chunks):
-> 3875 raise ValueError(
3876 "Attempt to save array to zarr with irregular "
3877 "chunking, please call arr.rechunk(...) first."
3878 )
3880 storage_options = storage_options or {}
3882 if storage_options:

ValueError: Attempt to save array to zarr with irregular chunking, please call arr.rechunk(...) first.

spatialdata version 0.3.0

[LucaMarconato]

Hi @Lem-P, this is likely due to this bug, now addressed Dask and in spatialdata (but unreleased): #861

Manually rechunking with numbers (not auto) after the query will fix the issue.

[Lem-P]

Hi @LucaMarconato, thank you for your answer.
As for manually rechunking, I have no idea how to do that.
I have checked the different issues linked to this PR, but I don't see how to manually rechunking my new sdata. (I am not familiar, at all, with Dask)

[LucaMarconato]

Please check the code available in this if block:

spatialdata/src/spatialdata/_core/query/_utils.py

Line 155 in 3adacc0

if chunks_still_irregular:

Let me know if it works 😊

[Lem-P]

Definitely out of my league.
I did it like this, is it ok?

CHUNK_SIZE = 1024
image = sdata.images['morphology_focus']
coords = list(image["scale0"].coords.keys())
rechunk_strategy = {c: CHUNK_SIZE for c in coords}
if "c" in coords:
 rechunk_strategy["c"] = image["scale0"]["image"].chunks[0][0]
image = image.chunk(rechunk_strategy)
sdata.images['morphology_focus']=image

[LucaMarconato]

Looks good yes! A release is coming soon btw, so the solution root problem will be available.

[LucaMarconato]

Release done, I close this.

[LouisK92]

@LucaMarconato
I still run into this issue.

Reproducible example:

import numpy as np
import spatialdata as sd
from spatialdata.models import Labels2DModel

labels_data = np.zeros((15000, 15000))
labels_data[:2, :] = 1
labels = Labels2DModel.parse(labels_data)

sdata = sd.SpatialData(labels={"labels": labels})


sdata_crop = sdata.query.bounding_box(
 axes=["y", "x"],
 min_coordinate=[2000,2000],
 max_coordinate=[10000,10000],
 target_coordinate_system="global",
)

sdata_crop.write("dataset.zarr")

[LucaMarconato]

@LouisK92 I just tried with spatialdata==0.6.1 and your code works (i.e. I cannot reproduce the issue). Can you try again? If still get the problem please open a new issue with OS details etc. Thanks.