Add RFC-5 Coordinate systems and transformations support

[thewtex]

Re: https://ngff.openmicroscopy.org/rfc/5/index.html

xref: InsightSoftwareConsortium/itkwidgets#757 (comment)

[jo-mueller]

@thewtex I have a sort of working version here :)

A few considerations that occurred to me along the way:

Frontend Metadata classes

I had to expand the functionalities of the introduced Transform instances. To use the full expressibility of transforms, users will have to instantiate those transform objects beforehand. Passing simple transforms like scale and translation beforehand as dicts like {"x": 2.0, "y": 2.0, "z": 2.0} is still supported, but converted to a new coordinateTransform metadata object under the hood. This takes from the simplicity of the current user interface a bit, but it seemed easier to me than trying to come up with a simple API that reflects what the metadata classes can already do.

Multiscales

So far, I got to a point where I can reproduce writing multiscales in a way that validation against the current version of the schemas of rfc5 check out. Great! 🥳

Writing spatial groups

Writing spatial groups (for lack of a better term) - collections of several multiscales with some sort of spatial relationship between them - requires a sort of nested structure. One way to do it would be to write each multiscales separately and then create the spatial collection structure on top, like suggested for the HCS dataset. That would require a small piece of code that would be able to create that top-level metadata (just like to_hcs_zarr)

Recursive writing

An idea that I was exploring a bit was recursive zarr writing. Essentially, the to_ngff_zarr function could be modified to accept a list of multiscales that would recursively call itself for each individual multiscale. The function would have to accept some additional metadata objects to tell it what to do with it, such as:

• coordinateTransformations: These could contain coordinate transformations between the frames of reference of each passed multiscale. It could be a requirement here that only multiscales with a coordinateSystem that is referenced by this transform are allowed to be written.
• Wells/Plates/HCS: Actually, this would also be a neat implementation of writing HCS data to disk. Multiscales in these structures could be passed as list of multiscales with some additional metadata that would allow a recursive writer to put each multiscale in a correct subgroup.

Requires: Writing recursively primarly requires writing multiscales to zarr groups rather than stores: When writing the arrays, the code uses zarr.create_array(StoreLike...) in a bunch of places. I think it could make sense to generally only write groups like (zarr.create_dataset(group...)). For "normal" multiscale images it shouldn't make much of a difference whether a dataset is written to a top-level group or a root store.

[thewtex]

@thewtex I have a sort of working version here :)

@jo-mueller awesome!! Could you please create a draft PR to facilitate review?

So far, I got to a point where I can reproduce writing multiscales in a way that validation against the current version of the schemas of rfc5 check out. Great! 🥳

Fantastic! 💯 🎉

Writing spatial groups (for lack of a better term) - collections of several multiscales with some sort of spatial relationship between them - requires a sort of nested structure. One way to do it would be to write each multiscales separately and then create the spatial collection structure on top, like #210. That would require a small piece of code that would be able to create that top-level metadata (just like to_hcs_zarr)

Yes. Ideally we are forwards-compatible with the OME-Zarr Collections RFC. @normanrz where can we find the latest and greatest Collections documentation?

Recursive writing

Sounds interesting! I do not have a good sense of the advantages and disadvantages of making a recursive refactor. My general impression is that it could simplify some cases but it would make other cases more complex and add more coupling and constraints? I think we want to continue to use stores instead of groups when possible to support de-coupled state and parallel writes. E.g. for HCS, the well images can ideally be written in parallel; even on different systems.

[jo-mueller]

@thewtex thanks for the thoughts on this! I agree that recursive stuff can often lead to...unforeseen consequences. Having looked at #212 just now it has become more clear to me how writing additional structures is intended to work. In the spirit of the newly added functions over there it will be more consistent to

• create the metadata structure as it is now (with the coordinateSystems and coordinateTransformations dataclasses)
• Write that top-level metadata structure to disk
• Add members to that spatial collection using a to_transform_zarr function that could directly check whether the images written to the collection share the coordinate systems that were previously defined at the top-level.

Here's a small MWE to write multiscales that embodies that idea already:

import ngff_zarr as nz
import numpy as np
from ngff_zarr.v06.zarr_metadata import Scale, CoordinateSystem, Axis

x_scale = 2
y_scale = 2

coordinate_system_physical = CoordinateSystem(
    name="physical",
    axes=[
        Axis(name="y", unit="micrometer", type="space"),
        Axis(name="x", unit="micrometer", type="space"),
    ],
)

coordinateTransformation = Scale(
    scale=[y_scale, x_scale],
    input='',
    output=coordinate_system_physical,
)

ngff_image = nz.to_ngff_image(
    np.random.random((1000, 1000)),
    coordinateTransformation=coordinateTransformation,
    name='image1'
    )

ngff_multiscales = nz.to_multiscales(
    ngff_image,
    scale_factors=[2, 4])

nz.to_ngff_zarr('multiscale_example.ome.zarr', ngff_multiscales, version='0.6.dev1')

[jo-mueller]

@thewtex small update/discussion point on this: The more I work on implementing RFC5/0.6 support in the ways of the current implementations, the more I find myself tangled up with deeply nested if/else checks and indexing into metadata of an unknown/diverse structure.

My conclusion at this point is that it's worth to put in some extra work in the metadata modelling to make sure that ngff-zarr can "absorb" larger, potentially breaking changes in the metadata in the future. This doesn't belong in the works for RFC5 implementation, hence I would propose the following roadmap:

• Implement converter routines for the existing dataclasses. We potentially need to do this for the Dataset dataclasses as well 🤔 For 0.5 metadata, this could look something like this

@dataclass
class Metadata:
    axes: List[Axis]
    datasets: List[Dataset]
    coordinateTransformations: Optional[List[Transform]]
    omero: Optional[Omero] = None
    name: str = "image"
    type: Optional[str] = None
    metadata: Optional[MethodMetadata] = None

    def to_version(self, version: str) -> Union["Metadata_v04", "Metadata"]:
        if version == "0.4":
            return self.to_v04()
        elif version == "0.5":
            return self
        else:
            raise ValueError(f"Unsupported version conversion: 0.5 -> {version}")
        
    @classmethod
    def from_version(cls, metadata: Union["Metadata_v04", "Metadata"]) -> "Metadata":
        from ..v04.zarr_metadata import Metadata as Metadata_v04
        
        if isinstance(metadata, Metadata_v04):
            return cls.from_v04(metadata)
        else:
            raise ValueError(f"Unsupported metadata type: {type(metadata)}")

    def to_v04(self) -> "Metadata_v04":
        from ..v04.zarr_metadata import Metadata as Metadata_v04
        
        metadata = Metadata_v04(
            axes=self.axes,
            datasets=self.datasets,
            coordinateTransformations=self.coordinateTransformations,
            name=self.name,
            metadata=self.metadata,
            type=self.type,
            omero=self.omero,
        )
        return metadata
    
    @classmethod
    def from_v04(cls, metadata_v04: "Metadata_v04"):
        from ..v04.zarr_metadata import Metadata as Metadata_v04
        
        metadata = Metadata(
            axes=metadata_v04.axes,
            datasets=metadata_v04.datasets,
            coordinateTransformations=metadata_v04.coordinateTransformations,
            name=metadata_v04.name,
            metadata=metadata_v04.metadata,
            type=metadata_v04.type,
            omero=metadata_v04.omero,
        )
        return metadata

It looks like a huge pile of boilerplate but I'm sure that it will go a long way in reducing complexity elsewhere quite a bit.

• Implement metadata parsing/dumping functionality. This should eat up a lot of the parsing/nested key checking in the from_ngff_zarr and to_ngff_zarr files.

At this point I think we are fairly close to re-inventing ome-zarr models, which I think is ok. Working on RFC5-related things I'm just convincing myself more and more that it's better to tidy the metadata handling first and then bring all the new features rather than the other way round.

Would love to hear your thoughts on this.

[thewtex]

@jo-mueller good idea! 🥇 I agree that adding conversion methods as proposed to the Metadata classes is a good refactor.

Side note: we probably want a string enum to use for supported versions to help prevent typos and to express what versions we support.