Add nonlinear X5 tests

Author: oesteban

CHANGES.rst

@@ -4,12 +4,14 @@ A new major release with critical updates.
 The new release includes a critical hotfix for 4D resamplings.
 The second major improvement is the inclusion of a first implementation of the X5 format (BIDS).
 The X5 implementation is currently restricted to reading/writing of linear transforms.
+It now supports nonlinear transforms as well.
 
 CHANGES
 -------
 * FIX: Broken 4D resampling by @oesteban in https://github.com/nipy/nitransforms/pull/247
 * ENH: Loading of X5 (linear) transforms by @oesteban in https://github.com/nipy/nitransforms/pull/243
 * ENH: Implement X5 representation and output to filesystem by @oesteban in https://github.com/nipy/nitransforms/pull/241
+* ENH: Support reading and writing of nonlinear transforms in X5
 * DOC: Fix references to ``os.PathLike`` by @oesteban in https://github.com/nipy/nitransforms/pull/242
 * MNT: Increase coverage by testing edge cases and adding docstrings by @oesteban in https://github.com/nipy/nitransforms/pull/248
 * MNT: Refactor io/lta to reduce one partial line by @oesteban in https://github.com/nipy/nitransforms/pull/246

nitransforms/nonlinear.py

@@ -9,6 +9,7 @@
 """Nonlinear transforms."""
 import warnings
 from functools import partial
+from collections import namedtuple
 import numpy as np
 
 from nitransforms import io
@@ -227,17 +228,54 @@ def __eq__(self, other):
             warnings.warn("Fields are equal, but references do not match.")
         return _eq
 
+    def to_x5(self, metadata=None):
+        """Return an :class:`~nitransforms.io.x5.X5Transform` representation."""
+        from ._version import __version__
+        from .io.x5 import X5Domain, X5Transform
+
+        metadata = {"WrittenBy": f"NiTransforms {__version__}"} | (metadata or {})
+
+        domain = None
+        if (reference := self.reference) is not None:
+            domain = X5Domain(
+                grid=True,
+                size=getattr(reference, "shape", (0, 0, 0)),
+                mapping=reference.affine,
+                coordinates="cartesian",
+            )
+
+        kinds = tuple("space" for _ in range(self.ndim)) + ("vector",)
+
+        return X5Transform(
+            type="nonlinear",
+            subtype="densefield",
+            representation="displacements",
+            metadata=metadata,
+            transform=self._deltas,
+            dimension_kinds=kinds,
+            domain=domain,
+        )
+
     @classmethod
     def from_filename(cls, filename, fmt="X5"):
         _factory = {
             "afni": io.afni.AFNIDisplacementsField,
             "itk": io.itk.ITKDisplacementsField,
             "fsl": io.fsl.FSLDisplacementsField,
+            "X5": None,
         }
-        if fmt not in _factory:
+        fmt = fmt.upper()
+        if fmt not in {k.upper() for k in _factory}:
             raise NotImplementedError(f"Unsupported format <{fmt}>")
 
-        return cls(_factory[fmt].from_filename(filename))
+        if fmt == "X5":
+            from .io.x5 import from_filename as load_x5
+            x5_xfm = load_x5(filename)[0]
+            Domain = namedtuple("Domain", "affine shape")
+            reference = Domain(x5_xfm.domain.mapping, x5_xfm.domain.size)
+            return cls(x5_xfm.transform, is_deltas=True, reference=reference)
+
+        return cls(_factory[fmt.lower()].from_filename(filename))
 
 
 load = DenseFieldTransform.from_filename
@@ -293,6 +331,39 @@ def to_field(self, reference=None, dtype="float32"):
             field.astype(dtype).reshape(*_ref.shape, -1), reference=_ref
         )
 
+    def to_x5(self, metadata=None):
+        """Return an :class:`~nitransforms.io.x5.X5Transform` representation."""
+        from ._version import __version__
+        from .io.x5 import X5Transform, X5Domain
+
+        metadata = {"WrittenBy": f"NiTransforms {__version__}"} | (metadata or {})
+
+        domain = None
+        if (reference := self.reference) is not None:
+            domain = X5Domain(
+                grid=True,
+                size=getattr(reference, "shape", (0, 0, 0)),
+                mapping=reference.affine,
+                coordinates="cartesian",
+            )
+
+        meta = metadata | {
+            "KnotsAffine": self._knots.affine.tolist(),
+            "KnotsShape": self._knots.shape,
+        }
+
+        kinds = tuple("space" for _ in range(self.ndim)) + ("vector",)
+
+        return X5Transform(
+            type="nonlinear",
+            subtype="bspline",
+            representation="coefficients",
+            metadata=meta,
+            transform=self._coeffs,
+            dimension_kinds=kinds,
+            domain=domain,
+        )
+
     def map(self, x, inverse=False):
         r"""
         Apply the transformation to a list of physical coordinate points.

nitransforms/tests/test_nonlinear.py

@@ -14,6 +14,7 @@
     BSplineFieldTransform,
     DenseFieldTransform,
 )
+from nitransforms import io
 from ..io.itk import ITKDisplacementsField
 
 
@@ -119,3 +120,47 @@ def test_bspline(tmp_path, testdata_path):
         ).mean()
         < 0.2
     )
+
+
+def test_densefield_x5_roundtrip(tmp_path):
+    """Ensure dense field transforms roundtrip via X5."""
+    ref = nb.Nifti1Image(np.zeros((2, 2, 2), dtype="uint8"), np.eye(4))
+    disp = nb.Nifti1Image(np.random.rand(2, 2, 2, 3).astype("float32"), np.eye(4))
+
+    xfm = DenseFieldTransform(disp, reference=ref)
+
+    node = xfm.to_x5(metadata={"GeneratedBy": "pytest"})
+    assert node.type == "nonlinear"
+    assert node.subtype == "densefield"
+    assert node.representation == "displacements"
+    assert node.domain.size == ref.shape
+    assert node.metadata["GeneratedBy"] == "pytest"
+
+    fname = tmp_path / "test.x5"
+    io.x5.to_filename(fname, [node])
+
+    xfm2 = DenseFieldTransform.from_filename(fname, fmt="X5")
+    diff = xfm2._deltas - xfm._deltas
+    coords = xfm.reference.ndcoords.T.reshape(xfm._deltas.shape)
+    assert np.allclose(diff, coords)
+    assert xfm2.reference.shape == ref.shape
+    assert np.allclose(xfm2.reference.affine, ref.affine)
+
+
+def test_bspline_to_x5(tmp_path):
+    """Check BSpline transforms export to X5."""
+    coeff = nb.Nifti1Image(np.zeros((2, 2, 2, 3), dtype="float32"), np.eye(4))
+    ref = nb.Nifti1Image(np.zeros((2, 2, 2), dtype="uint8"), np.eye(4))
+
+    xfm = BSplineFieldTransform(coeff, reference=ref)
+    node = xfm.to_x5(metadata={"tool": "pytest"})
+    assert node.type == "nonlinear"
+    assert node.subtype == "bspline"
+    assert node.representation == "coefficients"
+    assert node.metadata["tool"] == "pytest"
+
+    fname = tmp_path / "bspline.x5"
+    io.x5.to_filename(fname, [node])
+    node2 = io.x5.from_filename(fname)[0]
+    assert np.allclose(node2.transform, node.transform)
+    assert node2.metadata["tool"] == "pytest"