diff --git a/nibabel/pointset.py b/nibabel/pointset.py index 91b753140..c131b8131 100644 --- a/nibabel/pointset.py +++ b/nibabel/pointset.py @@ -1,34 +1,151 @@ -import operator as op -from functools import reduce +"""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.affines import apply_affine +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: - def __init__(self, coords): - self._coords = coords + """A collection of points described by coordinates. + + Parameters + ---------- + coords : array-like + 2-dimensional array with coordinates as rows + 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 + ndim = 2 + __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 shape(self) -> tuple[int, int]: + """The shape of the coordinate array""" + return self.coordinates.shape @property - def n_coords(self): + def n_coords(self) -> int: """Number of coordinates Subclasses should override with more efficient implementations. """ - return self.get_coords().shape[0] + 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 - def get_coords(self, name=None): - """Nx3 array of coordinates. - Parameters ---------- - name : :obj:`str` + 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. """ - return self._coords + 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 TriangularMesh(Pointset): @@ -65,14 +182,6 @@ def get_names(self): """ raise NotImplementedError - ## This method is called for by the BIAP, but it now seems simpler to wait to - ## provide it until there are any proposed implementations - # def decimate(self, *, n_coords=None, ratio=None): - # """ Return a TriangularMesh with a smaller number of vertices that - # preserves the geometry of the original """ - # # To be overridden when a format provides optimization opportunities - # raise NotImplementedError - class TriMeshFamily(TriangularMesh): def __init__(self, mapping, default=None): @@ -97,40 +206,51 @@ def get_coords(self, name=None): return self._coords[name] -class NdGrid(Pointset): - """ - Attributes - ---------- - shape : 3-tuple - number of coordinates in each dimension of grid +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. """ - def __init__(self, shape, affines): - self.shape = tuple(shape) - try: - self._affines = dict(affines) - except (TypeError, ValueError): - self._affines = {'world': np.array(affines)} - if 'voxels' not in self._affines: - self._affines['voxels'] = np.eye(4, dtype=np.uint8) - - def get_affine(self, name=None): - """4x4 array""" - if name is None: - name = next(iter(self._affines)) - return self._affines[name] + @classmethod + def from_image(cls, spatialimage: SpatialImage) -> Self: + return cls(coordinates=GridIndices(spatialimage.shape[:3]), affine=spatialimage.affine) - def get_coords(self, name=None): - if name is None: - name = next(iter(self._affines)) - aff = self.get_affine(name) - dt = np.result_type(*(np.min_scalar_type(dim) for dim in self.shape)) - # This is pretty wasteful; we almost certainly want instead an - # object that will retrieve a coordinate when indexed, but where - # np.array(obj) returns this - ijk_coords = np.array(list(np.ndindex(self.shape)), dtype=dt) - return apply_affine(aff, ijk_coords) + @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, + ) - @property - def n_coords(self): - return reduce(op.mul, self.shape) + def to_mask(self, shape=None) -> SpatialImage: + if shape is None: + shape = tuple(np.max(self.coordinates, axis=1)[: self.dim]) + mask_arr = np.zeros(shape, dtype='bool') + mask_arr[np.asanyarray(self.coordinates)[:, : self.dim]] = 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