transformnd 0.1.0__py3-none-any.whl

transformnd/__init__.py

@@ -0,0 +1,24 @@
+"""
+.. include:: ../../README.md
+"""
+
+from .base import Transform, TransformSequence, TransformWrapper
+from .util import SpaceRef, TransformSignature, check_ndim
+from . import transforms
+from . import adapters
+from .graph import TransformGraph
+from importlib.metadata import version as _version
+
+__version__ = _version("transformnd")
+
+__all__ = [
+    "Transform",
+    "TransformGraph",
+    "TransformSequence",
+    "TransformWrapper",
+    "TransformSignature",
+    "SpaceRef",
+    "check_ndim",
+    "transforms",
+    "adapters",
+]

transformnd/adapters/__init__.py

@@ -0,0 +1,34 @@
+"""Adapters for transforming objects which are not well-behaved numpy arrays.
+
+Adapter instances are callables which take the transform to be applied,
+the object to apply it to, and optionally some other arguments.
+The adapter knows how to get coordinates out of the object,
+and then create a new object with those transformed coordinates.
+
+Classes which compose over transformable objects can be adapted with the
+`AttrAdapter` class.
+See the `SimpleAdapter` or `FnAdapter` for wrapping simple adapting functions.
+Implement your own adapter by inheriting from `BaseAdapter`.
+
+See `.pandas.DataFrameAdapter` for an example of creating an adapter
+for an external type.
+
+"""
+
+from .base import (
+    AttrAdapter,
+    BaseAdapter,
+    FnAdapter,
+    NullAdapter,
+    ReshapeAdapter,
+    SimpleAdapter,
+)
+
+__all__ = [
+    "BaseAdapter",
+    "SimpleAdapter",
+    "NullAdapter",
+    "FnAdapter",
+    "AttrAdapter",
+    "ReshapeAdapter",
+]

transformnd/adapters/base.py

@@ -0,0 +1,167 @@
+"""Simple adapter cases."""
+
+from abc import ABC, abstractmethod
+from copy import deepcopy
+from functools import partial
+from collections.abc import Callable
+from typing import TypeVar, Any
+
+from array_api_compat import array_namespace
+
+from ..base import Transform, ArrayT
+
+ObjectT = TypeVar("ObjectT")
+
+
+class BaseAdapter[ObjectT, ArrayT](ABC):
+    @abstractmethod
+    def apply(self, transform: Transform[ArrayT], obj: ObjectT) -> ObjectT:
+        """Apply the given transformation to a non-array object.
+
+        Parameters
+        ----------
+        transform : Transform
+        obj : T
+
+        Returns
+        -------
+        T
+        """
+        pass
+
+    def partial(self, *args, **kwargs) -> Callable[..., ObjectT]:
+        """Create a partial function with frozen arguments.
+
+        Useful for applying the same transform to many objects,
+        or many transforms to the same object,
+        or for adapters with additional arguments,
+        using the same config repeatedly.
+
+        Returns
+        -------
+        Callable
+        """
+        return partial(self.apply, *args, **kwargs)
+
+
+class NullAdapter(BaseAdapter[ArrayT, ArrayT]):
+    """Adapter which simply applies the transform."""
+
+    def apply(self, transform: Transform[ArrayT], obj: ArrayT) -> ArrayT:
+        return transform.apply(obj)
+
+
+class FnAdapter(BaseAdapter[ObjectT, ArrayT]):
+    def __init__(self, fn: Callable[[Transform[ArrayT], ObjectT], ObjectT]):
+        """Adapter which simply wraps a function, for typing purposes.
+
+        Parameters
+        ----------
+        fn : Callable[[Transform, T], T]
+            Function which takes the object,
+            and applies the transformation to it.
+        """
+        self.fn = fn
+
+    def apply(self, transform: Transform[ArrayT], obj: ObjectT) -> ObjectT:
+        return self.fn(transform, obj)
+
+
+class AttrAdapter(BaseAdapter[ObjectT, ArrayT]):
+    def __init__(self, **kwargs: BaseAdapter[Any, ArrayT] | None) -> None:
+        """Adapter which transforms an object by applying transforms to its attributes.
+
+        Parameters
+        ----------
+        adapters : Dict[str, Optional[BaseAdapter]]
+            Keys are attribute names, values are adapters with which
+            to apply the transform to those attributes.
+            `None` is shorthand for `NullAdapter()`;
+            i.e. the attribute is an array and can be transformed
+            without being adapted.
+        """
+        self.adapters = {
+            k: NullAdapter[ArrayT]() if v is None else v for k, v in kwargs.items()
+        }
+
+    def apply(
+        self, transform: Transform[ArrayT], obj: ObjectT, in_place: bool = False
+    ) -> ObjectT:
+        """Apply the given transformation to the object, via its attributes.
+
+        Parameters
+        ----------
+        transform : Transform
+        obj : T
+        in_place : bool, optional
+            Whether to mutate the given object in place,
+            by default False (i.e. make a deep copy of it).
+
+        Returns
+        -------
+        T
+        """
+        if not in_place:
+            obj = deepcopy(obj)
+
+        for k, v in self.adapters.items():
+            member = getattr(obj, k)
+            try:
+                transformed = v.apply(transform, member, in_place=True)  # type: ignore
+            except TypeError as e:
+                if "got an unexpected keyword argument 'in_place'" in str(e):
+                    transformed = v.apply(transform, member)
+                else:
+                    raise e
+            setattr(obj, k, transformed)
+
+        return obj
+
+
+class SimpleAdapter(BaseAdapter[ObjectT, ArrayT], ABC):
+    """
+    Helper class for cases with simple conversion methods.
+    """
+
+    @abstractmethod
+    def _to_array(self, obj: ObjectT) -> ArrayT:
+        """Convert the object into an array of coordinates."""
+        pass
+
+    @abstractmethod
+    def _from_array(self, coords: ArrayT) -> ObjectT:
+        """Convert an array of coordinates into the correct type."""
+        pass
+
+    def apply(self, transform: Transform[ArrayT], obj: ObjectT) -> ObjectT:
+        coords: ArrayT = self._to_array(obj)
+        transformed = transform.apply(coords)
+        return self._from_array(transformed)
+
+
+class ReshapeAdapter(BaseAdapter[ArrayT, ArrayT]):
+    """Adapter which reshapes a numpy.ndarray"""
+
+    def __init__(self, dim_axis: int = -1) -> None:
+        """Adapt numpy arrays which are not of the correct shape.
+
+        Parameters
+        ----------
+        dim_axis : int, optional
+            Which axis contains the coordinates' dimensions,
+            by default -1 (last)
+        """
+        self.dim_axis: int = dim_axis
+
+    def apply(self, transform: Transform[ArrayT], obj: ArrayT) -> ArrayT:
+        xp = array_namespace(obj)
+        dim_axis = self.dim_axis
+        if self.dim_axis < 0:
+            dim_axis += xp.ndim(obj)
+
+        moved = xp.moveaxis(obj, dim_axis, -1)
+        m_shape = moved.shape
+
+        flattened = xp.reshape(moved, (-1, m_shape[-1]))
+        transformed = transform.apply(flattened)
+        return xp.moveaxis(xp.reshape(transformed, m_shape), -1, dim_axis)

transformnd/adapters/bounding_box.py

@@ -0,0 +1,16 @@
+from array_api_compat import array_namespace
+from transformnd.base import Transform
+from .base import BaseAdapter
+from ..extents.bounding_box import BoundingBox
+from ..base import ArrayT
+
+
+class BoundingBoxAdapter(BaseAdapter[BoundingBox[ArrayT], ArrayT]):
+    def apply(
+        self, transform: Transform[ArrayT], obj: BoundingBox[ArrayT]
+    ) -> BoundingBox[ArrayT]:
+        xp = array_namespace(obj.mins)
+        stacked = xp.stack([obj.mins, obj.maxes])
+        transformed = transform.apply(stacked)
+        mins, maxes = xp.unstack(transformed)
+        return BoundingBox(mins, maxes)

transformnd/adapters/pandas.py

@@ -0,0 +1,46 @@
+"""Adapt pandas DataFrames for transformation."""
+
+from collections.abc import Hashable
+
+import pandas as pd
+import numpy as np
+
+from ..base import Transform
+from .base import BaseAdapter
+
+
+class PandasAdapter(BaseAdapter[pd.DataFrame, np.ndarray]):
+    def __init__(self, columns: list[Hashable]):
+        """Adapt transformation for coordinates stored in a pandas DataFrame.
+
+        Parameters
+        ----------
+        columns : list of keys
+            Keys for columns containing coordinates, e.g. `["x", "y", "z"]`
+        """
+        self.columns = columns
+
+    def apply(
+        self, transform: Transform, df: pd.DataFrame, in_place: bool = False
+    ) -> pd.DataFrame:
+        """Transform the dataframe, optionally in-place.
+
+        Parameters
+        ----------
+        transform : Transform
+        df : pd.DataFrame
+
+        in_place : bool, optional
+            Whether to mutate the dataframe in place,
+            by default False (i.e. make a copy of it).
+
+        Returns
+        -------
+        pandas.DataFrame
+        """
+        coords = df[self.columns].to_numpy()
+        transformed = transform.apply(coords)
+        if not in_place:
+            df = df.copy()
+        df[self.columns] = transformed
+        return df

transformnd/adapters/polars.py

@@ -0,0 +1,44 @@
+"""Adapt polars DataFrames for transformation."""
+
+import polars as pl
+import numpy as np
+
+from ..base import Transform
+from .base import BaseAdapter
+
+
+class PolarsAdapter(BaseAdapter[pl.DataFrame, np.ndarray]):
+    def __init__(self, columns: list[str]):
+        """Adapt transformation for coordinates stored in a polars DataFrame.
+
+        Parameters
+        ----------
+        columns : list of keys
+            Keys for columns containing coordinates, e.g. `["x", "y", "z"]`
+        """
+        self.columns = columns
+
+    def apply(
+        self, transform: Transform, df: pl.DataFrame, in_place: bool = False
+    ) -> pl.DataFrame:
+        """Transform the dataframe, optionally in-place.
+
+        Parameters
+        ----------
+        transform : Transform
+        df : pl.DataFrame
+
+        in_place : bool, optional
+            Whether to mutate the dataframe in place,
+            by default False (i.e. make a copy of it).
+
+        Returns
+        -------
+        pandas.DataFrame
+        """
+        coords = df[self.columns].to_numpy()
+        transformed = transform.apply(coords)
+        if not in_place:
+            df = df.clone()
+        df[self.columns] = transformed
+        return df

transformnd/adapters/shapely.py

@@ -0,0 +1,124 @@
+import numpy as np
+from typing import Callable
+from shapely.geometry import (
+    LinearRing,
+    LineString,
+    MultiLineString,
+    MultiPoint,
+    MultiPolygon,
+    Point,
+    Polygon,
+    GeometryCollection,
+)
+from shapely.geometry.base import BaseGeometry, BaseMultipartGeometry
+from shapely.coords import CoordinateSequence
+
+from ..base import Transform, ArrayT
+from .base import BaseAdapter
+
+
+def as_numpy(coords: CoordinateSequence) -> np.ndarray:
+    return np.asarray(coords)
+
+
+class GeometryAdapter(BaseAdapter[BaseGeometry, ArrayT]):
+    """Transform shapely geometries.
+
+    As well as the generic `apply()`,
+    there are `apply_*()` methods for transforming different geometry subclasses.
+
+    N.B. some transforms may create invalid topologies
+    (incorrect winding, self-intersections etc.)
+    """
+
+    def __init__(
+        self,
+        array_fn: Callable[[CoordinateSequence], ArrayT] = as_numpy,  # type:ignore
+    ) -> None:
+        self.array_fn = array_fn
+
+    def apply_point(self, transform: Transform, point: Point) -> Point:
+        return Point(*transform.apply(self.array_fn(point.coords))[0])
+
+    def apply_linestring(
+        self, transform: Transform, linestring: LineString
+    ) -> LineString:
+        return LineString(transform.apply(self.array_fn(linestring.coords)))
+
+    def apply_linear_ring(
+        self, transform: Transform, linear_ring: LinearRing
+    ) -> LinearRing:
+        return LinearRing(transform.apply(self.array_fn(linear_ring.coords)))
+
+    def apply_polygon(self, transform: Transform, polygon: Polygon) -> Polygon:
+        return Polygon(
+            self.apply_linear_ring(transform, polygon.exterior),
+            [self.apply_linear_ring(transform, i) for i in polygon.interiors],
+        )
+
+    def apply_multipoint(self, transform: Transform, obj: MultiPoint) -> MultiPoint:
+        return MultiPoint([self.apply_point(transform, o) for o in obj.geoms])
+
+    def apply_multilinestring(
+        self, transform: Transform, obj: MultiLineString
+    ) -> MultiLineString:
+        return MultiLineString([self.apply_linestring(transform, o) for o in obj.geoms])
+
+    def apply_multipolygon(
+        self, transform: Transform, obj: MultiPolygon
+    ) -> MultiPolygon:
+        return MultiPolygon([self.apply_polygon(transform, o) for o in obj.geoms])
+
+    def apply_multipart(
+        self, transform: Transform, obj: BaseMultipartGeometry
+    ) -> BaseMultipartGeometry:
+        """Apply the transform to any shapely multipart geometry."""
+        if isinstance(obj, MultiPoint):
+            return self.apply_multipoint(transform, obj)
+        elif isinstance(obj, MultiLineString):
+            return self.apply_multilinestring(transform, obj)
+        elif isinstance(obj, MultiPolygon):
+            return self.apply_multipolygon(transform, obj)
+        elif isinstance(obj, GeometryCollection):
+            return self.apply_collection(transform, obj)
+        else:
+            raise ValueError(f"Unknown multipart geometry type {type(obj)}")
+
+    def apply_collection(
+        self, transform: Transform, obj: GeometryCollection
+    ) -> GeometryCollection:
+        return GeometryCollection([self.apply(transform, o) for o in obj.geoms])
+
+    def apply(
+        self,
+        transform: Transform,
+        obj: BaseGeometry,
+    ) -> BaseGeometry:
+        """Transform the shapely geometry.
+
+        See the other `apply_*` methods if you already know what type of geometry you're working with;
+        this may be a bit faster.
+
+        Parameters
+        ----------
+        transform : Transform
+        obj : BaseGeometry
+            Some shapely geometry in 2 or 3D
+
+        Returns
+        -------
+        BaseGeometry
+            An object of the same type as the input.
+        """
+        if isinstance(obj, BaseMultipartGeometry):
+            return self.apply_multipart(transform, obj)
+        elif isinstance(obj, Point):
+            return self.apply_point(transform, obj)
+        elif isinstance(obj, LineString):
+            return self.apply_linestring(transform, obj)
+        elif isinstance(obj, LinearRing):
+            return self.apply_linear_ring(transform, obj)
+        elif isinstance(obj, Polygon):
+            return self.apply_polygon(transform, obj)
+        else:
+            raise ValueError(f"Unknown geometry type {type(obj)}")