28class Transform[ArrayT](ABC):
 29    """Base class for transforms."""
 30
 31    def __init__(
 32        self,
 33        ndims: NDims,
 34        *,
 35        spaces: Spaces = Spaces(None, None),
 36    ):
 37        """
 38        Parameters
 39        ----------
 40        ndims
 41            Source and target dimensionality.
 42        spaces
 43            Optional source and target spaces
 44        """
 45        self.ndims: NDims = ndims
 46        self.spaces: Spaces = spaces
 47
 48    def is_identity(self) -> bool:
 49        """Whether this is a no-op transformation."""
 50        return False
 51
 52    def to_affine(self) -> Affine[ArrayT] | None:
 53        """Convert the transform into affine, if conversion is possible.
 54
 55        Returns
 56        -------
 57        Affine[ArrayT] | None
 58            The affine transformation, if conversion is possible.
 59            None otherwise.
 60        """
 61        return None
 62
 63    def _validate_coords(self, coords: ArrayT) -> ArrayT:
 64        """Check that input coordinates are of the correct shape.
 65
 66        Also ensure that coords is a 2Darray.
 67
 68        Parameters
 69        ----------
 70        coords
 71            NxD array of N D-dimensional coordinates.
 72
 73        Returns
 74        -------
 75        ArrayT
 76            The validated coordinates.
 77
 78        Raises
 79        ------
 80        ValueError
 81            If dimensions are not supported.
 82        """
 83        xp = array_namespace(coords)
 84        if xp.ndim(coords) != 2:
 85            raise ValueError("Coords must be a 2D array")
 86        dim = xp.shape(coords)[1]
 87        if xp.shape(coords)[1] != self.ndims.source:
 88            raise ValueError(
 89                f"Coords must have dimensionality {self.ndims.source}, got {dim}"
 90            )
 91        return coords
 92
 93    @abstractmethod
 94    def apply(self, coords: ArrayT) -> ArrayT:
 95        """Apply transformation.
 96
 97        Parameters
 98        ----------
 99        coords
100            NxD array of N D-dimensional coordinates.
101
102        Returns
103        -------
104        ArrayT
105            Transformed coordinates in the same shape.
106        """
107        pass
108
109    def invert(self) -> Transform | None:
110        """Invert the transformation, returning `None` if not possible."""
111        return None
112
113    def __invert__(self) -> Transform:
114        """Invert transformation if possible.
115
116        Returns `NotImplemented` otherwise (will raise `NotImplementedError`).
117
118        Returns
119        -------
120        Transform
121            Inverted transformation.
122        """
123        t = self.invert()
124        if t is None:
125            return NotImplemented
126        return t
127
128    def to_device(self, xp: ModuleType, device: str | None = None) -> Self:  # noqa: ARG002
129        """Return a copy of this transform with array parameters placed on the given device.
130
131        Useful for pre-allocating parameters on GPU before a tight apply() loop,
132        avoiding per-call host-to-device transfers.
133
134        Parameters
135        ----------
136        xp
137            The target array namespace (e.g. jax.numpy, torch).
138        device
139            Target device (e.g. from array_api_compat.device(array)).
140            If None, uses xp's default device.
141
142        Returns
143        -------
144        Self
145            A new transform instance with parameters on the target device,
146            or NotImplemented if the subclass does not support device placement.
147        """
148        return NotImplemented
149
150    def __or__(self, other: Transform[ArrayT]) -> TransformSequence[ArrayT]:
151        """Compose transformations into a sequence.
152
153        If other is a TransformSequence, prepend this transform to the others.
154
155        Parameters
156        ----------
157        other
158            The transform to compose with.
159
160        Returns
161        -------
162        TransformSequence[ArrayT]
163            The composed transform sequence.
164        """
165        if not isinstance(other, Transform):
166            return NotImplemented
167        transforms = as_transform_list(self) + as_transform_list(other)
168        return TransformSequence[ArrayT](
169            transforms,
170            spaces=Spaces(self.spaces.source, other.spaces.target),
171        )
172
173    def __ror__(self, other: Transform[ArrayT]) -> TransformSequence[ArrayT]:
174        """Compose transformations into a sequence.
175
176        If other is a TransformSequence, append this transform to the others.
177
178        Parameters
179        ----------
180        other
181            The transform to compose with.
182
183        Returns
184        -------
185        TransformSequence[ArrayT]
186            The composed transform sequence.
187        """
188        if not isinstance(other, Transform):
189            return NotImplemented
190        transforms = as_transform_list(other) + as_transform_list(self)
191        return TransformSequence(
192            transforms,
193            spaces=Spaces(other.spaces.source, self.spaces.target),
194        )
195
196    def __str__(self) -> str:
197        cls_name = type(self).__name__
198        src = space_str(self.spaces.source)
199        tgt = space_str(self.spaces.target)
200        return f"{cls_name}[{src}->{tgt}]"

 63class TransformGraph[ArrayT]:
 64    """Transform between any number of arbitrary spaces/ coordinate systems.
 65
 66    Finds the shortest path for transforming one space
 67    into another, via some intermediate spaces.
 68
 69    Populate with `my_transform_graph.add_transforms(my_transforms)`.
 70    """
 71
 72    def __init__(
 73        self,
 74    ):
 75        """Create an transform graph, optionally with some starting transforms.
 76
 77        See the `TransformGraph.add_transforms` documentation for restrictions on the
 78        given transforms.
 79        """
 80        self.graph = nx.MultiDiGraph()
 81        self.space_ndims: dict[SpaceRef, int] = dict()
 82
 83    def _update_spaces(
 84        self,
 85        transform: Transform[ArrayT],
 86        source: SpaceRef | None,
 87        target: SpaceRef | None,
 88    ) -> Spaces:
 89        """Check that the transform's spaces do not conflict with those given explicitly,
 90        that the source and target space is defined somewhere,
 91        and that the dimensionality of the spaces (inferred from the transforms)
 92        does not conflict with known spaces.
 93        """
 94        # check explicit spaces do not conflict with transform's spaces
 95        src = same_or_none(transform.spaces.source, source)
 96        tgt = same_or_none(transform.spaces.target, target)
 97
 98        # if the node already exists, make sure the dimensionality does not conflict
 99        self.space_ndims[src] = same_or_none(
100            self.space_ndims.get(src), transform.ndims.source
101        )
102        self.space_ndims[tgt] = same_or_none(
103            self.space_ndims.get(tgt), transform.ndims.target
104        )
105        return Spaces(src, tgt)
106
107    def _add_transform(
108        self,
109        transform: Transform[ArrayT],
110        source: SpaceRef | None,
111        target: SpaceRef | None,
112        edge_data: dict[str, Any] | None,
113    ) -> list[tuple[SpaceRef, SpaceRef]]:
114        """Clearing the get_sequence cache and splitting sequences and bijections should be handled outside this method."""
115        out = []
116
117        src, tgt = self._update_spaces(transform, source, target)
118
119        if edge_data is None:
120            edge_data = dict()
121
122        if TRANSFORM_KEY in edge_data:
123            raise ValueError(f"Must not use the key '{TRANSFORM_KEY}' in edge_data")
124
125        d = {TRANSFORM_KEY: transform, **edge_data}
126        self.graph.add_edge(src, tgt, **d)
127        out.append((src, tgt))
128        return out
129
130    def add_transform(
131        self,
132        transform: Transform[ArrayT],
133        source: SpaceRef | None = None,
134        target: SpaceRef | None = None,
135        *,
136        edge_data: dict[str, Any] | None = None,
137    ) -> list[tuple[SpaceRef, SpaceRef]]:
138        """Add a transform to the graph.
139
140        If the given transform is a `Bijection`,
141        only the forward component will be added as an independent edges.
142
143        This method will NOT overwrite intermediate edges.
144
145        N.B. Previously this method implicitly added inverse edges where possible.
146        Now these edges must be added explicitly by calling `add_transform(~transform)`.
147        Additionally, previously `TransformSequence`s would be split out into multiple edges
148        if any intermediate spaces were defined;
149        now these edges must be added explicitly with the `TransformSequence.split` method.
150
151        Parameters
152        ----------
153        transform
154            Transform to add to the graph as an edge.
155        source
156            May be omitted if `transform` has its source space defined.
157        target
158            May be omitted if `transform` has its target space defined.
159        edge_data
160            Dict of string keys to arbitrary values to associate with an edge.
161            Used during path-finding.
162            Must not have the `"_transform"` key.
163
164        Returns
165        -------
166        list[tuple[SpaceRef, SpaceRef]]
167            List of `(src, tgt)` edges added to the graph.
168        """
169        out: list[tuple[SpaceRef, SpaceRef]] = []
170        if isinstance(transform, Bijection):
171            out.extend(
172                self.add_transform(
173                    transform.forward,
174                    source,
175                    target,
176                    edge_data=edge_data,
177                )
178            )
179
180        else:
181            out.extend(self._add_transform(transform, source, target, edge_data))
182
183        if out:
184            self.get_sequence.cache_clear()
185
186        return out
187
188    @lru_cache()
189    def get_sequence(
190        self,
191        source_space: SpaceRef,
192        target_space: SpaceRef,
193        full: bool = False,
194        *,
195        weight: None | str | WeightFn = None,
196    ) -> TransformSequence[ArrayT]:
197        """Get the shortest TransformSequence for transforming between two spaces.
198
199        Parameters
200        ----------
201        source_space
202            The source coordinate space.
203        target_space
204            The target coordinate space.
205        full
206            By default, simplifies consecutive affines and drops bijections' inverse form.
207            If `full` is True, keeps each transformation as-is.
208        weight
209            str key in the `edge_data` dict given when an edge was added,
210            or a function to determine a weight from the args `src_space, tgt_space, edge_data`,
211            or None (all weights are 1).
212
213        Returns
214        -------
215        TransformSequence[ArrayT]
216            The shortest transform sequence between the spaces.
217        """
218        path = nx.shortest_path(self.graph, source_space, target_space, weight)  # type:ignore
219        transforms = []
220        wfn = normalise_edge_weight_fn(weight)
221
222        for src, tgt in pairwise(path):
223            edges = self.graph[src][tgt]
224            transforms.append(
225                min(edges.values(), key=lambda d: wfn(src, tgt, d))[TRANSFORM_KEY]
226            )
227
228        seq = TransformSequence(
229            transforms,
230            spaces=Spaces(source_space, target_space),
231        )
232        if not full:
233            seq = seq.simplify(drop_inverse=True)
234        return seq
235
236    def transform(
237        self,
238        source_space: SpaceRef,
239        target_space: SpaceRef,
240        coords: ArrayT,
241        *,
242        weight: None | str | WeightFn = None,
243    ) -> ArrayT:
244        """Transform coordinates from one space to another,
245        possibly via intermediates.
246
247        Parameters
248        ----------
249        source_space
250            The source coordinate space.
251        target_space
252            The target coordinate space.
253        coords
254            The coordinates to transform.
255        weight
256            str key in the `edge_data` dict given when an edge was added,
257            or a function to determine a weight from the args `src_space, tgt_space, edge_data`,
258            or None (all weights are 1).
259
260
261        Returns
262        -------
263        ArrayT
264            The transformed coordinates.
265        """
266        t = self.get_sequence(source_space, target_space, weight=weight)
267        return t.apply(coords)
268
269    def __iter__(self) -> Iterator[Transform[ArrayT]]:
270        """Iterate through the transforms present in the graph.
271
272        Includes inferred reverse transforms.
273
274        N.B. the `__iter__` method of some popular graph libraries like networkx
275        iterate through nodes, where this effectively iterates through edges.
276
277        Yields
278        ------
279        Transform[ArrayT]
280            The next transform in the graph.
281
282        Examples
283        --------
284        Create a new transform graph using another
285
286        >>> new_tgraph = TransformGraph([extra_transform, *old_tgraph])
287
288        """
289        for _, _, t in self.graph.edges.data(TRANSFORM_KEY):
290            yield t
291
292    def to_device(
293        self, xp: ModuleType, device: str | None = None
294    ) -> TransformGraph[ArrayT]:
295        result: TransformGraph[ArrayT] = TransformGraph()
296        for src, tgt, t in self.graph.edges.data(TRANSFORM_KEY):
297            result.graph.add_edge(src, tgt, transform=t.to_device(xp, device))
298        return result

276class TransformSequence(Transform[ArrayT], Sequence[Transform[ArrayT]]):
277    """Chain transforms, applying one after another."""
278
279    def __init__(
280        self,
281        transforms: Sequence[Transform[ArrayT]],
282        *,
283        spaces: Spaces = Spaces(None, None),
284    ) -> None:
285        """Combine transforms by chaining them.
286
287        Also checks for consistent dimensionality and space references,
288        inferring if None.
289
290        Parameters
291        ----------
292        transforms :
293            Items which are a TransformSequences
294            will each still be treated as a single transform.
295        spaces :
296            Optional source and target spaces.
297            Can also be inferred from the first and last transforms.
298
299        Raises
300        ------
301        ValueError
302            If spaces are incompatible.
303        """
304        ts = infer_spaces(transforms, *spaces)
305        if not ts:
306            raise ValueError("Empty transform sequence")
307
308        for idx, (t1, t2) in enumerate(pairwise(ts)):
309            if t1.ndims.target != t2.ndims.source:
310                raise ValueError(
311                    "Incompatible dimensionality: "
312                    f"transform {idx}'s target is {t1.ndims.target}D "
313                    f"and the next source is {t2.ndims.source}D"
314                )
315
316        spaces = Spaces(ts[0].spaces.source, ts[-1].spaces.target)
317        ndims = NDims(ts[0].ndims.source, ts[-1].ndims.target)
318
319        super().__init__(
320            ndims,
321            spaces=spaces,
322        )
323
324        self.transforms: list[Transform[ArrayT]] = ts
325
326    def __iter__(self) -> Iterator[Transform[ArrayT]]:
327        """Iterate through component transforms.
328
329        Yields
330        -------
331        Transform
332        """
333        yield from self.transforms
334
335    def __len__(self) -> int:
336        """Number of transforms.
337
338        Returns
339        -------
340        int
341        """
342        return len(self.transforms)
343
344    def invert(self) -> Transform[ArrayT] | None:
345        try:
346            transforms = [~t for t in reversed(self.transforms)]
347        except NotImplementedError:
348            return None
349        return type(self)(
350            transforms,
351            spaces=self.spaces.invert(),
352        )
353
354    def apply(self, coords: ArrayT) -> ArrayT:
355        for t in self.transforms:
356            coords = t.apply(coords)
357        return coords
358
359    def to_device(self, xp: ModuleType, device: str | None = None) -> Self:
360        result = copy(self)
361        result.transforms = [t.to_device(xp, device) for t in self.transforms]
362        return result
363
364    def list_spaces(self, skip_none: bool = False) -> list[SpaceRef]:
365        """List spaces in this transform.
366
367        Parameters
368        ----------
369        skip_none
370            Whether to skip undefined spaces, default False.
371
372        Returns
373        -------
374        list[SpaceRef]
375            The list of spaces.
376        """
377        spaces = [self.spaces.source] + [t.spaces.target for t in self.transforms]
378        if skip_none:
379            spaces = [s for s in spaces if s is not None]
380        return spaces
381
382    def split(self) -> Iterator[Transform[ArrayT]]:
383        """Split the sequence where an intermediate space is known."""
384        this_seq = []
385
386        for t in self.transforms:
387            if t.spaces.source is not None and t.spaces.target is not None:
388                yield t
389                continue
390
391            this_seq.append(t)
392            if t.spaces.target is not None:
393                yield type(self)(this_seq)
394                this_seq = []
395
396    def __str__(self) -> str:
397        cls_name = type(self).__name__
398        spaces_str = "->".join(space_str(s) for s in self.list_spaces())
399        return f"{cls_name}[{spaces_str}]"
400
401    def __getitem__(self, idx: slice | int):
402        if isinstance(idx, int):
403            return self.transforms[idx]
404        return type(self)(self.transforms[idx])
405
406    def is_identity(self) -> bool:
407        return all(t.is_identity() for t in self)
408
409    def flatten(self, drop_inverse: bool = True) -> Self:
410        """Flatten nested sequences."""
411        from .transforms.bijection import Bijection
412
413        out: list[Transform[ArrayT]] = []
414
415        for t in self.transforms:
416            if drop_inverse and isinstance(t, Bijection):
417                t = t.forward
418            if isinstance(t, TransformSequence):
419                out.extend(t.flatten())
420            else:
421                out.append(t)
422        return TransformSequence(out, spaces=self.spaces)  # type:ignore
423
424    def simplify(self, drop_inverse: bool = True):
425        """Reduce the number of transformations in this sequence if possible.
426
427        - Compose consecutive transformations which can be expressed as affines
428        - Drop trivial transforms (e.g. identity)
429        - Optionally drop explicit inverse transforms
430          (e.g. replace `Bijection`s with their `forward` transform)
431
432        Also drops all internal space tuples; only the sequence's remains.
433
434        Does not check whether transforms invert each other,
435        e.g. `Translation(1) | Translation(-1)`.
436        """
437        from .transforms import Identity
438
439        out: list[Transform[ArrayT]] = []
440        affine = None
441        for t in self.flatten(drop_inverse):
442            if t.is_identity():
443                continue
444
445            new_affine = t.to_affine()
446
447            if new_affine is None:
448                if affine is not None:
449                    add_to_output(affine, out)
450                    affine = None
451                add_to_output(t, out)
452                continue
453
454            if affine is None:
455                affine = new_affine
456            else:
457                affine = new_affine @ affine  # type: ignore[operator]
458
459        if affine is not None:
460            add_to_output(affine, out)
461
462        if not out:
463            out.append(Identity(self.ndims.source))
464
465        return type(self)(out, spaces=self.spaces)
466
467    def to_affine(self) -> Affine[ArrayT] | None:
468        simple = self.simplify(True)
469        if len(simple) != 1:
470            return None
471        return simple[0].to_affine()

203class TransformWrapper(Transform[ArrayT]):
204    """Wrapper around an arbitrary function which transforms coordinates."""
205
206    def __init__(
207        self,
208        fn: TransformSignature[ArrayT],
209        in_ndim: int,
210        out_ndim: int,
211        *,
212        spaces: Spaces = Spaces(None, None),
213    ):
214        """Wrapper around an arbitrary function.
215
216        `fn` should take and return an identically-shaped
217        NxD numpy array of N D-dimensional coordinates.
218
219        Parameters
220        ----------
221        fn
222            Callable.
223        in_ndim
224            Dimensionality of the input coordinates.
225        out_ndim
226            Dimensionality of the output coordinates.
227        spaces
228            Optional source and target spaces
229        """
230        super().__init__(NDims(in_ndim, out_ndim), spaces=spaces)
231        self.fn = fn
232
233    def apply(self, coords: ArrayT) -> ArrayT:
234        self._validate_coords(coords)
235        return self.fn(coords)