Review compatibility of MIP and import-ome-zarr task

[tcompa]

This comes up as part of #557 (issue: #521).

The high-level question is the following: our current FOV ROIs match with the original images, and also with the zarr YX chunks. Is this strictly required for the MIP task? Or could it work for any set of ROIs? And, more generally, why does it even depend on the concept of ROIs (given that #115 is still open)?

Obvious side question: what about illumination-correction task? In that case we must match ROIs to correction matrices, and then there is not much room for flexibility. But at least the illumination-correction processing scheme seems robust, with respect to the current issue (because we already work on a per-ROI level and because we only scan ROIs sequentially). Should we also use it for the MIP task?

---

1. MIP task hardcodes the FOV_ROI_table string, meaning it wouldn't work with an "imported" ome-zarr, unless we provide a table with that name. It's easy to make this more flexible (through a new task parameter), but see next point.
2. MIP task makes an important assumption on the ROIs in such table (be it FOV_ROI_table or an arbitrary table name that was passed as a task input parameter). That is, each ROI in the table must have the same size (measured in pixels) as the chunks of the zarr array at level 0. This is related to https://github.com/fractal-analytics-platform/fractal-tasks-core/issues/115, i.e. we currently rely on dask "smart" scheduling. I currently fail to understand why we added this constraint on chunk sizes in the past, but I'll review it a bit further.

[tcompa]

With 581d9a6 and following commits, I am experimenting with a more flexible MIP task.

Important note: the current issue does not represent a problem if the ome-zarr to be imported is made of a regular grid of FOVs and we know in advance their distribution (i.e. the YX grid shape). In that case, we can add this information as an import-ome-zarr parameter, so that we'll recreate the exact same ROI table as the original zarr would have.

Making this work with irregular FOVs would be more complex, as we would need to have an input parameter which is essentially the whole ROI table.

[jluethi]

hmm, good to review this. In general, I'd expect most OME-Zarrs we're importing to be (1) small images in comparison to our wells, e.g. not have the whole well stitched as one image, but just provide individual OME-Zarrs. Next, (2) I'd expect regular grids stitched. The (3) search-first stitched approach we also support for Yokogawa's certainly is an advanced use-case that I don't think will be very common.

As such, I think there will be some edge cases with scenario 2 & 3 in the import OME-Zarr task. But some of those can be solved as part of the import task (maybe not supported from the beginning and maybe a bit cumbersome).

Illumination correction task: Users need to be able to specify relevant ROIs. For the initial work, I'd limit this to having a well_ROI (covers everything for case 2) or regular grid ROIs that are specified during import (would cover 2). Covering scenario 3 for illumination correction would mean supporting importing an existing ROI table that the user would prepare beforehand.

On the MIP task: If the MIP task can work on the well_ROI_table / a total image ROI table as well, that should already cover many use cases. Whether it's worth making it more flexible in what ROIs to run over: not sure.

I can imagine we'll currently benefit from some memory efficiency & runtime benefits due to the ROIs being aligned with the chunks for grid scenarios. But for the search first, that doesn't hold anymore, so I'm not sure how valuable the constraint of equally sized ROIs actually is

[tcompa]

I can imagine we'll currently benefit from some memory efficiency & runtime benefits due to the ROIs being aligned with the chunks for grid scenarios. But for the search first, that doesn't hold anymore, so I'm not sure how valuable the constraint of equally sized ROIs actually is

A clarification: we do not use ROIs in the MIP task. The 3D->2D projection is like this

    # Load 0-th level
    data_czyx = da.from_zarr(zarrurl_old + "/0")
    num_channels = data_czyx.shape[0]
    # Loop over channels
    accumulate_chl = []
    for ind_ch in range(num_channels):
        # Perform MIP for each channel of level 0
        mip_yx = da.stack([da.max(data_czyx[ind_ch], axis=0)], axis=0)
        accumulate_chl.append(mip_yx)
    accumulated_array = da.stack(accumulate_chl, axis=0)

    # Write to disk (triggering execution)
    if accumulated_array.chunksize != chunksize:
        raise ValueError(
            f"ERROR\n{accumulated_array.chunksize=}\n{chunksize=}"
        )
    try:
        accumulated_array.to_zarr(
            f"{zarrurl_new}/0",
            overwrite=overwrite,
            dimension_separator="/",
            write_empty_chunks=False,
        )
    ...

That's why I mentioned that I need to review the reason for the preliminary check that we perform at the beginning of the task, which requires ROIs to correspond to zarr chunks. At the moment I cannot understand why it's needed.

[jluethi]

Hmm, that makes even less sense then! Because MIP does work for search-first scenarios, where the ROI sizes happen to match with chunk sizes as we set them at the moment. But they then don't match in their location.

Also, we do not want to have chunk-size related constraints hard-coded in our tasks if we can at all avoid it. Will become relevant for discussion around sharing & Zarr v3.

[tcompa]

If the MIP task can work on the well_ROI_table / a total image ROI table as well, that should already cover many use cases. Whether it's worth making it more flexible in what ROIs to run over: not sure.

The MIP task does work on the well_ROI_table level - see code snippet above. It just adds some more checks on the chunk sizes, but then it fully ignores the ROI table.

Because MIP does work for search-first scenarios,

It does work, by simply ignoring the existence of ROIs. It's making a 3D->2D projection of the whole zarr array, independently on whether it's a regular grid or only includes a few sparse FOV ROIs. That was the idea of #115, i.e. making the MIP task work per-ROI (to optimize it in the sparse case).

Also, we do not want to have chunk-size related constraints hard-coded in our tasks if we can at all avoid it.

I fully agree. Depending on work on sharding/V3, it may be useful to have a bit of flexibility in how large the chunks should be (e.g. "one chunk = one FOV", or maybe "one chunk = 2x2 FOVs"). But this should a be non-blocking optimization tool, rather than strict constraint.

[tcompa]

TLDR for MIP task

I suspect that the check on chunksizes does not serve any scope, and I'd like to remove it (after some additional review). Other than that, I see two options:

1. We proceed as in #115 and we make MIP task ROI-based. This may result in an optimized MIP task for search-first, but maybe it will be slower for typical use (regular grids) because we would loose a lot of the parallelization.

2. We move in the opposite direction, by fully removing ROIs from the MIP task. We accept a likely suboptimal performance for search-first MIP (which can be, however, mitigated by assigning larger SLURM resources - since dask likely will take advantage of it), but we have a more general MIP task. Also, the task code becomes quite simpler.

I'd be in favor of option 2, since it covers nicely the typical use case

TLDR for illumination-correction task

Independently on chunking, we always need ROIs with size that matches the correction matrix, and then it makes no sense to not use ROIs. Supporting import-ome-zarr + illumination-correction in a complex case (i.e. different from a regular grid) should require additional user's inputs (that is, a pre-made ROI table).

As of chunk-size checks: those are not needed as long as we stick with sequential ROI processing.

[jluethi]

I'd be in favor of option 2, since it covers nicely the typical use case

I also favor option 2. I'd rather have a generic MIP task as core. If there is a big need, there can be a custom, user-designed MIP task that works better for some edge-cases (like search-first). But the one in fractal-tasks-core should be a generic one that works as broadly as possible.

Also agree with the illumination-correction task summary