Adds affine transformation support

Adds support for affine transformations to the MLArray metadata.

This change introduces the `affine` property to the `MetaSpatial` class,
allowing users to specify a homogeneous affine matrix for spatial data.
This provides a more general way to represent spatial transformations
compared to separate spacing, origin, and direction parameters.

A new test case is added to verify that the affine transformation
is stored correctly in the metadata. A check prevents mixing `affine` with other
spatial parameters like `spacing`, `origin`, or `direction`.

Author: Karol-G

mlarray/meta.py

@@ -676,6 +676,7 @@ class MetaSpatial(BaseMeta):
         spacing: Per-dimension spacing values. Length must match ndims.
         origin: Per-dimension origin values. Length must match ndims.
         direction: Direction cosine matrix of shape [ndims, ndims].
+        affine: Homogeneous affine matrix of shape [ndims + 1, ndims + 1].
         shape: Array shape. Length must match (spatial + non-spatial) ndims.
         axis_labels: Per-axis labels or roles. Length must match ndims.
         axis_units: Per-axis units. Length must match ndims.
@@ -686,6 +687,7 @@ class MetaSpatial(BaseMeta):
     spacing: Optional[list[Union[int,float]]] = None
     origin: Optional[list[Union[int,float]]] = None
     direction: Optional[list[list[Union[int,float]]]] = None
+    affine: Optional[list[list[Union[int,float]]]] = None
     shape: Optional[list[int]] = None
     axis_labels: Optional[list[Union[str,AxisLabel]]] = None
     axis_units: Optional[list[str]] = None
@@ -716,6 +718,21 @@ def _validate_and_cast(self, *, ndims: Optional[int] = None, spatial_ndims: Opti
             self.direction = _cast_to_list(self.direction, "meta.spatial.direction")
             _validate_float_int_matrix(self.direction, "meta.spatial.direction", spatial_ndims)
 
+        if self.affine is not None:
+            self.affine = _cast_to_list(self.affine, "meta.spatial.affine")
+            if spatial_ndims is not None:
+                _validate_float_int_matrix(
+                    self.affine,
+                    "meta.spatial.affine",
+                    spatial_ndims + 1,
+                )
+            else:
+                _validate_float_int_matrix(self.affine, "meta.spatial.affine")
+                n_rows = len(self.affine)
+                for row in self.affine:
+                    if len(row) != n_rows:
+                        raise ValueError("meta.spatial.affine must be a square matrix")
+
         if self.shape is not None:
             self.shape = _cast_to_list(self.shape, "meta.spatial.shape")
             _validate_float_int_list(self.shape, "meta.spatial.shape", ndims)
@@ -915,7 +932,7 @@ class Meta(BaseMeta):
     Attributes:
         source: Source metadata from the original image source (JSON-serializable dict).
         extra: Additional metadata (JSON-serializable dict).
-        spatial: Spatial metadata (spacing, origin, direction, shape).
+        spatial: Spatial metadata (spacing, origin, direction, affine, shape).
         stats: Summary statistics.
         bbox: Bounding boxes.
         is_seg: Segmentation flag.

mlarray/mlarray.py

@@ -20,6 +20,7 @@ def __init__(
             spacing: Optional[Union[List, Tuple, np.ndarray]] = None,
             origin: Optional[Union[List, Tuple, np.ndarray]] = None,
             direction: Optional[Union[List, Tuple, np.ndarray]] = None,
+            affine: Optional[Union[List, Tuple, np.ndarray]] = None,
             meta: Optional[Union[Dict, Meta]] = None,
             axis_labels: Optional[List[Union[str, AxisLabel]]] = None,
             copy: Optional['MLArray'] = None,
@@ -49,6 +50,9 @@ def __init__(
             direction (Optional[Union[List, Tuple, np.ndarray]]): Direction
                 cosine matrix. Provide a 2D list/tuple/ndarray with shape
                 (ndims, ndims) for spatial dimensions.
+            affine (Optional[Union[List, Tuple, np.ndarray]]): Homogeneous
+                affine matrix. Provide a 2D list/tuple/ndarray with shape
+                (spatial_ndims + 1, spatial_ndims + 1).
             meta (Optional[Dict | Meta]): Free-form metadata dictionary or Meta
                 instance. Must be JSON-serializable when saving. 
                 If meta is passed as a Dict, it is internally converted into a
@@ -81,6 +85,7 @@ def __init__(
             spacing is not None
             or origin is not None
             or direction is not None
+            or affine is not None
             or meta is not None
             or axis_labels is not None
             or copy is not None
@@ -91,14 +96,23 @@ def __init__(
             or dparams is not None
         ):
             raise RuntimeError(
-                "Spacing, origin, direction, meta, axis_labels, copy, patch_size, "
+                "Spacing, origin, direction, affine, meta, axis_labels, copy, patch_size, "
                 "chunk_size, block_size, cparams or dparams cannot be set when "
                 "array is a filepath."
             )
         if isinstance(array, (str, Path)):
             self._load(array, compressed=compressed)
         else:
-            self._validate_and_add_meta(meta, spacing, origin, direction, axis_labels, False, validate=False)
+            self._validate_and_add_meta(
+                meta,
+                spacing=spacing,
+                origin=origin,
+                direction=direction,
+                affine=affine,
+                axis_labels=axis_labels,
+                has_array=False,
+                validate=False,
+            )
             if array is not None:
                 self._asarray(
                     array,
@@ -1112,6 +1126,8 @@ def affine(self) -> np.ndarray:
         """
         if self._store is None or self.meta._has_array.has_array == False:
             return None
+        if self.meta.spatial.affine is not None:
+            return self.meta.spatial.affine
         spacing  = np.array(self.spacing) if self.spacing is not None else np.ones(self.spatial_ndim)
         origin  = np.array(self.origin) if self.origin is not None else np.zeros(self.spatial_ndim)
         direction = np.array(self.direction) if self.direction is not None else np.eye(self.spatial_ndim)
@@ -1824,7 +1840,17 @@ def _comp_and_validate_blosc2_meta(self, meta_blosc2, patch_size, chunk_size, bl
         meta_blosc2._validate_and_cast(ndims=len(shape), spatial_ndims=num_spatial_axes)
         return meta_blosc2
 
-    def _validate_and_add_meta(self, meta, spacing=None, origin=None, direction=None, axis_labels=None, has_array=None, validate=True):
+    def _validate_and_add_meta(
+            self,
+            meta,
+            spacing=None,
+            origin=None,
+            direction=None,
+            affine=None,
+            axis_labels=None,
+            has_array=None,
+            validate=True,
+        ):
         """Validate and attach metadata to the MLArray instance.
 
         Args:
@@ -1836,6 +1862,8 @@ def _validate_and_add_meta(self, meta, spacing=None, origin=None, direction=None
                 spatial axis.
             direction (Optional[Union[List, Tuple, np.ndarray]]): Direction
                 cosine matrix with shape (ndims, ndims).
+            affine (Optional[Union[List, Tuple, np.ndarray]]): Homogeneous
+                affine matrix with shape (spatial_ndims + 1, spatial_ndims + 1).
             axis_labels (Optional[List[Union[str, AxisLabel]]]): Per-axis labels or roles. Length must match ndims.
             has_array (Optional[bool]): Explicitly set whether array data is
                 present. When True, metadata is validated with array-dependent
@@ -1853,12 +1881,22 @@ def _validate_and_add_meta(self, meta, spacing=None, origin=None, direction=None
             meta = Meta()
         self.meta = meta
         self.meta._mlarray_version = MLARRAY_VERSION
+
+        if affine is not None and (
+            spacing is not None or origin is not None or direction is not None
+        ):
+            raise ValueError(
+                "affine cannot be provided together with spacing, origin, or direction."
+            )
+
         if spacing is not None:
             self.meta.spatial.spacing = spacing
         if origin is not None:
             self.meta.spatial.origin = origin
         if direction is not None:
             self.meta.spatial.direction = direction
+        if affine is not None:
+            self.meta.spatial.affine = affine
         if axis_labels is not None:
             self.meta.spatial.axis_labels = axis_labels
         if has_array == True:

tests/test_usage.py

@@ -88,6 +88,32 @@ def test_metadata_inspection_and_manipulation(self):
             self.assertEqual(loaded.origin, [10.0, 10.0, 30.0])
             self.assertEqual(loaded.meta.source["study_id"], "study-001")
 
+    def test_affine_stored_in_meta_spatial(self):
+        array = _make_array(shape=(8, 8, 8))
+        affine = [
+            [1.0, 0.0, 0.0, 10.0],
+            [0.0, 2.0, 0.0, 20.0],
+            [0.0, 0.0, 3.0, 30.0],
+            [0.0, 0.0, 0.0, 1.0],
+        ]
+        image = MLArray(array, affine=affine)
+
+        self.assertEqual(image.meta.spatial.affine, affine)
+        self.assertIsNone(image.meta.spatial.spacing)
+        self.assertIsNone(image.meta.spatial.origin)
+        self.assertIsNone(image.meta.spatial.direction)
+
+    def test_affine_and_spacing_origin_direction_mix_raises(self):
+        array = _make_array(shape=(8, 8, 8))
+        affine = [
+            [1.0, 0.0, 0.0, 10.0],
+            [0.0, 1.0, 0.0, 20.0],
+            [0.0, 0.0, 1.0, 30.0],
+            [0.0, 0.0, 0.0, 1.0],
+        ]
+        with self.assertRaises(ValueError):
+            MLArray(array, spacing=(1.0, 1.0, 1.0), affine=affine)
+
     def test_copy_metadata_with_override(self):
         with tempfile.TemporaryDirectory() as tmpdir:
             array = _make_array()