transformnd 0.1.0__py3-none-any.whl

transformnd/transforms/map_axis.py

@@ -0,0 +1,67 @@
+from typing import Self
+from array_api_compat import array_namespace
+import numpy as np
+
+from ..base import Transform
+from ..util import ArrayT, SpaceTuple, to_single_ndim
+from ..transforms.affine import Affine
+
+
+class MapAxis(Transform[ArrayT]):
+    """Map coordinates from one axis to another.
+
+    For example, x -> y and y -> x"""
+
+    # ndim: Optional[Set[int]] = set(2)
+
+    def __init__(
+        self,
+        permutation: list[int],
+        *,
+        spaces: SpaceTuple = (None, None),
+    ):
+        """Base class for transformations.
+
+        Parameters
+        ----------
+        permutation: list[int]
+            New order of column axis. For example, [1, 0] means x -> y and y -> x.
+        spaces : tuple[SpaceRef, SpaceRef]
+            Optional source and target spaces
+        """
+        s_perm = sorted(permutation)
+        if any(a != b for a, b in enumerate(s_perm)):
+            raise ValueError(
+                "N-D permutation must contain all dimensions [0, N) exactly once"
+            )
+        self.permutation = permutation
+        self.ndim = {len(permutation)}
+        self.spaces = spaces
+
+    def is_identity(self) -> bool:
+        return all(a == b for a, b in enumerate(self.permutation))
+
+    def to_affine(self, ndim: int | None = None) -> Affine[ArrayT] | None:
+        ndim = to_single_ndim(ndim, self.ndim)
+        m = np.eye(ndim + 1)
+        perm = self.permutation + [ndim]
+        m = m[perm, :]
+        return Affine(m, spaces=self.spaces)
+
+    def apply(self, coords: ArrayT) -> ArrayT:
+        """Apply transformation to coordinates.
+
+        For example:
+        2-D with permutation [1, 0] will give you
+        [[x1, y1], [x2, y2]] -> [[y1, x1], [y2, x2]]
+        """
+
+        coords = self._validate_coords(coords)
+        xp = array_namespace(coords)
+        return xp.take(coords, self.permutation, 1)
+
+    def invert(self) -> Self | None:
+        return type(self)(
+            list(np.argsort(self.permutation)),
+            spaces=(self.spaces[1], self.spaces[0]),
+        )

transformnd/transforms/moving_least_squares.py

@@ -0,0 +1,66 @@
+"""@public
+Implementation of Moving Least Squares transformation.
+
+Requires the `movingleastsquares` extra.
+"""
+
+from array_api_compat import array_namespace
+import numpy as np
+from typing import Self
+from molesq.transform import Transformer as _Transformer
+
+from ..base import SpaceTuple, Transform
+from ..util import invert_spaces
+
+
+class MovingLeastSquares(Transform[np.ndarray]):
+    """Moving least squares transformation.
+
+    Deform based on a matched pairs of source and target control points; see <https://dl.acm.org/doi/10.1145/1141911.1141920>
+    """
+
+    def __init__(
+        self,
+        source_control_points: np.ndarray,
+        target_control_points: np.ndarray,
+        *,
+        spaces: SpaceTuple = (None, None),
+    ):
+        """Non-rigid transforms powered by molesq package.
+
+        Parameters
+        ----------
+        source_control_points : np.ndarray
+            NxD array of control point coordinates in the source space.
+        target_control_points : np.ndarray
+            NxD array of coordinates of the corresponding control points
+            in the target (deformed) space.
+        spaces : tuple[SpaceRef, SpaceRef]
+            Optional source and target spaces
+        """
+        super().__init__(spaces=spaces)
+        self._transformer = _Transformer(
+            np.asarray(source_control_points),
+            np.asarray(target_control_points),
+        )
+        self.ndim = {self._transformer.control_points.shape[1]}
+
+    def apply(self, coords: np.ndarray) -> np.ndarray:
+        coords = self._validate_coords(coords)
+        return self._transformer.transform(coords)
+
+    def is_identity(self) -> bool:
+        xp = array_namespace(self._transformer.control_points)
+        return xp.all(
+            xp.equal(
+                self._transformer.control_points,
+                self._transformer.deformed_control_points,
+            )
+        )
+
+    def invert(self) -> Self | None:
+        return type(self)(
+            self._transformer.deformed_control_points,
+            self._transformer.control_points,
+            spaces=invert_spaces(self.spaces),
+        )

transformnd/transforms/reflection.py

@@ -0,0 +1,207 @@
+from collections.abc import Sequence
+from copy import copy
+from typing import Self
+
+import numpy as np
+from numpy.typing import ArrayLike
+
+from ..base import SpaceTuple, Transform
+from ..util import is_square
+
+
+def proj(u: np.ndarray, v: np.ndarray) -> np.ndarray:
+    return (np.inner(u, v) / np.inner(u, u)) * u
+
+
+def gram_schmidt(vecs: np.ndarray) -> np.ndarray:
+    """
+    https://en.wikipedia.org/wiki/Gram%E2%80%93Schmidt_process
+    """
+    if not is_square(vecs):
+        raise ValueError("Wrong number of dimensions")
+
+    out: list[np.ndarray] = []
+    for v in vecs:
+        b = v.copy()
+
+        for u in out:
+            b -= proj(u, v)
+
+        out.append(b)
+    return np.array(out)
+
+
+def get_hyperplanes(
+    points: np.ndarray, unitise: bool = True, seed: int | None = None
+) -> tuple[np.ndarray, list[np.ndarray]]:
+    """
+    Reflective: point/line/.../hyperplane to be reflected around
+
+    Returns point-normal representation.
+    """
+    points = np.asarray(points)
+    if points.ndim <= 1:
+        points = np.expand_dims(points, -1)
+    elif points.ndim > 2:
+        raise ValueError("Points must be 2D array")
+
+    n_points, ndim = points.shape
+    n_reflections = ndim - n_points + 1
+
+    if n_reflections <= 0:
+        raise ValueError("Too many points given, must be hyperplane or lower-dim")
+
+    point = points[0]
+
+    if n_points == 1:
+        return point, list(np.eye(len(point)))
+
+    # non-orthogonal vectors spanning provided reflective
+    # transpose into row vectors for easier vectorisation
+    reflective_vecs = np.diff(points, axis=0)
+    rng = np.random.default_rng(seed)
+    randoms = rng.random((n_reflections, ndim))
+    non_orth = np.concatenate((reflective_vecs, randoms), 0)
+
+    basis = gram_schmidt(non_orth)
+    extras = basis[-n_reflections:]
+    if unitise:
+        extras /= np.linalg.norm(extras, axis=1)
+    return point, list(extras)
+
+
+def unitise(v: np.ndarray) -> np.ndarray:
+    return v / np.linalg.norm(v)
+
+
+def ensure_tuple(obj: int | Sequence[int]) -> tuple[int, ...]:
+    if isinstance(obj, int):
+        return (obj,)
+    return tuple(obj)
+
+
+class Reflect(Transform[np.ndarray]):
+    """Reflect coordinates about arbitrary planes."""
+
+    def __init__(
+        self,
+        normals: ArrayLike,
+        point: float | ArrayLike = 0.0,
+        *,
+        spaces: SpaceTuple = (None, None),
+    ):
+        """
+        Parameters
+        ----------
+        normals : sequence of arrays
+            Normal vectors to the planes of reflection.
+            Unitised internally.
+        point : float or array-like, optional
+            Intersection point of all reflection planes
+            (can be broadcast from scalar), by default 0 (i.e. the origin)
+        spaces : tuple[SpaceRef, SpaceRef]
+            Optional source and target spaces
+
+        Raises
+        ------
+        ValueError
+            Inconsistent dimensionality
+        """
+        super().__init__(spaces=spaces)
+        normals = np.asarray(normals)
+        if normals.ndim == 1:
+            normals = [normals]
+
+        n1 = normals[0]
+        if (
+            not np.isscalar(point)
+            and isinstance(point, Sequence)
+            and len(n1) != len(point)
+        ):
+            raise ValueError("Point and normals are not of the same dimensionality")
+        self.point: np.ndarray = np.asarray(point, dtype=float)
+        self.ndim = {len(n1)}
+        self.normals = [unitise(n) for n in normals]
+        # todo: matmul is associative, so turn this into an affine in 2/3D?
+
+    def apply(self, coords: np.ndarray) -> np.ndarray:
+        coords = self._validate_coords(coords)
+        out = coords - self.point
+        for n in self.normals:
+            # mul->sum vectorises dot product
+            # normals are unit, avoids unnecessary division by 1
+            out -= 2 * np.sum(coords * n, axis=1) * n
+        out += self.point
+        return out
+
+    @classmethod
+    def from_points(
+        cls,
+        points: ArrayLike,
+        *,
+        spaces: SpaceTuple = (None, None),
+    ):
+        """Infer a single plane of reflection from a minimal number of points on it.
+
+        Parameters
+        ----------
+        points : array-like
+            NxD array of N points in D dimensions. N == D
+        spaces : tuple[SpaceRef, SpaceRef]
+            Optional source and target spaces
+
+        Returns
+        -------
+        Reflection
+        """
+        point, normals = get_hyperplanes(np.asarray(points), unitise=False)
+        return cls(normals, point, spaces=spaces)
+
+    @classmethod
+    def from_axis(
+        cls,
+        axis: int | Sequence[int],
+        origin: ArrayLike,
+        *,
+        spaces: SpaceTuple = (None, None),
+    ):
+        """Reflect around hyperplane(s) parallel with axes.
+
+        Parameters
+        ----------
+        axis : int or sequence of int
+            Index (or indices) of axes in which to reflect.
+        origin : array-like
+            Point around which to reflect.
+        spaces : tuple[SpaceRef, SpaceRef]
+            Optional source and target spaces
+
+        Returns
+        -------
+        Reflection
+
+        Raises
+        ------
+        ValueError
+            Selected axis does not exist.
+        """
+        origin = np.asarray(origin)
+        axis = ensure_tuple(axis)
+
+        for a in axis:
+            if a >= len(axis):
+                raise ValueError(
+                    "Cannot reflect in axis which does not exist (too high)"
+                )
+
+        normals = []
+        for i in range(len(origin) - len(axis) + 1):
+            if i not in axis:
+                v = np.zeros_like(origin)
+                v[i] += 1
+                normals.append(v)
+
+        return cls(normals, origin, spaces=spaces)
+
+    def invert(self) -> Self | None:
+        return copy(self)

transformnd/transforms/simple.py

@@ -0,0 +1,172 @@
+"""
+Simple transformations like rigid translation and scaling.
+"""
+
+from copy import copy
+from typing import Self
+
+import numpy as np
+from numpy.typing import ArrayLike
+
+from array_api_compat import array_namespace
+from array_api_compat import device as xp_device
+from ..base import Transform
+from ..util import ArrayT, chain_or, SpaceTuple, invert_spaces, to_single_ndim
+from ..transforms.affine import Affine
+
+
+class Identity(Transform[ArrayT]):
+    """No-op transformation."""
+
+    def __init__(
+        self,
+        *,
+        spaces: SpaceTuple = (None, None),
+    ):
+        """
+        Transform which does nothing.
+
+        Parameters
+        ----------
+        spaces : tuple[SpaceRef, SpaceRef]
+            Optional source and target spaces
+
+        Raises
+        ------
+        ValueError
+            [description]
+        """
+        src = chain_or(*spaces, default=None)
+        tgt = chain_or(*spaces[::-1], default=None)
+        if src != tgt:
+            raise ValueError("Source and target spaces are different")
+        super().__init__(spaces=(src, src))
+
+    def __invert__(self) -> Transform[ArrayT]:
+        return self
+
+    def to_affine(self, ndim: int | None = None) -> Affine[ArrayT] | None:
+        ndim = to_single_ndim(ndim, self.ndim)
+        m = np.eye(ndim + 1)
+        return Affine(m, spaces=self.spaces)
+
+    def apply(self, coords: ArrayT) -> ArrayT:
+        xp = array_namespace(coords)
+        return xp.asarray(coords)
+
+
+class Translate(Transform[ArrayT]):
+    """Translate coordinates by addition."""
+
+    def __init__(
+        self,
+        translation: ArrayLike,
+        *,
+        spaces: SpaceTuple = (None, None),
+    ):
+        """Simple translation.
+
+        Parameters
+        ----------
+        translation : scalar or D-length array
+            Translation to apply in all dimensions, or each dimension.
+        spaces : tuple[SpaceRef, SpaceRef]
+            Optional source and target spaces
+
+        Raises
+        ------
+        ValueError
+            If the translation is the wrong shape
+        """
+        super().__init__(spaces=spaces)
+        self.translation = np.asarray(translation)
+        if self.translation.ndim > 1:
+            raise ValueError("Translation must be scalar or 1D")
+
+        if self.translation.shape not in [(), (1,)]:
+            self.ndim = {self.translation.shape[0]}
+        # otherwise, can be broadcast to anything
+
+    def to_affine(self, ndim: int | None = None) -> Affine[ArrayT] | None:
+        ndim = to_single_ndim(ndim, self.ndim)
+
+        m = np.eye(ndim + 1)
+        m[:-1, -1] = self.translation
+        return Affine(m, spaces=self.spaces)
+
+    def apply(self, coords: ArrayT) -> ArrayT:
+        coords = self._validate_coords(coords)
+        xp = array_namespace(coords)
+        d = xp_device(coords)
+        return coords + xp.asarray(self.translation, device=d)
+
+    def __invert__(self) -> Transform:
+        return type(self)(-self.translation, spaces=(self.spaces[1], self.spaces[0]))
+
+    def to_device(self, xp, device=None) -> Self:
+        result = copy(self)
+        result.translation = xp.asarray(self.translation, device=device)
+        return result
+
+
+class Scale(Transform[ArrayT]):
+    """Scale coordinates by multiplication."""
+
+    def __init__(
+        self,
+        scale: ArrayLike,
+        *,
+        spaces: SpaceTuple = (None, None),
+    ):
+        """Simple scale transform.
+
+        All points are scaled, i.e. distance from the origin may also change.
+
+        Parameters
+        ----------
+        scale : scalar or D-length array-like
+            Scaling to apply in all dimensions, or each dimension.
+        spaces : tuple[SpaceRef, SpaceRef]
+            Optional source and target spaces
+
+        Raises
+        ------
+        ValueError
+            If scale is the wrong shape.
+        """
+        super().__init__(spaces=spaces)
+        self.scale = np.asarray(scale)
+        if self.scale.ndim > 1:
+            raise ValueError("Scale must be scalar or 1D")
+
+        if self.scale.shape not in [(), (1,)]:
+            self.ndim = {self.scale.shape[0]}
+        # otherwise, can be broadcast to anything
+
+    def to_affine(self, ndim: int | None = None) -> Affine[ArrayT] | None:
+        ndim = to_single_ndim(ndim, self.ndim)
+
+        if np.ndim(self.scale) == 0:
+            scale_vec = np.full(ndim, self.scale)
+        else:
+            scale_vec = self.scale
+        m = np.eye(ndim + 1)
+        m[:-1, :-1] = np.diag(scale_vec)
+        return Affine(m, spaces=self.spaces)
+
+    def apply(self, coords: ArrayT) -> ArrayT:
+        coords = self._validate_coords(coords)
+        xp = array_namespace(coords)
+        d = xp_device(coords)
+        return coords * xp.asarray(self.scale, device=d)
+
+    def invert(self) -> Self | None:
+        return type(self)(
+            1 / self.scale,
+            spaces=invert_spaces(self.spaces),
+        )
+
+    def to_device(self, xp, device=None) -> Self:
+        result = copy(self)
+        result.scale = xp.asarray(self.scale, device=device)
+        return result

transformnd/transforms/thinplate.py

@@ -0,0 +1,83 @@
+"""@public
+Thin plate splines transformations.
+
+Requires the `thinplatesplines` extra.
+"""
+
+import logging
+
+import morphops as mops
+import numpy as np
+
+from ..base import SpaceTuple, Transform
+from ..util import check_ndim, invert_spaces
+
+logger = logging.getLogger(__name__)
+
+
+class ThinPlateSplines(Transform[np.ndarray]):
+    """Thin plate splines transforms.
+
+    Deform based on matched pairs of control points.
+    """
+
+    ndim = {2, 3}
+
+    def __init__(
+        self,
+        source_control_points: np.ndarray,
+        target_control_points: np.ndarray,
+        *,
+        spaces: SpaceTuple = (None, None),
+    ):
+        """Non-rigid control point based transforms in 2/3D.
+
+        Adapted from
+        https://github.com/schlegelp/navis/blob/master/navis/transforms/thinplate.py
+
+        Parameters
+        ----------
+        source_control_points : np.ndarray
+            NxD array of control point coordinates in the source space.
+        target_control_points : np.ndarray
+            NxD array of control point coordinates in the target (deformed) space.
+        spaces : tuple[SpaceRef, SpaceRef]
+            Optional source and target spaces
+
+        Raises
+        ------
+        ValueError
+            Invalid control points.
+        """
+        super().__init__(spaces=spaces)
+        self.source_control_points = np.asarray(source_control_points)
+        self.target_control_points = np.asarray(target_control_points)
+
+        if self.source_control_points.shape != self.target_control_points.shape:
+            raise ValueError("Control point arrays must be the same shape")
+
+        if self.source_control_points.ndim != 2:
+            raise ValueError("Control points array must be 2D")
+
+        ndim = self.source_control_points.shape[1]
+        check_ndim(ndim, self.ndim)
+        self.ndim = {ndim}
+
+        self.W, self.A = mops.tps_coefs(
+            self.source_control_points,
+            self.target_control_points,
+        )
+
+    def invert(self) -> Transform[np.ndarray] | None:
+        return type(self)(
+            self.target_control_points,
+            self.source_control_points,
+            spaces=invert_spaces(self.spaces),
+        )
+
+    def apply(self, coords: np.ndarray) -> np.ndarray:
+        coords = self._validate_coords(coords)
+        U = mops.K_matrix(coords, self.source_control_points)
+        P = mops.P_matrix(coords)
+        # The warped pts are the affine part + the non-uniform part
+        return P @ self.A + U @ self.W