diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index 1fc7efd0b..137aa4946 100644 --- a/.pre-commit-config.yaml +++ b/.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 diff --git a/nibabel/pointset.py b/nibabel/pointset.py new file mode 100644 index 000000000..b40449801 --- /dev/null +++ b/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 diff --git a/nibabel/tests/test_pointset.py b/nibabel/tests/test_pointset.py new file mode 100644 index 000000000..fb9a7c5c8 --- /dev/null +++ b/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) == '' + + 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) == '' + + 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))