ENH: Collapse linear and nonlinear transforms chains

nitransforms/io/__init__.py

@@ -1,6 +1,7 @@
 # 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
 
@@ -27,7 +28,37 @@
 
 
 def get_linear_factory(fmt, is_array=True):
-    """Return the type required by a given format."""
+    """
+    Return the type required by a given format.
+
+    Parameters
+    ----------
+    fmt : :obj:`str`
+        A format identifying string.
+    is_array : :obj:`bool`
+        Whether the array version of the class should be returned.
+
+    Returns
+    -------
+    type
+        The class object (not an instance) of the linear transfrom to be created
+        (for example, :obj:`~nitransforms.io.itk.ITKLinearTransform`).
+
+    Examples
+    --------
+    >>> get_linear_factory("itk")
+    <class 'nitransforms.io.itk.ITKLinearTransformArray'>
+    >>> get_linear_factory("itk", is_array=False)
+    <class 'nitransforms.io.itk.ITKLinearTransform'>
+    >>> get_linear_factory("fsl")
+    <class 'nitransforms.io.fsl.FSLLinearTransformArray'>
+    >>> get_linear_factory("fsl", is_array=False)
+    <class 'nitransforms.io.fsl.FSLLinearTransform'>
+    >>> get_linear_factory("fakepackage")  # doctest: +IGNORE_EXCEPTION_DETAIL
+    Traceback (most recent call last):
+    TypeError: Unsupported transform format <fakepackage>.
+
+    """
     if fmt.lower() not in _IO_TYPES:
         raise TypeError(f"Unsupported transform format <{fmt}>.")
 

nitransforms/io/lta.py

@@ -1,4 +1,5 @@
 """Read/write linear transforms."""
+
 import numpy as np
 from nibabel.volumeutils import Recoder
 from nibabel.affines import voxel_sizes, from_matvec
@@ -29,12 +30,12 @@ class VolumeGeometry(StringBasedStruct):
     template_dtype = np.dtype(
         [
             ("valid", "i4"),  # Valid values: 0, 1
-            ("volume", "i4", (3, )),  # width, height, depth
-            ("voxelsize", "f4", (3, )),  # xsize, ysize, zsize
+            ("volume", "i4", (3,)),  # width, height, depth
+            ("voxelsize", "f4", (3,)),  # xsize, ysize, zsize
             ("xras", "f8", (3, 1)),  # x_r, x_a, x_s
             ("yras", "f8", (3, 1)),  # y_r, y_a, y_s
             ("zras", "f8", (3, 1)),  # z_r, z_a, z_s
-            ("cras", "f8", (3, )),  # c_r, c_a, c_s
+            ("cras", "f8", (3,)),  # c_r, c_a, c_s
             ("filename", "U1024"),
         ]
     )  # Not conformant (may be >1024 bytes)
@@ -109,14 +110,19 @@ def from_string(cls, string):
             label, valstring = lines.pop(0).split(" =")
             assert label.strip() == key
 
-            val = ""
-            if valstring.strip():
-                parsed = np.genfromtxt(
+            parsed = (
+                np.genfromtxt(
                     [valstring.encode()], autostrip=True, dtype=cls.dtype[key]
                 )
-                if parsed.size:
-                    val = parsed.reshape(sa[key].shape)
-            sa[key] = val
+                if valstring.strip()
+                else None
+            )
+
+            if parsed is not None and parsed.size:
+                sa[key] = parsed.reshape(sa[key].shape)
+            else:  # pragma: no coverage
+                """Do not set sa[key]"""
+
         return volgeom
 
 
@@ -218,11 +224,15 @@ def to_ras(self, moving=None, reference=None):
     def to_string(self, partial=False):
         """Convert this transform to text."""
         sa = self.structarr
-        lines = [
-            "# LTA file created by NiTransforms",
-            "type      = {}".format(sa["type"]),
-            "nxforms   = 1",
-        ] if not partial else []
+        lines = (
+            [
+                "# LTA file created by NiTransforms",
+                "type      = {}".format(sa["type"]),
+                "nxforms   = 1",
+            ]
+            if not partial
+            else []
+        )
 
         # Standard preamble
         lines += [
@@ -232,10 +242,7 @@ def to_string(self, partial=False):
         ]
 
         # Format parameters matrix
-        lines += [
-            " ".join(f"{v:18.15e}" for v in sa["m_L"][i])
-            for i in range(4)
-        ]
+        lines += [" ".join(f"{v:18.15e}" for v in sa["m_L"][i]) for i in range(4)]
 
         lines += [
             "src volume info",
@@ -324,10 +331,7 @@ def __getitem__(self, idx):
     def to_ras(self, moving=None, reference=None):
         """Set type to RAS2RAS and return the new matrix."""
         self.structarr["type"] = 1
-        return [
-            xfm.to_ras(moving=moving, reference=reference)
-            for xfm in self.xforms
-        ]
+        return [xfm.to_ras(moving=moving, reference=reference) for xfm in self.xforms]
 
     def to_string(self):
         """Convert this LTA into text format."""
@@ -396,9 +400,11 @@ def from_ras(cls, ras, moving=None, reference=None):
         sa["type"] = 1
         sa["nxforms"] = ras.shape[0]
         for i in range(sa["nxforms"]):
-            lt._xforms.append(cls._inner_type.from_ras(
-                ras[i, ...], moving=moving, reference=reference
-            ))
+            lt._xforms.append(
+                cls._inner_type.from_ras(
+                    ras[i, ...], moving=moving, reference=reference
+                )
+            )
 
         sa["subject"] = "unset"
         sa["fscale"] = 0.0
@@ -407,8 +413,10 @@ def from_ras(cls, ras, moving=None, reference=None):
 
 def _drop_comments(string):
     """Drop comments."""
-    return "\n".join([
-        line.split("#")[0].strip()
-        for line in string.splitlines()
-        if line.split("#")[0].strip()
-    ])
+    return "\n".join(
+        [
+            line.split("#")[0].strip()
+            for line in string.splitlines()
+            if line.split("#")[0].strip()
+        ]
+    )

nitransforms/io/x5.py

@@ -136,3 +136,53 @@ def to_filename(fname: str | Path, x5_list: List[X5Transform]):
             #         "AdditionalParameters", data=node.additional_parameters
             #     )
     return fname
+
+
+def from_filename(fname: str | Path) -> List[X5Transform]:
+    """Read a list of :class:`X5Transform` objects from an X5 HDF5 file."""
+    try:
+        with h5py.File(str(fname), "r") as in_file:
+            if in_file.attrs.get("Format") != "X5":
+                raise TypeError("Input file is not in X5 format")
+
+            tg = in_file["TransformGroup"]
+            return [
+                _read_x5_group(node)
+                for _, node in sorted(tg.items(), key=lambda kv: int(kv[0]))
+            ]
+    except OSError as err:
+        if "file signature not found" in err.args[0]:
+            raise TypeError("Input file is not HDF5.")
+
+        raise  # pragma: no cover
+
+
+def _read_x5_group(node) -> X5Transform:
+    x5 = X5Transform(
+        type=node.attrs["Type"],
+        transform=np.asarray(node["Transform"]),
+        subtype=node.attrs.get("SubType"),
+        representation=node.attrs.get("Representation"),
+        metadata=json.loads(node.attrs["Metadata"])
+        if "Metadata" in node.attrs
+        else None,
+        dimension_kinds=[
+            k.decode() if isinstance(k, bytes) else k
+            for k in node["DimensionKinds"][()]
+        ],
+        domain=None,
+        inverse=np.asarray(node["Inverse"]) if "Inverse" in node else None,
+        jacobian=np.asarray(node["Jacobian"]) if "Jacobian" in node else None,
+        array_length=int(node.attrs.get("ArrayLength", 1)),
+    )
+
+    if "Domain" in node:
+        dgrp = node["Domain"]
+        x5.domain = X5Domain(
+            grid=bool(int(np.asarray(dgrp["Grid"]))),
+            size=tuple(np.asarray(dgrp["Size"])),
+            mapping=np.asarray(dgrp["Mapping"]),
+            coordinates=dgrp.attrs.get("Coordinates"),
+        )
+
+    return x5

nitransforms/linear.py

@@ -9,6 +9,7 @@
 """Linear transforms."""
 
 import warnings
+from collections import namedtuple
 import numpy as np
 from pathlib import Path
 
@@ -27,7 +28,12 @@
     EQUALITY_TOL,
 )
 from nitransforms.io import get_linear_factory, TransformFileError
-from nitransforms.io.x5 import X5Transform, X5Domain, to_filename as save_x5
+from nitransforms.io.x5 import (
+    X5Transform,
+    X5Domain,
+    to_filename as save_x5,
+    from_filename as load_x5,
+)
 
 
 class Affine(TransformBase):
@@ -149,19 +155,17 @@ def __matmul__(self, b):
         True
 
         >>> xfm1 = Affine([[1, 0, 0, 4], [0, 1, 0, 0], [0, 0, 1, 0], [0, 0, 0, 1]])
-        >>> xfm1 @ np.eye(4) == xfm1
+        >>> xfm1 @ Affine() == xfm1
         True
 
         """
-        if not isinstance(b, self.__class__):
-            _b = self.__class__(b)
-        else:
-            _b = b
+        if isinstance(b, self.__class__):
+            return self.__class__(
+                b.matrix @ self.matrix,
+                reference=b.reference,
+            )
 
-        retval = self.__class__(self.matrix.dot(_b.matrix))
-        if _b.reference:
-            retval.reference = _b.reference
-        return retval
+        return b @ self
 
     @property
     def matrix(self):
@@ -174,8 +178,29 @@ def ndim(self):
         return self._matrix.ndim + 1
 
     @classmethod
-    def from_filename(cls, filename, fmt=None, reference=None, moving=None):
+    def from_filename(
+        cls, filename, fmt=None, reference=None, moving=None, x5_position=0
+    ):
         """Create an affine from a transform file."""
+
+        if fmt and fmt.upper() == "X5":
+            x5_xfm = load_x5(filename)[x5_position]
+            Transform = cls if x5_xfm.array_length == 1 else LinearTransformsMapping
+            if (
+                x5_xfm.domain
+                and not x5_xfm.domain.grid
+                and len(x5_xfm.domain.size) == 3
+            ):  # pragma: no cover
+                raise NotImplementedError(
+                    "Only 3D regularly gridded domains are supported"
+                )
+            elif x5_xfm.domain:
+                # Override reference
+                Domain = namedtuple("Domain", "affine shape")
+                reference = Domain(x5_xfm.domain.mapping, x5_xfm.domain.size)
+
+            return Transform(x5_xfm.transform, reference=reference)
+
         fmtlist = [fmt] if fmt is not None else ("itk", "lta", "afni", "fsl")
 
         if fmt is not None and not Path(filename).exists():
@@ -265,7 +290,9 @@ def to_filename(self, filename, fmt="X5", moving=None, x5_inverse=False):
         if fmt.upper() == "X5":
             return save_x5(filename, [self.to_x5(store_inverse=x5_inverse)])
 
-        writer = get_linear_factory(fmt, is_array=isinstance(self, LinearTransformsMapping))
+        writer = get_linear_factory(
+            fmt, is_array=isinstance(self, LinearTransformsMapping)
+        )
 
         if fmt.lower() in ("itk", "ants", "elastix"):
             writer.from_ras(self.matrix).to_filename(filename)
@@ -348,11 +375,6 @@ def __init__(self, transforms, reference=None):
         )
         self._inverse = np.linalg.inv(self._matrix)
 
-    def __iter__(self):
-        """Enable iterating over the series of transforms."""
-        for _m in self.matrix:
-            yield Affine(_m, reference=self._reference)
-
     def __getitem__(self, i):
         """Enable indexed access to the series of matrices."""
         return Affine(self.matrix[i, ...], reference=self._reference)

nitransforms/manip.py

@@ -8,7 +8,6 @@
 ### ### ### ### ### ### ### ### ### ### ### ### ### ### ### ### ### ### ### ##
 """Common interface for transforms."""
 from collections.abc import Iterable
-import numpy as np
 
 from .base import (
     TransformBase,
@@ -145,17 +144,17 @@ def map(self, x, inverse=False):
 
         return x
 
-    def asaffine(self, indices=None):
+    def collapse(self):
         """
-        Combine a succession of linear transforms into one.
+        Combine a succession of transforms into one.
 
         Example
         ------
         >>> chain = TransformChain(transforms=[
         ...     Affine.from_matvec(vec=(2, -10, 3)),
         ...     Affine.from_matvec(vec=(-2, 10, -3)),
         ... ])
-        >>> chain.asaffine()
+        >>> chain.collapse()
         array([[1., 0., 0., 0.],
                [0., 1., 0., 0.],
                [0., 0., 1., 0.],
@@ -165,15 +164,15 @@ def asaffine(self, indices=None):
         ...     Affine.from_matvec(vec=(1, 2, 3)),
         ...     Affine.from_matvec(mat=[[0, 1, 0], [0, 0, 1], [1, 0, 0]]),
         ... ])
-        >>> chain.asaffine()
+        >>> chain.collapse()
         array([[0., 1., 0., 2.],
                [0., 0., 1., 3.],
                [1., 0., 0., 1.],
                [0., 0., 0., 1.]])
 
         >>> np.allclose(
         ...     chain.map((4, -2, 1)),
-        ...     chain.asaffine().map((4, -2, 1)),
+        ...     chain.collapse().map((4, -2, 1)),
         ... )
         True
 
@@ -183,9 +182,8 @@ def asaffine(self, indices=None):
             The indices of the values to extract.
 
         """
-        affines = self.transforms if indices is None else np.take(self.transforms, indices)
-        retval = affines[0]
-        for xfm in affines[1:]:
+        retval = self.transforms[-1]
+        for xfm in reversed(self.transforms[:-1]):
             retval = xfm @ retval
         return retval