transformnd.adapters

17class BaseAdapter[ObjectT, ArrayT](ABC):
18    """Base class for adapters that transform non-array objects."""
19
20    @abstractmethod
21    def apply(self, transform: Transform[ArrayT], obj: ObjectT) -> ObjectT:
22        """Apply the given transformation to a non-array object.
23
24        Parameters
25        ----------
26        transform
27            The transformation to apply.
28        obj
29            The object to transform.
30
31        Returns
32        -------
33        ObjectT
34            The transformed object.
35        """
36        pass
37
38    def partial(self, *args: Any, **kwargs: Any) -> Callable[..., ObjectT]:
39        """Create a partial function with frozen arguments.
40
41        Useful for applying the same transform to many objects,
42        or many transforms to the same object,
43        or for adapters with additional arguments,
44        using the same config repeatedly.
45
46        Parameters
47        ----------
48        *args
49            Positional arguments to freeze.
50        **kwargs
51            Keyword arguments to freeze.
52
53        Returns
54        -------
55        Callable[..., ObjectT]
56            A partial function with the given arguments frozen.
57        """
58        return partial(self.apply, *args, **kwargs)

145class SimpleAdapter(BaseAdapter[ObjectT, ArrayT], ABC):
146    """
147    Helper class for cases with simple conversion methods.
148    """
149
150    @abstractmethod
151    def _to_array(self, obj: ObjectT) -> ArrayT:
152        """Convert the object into an array of coordinates."""
153        pass
154
155    @abstractmethod
156    def _from_array(self, coords: ArrayT) -> ObjectT:
157        """Convert an array of coordinates into the correct type."""
158        pass
159
160    def apply(self, transform: Transform[ArrayT], obj: ObjectT) -> ObjectT:
161        coords: ArrayT = self._to_array(obj)
162        transformed = transform.apply(coords)
163        return self._from_array(transformed)

61class NullAdapter(BaseAdapter[ArrayT, ArrayT]):
62    """Adapter which simply applies the transform."""
63
64    def apply(self, transform: Transform[ArrayT], obj: ArrayT) -> ArrayT:
65        return transform.apply(obj)

68class FnAdapter(BaseAdapter[ObjectT, ArrayT]):
69    """Adapter which simply wraps a function, for typing purposes.
70
71    Parameters
72    ----------
73    fn
74        Function which takes the object,
75        and applies the transformation to it.
76    """
77
78    def __init__(self, fn: Callable[[Transform[ArrayT], ObjectT], ObjectT]):
79        self.fn = fn
80
81    def apply(self, transform: Transform[ArrayT], obj: ObjectT) -> ObjectT:
82        return self.fn(transform, obj)

 85class AttrAdapter(BaseAdapter[ObjectT, ArrayT]):
 86    """Adapter which transforms an object by applying transforms to its attributes.
 87
 88    Parameters
 89    ----------
 90    **kwargs
 91        Keys are attribute names, values are adapters with which
 92        to apply the transform to those attributes.
 93        `None` is shorthand for `NullAdapter()`;
 94        i.e. the attribute is an array and can be transformed
 95        without being adapted.
 96    """
 97
 98    def __init__(self, **kwargs: BaseAdapter[Any, ArrayT] | None) -> None:
 99        self.adapters = {
100            k: NullAdapter[ArrayT]() if v is None else v for k, v in kwargs.items()
101        }
102
103    def apply(
104        self, transform: Transform[ArrayT], obj: ObjectT, in_place: bool = False
105    ) -> ObjectT:
106        """Apply the given transformation to the object, via its attributes.
107
108        Parameters
109        ----------
110        transform
111            The transformation to apply.
112        obj
113            The object to transform.
114        in_place
115            Whether to mutate the given object in place,
116            by default False (i.e. make a deep copy of it).
117
118        Returns
119        -------
120        ObjectT
121            The transformed object.
122
123        Raises
124        ------
125        TypeError
126            If the adapter does not support the in_place argument.
127        """
128        if not in_place:
129            obj = deepcopy(obj)
130
131        for k, v in self.adapters.items():
132            member = getattr(obj, k)
133            try:
134                transformed = v.apply(transform, member, in_place=True)  # type: ignore
135            except TypeError as e:
136                if "got an unexpected keyword argument 'in_place'" in str(e):
137                    transformed = v.apply(transform, member)
138                else:
139                    raise e
140            setattr(obj, k, transformed)
141
142        return obj

166class ReshapeAdapter(BaseAdapter[ArrayT, ArrayT]):
167    """Adapter which reshapes a numpy.ndarray"""
168
169    def __init__(self, dim_axis: int = -1) -> None:
170        """Adapt numpy arrays which are not of the correct shape.
171
172        Parameters
173        ----------
174        dim_axis
175            Which axis contains the coordinates' dimensions,
176            by default -1 (last)
177        """
178        self.dim_axis: int = dim_axis
179
180    def apply(self, transform: Transform[ArrayT], obj: ArrayT) -> ArrayT:
181        xp = array_namespace(obj)
182        dim_axis = self.dim_axis
183        if self.dim_axis < 0:
184            dim_axis += xp.ndim(obj)
185
186        moved = xp.moveaxis(obj, dim_axis, -1)
187        m_shape = moved.shape
188
189        flattened = xp.reshape(moved, (-1, m_shape[-1]))
190        transformed = transform.apply(flattened)
191        return xp.moveaxis(xp.reshape(transformed, m_shape), -1, dim_axis)

16class PandasAdapter(BaseAdapter["pd.DataFrame", np.ndarray]):
17    def __init__(self, columns: list[Hashable]):
18        """Adapt transformation for coordinates stored in a pandas DataFrame.
19
20        Parameters
21        ----------
22        columns
23            Keys for columns containing coordinates, e.g. `["x", "y", "z"]`
24        """
25        self.columns = columns
26
27    def apply(
28        self, transform: Transform, df: "pd.DataFrame", in_place: bool = False
29    ) -> "pd.DataFrame":
30        """Transform the dataframe, optionally in-place.
31
32        Parameters
33        ----------
34        transform
35            The transformation to apply.
36        df
37            The DataFrame to transform.
38        in_place
39            Whether to mutate the dataframe in place,
40            by default False (i.e. make a copy of it).
41
42        Returns
43        -------
44        pd.DataFrame
45            The transformed DataFrame.
46        """
47        coords = df[self.columns].to_numpy()
48        transformed = transform.apply(coords)
49        if not in_place:
50            df = df.copy()
51        df[self.columns] = transformed
52        return df

14class PolarsAdapter(BaseAdapter["pl.DataFrame", np.ndarray]):
15    def __init__(self, columns: list[str]):
16        """Adapt transformation for coordinates stored in a polars DataFrame.
17
18        Parameters
19        ----------
20        columns
21            Keys for columns containing coordinates, e.g. `["x", "y", "z"]`
22        """
23        self.columns = columns
24
25    def apply(
26        self, transform: Transform, df: "pl.DataFrame", in_place: bool = False
27    ) -> "pl.DataFrame":
28        """Transform the dataframe, optionally in-place.
29
30        Parameters
31        ----------
32        transform
33            The transformation to apply.
34        df
35            The DataFrame to transform.
36        in_place
37            Whether to mutate the dataframe in place,
38            by default False (i.e. make a copy of it).
39
40        Returns
41        -------
42        pl.DataFrame
43            The transformed DataFrame.
44        """
45        coords = df[self.columns].to_numpy()
46        transformed = transform.apply(coords)
47        if not in_place:
48            df = df.clone()
49        df[self.columns] = transformed
50        return df

22class ShapelyAdapter(BaseAdapter["BaseGeometry", ArrayT]):
23    """Transform shapely geometries.
24
25    As well as the generic `apply()`,
26    there are `apply_*()` methods for transforming different geometry subclasses.
27
28    N.B. some transforms may create invalid topologies
29    (incorrect winding, self-intersections etc.).
30
31    N.B. shapely geometries' coordinates are in `XY(Z)` order
32    """
33
34    def apply[T: "BaseGeometry"](
35        self,
36        transform: Transform,
37        obj: T,
38        *,
39        include_z: bool | None = None,
40    ) -> T:
41        """Transform the shapely geometry.
42
43        Parameters
44        ----------
45        transform
46            The transformation to apply.
47        obj
48            Some shapely geometry in 2 or 3D
49        include_z
50            Force inclusion/ exclusion of Z coordinate.
51            By default (None), checks whether the given geometry has Z coordinates.
52
53        Returns
54        -------
55        T
56            An object of the same type as the input.
57        """
58        import shapely
59
60        def fn(coords: np.ndarray) -> np.ndarray:
61            c = coords.copy()
62            return transform.apply(c)
63
64        if include_z is None:
65            inc_z = bool(shapely.has_z(obj))
66        else:
67            inc_z = include_z
68
69        return shapely.transform(obj, fn, include_z=inc_z)