transformnd 0.2.1__tar.gz → 0.4.0__tar.gz

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (31) hide show
  1. {transformnd-0.2.1 → transformnd-0.4.0}/PKG-INFO +22 -13
  2. {transformnd-0.2.1 → transformnd-0.4.0}/README.md +17 -12
  3. {transformnd-0.2.1 → transformnd-0.4.0}/pyproject.toml +22 -2
  4. {transformnd-0.2.1 → transformnd-0.4.0}/src/transformnd/adapters/base.py +52 -29
  5. {transformnd-0.2.1 → transformnd-0.4.0}/src/transformnd/adapters/pandas.py +8 -6
  6. {transformnd-0.2.1 → transformnd-0.4.0}/src/transformnd/adapters/polars.py +8 -6
  7. transformnd-0.4.0/src/transformnd/adapters/shapely.py +64 -0
  8. {transformnd-0.2.1 → transformnd-0.4.0}/src/transformnd/base.py +35 -18
  9. transformnd-0.4.0/src/transformnd/graph.py +376 -0
  10. {transformnd-0.2.1 → transformnd-0.4.0}/src/transformnd/transforms/affine.py +75 -72
  11. {transformnd-0.2.1 → transformnd-0.4.0}/src/transformnd/transforms/bijection.py +18 -2
  12. {transformnd-0.2.1 → transformnd-0.4.0}/src/transformnd/transforms/by_dimension.py +8 -3
  13. {transformnd-0.2.1 → transformnd-0.4.0}/src/transformnd/transforms/map_axis.py +7 -2
  14. {transformnd-0.2.1 → transformnd-0.4.0}/src/transformnd/transforms/moving_least_squares.py +3 -3
  15. {transformnd-0.2.1 → transformnd-0.4.0}/src/transformnd/transforms/reflection.py +12 -12
  16. {transformnd-0.2.1 → transformnd-0.4.0}/src/transformnd/transforms/simple.py +7 -6
  17. {transformnd-0.2.1 → transformnd-0.4.0}/src/transformnd/transforms/thinplate.py +3 -3
  18. transformnd-0.4.0/src/transformnd/transforms/vector_field.py +229 -0
  19. {transformnd-0.2.1 → transformnd-0.4.0}/src/transformnd/util.py +89 -25
  20. transformnd-0.2.1/src/transformnd/adapters/shapely.py +0 -124
  21. transformnd-0.2.1/src/transformnd/graph.py +0 -251
  22. {transformnd-0.2.1 → transformnd-0.4.0}/src/transformnd/__init__.py +0 -0
  23. {transformnd-0.2.1 → transformnd-0.4.0}/src/transformnd/adapters/__init__.py +0 -0
  24. {transformnd-0.2.1 → transformnd-0.4.0}/src/transformnd/adapters/bounding_box.py +0 -0
  25. {transformnd-0.2.1 → transformnd-0.4.0}/src/transformnd/constants.py +0 -0
  26. {transformnd-0.2.1 → transformnd-0.4.0}/src/transformnd/extents/__init__.py +0 -0
  27. {transformnd-0.2.1 → transformnd-0.4.0}/src/transformnd/extents/base.py +0 -0
  28. {transformnd-0.2.1 → transformnd-0.4.0}/src/transformnd/extents/bounding_box.py +0 -0
  29. {transformnd-0.2.1 → transformnd-0.4.0}/src/transformnd/py.typed +0 -0
  30. {transformnd-0.2.1 → transformnd-0.4.0}/src/transformnd/transforms/__init__.py +0 -0
  31. {transformnd-0.2.1 → transformnd-0.4.0}/src/transformnd/types.py +0 -0
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.4
2
2
  Name: transformnd
3
- Version: 0.2.1
3
+ Version: 0.4.0
4
4
  Summary: ND coordinate transformations
5
5
  Author: Chris Barnes
6
6
  Author-email: Chris Barnes <chris.barnes@gerbi-gmb.de>
@@ -20,11 +20,14 @@ Requires-Dist: numpy>=2
20
20
  Requires-Dist: networkx>=3
21
21
  Requires-Dist: array-api-compat>=1.14
22
22
  Requires-Dist: typing-extensions>=4.15.0
23
+ Requires-Dist: dask[dask]>=2026.3.0
24
+ Requires-Dist: dask-image[dask]>=2026.5.0
23
25
  Requires-Dist: molesq>=0.4.0 ; extra == 'movingleastsquares'
24
26
  Requires-Dist: pandas>=3.0.2 ; extra == 'pandas'
25
27
  Requires-Dist: polars>=1.40.1 ; extra == 'polars'
26
28
  Requires-Dist: shapely>=2.1.2 ; extra == 'shapely'
27
29
  Requires-Dist: morphops>=0.1.13 ; extra == 'thinplatesplines'
30
+ Requires-Dist: scipy>=1.17.1 ; extra == 'vectorfield'
28
31
  Requires-Python: >=3.12, <4.0
29
32
  Project-URL: documentation, https://transformnd.readthedocs.io/en/latest/
30
33
  Project-URL: source, https://github.com/clbarnes/transformnd
@@ -34,6 +37,7 @@ Provides-Extra: pandas
34
37
  Provides-Extra: polars
35
38
  Provides-Extra: shapely
36
39
  Provides-Extra: thinplatesplines
40
+ Provides-Extra: vectorfield
37
41
  Description-Content-Type: text/markdown
38
42
 
39
43
  # transformnd
@@ -71,18 +75,21 @@ Open it with `uv run --group tutorial marimo edit examples/tutorial.py`.
71
75
 
72
76
  ## Implemented transforms
73
77
 
74
- - Identity (`transformnd.transforms.Identity`)
75
- - Translation (`transformnd.transforms.Translate`)
76
- - Scale (`transformnd.transforms.Scale`)
77
- - Reflection (`transformnd.transforms.Reflect`)
78
- - Affine (`transformnd.transforms.Affine`)
79
- - Can be composed efficiently with `@` operator; the right hand operand is effectively applied first
80
- - MapAxis (`transformnd.transforms.MapAxis`): permute coordinate axes
81
- - ByDimension (`transformnd.transforms.ByDimension`): apply transformations to subsets of coordinate axes
82
- - Moving Least Squares, affine (`transformnd.transforms.moving_least_squares.MovingLeastSquares`)
83
- - uses `movingleastsquares` extra
84
- - Thin Plate Splines (`transformnd.transforms.thinplate.ThinPlateSplines`)
85
- - uses `thinplatesplines` extra
78
+ All transforms are accessed under the `transformnd.transforms` subpackage.
79
+
80
+ | Transform | Extra | Description |
81
+ | --------- | ----- | ----------- |
82
+ | `Identity` | | No-op transformation |
83
+ | `Translation` | | Add a constant translation to the input coordinates |
84
+ | `Scale` | | Multiply the input coordinates by constant scale factor |
85
+ | `Reflection` | | Reflect coordinates about arbitrary planes |
86
+ | `MapAxis` | | Rearrange axes of the input coordinates |
87
+ | `Affine` | | Multiply augmented coordinates by an affine transformation matrix. Can represent all of the above transformations. Can be composed with matrix multiplication `aff2 @ aff1`. |
88
+ | `ByDimension` | | Apply different transformations to subsets of the input coordinates' dimensions |
89
+ | `moving_least_squares.MovingLeastSquares` | `movingleastsquares` | Landmark-based transformation. |
90
+ | `thin_plate_splines.ThinPlateSplines` | `thinplatesplines` | Landmark-based transformation. |
91
+ | `vector_field.Coordinates` | `vectorfield` | Look up output coordinates in a vector field indexed by the input coordinates |
92
+ | `vector_field.Displacements` | `vectorfield` | Look up translations in a vector field indexed by the input coordinates, and add them to input coordinates |
86
93
 
87
94
  Arbitrary transforms can be composed into a `TransformSequence` with `transform1 | transform2`.
88
95
  A graph of transforms between defined spaces can be traversed using the `TransformGraph`.
@@ -135,12 +142,14 @@ Methods which SHOULD be implemented if applicable:
135
142
  - `just` to list commands.
136
143
  - Docs are generated with [`pdoc`](https://pdoc.dev/) (use `just doc`) and hosted on ReadTheDocs
137
144
  - `just bump` bumps the version, commits, and tags (but does not push); depends on [`schpet/changelog`](https://github.com/schpet/changelog)
145
+ - `just repl` starts an IPython shell with all dependencies installed
138
146
 
139
147
  ## Thanks
140
148
 
141
149
  Thanks to contributors
142
150
 
143
151
  - [Francesca Drummer](https://github.com/FrancescaDr)
152
+ - [Johannes Soltwedel](https://github.com/jo-mueller)
144
153
  - [Lorenzo Cerrone](https://github.com/lorenzocerrone)
145
154
  - [Maks Hess](https://github.com/MaksHess)
146
155
  - [Silvia Maria Macrì](https://github.com/SilviaMariaMacri)
@@ -33,18 +33,21 @@ Open it with `uv run --group tutorial marimo edit examples/tutorial.py`.
33
33
 
34
34
  ## Implemented transforms
35
35
 
36
- - Identity (`transformnd.transforms.Identity`)
37
- - Translation (`transformnd.transforms.Translate`)
38
- - Scale (`transformnd.transforms.Scale`)
39
- - Reflection (`transformnd.transforms.Reflect`)
40
- - Affine (`transformnd.transforms.Affine`)
41
- - Can be composed efficiently with `@` operator; the right hand operand is effectively applied first
42
- - MapAxis (`transformnd.transforms.MapAxis`): permute coordinate axes
43
- - ByDimension (`transformnd.transforms.ByDimension`): apply transformations to subsets of coordinate axes
44
- - Moving Least Squares, affine (`transformnd.transforms.moving_least_squares.MovingLeastSquares`)
45
- - uses `movingleastsquares` extra
46
- - Thin Plate Splines (`transformnd.transforms.thinplate.ThinPlateSplines`)
47
- - uses `thinplatesplines` extra
36
+ All transforms are accessed under the `transformnd.transforms` subpackage.
37
+
38
+ | Transform | Extra | Description |
39
+ | --------- | ----- | ----------- |
40
+ | `Identity` | | No-op transformation |
41
+ | `Translation` | | Add a constant translation to the input coordinates |
42
+ | `Scale` | | Multiply the input coordinates by constant scale factor |
43
+ | `Reflection` | | Reflect coordinates about arbitrary planes |
44
+ | `MapAxis` | | Rearrange axes of the input coordinates |
45
+ | `Affine` | | Multiply augmented coordinates by an affine transformation matrix. Can represent all of the above transformations. Can be composed with matrix multiplication `aff2 @ aff1`. |
46
+ | `ByDimension` | | Apply different transformations to subsets of the input coordinates' dimensions |
47
+ | `moving_least_squares.MovingLeastSquares` | `movingleastsquares` | Landmark-based transformation. |
48
+ | `thin_plate_splines.ThinPlateSplines` | `thinplatesplines` | Landmark-based transformation. |
49
+ | `vector_field.Coordinates` | `vectorfield` | Look up output coordinates in a vector field indexed by the input coordinates |
50
+ | `vector_field.Displacements` | `vectorfield` | Look up translations in a vector field indexed by the input coordinates, and add them to input coordinates |
48
51
 
49
52
  Arbitrary transforms can be composed into a `TransformSequence` with `transform1 | transform2`.
50
53
  A graph of transforms between defined spaces can be traversed using the `TransformGraph`.
@@ -97,12 +100,14 @@ Methods which SHOULD be implemented if applicable:
97
100
  - `just` to list commands.
98
101
  - Docs are generated with [`pdoc`](https://pdoc.dev/) (use `just doc`) and hosted on ReadTheDocs
99
102
  - `just bump` bumps the version, commits, and tags (but does not push); depends on [`schpet/changelog`](https://github.com/schpet/changelog)
103
+ - `just repl` starts an IPython shell with all dependencies installed
100
104
 
101
105
  ## Thanks
102
106
 
103
107
  Thanks to contributors
104
108
 
105
109
  - [Francesca Drummer](https://github.com/FrancescaDr)
110
+ - [Johannes Soltwedel](https://github.com/jo-mueller)
106
111
  - [Lorenzo Cerrone](https://github.com/lorenzocerrone)
107
112
  - [Maks Hess](https://github.com/MaksHess)
108
113
  - [Silvia Maria Macrì](https://github.com/SilviaMariaMacri)
@@ -1,6 +1,6 @@
1
1
  [project]
2
2
  name = "transformnd"
3
- version = "0.2.1"
3
+ version = "0.4.0"
4
4
  description = "ND coordinate transformations"
5
5
  readme = "README.md"
6
6
  authors = [{ name = "Chris Barnes", email = "chris.barnes@gerbi-gmb.de" }]
@@ -10,6 +10,8 @@ dependencies = [
10
10
  "networkx>=3",
11
11
  "array_api_compat>=1.14",
12
12
  "typing-extensions>=4.15.0",
13
+ "dask[dask]>=2026.3.0",
14
+ "dask-image[dask]>=2026.5.0",
13
15
  ]
14
16
  license = "MIT"
15
17
  classifiers = [
@@ -41,6 +43,9 @@ shapely = ["shapely>=2.1.2"]
41
43
  polars = [
42
44
  "polars>=1.40.1",
43
45
  ]
46
+ vectorfield = [
47
+ "scipy>=1.17.1",
48
+ ]
44
49
 
45
50
  [dependency-groups]
46
51
  dev = [
@@ -50,12 +55,17 @@ dev = [
50
55
  {include-group = "tutorial"},
51
56
  {include-group = "bench"},
52
57
  ]
53
- test = ["pytest", "jax"]
58
+ test = [
59
+ "pytest",
60
+ "jax",
61
+ "dask>=2026.3.0",
62
+ ]
54
63
  lint = [
55
64
  "ruff",
56
65
  "mypy",
57
66
  "prek>=0.3.9",
58
67
  "types-shapely>=2.1.0.20260408",
68
+ "pydoclint>=0.8.6",
59
69
  ]
60
70
  doc = [
61
71
  "pdoc>=16.0.0",
@@ -80,3 +90,13 @@ check_untyped_defs = true
80
90
  [tool.pytest]
81
91
  testpaths = ["tests", "bench"]
82
92
  addopts = ["--benchmark-skip"]
93
+
94
+ [tool.pydoclint]
95
+ arg-type-hints-in-docstring = false
96
+ allow-init-docstring = true
97
+
98
+ [tool.ruff]
99
+ lint.external = [
100
+ "DOC", # pydoclint
101
+ "D", # pydocstyle
102
+ ]
@@ -14,22 +14,27 @@ ObjectT = TypeVar("ObjectT")
14
14
 
15
15
 
16
16
  class BaseAdapter[ObjectT, ArrayT](ABC):
17
+ """Base class for adapters that transform non-array objects."""
18
+
17
19
  @abstractmethod
18
20
  def apply(self, transform: Transform[ArrayT], obj: ObjectT) -> ObjectT:
19
21
  """Apply the given transformation to a non-array object.
20
22
 
21
23
  Parameters
22
24
  ----------
23
- transform : Transform
24
- obj : T
25
+ transform
26
+ The transformation to apply.
27
+ obj
28
+ The object to transform.
25
29
 
26
30
  Returns
27
31
  -------
28
- T
32
+ ObjectT
33
+ The transformed object.
29
34
  """
30
35
  pass
31
36
 
32
- def partial(self, *args, **kwargs) -> Callable[..., ObjectT]:
37
+ def partial(self, *args: Any, **kwargs: Any) -> Callable[..., ObjectT]:
33
38
  """Create a partial function with frozen arguments.
34
39
 
35
40
  Useful for applying the same transform to many objects,
@@ -37,9 +42,17 @@ class BaseAdapter[ObjectT, ArrayT](ABC):
37
42
  or for adapters with additional arguments,
38
43
  using the same config repeatedly.
39
44
 
45
+ Parameters
46
+ ----------
47
+ *args
48
+ Positional arguments to freeze.
49
+ **kwargs
50
+ Keyword arguments to freeze.
51
+
40
52
  Returns
41
53
  -------
42
- Callable
54
+ Callable[..., ObjectT]
55
+ A partial function with the given arguments frozen.
43
56
  """
44
57
  return partial(self.apply, *args, **kwargs)
45
58
 
@@ -52,15 +65,16 @@ class NullAdapter(BaseAdapter[ArrayT, ArrayT]):
52
65
 
53
66
 
54
67
  class FnAdapter(BaseAdapter[ObjectT, ArrayT]):
55
- def __init__(self, fn: Callable[[Transform[ArrayT], ObjectT], ObjectT]):
56
- """Adapter which simply wraps a function, for typing purposes.
68
+ """Adapter which simply wraps a function, for typing purposes.
57
69
 
58
- Parameters
59
- ----------
60
- fn : Callable[[Transform, T], T]
61
- Function which takes the object,
62
- and applies the transformation to it.
63
- """
70
+ Parameters
71
+ ----------
72
+ fn
73
+ Function which takes the object,
74
+ and applies the transformation to it.
75
+ """
76
+
77
+ def __init__(self, fn: Callable[[Transform[ArrayT], ObjectT], ObjectT]):
64
78
  self.fn = fn
65
79
 
66
80
  def apply(self, transform: Transform[ArrayT], obj: ObjectT) -> ObjectT:
@@ -68,18 +82,19 @@ class FnAdapter(BaseAdapter[ObjectT, ArrayT]):
68
82
 
69
83
 
70
84
  class AttrAdapter(BaseAdapter[ObjectT, ArrayT]):
71
- def __init__(self, **kwargs: BaseAdapter[Any, ArrayT] | None) -> None:
72
- """Adapter which transforms an object by applying transforms to its attributes.
85
+ """Adapter which transforms an object by applying transforms to its attributes.
86
+
87
+ Parameters
88
+ ----------
89
+ **kwargs
90
+ Keys are attribute names, values are adapters with which
91
+ to apply the transform to those attributes.
92
+ `None` is shorthand for `NullAdapter()`;
93
+ i.e. the attribute is an array and can be transformed
94
+ without being adapted.
95
+ """
73
96
 
74
- Parameters
75
- ----------
76
- adapters : Dict[str, Optional[BaseAdapter]]
77
- Keys are attribute names, values are adapters with which
78
- to apply the transform to those attributes.
79
- `None` is shorthand for `NullAdapter()`;
80
- i.e. the attribute is an array and can be transformed
81
- without being adapted.
82
- """
97
+ def __init__(self, **kwargs: BaseAdapter[Any, ArrayT] | None) -> None:
83
98
  self.adapters = {
84
99
  k: NullAdapter[ArrayT]() if v is None else v for k, v in kwargs.items()
85
100
  }
@@ -91,15 +106,23 @@ class AttrAdapter(BaseAdapter[ObjectT, ArrayT]):
91
106
 
92
107
  Parameters
93
108
  ----------
94
- transform : Transform
95
- obj : T
96
- in_place : bool, optional
109
+ transform
110
+ The transformation to apply.
111
+ obj
112
+ The object to transform.
113
+ in_place
97
114
  Whether to mutate the given object in place,
98
115
  by default False (i.e. make a deep copy of it).
99
116
 
100
117
  Returns
101
118
  -------
102
- T
119
+ ObjectT
120
+ The transformed object.
121
+
122
+ Raises
123
+ ------
124
+ TypeError
125
+ If the adapter does not support the in_place argument.
103
126
  """
104
127
  if not in_place:
105
128
  obj = deepcopy(obj)
@@ -147,7 +170,7 @@ class ReshapeAdapter(BaseAdapter[ArrayT, ArrayT]):
147
170
 
148
171
  Parameters
149
172
  ----------
150
- dim_axis : int, optional
173
+ dim_axis
151
174
  Which axis contains the coordinates' dimensions,
152
175
  by default -1 (last)
153
176
  """
@@ -15,7 +15,7 @@ class PandasAdapter(BaseAdapter[pd.DataFrame, np.ndarray]):
15
15
 
16
16
  Parameters
17
17
  ----------
18
- columns : list of keys
18
+ columns
19
19
  Keys for columns containing coordinates, e.g. `["x", "y", "z"]`
20
20
  """
21
21
  self.columns = columns
@@ -27,16 +27,18 @@ class PandasAdapter(BaseAdapter[pd.DataFrame, np.ndarray]):
27
27
 
28
28
  Parameters
29
29
  ----------
30
- transform : Transform
31
- df : pd.DataFrame
32
-
33
- in_place : bool, optional
30
+ transform
31
+ The transformation to apply.
32
+ df
33
+ The DataFrame to transform.
34
+ in_place
34
35
  Whether to mutate the dataframe in place,
35
36
  by default False (i.e. make a copy of it).
36
37
 
37
38
  Returns
38
39
  -------
39
- pandas.DataFrame
40
+ pd.DataFrame
41
+ The transformed DataFrame.
40
42
  """
41
43
  coords = df[self.columns].to_numpy()
42
44
  transformed = transform.apply(coords)
@@ -13,7 +13,7 @@ class PolarsAdapter(BaseAdapter[pl.DataFrame, np.ndarray]):
13
13
 
14
14
  Parameters
15
15
  ----------
16
- columns : list of keys
16
+ columns
17
17
  Keys for columns containing coordinates, e.g. `["x", "y", "z"]`
18
18
  """
19
19
  self.columns = columns
@@ -25,16 +25,18 @@ class PolarsAdapter(BaseAdapter[pl.DataFrame, np.ndarray]):
25
25
 
26
26
  Parameters
27
27
  ----------
28
- transform : Transform
29
- df : pl.DataFrame
30
-
31
- in_place : bool, optional
28
+ transform
29
+ The transformation to apply.
30
+ df
31
+ The DataFrame to transform.
32
+ in_place
32
33
  Whether to mutate the dataframe in place,
33
34
  by default False (i.e. make a copy of it).
34
35
 
35
36
  Returns
36
37
  -------
37
- pandas.DataFrame
38
+ pl.DataFrame
39
+ The transformed DataFrame.
38
40
  """
39
41
  coords = df[self.columns].to_numpy()
40
42
  transformed = transform.apply(coords)
@@ -0,0 +1,64 @@
1
+ import logging
2
+
3
+ import numpy as np
4
+ import shapely
5
+ from shapely.geometry.base import BaseGeometry
6
+ from shapely.coords import CoordinateSequence
7
+
8
+ from ..base import Transform, ArrayT
9
+ from .base import BaseAdapter
10
+
11
+ logger = logging.getLogger(__name__)
12
+
13
+
14
+ def as_numpy(coords: CoordinateSequence) -> np.ndarray:
15
+ return np.asarray(coords)
16
+
17
+
18
+ class GeometryAdapter(BaseAdapter[BaseGeometry, ArrayT]):
19
+ """Transform shapely geometries.
20
+
21
+ As well as the generic `apply()`,
22
+ there are `apply_*()` methods for transforming different geometry subclasses.
23
+
24
+ N.B. some transforms may create invalid topologies
25
+ (incorrect winding, self-intersections etc.).
26
+
27
+ N.B. shapely geometries' coordinates are in `XY(Z)` order
28
+ """
29
+
30
+ def apply[T: BaseGeometry](
31
+ self,
32
+ transform: Transform,
33
+ obj: T,
34
+ *,
35
+ include_z: bool | None = None,
36
+ ) -> T:
37
+ """Transform the shapely geometry.
38
+
39
+ Parameters
40
+ ----------
41
+ transform
42
+ The transformation to apply.
43
+ obj
44
+ Some shapely geometry in 2 or 3D
45
+ include_z
46
+ Force inclusion/ exclusion of Z coordinate.
47
+ By default (None), checks whether the given geometry has Z coordinates.
48
+
49
+ Returns
50
+ -------
51
+ T
52
+ An object of the same type as the input.
53
+ """
54
+
55
+ def fn(coords: np.ndarray) -> np.ndarray:
56
+ c = coords.copy()
57
+ return transform.apply(c)
58
+
59
+ if include_z is None:
60
+ inc_z = bool(shapely.has_z(obj))
61
+ else:
62
+ inc_z = include_z
63
+
64
+ return shapely.transform(obj, fn, include_z=inc_z)
@@ -6,6 +6,7 @@ from abc import ABC, abstractmethod
6
6
  from collections.abc import Iterator, Sequence
7
7
  from copy import copy
8
8
  from typing import Self, TYPE_CHECKING
9
+ from types import ModuleType
9
10
 
10
11
  from array_api_compat import array_namespace
11
12
 
@@ -35,6 +36,8 @@ class Transform[ArrayT](ABC):
35
36
  """
36
37
  Parameters
37
38
  ----------
39
+ ndims
40
+ Source and target dimensionality.
38
41
  spaces
39
42
  Optional source and target spaces
40
43
  """
@@ -50,7 +53,7 @@ class Transform[ArrayT](ABC):
50
53
 
51
54
  Returns
52
55
  -------
53
- Transform | None:
56
+ Affine[ArrayT] | None
54
57
  The affine transformation, if conversion is possible.
55
58
  None otherwise.
56
59
  """
@@ -63,9 +66,14 @@ class Transform[ArrayT](ABC):
63
66
 
64
67
  Parameters
65
68
  ----------
66
- coords : ArrayT
69
+ coords
67
70
  NxD array of N D-dimensional coordinates.
68
71
 
72
+ Returns
73
+ -------
74
+ ArrayT
75
+ The validated coordinates.
76
+
69
77
  Raises
70
78
  ------
71
79
  ValueError
@@ -87,12 +95,12 @@ class Transform[ArrayT](ABC):
87
95
 
88
96
  Parameters
89
97
  ----------
90
- coords : ArrayT
98
+ coords
91
99
  NxD array of N D-dimensional coordinates.
92
100
 
93
101
  Returns
94
102
  -------
95
- np.ndarray
103
+ ArrayT
96
104
  Transformed coordinates in the same shape.
97
105
  """
98
106
  pass
@@ -116,7 +124,7 @@ class Transform[ArrayT](ABC):
116
124
  return NotImplemented
117
125
  return t
118
126
 
119
- def to_device(self, xp, device=None) -> Self: # noqa: ARG002
127
+ def to_device(self, xp: ModuleType, device: str | None = None) -> Self: # noqa: ARG002
120
128
  """Return a copy of this transform with array parameters placed on the given device.
121
129
 
122
130
  Useful for pre-allocating parameters on GPU before a tight apply() loop,
@@ -124,15 +132,15 @@ class Transform[ArrayT](ABC):
124
132
 
125
133
  Parameters
126
134
  ----------
127
- xp : array namespace
135
+ xp
128
136
  The target array namespace (e.g. jax.numpy, torch).
129
- device : device object, optional
137
+ device
130
138
  Target device (e.g. from array_api_compat.device(array)).
131
139
  If None, uses xp's default device.
132
140
 
133
141
  Returns
134
142
  -------
135
- Transform
143
+ Self
136
144
  A new transform instance with parameters on the target device,
137
145
  or NotImplemented if the subclass does not support device placement.
138
146
  """
@@ -145,11 +153,13 @@ class Transform[ArrayT](ABC):
145
153
 
146
154
  Parameters
147
155
  ----------
148
- other : Transform
156
+ other
157
+ The transform to compose with.
149
158
 
150
159
  Returns
151
160
  -------
152
- TransformSequence
161
+ TransformSequence[ArrayT]
162
+ The composed transform sequence.
153
163
  """
154
164
  if not isinstance(other, Transform):
155
165
  return NotImplemented
@@ -166,11 +176,13 @@ class Transform[ArrayT](ABC):
166
176
 
167
177
  Parameters
168
178
  ----------
169
- other : Transform
179
+ other
180
+ The transform to compose with.
170
181
 
171
182
  Returns
172
183
  -------
173
- TransformSequence
184
+ TransformSequence[ArrayT]
185
+ The composed transform sequence.
174
186
  """
175
187
  if not isinstance(other, Transform):
176
188
  return NotImplemented
@@ -205,9 +217,13 @@ class TransformWrapper(Transform[ArrayT]):
205
217
 
206
218
  Parameters
207
219
  ----------
208
- fn : TransformSignature
220
+ fn
209
221
  Callable.
210
- spaces : tuple[SpaceRef, SpaceRef]
222
+ in_ndim
223
+ Dimensionality of the input coordinates.
224
+ out_ndim
225
+ Dimensionality of the output coordinates.
226
+ spaces
211
227
  Optional source and target spaces
212
228
  """
213
229
  super().__init__(NDims(in_ndim, out_ndim), spaces=spaces)
@@ -331,22 +347,23 @@ class TransformSequence(Transform[ArrayT], Sequence[Transform[ArrayT]]):
331
347
  coords = t.apply(coords)
332
348
  return coords
333
349
 
334
- def to_device(self, xp, device=None) -> Self:
350
+ def to_device(self, xp: ModuleType, device: str | None = None) -> Self:
335
351
  result = copy(self)
336
352
  result.transforms = [t.to_device(xp, device) for t in self.transforms]
337
353
  return result
338
354
 
339
- def list_spaces(self, skip_none=False) -> list[SpaceRef]:
355
+ def list_spaces(self, skip_none: bool = False) -> list[SpaceRef]:
340
356
  """List spaces in this transform.
341
357
 
342
358
  Parameters
343
359
  ----------
344
- skip_none : bool, optional
360
+ skip_none
345
361
  Whether to skip undefined spaces, default False.
346
362
 
347
363
  Returns
348
364
  -------
349
- List[SpaceRef]
365
+ list[SpaceRef]
366
+ The list of spaces.
350
367
  """
351
368
  spaces = [self.spaces.source] + [t.spaces.target for t in self.transforms]
352
369
  if skip_none: