fix: improve loading and equalities

Author: oesteban

nitransforms/linear.py

@@ -122,6 +122,9 @@ def __eq__(self, other):
         True
 
         """
+        if not hasattr(other, "matrix"):
+            return False
+
         _eq = np.allclose(self.matrix, other.matrix, rtol=EQUALITY_TOL)
         if _eq and self._reference != other._reference:
             warnings.warn("Affines are equal, but references do not match.")

nitransforms/manip.py

@@ -208,10 +208,16 @@ def from_filename(cls, filename, fmt="X5", reference=None, moving=None, x5_chain
 
         retval = []
         if fmt and fmt.upper() == "X5":
-            xfm_list = load_x5(filename)
+            # Get list of X5 nodes and generate transforms
+            xfm_list = [
+                globals()[f"{node.type}_from_x5"]([node]) for node in load_x5(filename)
+            ]
             if not xfm_list:
                 raise TransformError("Empty transform group")
 
+            if x5_chain is None:
+                return xfm_list
+
             with h5py.File(str(filename), "r") as f:
                 chain_grp = f.get("TransformChain")
                 if chain_grp is None:
@@ -221,12 +227,7 @@ def from_filename(cls, filename, fmt="X5", reference=None, moving=None, 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)
+            return TransformChain([xfm_list[int(idx)] for idx in chain_path.split("/")])
 
         if str(filename).endswith(".h5"):
             reference = None
@@ -247,9 +248,14 @@ def to_filename(self, filename, fmt="X5"):
         if fmt.upper() != "X5":
             raise NotImplementedError("Only X5 format is supported for chains")
 
-        existing = load_x5(filename) if os.path.exists(filename) else []
+        existing = (
+            self.from_filename(filename, x5_chain=None)
+            if os.path.exists(filename)
+            else []
+        )
+
         xfm_chain = []
-        new_nodes = []
+        new_xfms = []
         next_xfm_index = len(existing)
         for xfm in self.transforms:
             for eidx, existing_xfm in enumerate(existing):
@@ -258,7 +264,7 @@ def to_filename(self, filename, fmt="X5"):
                     break
             else:
                 xfm_chain.append(next_xfm_index)
-                new_nodes.append((next_xfm_index, xfm))
+                new_xfms.append((next_xfm_index, xfm))
                 existing.append(xfm)
                 next_xfm_index += 1
 
@@ -269,9 +275,9 @@ def to_filename(self, filename, fmt="X5"):
                 f.attrs["Version"] = np.uint16(1)
 
             tg = f.require_group("TransformGroup")
-            for idx, node in new_nodes:
+            for idx, node in new_xfms:
                 g = tg.create_group(str(idx))
-                x5io._write_x5_group(g, node)
+                x5io._write_x5_group(g, node.to_x5())
 
             cg = f.require_group("TransformChain")
             cg.create_dataset(str(len(cg)), data="/".join(str(i) for i in xfm_chain))

nitransforms/nonlinear.py

@@ -247,6 +247,9 @@ def __eq__(self, other):
         True
 
         """
+        if not hasattr(other, "_field"):
+            return False
+
         _eq = np.allclose(self._field, other._field)
         if _eq and self._reference != other._reference:
             warnings.warn("Fields are equal, but references do not match.")
@@ -322,6 +325,28 @@ def __init__(self, coefficients, reference=None, order=3):
                     "not match the number of dimensions"
                 )
 
+    def __eq__(self, other):
+        """
+        Overload equals operator.
+
+        Examples
+        --------
+        >>> xfm1 = BSplineFieldTransform(test_dir / "someones_bspline_coefficients.nii.gz")
+        >>> xfm2 = BSplineFieldTransform(test_dir / "someones_bspline_coefficients.nii.gz")
+        >>> xfm1 == xfm2
+        True
+
+        """
+        if not hasattr(other, "_coeffs"):
+            return False
+
+        _eq = np.allclose(self._coeffs, other._coeffs)
+        _eq = _eq and self._order == other._order
+
+        if _eq and self._reference != other._reference:
+            warnings.warn("Coefficients are equal, but references do not match.")
+        return _eq
+
     @property
     def ndim(self):
         """Get the dimensions of the transform."""

nitransforms/tests/test_manip.py

@@ -59,11 +59,14 @@ def test_transformchain_x5_roundtrip(tmp_path):
 
     fname = tmp_path / "chain.x5"
     chain.to_filename(fname)
+
+    with h5py.File(fname) as f:
+        assert len(f["TransformGroup"]) == 2
+
     chain.to_filename(fname)  # append again, should not duplicate transforms
 
     with h5py.File(fname) as f:
         assert len(f["TransformGroup"]) == 2
-        assert len(f["TransformChain"]) == 2
 
     loaded0 = TransformChain.from_filename(fname, fmt="X5", x5_chain=0)
     loaded1 = TransformChain.from_filename(fname, fmt="X5", x5_chain=1)

nitransforms/tests/test_resampling.py

@@ -284,7 +284,7 @@ def test_apply_transformchain(tmp_path, testdata_path):
         / "ds-005_sub-01_from-T1w_to-MNI152NLin2009cAsym_mode-image_xfm.h5"
     )
 
-    xfm = nitm.load(xfm_fname)
+    xfm = nitm.load(xfm_fname, fmt="itk")
 
     assert len(xfm) == 2