Merge pull request #1251 from effigies/enh/pointset

ENH: Add pointset data structures [BIAP9]

Author: effigies

.pre-commit-config.yaml

@@ -1,7 +1,7 @@
 exclude: ".*/data/.*"
 repos:
   - repo: https://github.com/pre-commit/pre-commit-hooks
-    rev: v4.1.0
+    rev: v4.4.0
     hooks:
       - id: trailing-whitespace
       - id: end-of-file-fixer
@@ -21,12 +21,12 @@ repos:
     hooks:
       - id: isort
   - repo: https://github.com/pycqa/flake8
-    rev: 6.0.0
+    rev: 6.1.0
     hooks:
       - id: flake8
         exclude: "^(doc|nisext|tools)/"
   - repo: https://github.com/pre-commit/mirrors-mypy
-    rev: v0.991
+    rev: v1.5.1
     hooks:
       - id: mypy
         # Sync with project.optional-dependencies.typing

nibabel/pointset.py

@@ -0,0 +1,194 @@
+"""Point-set structures
+
+Imaging data are sampled at points in space, and these points
+can be described by coordinates.
+These structures are designed to enable operations on sets of
+points, as opposed to the data sampled at those points.
+
+Abstractly, a point set is any collection of points, but there are
+two types that warrant special consideration in the neuroimaging
+context: grids and meshes.
+
+A *grid* is a collection of regularly-spaced points. The canonical
+examples of grids are the indices of voxels and their affine
+projection into a reference space.
+
+A *mesh* is a collection of points and some structure that enables
+adjacent points to be identified. A *triangular mesh* in particular
+uses triplets of adjacent vertices to describe faces.
+"""
+from __future__ import annotations
+
+import math
+import typing as ty
+from dataclasses import dataclass, replace
+
+import numpy as np
+
+from nibabel.casting import able_int_type
+from nibabel.fileslice import strided_scalar
+from nibabel.spatialimages import SpatialImage
+
+if ty.TYPE_CHECKING:  # pragma: no cover
+    from typing_extensions import Self
+
+    _DType = ty.TypeVar('_DType', bound=np.dtype[ty.Any])
+
+
+class CoordinateArray(ty.Protocol):
+    ndim: int
+    shape: tuple[int, int]
+
+    @ty.overload
+    def __array__(self, dtype: None = ..., /) -> np.ndarray[ty.Any, np.dtype[ty.Any]]:
+        ...  # pragma: no cover
+
+    @ty.overload
+    def __array__(self, dtype: _DType, /) -> np.ndarray[ty.Any, _DType]:
+        ...  # pragma: no cover
+
+
+@dataclass
+class Pointset:
+    """A collection of points described by coordinates.
+
+    Parameters
+    ----------
+    coords : array-like
+      (*N*, *n*) array with *N* being points and columns their *n*-dimensional coordinates
+    affine : :class:`numpy.ndarray`
+      Affine transform to be applied to coordinates array
+    homogeneous : :class:`bool`
+      Indicate whether the provided coordinates are homogeneous,
+      i.e., homogeneous 3D coordinates have the form ``(x, y, z, 1)``
+    """
+
+    coordinates: CoordinateArray
+    affine: np.ndarray
+    homogeneous: bool = False
+
+    # Force use of __rmatmul__ with numpy arrays
+    __array_priority__ = 99
+
+    def __init__(
+        self,
+        coordinates: CoordinateArray,
+        affine: np.ndarray | None = None,
+        homogeneous: bool = False,
+    ):
+        self.coordinates = coordinates
+        self.homogeneous = homogeneous
+
+        if affine is None:
+            self.affine = np.eye(self.dim + 1)
+        else:
+            self.affine = np.asanyarray(affine)
+
+        if self.affine.shape != (self.dim + 1,) * 2:
+            raise ValueError(f'Invalid affine for {self.dim}D coordinates:\n{self.affine}')
+        if np.any(self.affine[-1, :-1] != 0) or self.affine[-1, -1] != 1:
+            raise ValueError(f'Invalid affine matrix:\n{self.affine}')
+
+    @property
+    def n_coords(self) -> int:
+        """Number of coordinates
+
+        Subclasses should override with more efficient implementations.
+        """
+        return self.coordinates.shape[0]
+
+    @property
+    def dim(self) -> int:
+        """The dimensionality of the space the coordinates are in"""
+        return self.coordinates.shape[1] - self.homogeneous
+
+    def __rmatmul__(self, affine: np.ndarray) -> Self:
+        """Apply an affine transformation to the pointset
+
+        This will return a new pointset with an updated affine matrix only.
+        """
+        return replace(self, affine=np.asanyarray(affine) @ self.affine)
+
+    def _homogeneous_coords(self):
+        if self.homogeneous:
+            return np.asanyarray(self.coordinates)
+
+        ones = strided_scalar(
+            shape=(self.coordinates.shape[0], 1),
+            scalar=np.array(1, dtype=self.coordinates.dtype),
+        )
+        return np.hstack((self.coordinates, ones))
+
+    def get_coords(self, *, as_homogeneous: bool = False):
+        """Retrieve the coordinates
+
+        Parameters
+        ----------
+        as_homogeneous : :class:`bool`
+            Return homogeneous coordinates if ``True``, or Cartesian
+            coordiantes if ``False``.
+
+        name : :class:`str`
+            Select a particular coordinate system if more than one may exist.
+            By default, `None` is equivalent to `"world"` and corresponds to
+            an RAS+ coordinate system.
+        """
+        ident = np.allclose(self.affine, np.eye(self.affine.shape[0]))
+        if self.homogeneous == as_homogeneous and ident:
+            return np.asanyarray(self.coordinates)
+        coords = self._homogeneous_coords()
+        if not ident:
+            coords = (self.affine @ coords.T).T
+        if not as_homogeneous:
+            coords = coords[:, :-1]
+        return coords
+
+
+class Grid(Pointset):
+    r"""A regularly-spaced collection of coordinates
+
+    This class provides factory methods for generating Pointsets from
+    :class:`~nibabel.spatialimages.SpatialImage`\s and generating masks
+    from coordinate sets.
+    """
+
+    @classmethod
+    def from_image(cls, spatialimage: SpatialImage) -> Self:
+        return cls(coordinates=GridIndices(spatialimage.shape[:3]), affine=spatialimage.affine)
+
+    @classmethod
+    def from_mask(cls, mask: SpatialImage) -> Self:
+        mask_arr = np.bool_(mask.dataobj)
+        return cls(
+            coordinates=np.c_[np.nonzero(mask_arr)].astype(able_int_type(mask.shape)),
+            affine=mask.affine,
+        )
+
+    def to_mask(self, shape=None) -> SpatialImage:
+        if shape is None:
+            shape = tuple(np.max(self.coordinates, axis=0)[: self.dim] + 1)
+        mask_arr = np.zeros(shape, dtype='bool')
+        mask_arr[tuple(np.asanyarray(self.coordinates)[:, : self.dim].T)] = True
+        return SpatialImage(mask_arr, self.affine)
+
+
+class GridIndices:
+    """Class for generating indices just-in-time"""
+
+    __slots__ = ('gridshape', 'dtype', 'shape')
+    ndim = 2
+
+    def __init__(self, shape, dtype=None):
+        self.gridshape = shape
+        self.dtype = dtype or able_int_type(shape)
+        self.shape = (math.prod(self.gridshape), len(self.gridshape))
+
+    def __repr__(self):
+        return f'<{self.__class__.__name__}{self.gridshape}>'
+
+    def __array__(self, dtype=None):
+        if dtype is None:
+            dtype = self.dtype
+
+        axes = [np.arange(s, dtype=dtype) for s in self.gridshape]
+        return np.reshape(np.meshgrid(*axes, copy=False, indexing='ij'), (len(axes), -1)).T

nibabel/tests/test_pointset.py

@@ -0,0 +1,184 @@
+from math import prod
+from pathlib import Path
+from unittest import skipUnless
+
+import numpy as np
+import pytest
+
+from nibabel import pointset as ps
+from nibabel.affines import apply_affine
+from nibabel.arrayproxy import ArrayProxy
+from nibabel.fileslice import strided_scalar
+from nibabel.onetime import auto_attr
+from nibabel.optpkg import optional_package
+from nibabel.spatialimages import SpatialImage
+from nibabel.tests.nibabel_data import get_nibabel_data
+
+h5, has_h5py, _ = optional_package('h5py')
+
+FS_DATA = Path(get_nibabel_data()) / 'nitest-freesurfer'
+
+
+class TestPointsets:
+    rng = np.random.default_rng()
+
+    @pytest.mark.parametrize('shape', [(5, 2), (5, 3), (5, 4)])
+    @pytest.mark.parametrize('homogeneous', [True, False])
+    def test_init(self, shape, homogeneous):
+        coords = self.rng.random(shape)
+
+        if homogeneous:
+            coords = np.column_stack([coords, np.ones(shape[0])])
+
+        points = ps.Pointset(coords, homogeneous=homogeneous)
+        assert np.allclose(points.affine, np.eye(shape[1] + 1))
+        assert points.homogeneous is homogeneous
+        assert (points.n_coords, points.dim) == shape
+
+        points = ps.Pointset(coords, affine=np.diag([2] * shape[1] + [1]), homogeneous=homogeneous)
+        assert np.allclose(points.affine, np.diag([2] * shape[1] + [1]))
+        assert points.homogeneous is homogeneous
+        assert (points.n_coords, points.dim) == shape
+
+        # Badly shaped affine
+        with pytest.raises(ValueError):
+            ps.Pointset(coords, affine=[0, 1])
+
+        # Badly valued affine
+        with pytest.raises(ValueError):
+            ps.Pointset(coords, affine=np.ones((shape[1] + 1, shape[1] + 1)))
+
+    @pytest.mark.parametrize('shape', [(5, 2), (5, 3), (5, 4)])
+    @pytest.mark.parametrize('homogeneous', [True, False])
+    def test_affines(self, shape, homogeneous):
+        orig_coords = coords = self.rng.random(shape)
+
+        if homogeneous:
+            coords = np.column_stack([coords, np.ones(shape[0])])
+
+        points = ps.Pointset(coords, homogeneous=homogeneous)
+        assert np.allclose(points.get_coords(), orig_coords)
+
+        # Apply affines
+        scaler = np.diag([2] * shape[1] + [1])
+        scaled = scaler @ points
+        assert np.array_equal(scaled.coordinates, points.coordinates)
+        assert np.array_equal(scaled.affine, scaler)
+        assert np.allclose(scaled.get_coords(), 2 * orig_coords)
+
+        flipper = np.eye(shape[1] + 1)
+        # [[1, 0, 0], [0, 1, 0], [0, 0, 1]] becomes [[0, 1, 0], [1, 0, 0], [0, 0, 1]]
+        flipper[:-1] = flipper[-2::-1]
+        flipped = flipper @ points
+        assert np.array_equal(flipped.coordinates, points.coordinates)
+        assert np.array_equal(flipped.affine, flipper)
+        assert np.allclose(flipped.get_coords(), orig_coords[:, ::-1])
+
+        # Concatenate affines, with any associativity
+        for doubledup in [(scaler @ flipper) @ points, scaler @ (flipper @ points)]:
+            assert np.array_equal(doubledup.coordinates, points.coordinates)
+            assert np.allclose(doubledup.affine, scaler @ flipper)
+            assert np.allclose(doubledup.get_coords(), 2 * orig_coords[:, ::-1])
+
+    def test_homogeneous_coordinates(self):
+        ccoords = self.rng.random((5, 3))
+        hcoords = np.column_stack([ccoords, np.ones(5)])
+
+        cartesian = ps.Pointset(ccoords)
+        homogeneous = ps.Pointset(hcoords, homogeneous=True)
+
+        for points in (cartesian, homogeneous):
+            assert np.array_equal(points.get_coords(), ccoords)
+            assert np.array_equal(points.get_coords(as_homogeneous=True), hcoords)
+
+        affine = np.diag([2, 3, 4, 1])
+        cart2 = affine @ cartesian
+        homo2 = affine @ homogeneous
+
+        exp_c = apply_affine(affine, ccoords)
+        exp_h = (affine @ hcoords.T).T
+        for points in (cart2, homo2):
+            assert np.array_equal(points.get_coords(), exp_c)
+            assert np.array_equal(points.get_coords(as_homogeneous=True), exp_h)
+
+
+def test_GridIndices():
+    # 2D case
+    shape = (2, 3)
+    gi = ps.GridIndices(shape)
+
+    assert gi.dtype == np.dtype('u1')
+    assert gi.shape == (6, 2)
+    assert repr(gi) == '<GridIndices(2, 3)>'
+
+    gi_arr = np.asanyarray(gi)
+    assert gi_arr.dtype == np.dtype('u1')
+    assert gi_arr.shape == (6, 2)
+    # Tractable to write out
+    assert np.array_equal(gi_arr, [[0, 0], [0, 1], [0, 2], [1, 0], [1, 1], [1, 2]])
+
+    shape = (2, 3, 4)
+    gi = ps.GridIndices(shape)
+
+    assert gi.dtype == np.dtype('u1')
+    assert gi.shape == (24, 3)
+    assert repr(gi) == '<GridIndices(2, 3, 4)>'
+
+    gi_arr = np.asanyarray(gi)
+    assert gi_arr.dtype == np.dtype('u1')
+    assert gi_arr.shape == (24, 3)
+    # Separate implementation
+    assert np.array_equal(gi_arr, np.mgrid[:2, :3, :4].reshape(3, -1).T)
+
+
+class TestGrids(TestPointsets):
+    @pytest.mark.parametrize('shape', [(5, 5, 5), (5, 5, 5, 5), (5, 5, 5, 5, 5)])
+    def test_from_image(self, shape):
+        # Check image is generates voxel coordinates
+        affine = np.diag([2, 3, 4, 1])
+        img = SpatialImage(strided_scalar(shape), affine)
+        grid = ps.Grid.from_image(img)
+        grid_coords = grid.get_coords()
+
+        assert grid.n_coords == prod(shape[:3])
+        assert grid.dim == 3
+        assert np.allclose(grid.affine, affine)
+
+        assert np.allclose(grid_coords[0], [0, 0, 0])
+        # Final index is [4, 4, 4], scaled by affine
+        assert np.allclose(grid_coords[-1], [8, 12, 16])
+
+    def test_from_mask(self):
+        affine = np.diag([2, 3, 4, 1])
+        mask = np.zeros((3, 3, 3))
+        mask[1, 1, 1] = 1
+        img = SpatialImage(mask, affine)
+
+        grid = ps.Grid.from_mask(img)
+        grid_coords = grid.get_coords()
+
+        assert grid.n_coords == 1
+        assert grid.dim == 3
+        assert np.array_equal(grid_coords, [[2, 3, 4]])
+
+    def test_to_mask(self):
+        coords = np.array([[1, 1, 1]])
+
+        grid = ps.Grid(coords)
+
+        mask_img = grid.to_mask()
+        assert mask_img.shape == (2, 2, 2)
+        assert np.array_equal(mask_img.get_fdata(), [[[0, 0], [0, 0]], [[0, 0], [0, 1]]])
+        assert np.array_equal(mask_img.affine, np.eye(4))
+
+        mask_img = grid.to_mask(shape=(3, 3, 3))
+        assert mask_img.shape == (3, 3, 3)
+        assert np.array_equal(
+            mask_img.get_fdata(),
+            [
+                [[0, 0, 0], [0, 0, 0], [0, 0, 0]],
+                [[0, 0, 0], [0, 1, 0], [0, 0, 0]],
+                [[0, 0, 0], [0, 0, 0], [0, 0, 0]],
+            ],
+        )
+        assert np.array_equal(mask_img.affine, np.eye(4))