Add array_length support for X5 and adjust mapping

Author: oesteban

nitransforms/io/__init__.py

@@ -1,7 +1,6 @@
 # emacs: -*- mode: python-mode; py-indent-offset: 4; indent-tabs-mode: nil -*-
 # vi: set ft=python sts=4 ts=4 sw=4 et:
 """Read and write transforms."""
-
 from nitransforms.io import afni, fsl, itk, lta, x5
 from nitransforms.io.base import TransformIOError, TransformFileError
 
@@ -10,6 +9,7 @@
     "fsl",
     "itk",
     "lta",
+    "x5",
     "get_linear_factory",
     "TransformFileError",
     "TransformIOError",
@@ -23,7 +23,6 @@
     "fs": (lta, "FSLinearTransform"),
     "fsl": (fsl, "FSLLinearTransform"),
     "afni": (afni, "AFNILinearTransform"),
-    "x5": (x5, "X5LinearTransform"),
 }
 
 

nitransforms/io/x5.py

@@ -1,39 +1,80 @@
-# emacs: -*- mode: python-mode; py-indent-offset: 4; indent-tabs-mode: nil -*-
-# vi: set ft=python sts=4 ts=4 sw=4 et:
-### ### ### ### ### ### ### ### ### ### ### ### ### ### ### ### ### ### ### ##
-#
-#   See COPYING file distributed along with the NiBabel package for the
-#   copyright and license terms.
-#
-### ### ### ### ### ### ### ### ### ### ### ### ### ### ### ### ### ### ### ##
-"""Read/write x5 transforms."""
+"""Data structures for the X5 transform format."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Any, Dict, Optional, Sequence, List
+
+import json
+import h5py
 
 import numpy as np
-from h5py import File as H5File
 
 
-def to_filename(fname, xfm, moving=None):
-    """Store the transform in BIDS-Transforms HDF5 file format (.x5)."""
-    with H5File(fname, "w") as out_file:
-        out_file.attrs["Format"] = "X5"
-        out_file.attrs["Version"] = np.uint16(1)
-        x5_root = out_file.create_group("/0")
+@dataclass
+class X5Domain:
+    """Domain information of a transform."""
 
-        # Serialize this object into the x5 file format.
-        transform_group = x5_root.create_group("TransformGroup")
+    grid: bool
+    size: Sequence[int]
+    mapping: np.ndarray
+    coordinates: Optional[str] = None
 
-        # Group '0' containing Affine transform
-        transform_0 = transform_group.create_group("0")
 
-        transform_0.attrs["Type"] = "Affine"
-        transform_0.create_dataset("Transform", data=xfm.matrix)
-        transform_0.create_dataset("Inverse", data=~xfm)
+@dataclass
+class X5Transform:
+    """Represent one transform entry of an X5 file."""
 
-        metadata = {"key": "value"}
-        transform_0.attrs["Metadata"] = str(metadata)
+    type: str
+    transform: np.ndarray
+    dimension_kinds: Sequence[str]
+    array_length: int = 1
+    domain: Optional[X5Domain] = None
+    subtype: Optional[str] = None
+    representation: Optional[str] = None
+    metadata: Optional[Dict[str, Any]] = None
+    inverse: Optional[np.ndarray] = None
+    jacobian: Optional[np.ndarray] = None
+    additional_parameters: Optional[np.ndarray] = None
+
+
+def to_filename(fname: str, x5_list: List[X5Transform]):
+    """Write a list of :class:`X5Transform` objects to an X5 HDF5 file."""
+    with h5py.File(str(fname), "w") as out_file:
+        out_file.attrs["Format"] = "X5"
+        out_file.attrs["Version"] = np.uint16(1)
+        tg = out_file.create_group("TransformGroup")
+        for i, node in enumerate(x5_list):
+            g = tg.create_group(str(i))
+            g.attrs["Type"] = node.type
+            g.attrs["ArrayLength"] = node.array_length
+            if node.subtype is not None:
+                g.attrs["SubType"] = node.subtype
+            if node.representation is not None:
+                g.attrs["Representation"] = node.representation
+            if node.metadata is not None:
+                g.attrs["Metadata"] = json.dumps(node.metadata)
+            g.create_dataset("Transform", data=node.transform)
+            g.create_dataset(
+                "DimensionKinds",
+                data=np.asarray(node.dimension_kinds, dtype="S"),
+            )
+            if node.domain is not None:
+                dgrp = g.create_group("Domain")
+                dgrp.create_dataset(
+                    "Grid", data=np.uint8(1 if node.domain.grid else 0)
+                )
+                dgrp.create_dataset("Size", data=np.asarray(node.domain.size))
+                dgrp.create_dataset("Mapping", data=node.domain.mapping)
+                if node.domain.coordinates is not None:
+                    dgrp.attrs["Coordinates"] = node.domain.coordinates
 
-        # sub-group 'Domain' contained within group '0'
-        transform_0.create_group("Domain")
-        # domain_group.attrs["Grid"] = self._grid
-        # domain_group.create_dataset("Size", data=_as_homogeneous(self._reference.shape))
-        # domain_group.create_dataset("Mapping", data=self.mapping)
+            if node.inverse is not None:
+                g.create_dataset("Inverse", data=node.inverse)
+            if node.jacobian is not None:
+                g.create_dataset("Jacobian", data=node.jacobian)
+            if node.additional_parameters is not None:
+                g.create_dataset(
+                    "AdditionalParameters", data=node.additional_parameters
+                )
+    return str(fname)

nitransforms/linear.py

@@ -20,6 +20,7 @@
     EQUALITY_TOL,
 )
 from nitransforms.io import get_linear_factory, TransformFileError
+from nitransforms.io.x5 import X5Transform, X5Domain
 
 
 class Affine(TransformBase):
@@ -276,6 +277,25 @@ def __repr__(self):
         """
         return repr(self.matrix)
 
+    def to_x5(self):
+        """Return an :class:`~nitransforms.io.x5.X5Transform` representation."""
+        domain = None
+        if self._reference is not None:
+            domain = X5Domain(
+                grid=True,
+                size=self.reference.shape,
+                mapping=self.reference.affine,
+            )
+        kinds = tuple("space" for _ in range(self.ndim)) + ("vector",)
+        return X5Transform(
+            type="linear",
+            subtype="affine",
+            transform=self.matrix,
+            dimension_kinds=kinds,
+            domain=domain,
+            inverse=(~self).matrix,
+        )
+
 
 class LinearTransformsMapping(Affine):
     """Represents a series of linear transforms."""
@@ -330,6 +350,26 @@ def __getitem__(self, i):
         """Enable indexed access to the series of matrices."""
         return Affine(self.matrix[i, ...], reference=self._reference)
 
+    def to_x5(self):
+        """Return an :class:`~nitransforms.io.x5.X5Transform` object."""
+        domain = None
+        if self._reference is not None:
+            domain = X5Domain(
+                grid=True,
+                size=self.reference.shape,
+                mapping=self.reference.affine,
+            )
+        kinds = tuple("space" for _ in range(self.ndim - 1)) + ("vector",)
+        return X5Transform(
+            type="linear",
+            subtype="affine",
+            transform=self.matrix,
+            dimension_kinds=kinds,
+            domain=domain,
+            inverse=self._inverse,
+            array_length=len(self),
+        )
+
     def map(self, x, inverse=False):
         r"""
         Apply :math:`y = f(x)`.

nitransforms/tests/test_linear.py

@@ -6,8 +6,11 @@
 import numpy as np
 import h5py
 
+from pathlib import Path
+
 from nibabel.eulerangles import euler2mat
 from nibabel.affines import from_matvec
+import nibabel as nb
 from nitransforms import linear as nitl
 from nitransforms import io
 from .utils import assert_affines_by_filename
@@ -243,16 +246,35 @@ def test_linear_save(tmpdir, data_path, get_testdata, image_orientation, sw_tool
     assert_affines_by_filename(xfm_fname1, xfm_fname2)
 
 
-def test_Affine_to_x5(tmpdir, testdata_path):
+def test_Affine_to_x5(tmpdir):
     """Test affine's operations."""
     tmpdir.chdir()
     aff = nitl.Affine()
-    with h5py.File("xfm.x5", "w") as f:
-        aff._to_hdf5(f.create_group("Affine"))
-
-    aff.reference = testdata_path / "someones_anatomy.nii.gz"
-    with h5py.File("withref-xfm.x5", "w") as f:
-        aff._to_hdf5(f.create_group("Affine"))
+    node = aff.to_x5()
+    assert node.type == "linear"
+    assert node.domain is None
+    assert node.transform.shape == (4, 4)
+    assert node.array_length == 1
+
+    img = nb.Nifti1Image(np.zeros((2, 2, 2), dtype="float32"), np.eye(4))
+    img_path = Path(tmpdir) / "ref.nii.gz"
+    img.to_filename(str(img_path))
+
+    aff.reference = img_path
+    node = aff.to_x5()
+    assert node.domain.grid
+    assert node.domain.size == aff.reference.shape
+
+
+def test_mapping_to_x5():
+    mats = [
+        np.eye(4),
+        np.array([[1, 0, 0, 1], [0, 1, 0, 2], [0, 0, 1, 3], [0, 0, 0, 1]]),
+    ]
+    mapping = nitl.LinearTransformsMapping(mats)
+    node = mapping.to_x5()
+    assert node.array_length == 2
+    assert node.transform.shape == (2, 4, 4)
 
 
 def test_mulmat_operator(testdata_path):

nitransforms/tests/test_x5.py

@@ -0,0 +1,40 @@
+import numpy as np
+from h5py import File as H5File
+
+from ..io.x5 import X5Transform, X5Domain, to_filename
+
+
+def test_x5_transform_defaults():
+    xf = X5Transform(
+        type="linear",
+        transform=np.eye(4),
+        dimension_kinds=("space", "space", "space", "vector"),
+    )
+    assert xf.domain is None
+    assert xf.subtype is None
+    assert xf.representation is None
+    assert xf.metadata is None
+    assert xf.inverse is None
+    assert xf.jacobian is None
+    assert xf.additional_parameters is None
+    assert xf.array_length == 1
+
+
+def test_to_filename(tmp_path):
+    domain = X5Domain(grid=True, size=(10, 10, 10), mapping=np.eye(4))
+    node = X5Transform(
+        type="linear",
+        transform=np.eye(4),
+        dimension_kinds=("space", "space", "space", "vector"),
+        domain=domain,
+    )
+    fname = tmp_path / "test.x5"
+    to_filename(fname, [node])
+
+    with H5File(fname, "r") as f:
+        assert f.attrs["Format"] == "X5"
+        assert f.attrs["Version"] == 1
+        grp = f["TransformGroup"]
+        assert "0" in grp
+        assert grp["0"].attrs["Type"] == "linear"
+        assert grp["0"].attrs["ArrayLength"] == 1