fix: general re-design of the PR

Author: oesteban

nitransforms/io/x5.py

@@ -26,7 +26,7 @@
 import numpy as np
 
 
-@dataclass
+@dataclass(eq=True)
 class X5Domain:
     """Domain information of a transform representing reference/moving spaces."""
 

nitransforms/linear.py

@@ -186,22 +186,7 @@ def from_filename(
         """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)
+            return from_x5(load_x5(filename), x5_position)
 
         fmtlist = [fmt] if fmt is not None else ("itk", "lta", "afni", "fsl")
 
@@ -458,3 +443,20 @@ def load(filename, fmt=None, reference=None, moving=None):
         xfm = xfm[0]
 
     return xfm
+
+
+def from_x5(x5_list, x5_position=0):
+    """Create an affine from a list of :class:`~nitransforms.io.x5.X5Transform` objects."""
+
+    x5_xfm = x5_list[x5_position]
+    Transform = Affine 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)

nitransforms/manip.py

@@ -7,15 +7,26 @@
 #
 ### ### ### ### ### ### ### ### ### ### ### ### ### ### ### ### ### ### ### ##
 """Common interface for transforms."""
+
+import os
 from collections.abc import Iterable
 import numpy as np
 
-from .base import (
+import h5py
+from nitransforms.base import (
     TransformBase,
     TransformError,
 )
-from .linear import Affine, LinearTransformsMapping
-from .nonlinear import DenseFieldTransform
+from nitransforms.io import itk, x5 as x5io
+from nitransforms.io.x5 import from_filename as load_x5
+from nitransforms.linear import (
+    Affine,
+    from_x5 as linear_from_x5,  # noqa: F401
+)
+from nitransforms.nonlinear import (
+    DenseFieldTransform,
+    from_x5 as nonlinear_from_x5,  # noqa: F401
+)
 
 
 class TransformChain(TransformBase):
@@ -183,7 +194,9 @@ 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)
+        affines = (
+            self.transforms if indices is None else np.take(self.transforms, indices)
+        )
         retval = affines[0]
         for xfm in affines[1:]:
             retval = xfm @ retval
@@ -192,51 +205,28 @@ def asaffine(self, indices=None):
     @classmethod
     def from_filename(cls, filename, fmt="X5", reference=None, moving=None, x5_chain=0):
         """Load a transform file."""
-        from .io import itk, x5 as x5io
-        import h5py
-        import nibabel as nb
-        from collections import namedtuple
 
         retval = []
         if fmt and fmt.upper() == "X5":
+            xfm_list = load_x5(filename)
+            if not xfm_list:
+                raise TransformError("Empty transform group")
+
             with h5py.File(str(filename), "r") as f:
-                if f.attrs.get("Format") == "X5":
-                    tg = [
-                        x5io._read_x5_group(node)
-                        for _, node in sorted(f["TransformGroup"].items(), key=lambda kv: int(kv[0]))
-                    ]
-                    chain_grp = f.get("TransformChain")
-                    if chain_grp is None:
-                        raise TransformError("X5 file contains no TransformChain")
-
-                    chain_path = chain_grp[str(x5_chain)][()]
-                    if isinstance(chain_path, bytes):
-                        chain_path = chain_path.decode()
-                    indices = [int(idx) for idx in chain_path.split("/") if idx]
-
-                    Domain = namedtuple("Domain", "affine shape")
-                    for idx in indices:
-                        node = tg[idx]
-                        if node.type == "linear":
-                            Transform = Affine if node.array_length == 1 else LinearTransformsMapping
-                            reference = None
-                            if node.domain is not None:
-                                reference = Domain(node.domain.mapping, node.domain.size)
-                            retval.append(Transform(node.transform, reference=reference))
-                        elif node.type == "nonlinear":
-                            reference = Domain(node.domain.mapping, node.domain.size)
-                            field = nb.Nifti1Image(node.transform, reference.affine)
-                            retval.append(
-                                DenseFieldTransform(
-                                    field,
-                                    is_deltas=node.representation == "displacements",
-                                    reference=reference,
-                                )
-                            )
-                        else:  # pragma: no cover - unsupported type
-                            raise NotImplementedError(f"Unsupported transform type {node.type}")
-
-                    return TransformChain(retval)
+                chain_grp = f.get("TransformChain")
+                if chain_grp is None:
+                    raise TransformError("X5 file contains no TransformChain")
+
+                chain_path = chain_grp[str(x5_chain)][()]
+                if isinstance(chain_path, bytes):
+                    chain_path = chain_path.decode()
+
+            for idx in chain_path.split("/"):
+                node = x5io._read_x5_group(xfm_list[int(idx)])
+                from_x5 = globals()[f"{node.type}_from_x5"]
+                retval.append(from_x5([node]))
+
+            return TransformChain(retval)
 
         if str(filename).endswith(".h5"):
             reference = None
@@ -253,57 +243,24 @@ def from_filename(cls, filename, fmt="X5", reference=None, moving=None, x5_chain
 
     def to_filename(self, filename, fmt="X5"):
         """Store the transform chain in X5 format."""
-        from .io import x5 as x5io
-        import os
-        import h5py
 
         if fmt.upper() != "X5":
             raise NotImplementedError("Only X5 format is supported for chains")
 
-        if os.path.exists(filename):
-            with h5py.File(str(filename), "r") as f:
-                existing = [
-                    x5io._read_x5_group(node)
-                    for _, node in sorted(f["TransformGroup"].items(), key=lambda kv: int(kv[0]))
-                ]
-        else:
-            existing = []
-
-        # convert to objects for equality check
-        from collections import namedtuple
-        import nibabel as nb
-
-        def _as_transform(x5node):
-            Domain = namedtuple("Domain", "affine shape")
-            if x5node.type == "linear":
-                Transform = Affine if x5node.array_length == 1 else LinearTransformsMapping
-                ref = None
-                if x5node.domain is not None:
-                    ref = Domain(x5node.domain.mapping, x5node.domain.size)
-                return Transform(x5node.transform, reference=ref)
-            reference = Domain(x5node.domain.mapping, x5node.domain.size)
-            field = nb.Nifti1Image(x5node.transform, reference.affine)
-            return DenseFieldTransform(
-                field,
-                is_deltas=x5node.representation == "displacements",
-                reference=reference,
-            )
-
-        existing_objs = [_as_transform(n) for n in existing]
-        path_indices = []
+        existing = load_x5(filename) if os.path.exists(filename) else []
+        xfm_chain = []
         new_nodes = []
+        next_xfm_index = len(existing)
         for xfm in self.transforms:
-            # find existing
-            idx = None
-            for i, obj in enumerate(existing_objs):
-                if type(xfm) is type(obj) and xfm == obj:
-                    idx = i
+            for eidx, existing_xfm in enumerate(existing):
+                if xfm == existing_xfm:
+                    xfm_chain.append(eidx)
                     break
-            if idx is None:
-                idx = len(existing_objs)
-                new_nodes.append((idx, xfm.to_x5()))
-                existing_objs.append(xfm)
-            path_indices.append(idx)
+            else:
+                xfm_chain.append(next_xfm_index)
+                new_nodes.append((next_xfm_index, xfm))
+                existing.append(xfm)
+                next_xfm_index += 1
 
         mode = "r+" if os.path.exists(filename) else "w"
         with h5py.File(str(filename), mode) as f:
@@ -317,7 +274,7 @@ def _as_transform(x5node):
                 x5io._write_x5_group(g, node)
 
             cg = f.require_group("TransformChain")
-            cg.create_dataset(str(len(cg)), data="/".join(str(i) for i in path_indices))
+            cg.create_dataset(str(len(cg)), data="/".join(str(i) for i in xfm_chain))
 
         return filename
 

nitransforms/nonlinear.py

@@ -278,7 +278,7 @@ def to_x5(self, metadata=None):
         )
 
     @classmethod
-    def from_filename(cls, filename, fmt="X5"):
+    def from_filename(cls, filename, fmt="X5", x5_position=0):
         _factory = {
             "afni": io.afni.AFNIDisplacementsField,
             "itk": io.itk.ITKDisplacementsField,
@@ -290,15 +290,7 @@ def from_filename(cls, filename, fmt="X5"):
             raise NotImplementedError(f"Unsupported format <{fmt}>")
 
         if fmt == "X5":
-            x5_xfm = load_x5(filename)[0]
-            Domain = namedtuple("Domain", "affine shape")
-            reference = Domain(x5_xfm.domain.mapping, x5_xfm.domain.size)
-            field = nb.Nifti1Image(x5_xfm.transform, reference.affine)
-            return cls(
-                field,
-                is_deltas=x5_xfm.representation == "displacements",
-                reference=reference,
-            )
+            return from_x5(load_x5(filename), x5_position=x5_position)
 
         return cls(_factory[fmt.lower()].from_filename(filename))
 
@@ -336,21 +328,15 @@ def ndim(self):
         return self._coeffs.ndim - 1
 
     @classmethod
-    def from_filename(cls, filename, fmt="X5"):
+    def from_filename(cls, filename, fmt="X5", x5_position=0):
         _factory = {
             "X5": None,
         }
         fmt = fmt.upper()
         if fmt not in {k.upper() for k in _factory}:
             raise NotImplementedError(f"Unsupported format <{fmt}>")
 
-        x5_xfm = load_x5(filename)[0]
-        Domain = namedtuple("Domain", "affine shape")
-        reference = Domain(x5_xfm.domain.mapping, x5_xfm.domain.size)
-
-        coefficients = nb.Nifti1Image(x5_xfm.transform, x5_xfm.additional_parameters)
-        return cls(coefficients, reference=reference)
-
+        return from_x5(load_x5(filename), x5_position=x5_position)
         # return cls(_factory[fmt.lower()].from_filename(filename))
 
     def to_field(self, reference=None, dtype="float32"):
@@ -440,6 +426,31 @@ def map(self, x, inverse=False):
         return np.array([vfunc(_x).tolist() for _x in np.atleast_2d(x)])
 
 
+def from_x5(x5_list, x5_position=0):
+    """Create a transform from a list of :class:`~nitransforms.io.x5.X5Transform` objects."""
+
+    x5_xfm = x5_list[x5_position]
+
+    Transform = (
+        BSplineFieldTransform if x5_xfm.subtype == "bspline" else DenseFieldTransform
+    )
+    Domain = namedtuple("Domain", "affine shape")
+    reference = Domain(x5_xfm.domain.mapping, x5_xfm.domain.size)
+    xfm_params = (
+        nb.Nifti1Image(x5_xfm.transform, x5_xfm.additional_parameters)
+        if x5_xfm.subtype == "bspline"
+        else x5_xfm.transform
+    )
+
+    xfm_kwargs = (
+        {}
+        if x5_xfm.subtype == "bspline"
+        else {"is_deltas": x5_xfm.representation == "displacements"}
+    )
+
+    return Transform(xfm_params, reference=reference, **xfm_kwargs)
+
+
 def _map_xyz(x, reference, knots, coeffs):
     """Apply the transformation to just one coordinate."""
     ndim = len(x)