diff --git a/docs/_api/io.rst b/docs/_api/io.rst index 7666253b..1a3b6865 100644 --- a/docs/_api/io.rst +++ b/docs/_api/io.rst @@ -20,6 +20,18 @@ AFNI .. automodule:: nitransforms.io.afni :members: +^^^^^^^^ +BIDS' X5 +^^^^^^^^ +.. automodule:: nitransforms.io.x5 + :members: + +^^^^^^^^^^^^^^ +FreeSurfer/LTA +^^^^^^^^^^^^^^ +.. automodule:: nitransforms.io.lta + :members: + ^^^ FSL ^^^ @@ -31,9 +43,3 @@ ITK ^^^ .. automodule:: nitransforms.io.itk :members: - -^^^^^^^^^^^^^^ -FreeSurfer/LTA -^^^^^^^^^^^^^^ -.. automodule:: nitransforms.io.lta - :members: diff --git a/nitransforms/io/__init__.py b/nitransforms/io/__init__.py index c38d11c2..f9030724 100644 --- a/nitransforms/io/__init__.py +++ b/nitransforms/io/__init__.py @@ -1,7 +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 +from nitransforms.io import afni, fsl, itk, lta, x5 from nitransforms.io.base import TransformIOError, TransformFileError __all__ = [ @@ -9,6 +9,7 @@ "fsl", "itk", "lta", + "x5", "get_linear_factory", "TransformFileError", "TransformIOError", diff --git a/nitransforms/io/x5.py b/nitransforms/io/x5.py new file mode 100644 index 00000000..f4c14113 --- /dev/null +++ b/nitransforms/io/x5.py @@ -0,0 +1,138 @@ +# 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. +# +### ### ### ### ### ### ### ### ### ### ### ### ### ### ### ### ### ### ### ## +""" +Data structures for the X5 transform format. + +Implements what's drafted in the `BIDS X5 specification draft +`__. + +""" + +from __future__ import annotations + +from dataclasses import dataclass +from pathlib import Path +from typing import Any, Dict, Optional, Sequence, List + +import json +import h5py + +import numpy as np + + +@dataclass +class X5Domain: + """Domain information of a transform representing reference/moving spaces.""" + + grid: bool + """Whether sampling locations in the manifold are located regularly.""" + size: Sequence[int] + """The number of sampling locations per dimension (or total if not a grid).""" + mapping: Optional[np.ndarray] + """A mapping to go from samples (pixel/voxel coordinates, indices) to space coordinates.""" + coordinates: Optional[str] = None + """Indexing type of the Mapping field (for example, "cartesian", "spherical" or "index").""" + + +@dataclass +class X5Transform: + """Represent one transform entry of an X5 file.""" + + type: str + """A REQUIRED unicode string with possible values: "linear", "nonlinear", "composite".""" + transform: np.ndarray + """A REQUIRED array of parameters (e.g., affine matrix, or dense displacements field).""" + subtype: Optional[str] = None + """An OPTIONAL extension of type to drive the interpretation of AdditionalParameters.""" + representation: Optional[str] = None + """ + A string specifiying the transform representation or model, REQUIRED only for nonlinear Type. + """ + metadata: Optional[Dict[str, Any]] = None + """An OPTIONAL string (JSON) to embed metadata.""" + dimension_kinds: Optional[Sequence[str]] = None + """Identifies what "kind" of information is represented by the samples along each axis.""" + domain: Optional[X5Domain] = None + """ + A dataset specifying the reference manifold for the transform (either + a regularly gridded 3D space or a surface/sphere). + REQUIRED for nonlinear Type, RECOMMENDED for linear Type. + """ + inverse: Optional[np.ndarray] = None + """Placeholder to pre-calculated inverses.""" + jacobian: Optional[np.ndarray] = None + """ + A RECOMMENDED data array to keep cached the determinant of Jacobian of the transform + in case tools have calculated it. + For parametric models it is generally possible to obtain it analytically, so this dataset + could not be as useful in that case. + """ + # additional_parameters: Optional[np.ndarray] = None + # AdditionalParameters is empty in the draft spec - ignore for now. + # Only documentation ATM is for SubType: + # The SubType setting enables setting the additional parameters on a dataset called + # "AdditionalParameters" that hangs directly from this transform node. + array_length: int = 1 + """Undocumented field in the draft to enable a single transform group for 4D transforms.""" + + +def to_filename(fname: str | Path, x5_list: List[X5Transform]): + """ + Write a list of :class:`X5Transform` objects to an X5 HDF5 file. + + Parameters + ---------- + fname : :obj:`os.pathlike` + The file name (preferably with the ".x5" extension) in which transforms will be stored. + x5_list : :obj:`list` + The list of transforms to be stored in the output dataset. + + Returns + ------- + fname : :obj:`os.pathlike` + File containing the transform(s). + + """ + 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 + + 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) + # Disabled until we need SubType and AdditionalParameters + # if node.additional_parameters is not None: + # g.create_dataset( + # "AdditionalParameters", data=node.additional_parameters + # ) + return fname diff --git a/nitransforms/linear.py b/nitransforms/linear.py index 71df6a16..cf8f8465 100644 --- a/nitransforms/linear.py +++ b/nitransforms/linear.py @@ -7,12 +7,19 @@ # ### ### ### ### ### ### ### ### ### ### ### ### ### ### ### ### ### ### ### ## """Linear transforms.""" + import warnings import numpy as np from pathlib import Path from nibabel.affines import from_matvec +# Avoids circular imports +try: + from nitransforms._version import __version__ +except ModuleNotFoundError: # pragma: no cover + __version__ = "0+unknown" + from nitransforms.base import ( ImageGrid, TransformBase, @@ -20,6 +27,7 @@ EQUALITY_TOL, ) from nitransforms.io import get_linear_factory, TransformFileError +from nitransforms.io.x5 import X5Transform, X5Domain, to_filename as save_x5 class Affine(TransformBase): @@ -79,6 +87,23 @@ def __init__(self, matrix=None, reference=None): self._matrix[3, :] = (0, 0, 0, 1) self._inverse = np.linalg.inv(self._matrix) + def __repr__(self): + """ + Change representation to the internal matrix. + + Example + ------- + >>> Affine([ + ... [1, 0, 0, 4], [0, 1, 0, 0], [0, 0, 1, 0], [0, 0, 0, 1] + ... ]) # doctest: +NORMALIZE_WHITESPACE + array([[1, 0, 0, 4], + [0, 1, 0, 0], + [0, 0, 1, 0], + [0, 0, 0, 1]]) + + """ + return repr(self.matrix) + def __eq__(self, other): """ Overload equals operator. @@ -148,62 +173,6 @@ def ndim(self): """Access the internal representation of this affine.""" return self._matrix.ndim + 1 - def map(self, x, inverse=False): - r""" - Apply :math:`y = f(x)`. - - Parameters - ---------- - x : N x D numpy.ndarray - Input RAS+ coordinates (i.e., physical coordinates). - inverse : bool - If ``True``, apply the inverse transform :math:`x = f^{-1}(y)`. - - Returns - ------- - y : N x D numpy.ndarray - Transformed (mapped) RAS+ coordinates (i.e., physical coordinates). - - Examples - -------- - >>> xfm = Affine([[1, 0, 0, 1], [0, 1, 0, 2], [0, 0, 1, 3], [0, 0, 0, 1]]) - >>> xfm.map((0,0,0)) - array([[1., 2., 3.]]) - - >>> xfm.map((0,0,0), inverse=True) - array([[-1., -2., -3.]]) - - """ - affine = self._matrix - coords = _as_homogeneous(x, dim=affine.shape[0] - 1).T - if inverse is True: - affine = self._inverse - return affine.dot(coords).T[..., :-1] - - def _to_hdf5(self, x5_root): - """Serialize this object into the x5 file format.""" - xform = x5_root.create_dataset("Transform", data=[self._matrix]) - xform.attrs["Type"] = "affine" - x5_root.create_dataset("Inverse", data=[(~self).matrix]) - - if self._reference: - self.reference._to_hdf5(x5_root.create_group("Reference")) - - def to_filename(self, filename, fmt="X5", moving=None): - """Store the transform in the requested output format.""" - writer = get_linear_factory(fmt, is_array=False) - - if fmt.lower() in ("itk", "ants", "elastix"): - writer.from_ras(self.matrix).to_filename(filename) - else: - # Rest of the formats peek into moving and reference image grids - writer.from_ras( - self.matrix, - reference=self.reference, - moving=ImageGrid(moving) if moving is not None else self.reference, - ).to_filename(filename) - return filename - @classmethod def from_filename(cls, filename, fmt=None, reference=None, moving=None): """Create an affine from a transform file.""" @@ -259,22 +228,80 @@ def from_matvec(cls, mat=None, vec=None, reference=None): vec = vec if vec is not None else np.zeros((3,)) return cls(from_matvec(mat, vector=vec), reference=reference) - def __repr__(self): - """ - Change representation to the internal matrix. + def map(self, x, inverse=False): + r""" + Apply :math:`y = f(x)`. - Example + Parameters + ---------- + x : N x D numpy.ndarray + Input RAS+ coordinates (i.e., physical coordinates). + inverse : bool + If ``True``, apply the inverse transform :math:`x = f^{-1}(y)`. + + Returns ------- - >>> Affine([ - ... [1, 0, 0, 4], [0, 1, 0, 0], [0, 0, 1, 0], [0, 0, 0, 1] - ... ]) # doctest: +NORMALIZE_WHITESPACE - array([[1, 0, 0, 4], - [0, 1, 0, 0], - [0, 0, 1, 0], - [0, 0, 0, 1]]) + y : N x D numpy.ndarray + Transformed (mapped) RAS+ coordinates (i.e., physical coordinates). + + Examples + -------- + >>> xfm = Affine([[1, 0, 0, 1], [0, 1, 0, 2], [0, 0, 1, 3], [0, 0, 0, 1]]) + >>> xfm.map((0,0,0)) + array([[1., 2., 3.]]) + + >>> xfm.map((0,0,0), inverse=True) + array([[-1., -2., -3.]]) """ - return repr(self.matrix) + affine = self._matrix + coords = _as_homogeneous(x, dim=affine.shape[0] - 1).T + if inverse is True: + affine = self._inverse + return affine.dot(coords).T[..., :-1] + + def to_filename(self, filename, fmt="X5", moving=None, x5_inverse=False): + """Store the transform in the requested output format.""" + 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)) + + if fmt.lower() in ("itk", "ants", "elastix"): + writer.from_ras(self.matrix).to_filename(filename) + else: + # Rest of the formats peek into moving and reference image grids + writer.from_ras( + self.matrix, + reference=self.reference, + moving=ImageGrid(moving) if moving is not None else self.reference, + ).to_filename(filename) + return filename + + def to_x5(self, store_inverse=False, metadata=None): + """Return an :class:`~nitransforms.io.x5.X5Transform` representation.""" + metadata = {"WrittenBy": f"NiTransforms {__version__}"} | (metadata or {}) + + domain = None + if (reference := self.reference) is not None: + domain = X5Domain( + grid=True, + size=getattr(reference or {}, "shape", (0, 0, 0)), + mapping=reference.affine, + coordinates="cartesian", + ) + kinds = tuple("space" for _ in range(self.ndim)) + ("vector",) + return X5Transform( + type="linear", + subtype="affine", + representation="matrix", + metadata=metadata, + transform=self.matrix, + dimension_kinds=kinds, + domain=domain, + inverse=(~self).matrix if store_inverse else None, + array_length=len(self), + ) class LinearTransformsMapping(Affine): @@ -383,21 +410,6 @@ def map(self, x, inverse=False): affine = self._inverse return np.swapaxes(affine.dot(coords), 1, 2) - def to_filename(self, filename, fmt="X5", moving=None): - """Store the transform in the requested output format.""" - writer = get_linear_factory(fmt, is_array=True) - - if fmt.lower() in ("itk", "ants", "elastix"): - writer.from_ras(self.matrix).to_filename(filename) - else: - # Rest of the formats peek into moving and reference image grids - writer.from_ras( - self.matrix, - reference=self.reference, - moving=ImageGrid(moving) if moving is not None else self.reference, - ).to_filename(filename) - return filename - def load(filename, fmt=None, reference=None, moving=None): """ diff --git a/nitransforms/tests/test_linear.py b/nitransforms/tests/test_linear.py index 31627159..32634c61 100644 --- a/nitransforms/tests/test_linear.py +++ b/nitransforms/tests/test_linear.py @@ -4,10 +4,12 @@ import pytest 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 +245,50 @@ 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): +@pytest.mark.parametrize("store_inverse", [True, False]) +def test_linear_to_x5(tmpdir, store_inverse): """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")) + # Test base operations + aff = nitl.Affine() + node = aff.to_x5( + metadata={"GeneratedBy": "FreeSurfer 8"}, store_inverse=store_inverse + ) + assert node.type == "linear" + assert node.subtype == "affine" + assert node.representation == "matrix" + assert node.domain is None + assert node.transform.shape == (4, 4) + assert node.array_length == 1 + assert (node.metadata or {}).get("GeneratedBy") == "FreeSurfer 8" + + aff.to_filename("export1.x5", x5_inverse=store_inverse) + + # Test with Domain + 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 + aff.to_filename("export2.x5", x5_inverse=store_inverse) + + # Test with Jacobian + node.jacobian = np.zeros((2, 2, 2), dtype="float32") + io.x5.to_filename("export3.x5", [node]) + + +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): diff --git a/nitransforms/tests/test_x5.py b/nitransforms/tests/test_x5.py new file mode 100644 index 00000000..8502a387 --- /dev/null +++ b/nitransforms/tests/test_x5.py @@ -0,0 +1,41 @@ +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.array_length == 1 + # Disabled for now + # assert xf.additional_parameters is None + + +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