SurfaceTransform class

nitransforms/surface.py

@@ -0,0 +1,124 @@
+# 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.
+#
+### ### ### ### ### ### ### ### ### ### ### ### ### ### ### ### ### ### ### ##
+"""Surface transforms."""
+
+import h5py
+import numpy as np
+import scipy.sparse as sparse
+
+from nitransforms.base import TransformBase
+
+
+class SurfaceTransform(TransformBase):
+    """Represents transforms between surface spaces."""
+
+    __slots__ = ("mat",)
+
+    def __init__(self, mat):
+        """Initialize the transform.
+
+        Parameters
+        ----------
+        mat : array-like, shape (nv1, nv2)
+            Sparse matrix representing the transform.
+        """
+        super().__init__()
+        if isinstance(mat, sparse.csr_array):
+            self.mat = mat
+        else:
+            self.mat = sparse.csr_array(mat)
+
+    def apply(self, x, inverse=False, normalize="element"):
+        """Apply the transform to surface data.
+
+        Parameters
+        ----------
+        x : array-like, shape (..., nv1)
+            Data to transform.
+        inverse : bool, default=False
+            Whether to apply the inverse transform. If True, ``x`` has shape
+            (..., nv2), and the output will have shape (..., nv1).
+        normalize : {"element", "sum", "none"}, default="element"
+            Normalization strategy. If "element", the scale of each value in
+            the output is comparable to each value of the input. If "sum", the
+            sum of the output is comparable to the sum of the input. If
+            "none", no normalization is applied.
+
+        Returns
+        -------
+        y : array-like, shape (..., nv2)
+            Transformed data.
+        """
+        if normalize not in ("element", "sum", "none"):
+            raise ValueError("Invalid normalization strategy.")
+
+        mat = self.mat.T if inverse else self.mat
+
+        if normalize == "element":
+            sum_ = mat.sum(axis=0)
+            scale = np.zeros_like(sum_)
+            mask = sum_ != 0
+            scale[mask] = 1.0 / sum_[mask]
+            mat = mat @ sparse.diags(scale)
+        elif normalize == "sum":
+            sum_ = mat.sum(axis=1)
+            scale = np.zeros_like(sum_)
+            mask = sum_ != 0
+            scale[mask] = 1.0 / sum_[mask]
+            mat = sparse.diags(scale) @ mat
+
+        y = x @ mat
+        return y
+
+    def _to_hdf5(self, x5_root):
+        """Write transform to HDF5 file."""
+        xform = x5_root.create_group("Transform")
+        xform.attrs["Type"] = "surface"
+        xform.create_dataset("data", data=self.mat.data)
+        xform.create_dataset("indices", data=self.mat.indices)
+        xform.create_dataset("indptr", data=self.mat.indptr)
+        xform.create_dataset("shape", data=self.mat.shape)
+
+    def to_filename(self, filename, fmt=None):
+        """Store the transform."""
+        if fmt is None:
+            fmt = "npz" if filename.endswith(".npz") else "X5"
+
+        if fmt == "npz":
+            sparse.save_npz(filename, self.mat)
+            return filename
+
+        with h5py.File(filename, "w") as out_file:
+            out_file.attrs["Format"] = "X5"
+            out_file.attrs["Version"] = np.uint16(1)
+            root = out_file.create_group("/0")
+            self._to_hdf5(root)
+
+        return filename
+
+    @classmethod
+    def from_filename(cls, filename, fmt=None):
+        """Load transform from file."""
+        if fmt is None:
+            fmt = "npz" if filename.endswith(".npz") else "X5"
+
+        if fmt == "npz":
+            return cls(sparse.load_npz(filename))
+
+        if fmt != "X5":
+            raise ValueError("Only npz and X5 formats are supported.")
+
+        with h5py.File(filename, "r") as f:
+            assert f.attrs["Format"] == "X5"
+            xform = f["/0/Transform"]
+            mat = sparse.csr_matrix(
+                (xform["data"][()], xform["indices"][()], xform["indptr"][()]),
+                shape=xform["shape"][()],
+            )
+        return cls(mat)

nitransforms/tests/test_surface.py

@@ -0,0 +1,58 @@
+import os
+import tempfile
+
+import numpy as np
+import scipy.sparse as sparse
+
+from nitransforms.surface import SurfaceTransform
+
+
+def test_surface_transform_x5():
+    mat = sparse.random(10, 10, density=0.5)
+    xfm = SurfaceTransform(mat)
+    fn = tempfile.mktemp(suffix=".h5")
+    print(fn)
+    xfm.to_filename(fn)
+
+    xfm2 = SurfaceTransform.from_filename(fn)
+    try:
+        assert xfm.mat.shape == xfm2.mat.shape
+        np.testing.assert_array_equal(xfm.mat.data, xfm2.mat.data)
+        np.testing.assert_array_equal(xfm.mat.indices, xfm2.mat.indices)
+        np.testing.assert_array_equal(xfm.mat.indptr, xfm2.mat.indptr)
+    except Exception:
+        os.remove(fn)
+        raise
+    os.remove(fn)
+
+
+def test_surface_transform_npz():
+    mat = sparse.random(10, 10, density=0.5)
+    xfm = SurfaceTransform(mat)
+    fn = tempfile.mktemp(suffix=".npz")
+    print(fn)
+    xfm.to_filename(fn)
+
+    xfm2 = SurfaceTransform.from_filename(fn)
+    try:
+        assert xfm.mat.shape == xfm2.mat.shape
+        np.testing.assert_array_equal(xfm.mat.data, xfm2.mat.data)
+        np.testing.assert_array_equal(xfm.mat.indices, xfm2.mat.indices)
+        np.testing.assert_array_equal(xfm.mat.indptr, xfm2.mat.indptr)
+    except Exception:
+        os.remove(fn)
+        raise
+    os.remove(fn)
+
+
+def test_surface_transform_normalization():
+    mat = np.random.uniform(size=(20, 10))
+    xfm = SurfaceTransform(mat)
+    x = np.random.uniform(size=(5, 20))
+    y_element = xfm.apply(x, normalize="element")
+    np.testing.assert_array_less(y_element.sum(axis=1), x.sum(axis=1))
+    y_sum = xfm.apply(x, normalize="sum")
+    np.testing.assert_allclose(y_sum.sum(axis=1), x.sum(axis=1))
+    y_none = xfm.apply(x, normalize="none")
+    assert y_none.sum() != y_element.sum()
+    assert y_none.sum() != y_sum.sum()