NGFF v06 transformations - #1182
Conversation
Cleans up and modernizes the NGff classes so that they have stronger invariants and guarantee that they always have input and output. Reading and writing to zarr is done via ome-zarr-models-py. Uses graph module as part of io_raster.py::try_read_ngff06_multiscale to interpret Ngff trnasformations and produce an output that could be added to the new graph implementation.
0dfccb5 to
f57dc1d
Compare
Codecov Report❌ Patch coverage is Additional details and impacted files@@ Coverage Diff @@
## transformation_manager #1182 +/- ##
=========================================================
Coverage ? 88.86%
=========================================================
Files ? 55
Lines ? 8595
Branches ? 0
=========================================================
Hits ? 7638
Misses ? 957
Partials ? 0
🚀 New features to boost your workflow:
|
| "numpy", | ||
| "ome_zarr>=0.16.0", | ||
| # "ome_zarr>=0.16.0", | ||
| "ome_zarr @ git+https://github.com/ome/ome-zarr-py/@e859b55425b740335309876bb023fe52977a3045", |
There was a problem hiding this comment.
this commit is ~40 commits old wrt main; I'd update already.
| multiscale = oz.OMEZarrMultiscale.from_ome_zarr(str(store)) | ||
| assert isinstance(multiscale, oz.OMEZarrMultiscale) # disambiguate from OMEZarrLabel | ||
|
|
||
| name_to_cs: dict[str, CoordSystem] = {} |
There was a problem hiding this comment.
Reminder: later when we work with Scenes, the coordinate system name is not enough for uniquely identifying a CS. It will be the combniation of path where the cs is defined, and the name.
There was a problem hiding this comment.
I saw later in that in vert.py there is already some logic needed for this (the CoordinateSystemIndentifier usage.
| continue | ||
| coords = coords.merge( | ||
| xr.Coordinates.from_xindex( | ||
| RangeIndex.linspace( |
| assert transf.input.path is not None | ||
| in_cs = CoordSystem(name=str(transf.input.name), axes=[Axis(name=ax.name, type=ax.type) for ax in out_cs.axes]) | ||
|
|
||
| ozm_seq = ozm06trans.Sequence(transformations=ds_md.coordinateTransformations) |
There was a problem hiding this comment.
IIRC, according to the specs, here we either have a Sequence (of Scale + Translation) or a Scale. We need to double check.
There was a problem hiding this comment.
Anyway the logic of doing .transform_points() is correct in both cases
|
|
||
| data_tree[f"scale{scale_idx}"] = xr.Dataset( | ||
| { | ||
| "image": xr.DataArray( |
There was a problem hiding this comment.
Not for this PR (we can "resolve the conversation"), but a reminder. This string and the name="image" below needs to be documented somewhere, e.g. in docstrings of model classes/contribution guide for developers using spaitaldata/in-memory design doc.
E.g. will the users expect to have always image or any string, but always a len(dataset) == 1? The syntax to retrieve the DataArray from the DataTree will change.
CC @jan-glx
| class AxisParsingException(Exception): | ||
| pass |
There was a problem hiding this comment.
I like it here. Just for awareness, in the transformation manager codebase the exceptions are in a separate exceptions.py file. https://github.com/scverse/spatialdata/pull/1164/changes
We can choose to keep things local unless there are many exceptions, in that case it is better to place in a dedicated files.
| if not isinstance(model.unit, str): | ||
| raise AxisParsingException("Can't handle axis unit") |
There was a problem hiding this comment.
Something to keep in mind for when we'll try to read legacy data. In the old CS we were settings unit="unit". This is not allowed by the specs (having "unit" is a SHOULD, but if they are specified they need to be real units).
There was a problem hiding this comment.
But also here, shouldn't we allow for model.unit to be None?
| name=name, | ||
| type=model.type, | ||
| unit=model.unit, | ||
| long_name=model.longName, |
There was a problem hiding this comment.
why is here camel case, and elsewhere with underscore?
| class LegacyAxes: | ||
| pass |
There was a problem hiding this comment.
unused for now; placeholder for later?
| if len(self.axes) != len({axis.name for axis in self.axes}): | ||
| raise ValueError("Axes names must be unique") |
There was a problem hiding this comment.
Can we delegate all the validation to ozm, including this one?
There was a problem hiding this comment.
Another issue with delegating to ozm is that it can only validate based on their own models, so we'd have to convert e.g. our Axis type into their ozm.Axis to do the validation, which is also kinda clunky 🤔
| ) | ||
|
|
||
| @classmethod | ||
| def try_from_model_or_default[T](cls, model: ozi.CoordinateSystem | None, *, default: T) -> CoordSystem | T: |
There was a problem hiding this comment.
Who is going to use this/how?
There was a problem hiding this comment.
The logic for ensuring that default has unique names is something that would live outside this class. I wonder if then even the whole 3-line implementation of this function should live outside.
| from collections.abc import Sequence | ||
| from typing import Final, Literal | ||
|
|
||
| import ome_zarr.classes.image as ozi |
There was a problem hiding this comment.
How about being more explicit and import directly from ome-zarr-models-py? WDYT?
| input = self.to_model() | ||
| if input is None: | ||
| return None | ||
| return ozi.CoordinateSystemIdentifier(name=input.name) |
There was a problem hiding this comment.
ok good, this is already the foundation for when we'll deal with Scenes
| from spatialdata.transformations.graph.vert import Axis, CoordSystem | ||
|
|
||
|
|
||
| class GarbledInput(Exception): |
There was a problem hiding this comment.
good to have the exception
|
|
||
| @abstractmethod | ||
| def to_affine(self) -> AffineEdge: | ||
| """Convert the transformation to an affine transformation, whenever the conversion can be made.""" |
There was a problem hiding this comment.
Up to you if doing it now or later, but I think it is going to be very beneficial to already play around with the displacement class to test code paths where a walk in the transformation graph contains a "non-affinable-transformation".
| def to_affine(self) -> AffineEdge: | ||
| """Convert the transformation to an affine transformation, whenever the conversion can be made.""" | ||
|
|
||
| def _validate_transform_points_shapes(self, points: xr.DataArray | xr.DataTree | ArrayLike) -> None: |
There was a problem hiding this comment.
why datatree and dataarray here? I think datatree should not be here, and dataarray should not be used for images. Maybe it is to use this as an arraylike with annotated coordinates (e.g. n x dims, with dims for example ['x', 'y', 'k', 'u'])?
|
|
||
| def _validate_transform_points_shapes(self, points: xr.DataArray | xr.DataTree | ArrayLike) -> None: | ||
| """ | ||
| Validate if the shape of the points (coordinats to be transformed) are consistent with the input size of the |
There was a problem hiding this comment.
coordinats -> coordinates
| f"points must be a tensor of shape (n, d), where n is the number of points and d is the " | ||
| f"the number of spatial dimensions. Points shape: {points.shape}, input size: {input_size}" |
There was a problem hiding this comment.
NGFF is more general: the matrix also contains the non-spatial dimensions. The dimensionality constraint to check is the same, but we should not tell the user about spatial dimensions, just dimensions
| # order of the composition: self is applied first, then the transformation passed as argument | ||
| def compose_with(self, transformation: BaseTransfEdge) -> BaseTransfEdge: | ||
| """ | ||
| Compose the transfomation object with another transformation | ||
|
|
||
| Parameters | ||
| ---------- | ||
| transformation | ||
| The transformation to compose with. | ||
|
|
||
| Returns | ||
| ------- | ||
| The compoesed transformation. | ||
|
|
||
| Notes | ||
| ------- | ||
| Self is applied first, then the transformation passed as argument. | ||
| """ | ||
| return SequenceEdge([self, transformation], name=None) # FIXME: no name? |
There was a problem hiding this comment.
I think I had this in the old NGFF classes but never used. Also, nitpick, this would lead to a transformation like [[[t0, t1], t2], t3]; which is fine, but composing in one go from outside may be more appealing [t0, t1, t2, t3].
Anyway, no big deal, with can keep.
There was a problem hiding this comment.
Regarding the name: it is optional, so we can keep it optional (or let the user pass it as optional argument to compose_with())
|
|
||
| def __init__( | ||
| self, | ||
| name: str | None, |
There was a problem hiding this comment.
after all the coordinate system name is crucial, the transformation name is not a big deal
| linear: ArrayLike, | ||
| translation: ArrayLike | None = None, |
There was a problem hiding this comment.
This removes all the ambiguity of having to pass the homogenous affine matrix vs not.
| input: CoordSystem, | ||
| output: CoordSystem, | ||
| ) -> AffineEdge: | ||
| """Creates an AffineEdge from a raw affine matrix |
There was a problem hiding this comment.
I'd add "in homogenous coordinates"
| return SequenceEdge([self, transformation], name=None) # FIXME: no name? | ||
|
|
||
| @abstractmethod | ||
| def to_model(self) -> ozm06trans.AnyTransform: |
There was a problem hiding this comment.
shall we add try_from_model as abstractmethod as well? this would mirror what done in vert.py
There was a problem hiding this comment.
It would look nice and symmetrical, but it's not sound from a typing perspective; If we make it an abstract class method, then all child methods have to have a compatible signature, but each class deserializes from a specific model (e.g.: ScaleEdge deserializes from omz.Scale, IdentityEdge deserializes from omz.Identity, etc.), but child methods can't have stricter parameters than their parents.
Serializing in to_model is fine, though, because the child methods are free to have more precise return types than the parent method.
| if len(set(output_to_input.values())) != len(output_to_input): | ||
| raise ValueError("input_to_output must map unique inputs to unique outputs") |
There was a problem hiding this comment.
Do we want to do this validation here instead of relying on ozmp? Added discussion in Zulip. Let's follow up there.
| def __repr__(self) -> str: | ||
| s = super().__repr__() + "\n" | ||
| s += "\n".join(f" {out} <- {inp}\n" for out, inp in self.output_to_input.items()) | ||
| return s |
There was a problem hiding this comment.
Rough edge: MapAxis is basically an affine (made of zeros and a few ones to permute the data). When we do repr on the affine we print the matrix and not the axes (which will be handled by the coordinate systems). Here we do the opposite, we print the axes and don't care about the affine.
It's not incorrect, but I see that this goes more towards the philosophy of spatialdata transformations, while the affine stays "true" to the NGFF transformations.
For clarify we could actually always print both: the "index-based/matrix representation" (in this case index 0 goes to 1, index 3 goes to 4 etc; in the affine case the matrix); and then the axes annotation (x -> y, y -> z here; in the case of affine the input and output axes, as we do when we print spatialdata Affine transformations).
| ) | ||
|
|
||
|
|
||
| class TranslationEdge(BaseTransfEdge): |
There was a problem hiding this comment.
Up to you if doing it here or in a follow-up PR. projectAxis is a crucial transform and we should start using it/play around with it soon.
| ) | ||
|
|
||
| def inverse(self) -> BaseTransfEdge: | ||
| inv = np.linalg.inv(self.affine) |
There was a problem hiding this comment.
worth adding a check on the stability by inspecting the conditioning number. We can reuse the code of decompose affine.
There was a problem hiding this comment.
In general inverting matrices is not a good idea, but I think the matrices coming from spatial alignments should be generally fine to be inverted.
| raise ValueError("input and output should have the same numbe rof axes") | ||
| if not np.isclose(np.linalg.det(linear_matrix), 1.0): | ||
| raise ValueError("det(linear_matrix) should be ~= 1") | ||
| linear_matrix.flags.writeable = False |
There was a problem hiding this comment.
Even though self.rotation is Final, there is nothing preventing it from being modified internally (potentially through another reference that was kept outside of __init__). This flags raises an exception if anything tries writing to the array.
Still, for extra paranoia we should probably make a copy of linear_matrix so we don't accidentally freeze an array that a user might be meaning to recycle 🤔
| for current_transf in transformations[1:]: | ||
| if previous_transf.output != current_transf.input: | ||
| raise ValueError(f"Mismatched input/output from {previous_transf} to {current_transf}") | ||
| previous_transf = current_transf | ||
| self.transformations = transformations |
There was a problem hiding this comment.
NGFF transformations in a NGFF sequence may not specify the input and output (sometimes). So we should relax this check.
There was a problem hiding this comment.
They don't need to specify input and output in their serialized, ome-zarr-models-py form, but in our runtime representation (*Edge) we gurantee that this is always the case (the type of BaseTransfEdge.input is a non-optional CoordSystem). If the inner transformations are being read in from a file, then they will be given virtual coordinate systems for input and output.
2 participants
Mirrors Ngff* classes in the
graphmoduleCoordSystem(based onNgffCoordinateSystem)__hash__(needed for usage in graph, but we might still have the identity issue when searching for verts)*Edge(based onNgffTransformation)omz-models-pyclasses__init__and nowhere elseio_raster.py
RangeIndexto the resulting xarraysTODO