transformnd 0.3.0__tar.gz → 0.4.1__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 (30) hide show
  1. {transformnd-0.3.0 → transformnd-0.4.1}/PKG-INFO +35 -13
  2. {transformnd-0.3.0 → transformnd-0.4.1}/README.md +17 -12
  3. {transformnd-0.3.0 → transformnd-0.4.1}/pyproject.toml +21 -2
  4. transformnd-0.4.1/src/transformnd/adapters/shapely.py +64 -0
  5. {transformnd-0.3.0 → transformnd-0.4.1}/src/transformnd/graph.py +0 -36
  6. transformnd-0.4.1/src/transformnd/transforms/vector_field.py +229 -0
  7. {transformnd-0.3.0 → transformnd-0.4.1}/src/transformnd/util.py +65 -8
  8. transformnd-0.3.0/src/transformnd/adapters/shapely.py +0 -130
  9. {transformnd-0.3.0 → transformnd-0.4.1}/src/transformnd/__init__.py +0 -0
  10. {transformnd-0.3.0 → transformnd-0.4.1}/src/transformnd/adapters/__init__.py +0 -0
  11. {transformnd-0.3.0 → transformnd-0.4.1}/src/transformnd/adapters/base.py +0 -0
  12. {transformnd-0.3.0 → transformnd-0.4.1}/src/transformnd/adapters/bounding_box.py +0 -0
  13. {transformnd-0.3.0 → transformnd-0.4.1}/src/transformnd/adapters/pandas.py +0 -0
  14. {transformnd-0.3.0 → transformnd-0.4.1}/src/transformnd/adapters/polars.py +0 -0
  15. {transformnd-0.3.0 → transformnd-0.4.1}/src/transformnd/base.py +0 -0
  16. {transformnd-0.3.0 → transformnd-0.4.1}/src/transformnd/constants.py +0 -0
  17. {transformnd-0.3.0 → transformnd-0.4.1}/src/transformnd/extents/__init__.py +0 -0
  18. {transformnd-0.3.0 → transformnd-0.4.1}/src/transformnd/extents/base.py +0 -0
  19. {transformnd-0.3.0 → transformnd-0.4.1}/src/transformnd/extents/bounding_box.py +0 -0
  20. {transformnd-0.3.0 → transformnd-0.4.1}/src/transformnd/py.typed +0 -0
  21. {transformnd-0.3.0 → transformnd-0.4.1}/src/transformnd/transforms/__init__.py +0 -0
  22. {transformnd-0.3.0 → transformnd-0.4.1}/src/transformnd/transforms/affine.py +0 -0
  23. {transformnd-0.3.0 → transformnd-0.4.1}/src/transformnd/transforms/bijection.py +0 -0
  24. {transformnd-0.3.0 → transformnd-0.4.1}/src/transformnd/transforms/by_dimension.py +0 -0
  25. {transformnd-0.3.0 → transformnd-0.4.1}/src/transformnd/transforms/map_axis.py +0 -0
  26. {transformnd-0.3.0 → transformnd-0.4.1}/src/transformnd/transforms/moving_least_squares.py +0 -0
  27. {transformnd-0.3.0 → transformnd-0.4.1}/src/transformnd/transforms/reflection.py +0 -0
  28. {transformnd-0.3.0 → transformnd-0.4.1}/src/transformnd/transforms/simple.py +0 -0
  29. {transformnd-0.3.0 → transformnd-0.4.1}/src/transformnd/transforms/thinplate.py +0 -0
  30. {transformnd-0.3.0 → transformnd-0.4.1}/src/transformnd/types.py +0 -0
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.4
2
2
  Name: transformnd
3
- Version: 0.3.0
3
+ Version: 0.4.1
4
4
  Summary: ND coordinate transformations
5
5
  Author: Chris Barnes
6
6
  Author-email: Chris Barnes <chris.barnes@gerbi-gmb.de>
@@ -20,20 +20,37 @@ 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: transformnd[pandas] ; extra == 'adapters'
24
+ Requires-Dist: transformnd[shapely] ; extra == 'adapters'
25
+ Requires-Dist: transformnd[polars] ; extra == 'adapters'
26
+ Requires-Dist: transformnd[transforms] ; extra == 'all'
27
+ Requires-Dist: transformnd[adapters] ; extra == 'all'
23
28
  Requires-Dist: molesq>=0.4.0 ; extra == 'movingleastsquares'
24
29
  Requires-Dist: pandas>=3.0.2 ; extra == 'pandas'
25
30
  Requires-Dist: polars>=1.40.1 ; extra == 'polars'
26
31
  Requires-Dist: shapely>=2.1.2 ; extra == 'shapely'
27
32
  Requires-Dist: morphops>=0.1.13 ; extra == 'thinplatesplines'
33
+ Requires-Dist: transformnd[thinplatesplines] ; extra == 'transforms'
34
+ Requires-Dist: transformnd[movingleastsquares] ; extra == 'transforms'
35
+ Requires-Dist: transformnd[vectorfield] ; extra == 'transforms'
36
+ Requires-Dist: transformnd[vectorfield-dask] ; extra == 'transforms'
37
+ Requires-Dist: scipy>=1.17.1 ; extra == 'vectorfield'
38
+ Requires-Dist: dask>=2026.3.0 ; extra == 'vectorfield-dask'
39
+ Requires-Dist: dask-image>=2026.5.0 ; extra == 'vectorfield-dask'
28
40
  Requires-Python: >=3.12, <4.0
29
41
  Project-URL: documentation, https://transformnd.readthedocs.io/en/latest/
30
42
  Project-URL: source, https://github.com/clbarnes/transformnd
31
43
  Project-URL: issues, https://github.com/clbarnes/transformnd/issues
44
+ Provides-Extra: adapters
45
+ Provides-Extra: all
32
46
  Provides-Extra: movingleastsquares
33
47
  Provides-Extra: pandas
34
48
  Provides-Extra: polars
35
49
  Provides-Extra: shapely
36
50
  Provides-Extra: thinplatesplines
51
+ Provides-Extra: transforms
52
+ Provides-Extra: vectorfield
53
+ Provides-Extra: vectorfield-dask
37
54
  Description-Content-Type: text/markdown
38
55
 
39
56
  # transformnd
@@ -71,18 +88,21 @@ Open it with `uv run --group tutorial marimo edit examples/tutorial.py`.
71
88
 
72
89
  ## Implemented transforms
73
90
 
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
91
+ All transforms are accessed under the `transformnd.transforms` subpackage.
92
+
93
+ | Transform | Extra | Description |
94
+ | --------- | ----- | ----------- |
95
+ | `Identity` | | No-op transformation |
96
+ | `Translation` | | Add a constant translation to the input coordinates |
97
+ | `Scale` | | Multiply the input coordinates by constant scale factor |
98
+ | `Reflection` | | Reflect coordinates about arbitrary planes |
99
+ | `MapAxis` | | Rearrange axes of the input coordinates |
100
+ | `Affine` | | Multiply augmented coordinates by an affine transformation matrix. Can represent all of the above transformations. Can be composed with matrix multiplication `aff2 @ aff1`. |
101
+ | `ByDimension` | | Apply different transformations to subsets of the input coordinates' dimensions |
102
+ | `moving_least_squares.MovingLeastSquares` | `movingleastsquares` | Landmark-based transformation. |
103
+ | `thin_plate_splines.ThinPlateSplines` | `thinplatesplines` | Landmark-based transformation. |
104
+ | `vector_field.Coordinates` | `vectorfield` for in-memory, `vectorfield-dask` for chunked | Look up output coordinates in a vector field indexed by the input coordinates |
105
+ | `vector_field.Displacements` | `vectorfield`, `vectorfield-dask` for chunked | Look up translations in a vector field indexed by the input coordinates, and add them to input coordinates |
86
106
 
87
107
  Arbitrary transforms can be composed into a `TransformSequence` with `transform1 | transform2`.
88
108
  A graph of transforms between defined spaces can be traversed using the `TransformGraph`.
@@ -135,12 +155,14 @@ Methods which SHOULD be implemented if applicable:
135
155
  - `just` to list commands.
136
156
  - Docs are generated with [`pdoc`](https://pdoc.dev/) (use `just doc`) and hosted on ReadTheDocs
137
157
  - `just bump` bumps the version, commits, and tags (but does not push); depends on [`schpet/changelog`](https://github.com/schpet/changelog)
158
+ - `just repl` starts an IPython shell with all dependencies installed
138
159
 
139
160
  ## Thanks
140
161
 
141
162
  Thanks to contributors
142
163
 
143
164
  - [Francesca Drummer](https://github.com/FrancescaDr)
165
+ - [Johannes Soltwedel](https://github.com/jo-mueller)
144
166
  - [Lorenzo Cerrone](https://github.com/lorenzocerrone)
145
167
  - [Maks Hess](https://github.com/MaksHess)
146
168
  - [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` for in-memory, `vectorfield-dask` for chunked | Look up output coordinates in a vector field indexed by the input coordinates |
50
+ | `vector_field.Displacements` | `vectorfield`, `vectorfield-dask` for chunked | 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.3.0"
3
+ version = "0.4.1"
4
4
  description = "ND coordinate transformations"
5
5
  readme = "README.md"
6
6
  authors = [{ name = "Chris Barnes", email = "chris.barnes@gerbi-gmb.de" }]
@@ -41,6 +41,21 @@ shapely = ["shapely>=2.1.2"]
41
41
  polars = [
42
42
  "polars>=1.40.1",
43
43
  ]
44
+ vectorfield = [
45
+ "scipy>=1.17.1",
46
+ ]
47
+ vectorfield-dask = [
48
+ "dask>=2026.3.0",
49
+ "dask-image>=2026.5.0",
50
+ ]
51
+ transforms = [
52
+ "transformnd[thinplatesplines]",
53
+ "transformnd[movingleastsquares]",
54
+ "transformnd[vectorfield]",
55
+ "transformnd[vectorfield-dask]",
56
+ ]
57
+ adapters = ["transformnd[pandas]", "transformnd[shapely]", "transformnd[polars]"]
58
+ all = ["transformnd[transforms]", "transformnd[adapters]"]
44
59
 
45
60
  [dependency-groups]
46
61
  dev = [
@@ -50,7 +65,11 @@ dev = [
50
65
  {include-group = "tutorial"},
51
66
  {include-group = "bench"},
52
67
  ]
53
- test = ["pytest", "jax"]
68
+ test = [
69
+ "pytest",
70
+ "jax",
71
+ "dask>=2026.3.0",
72
+ ]
54
73
  lint = [
55
74
  "ruff",
56
75
  "mypy",
@@ -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)
@@ -1,7 +1,6 @@
1
1
  """Bridging transforms between known spaces."""
2
2
 
3
3
  from __future__ import annotations
4
- from dataclasses import dataclass
5
4
  from functools import lru_cache
6
5
  from collections.abc import Iterable, Iterator
7
6
  import logging
@@ -47,41 +46,6 @@ def split_sequence(seq: TransformSequence[ArrayT]) -> Iterator[Transform[ArrayT]
47
46
  this_seq = []
48
47
 
49
48
 
50
- @dataclass(frozen=True, eq=True)
51
- class SimplifyConfig:
52
- ndim: int | None = None
53
- """Force specific dimensionality, allowing conversion to affines."""
54
-
55
- drop_inverse: bool = False
56
- """Drop explicit inverses in bijection transformations."""
57
-
58
-
59
- class NDimRegistries:
60
- def __init__(self, perm: dict[SpaceRef, int]) -> None:
61
- self.perm = perm
62
- self.temp: dict[SpaceRef, int] = dict()
63
-
64
- def _check_inner(
65
- self, space: SpaceRef, ndim: int, reg: dict[SpaceRef, int], add: bool = False
66
- ):
67
- val = reg.get(space)
68
- if val is None:
69
- if add:
70
- reg[space] = ndim
71
- return None
72
- if val != ndim:
73
- raise ValueError(
74
- f"New transform implies space {space} is {ndim}D, but it is already registered as {val}D"
75
- )
76
-
77
- def merge(self):
78
- self.perm.update(self.temp)
79
-
80
- def check(self, space: SpaceRef, ndim: int):
81
- self._check_inner(space, ndim, self.perm)
82
- self._check_inner(space, ndim, self.temp, True)
83
-
84
-
85
49
  class TransformGraph[ArrayT]:
86
50
  """Transform between any number of arbitrary spaces/ coordinate systems.
87
51
 
@@ -0,0 +1,229 @@
1
+ from abc import ABC
2
+ from collections.abc import Iterable
3
+ from types import ModuleType
4
+ from typing import Self
5
+
6
+ import numpy as np
7
+ from array_api_compat import array_namespace, is_dask_array
8
+
9
+ from ..types import NDims, Spaces
10
+ from ..base import Transform, ArrayT
11
+ from ..util import set_scipy_array_api, as_floats
12
+
13
+ set_scipy_array_api()
14
+
15
+ __all__ = ["Coordinates", "Displacements"]
16
+
17
+
18
+ class BaseVectorField(Transform[ArrayT], ABC):
19
+ def __init__(
20
+ self,
21
+ vector_field: ArrayT,
22
+ index_transform: Transform[ArrayT] | None = None,
23
+ interpolation_order: int = 3,
24
+ vector_axis: int = -1,
25
+ *,
26
+ spaces: Spaces = Spaces(None, None),
27
+ ):
28
+ """Look up a vector in array.
29
+
30
+ Parameters
31
+ ----------
32
+ vector_field
33
+ Array with `Di + 1` dimensions, where `Di` is the input dimensionality.
34
+ index_transform
35
+ Transformation from source coordinate to array indices.
36
+ interpolation_order
37
+ Order of the spline interpolation used for coordinates which are not integer array indices.
38
+ vector_axis
39
+ Which axis of the `vector_field` contains the vector values; defaults to the last (`-1`).
40
+ spaces
41
+ References for source and target spaces
42
+
43
+ Raises
44
+ ------
45
+ ValueError
46
+ If index_transform's output dimensionality is not exactly one less than the vector field's number of dimensions.
47
+ """
48
+ self.vector_field: ArrayT = as_floats(vector_field) # type: ignore
49
+ xp = array_namespace(vector_field)
50
+ sh = xp.shape(vector_field)
51
+ in_ndim = len(sh) - 1
52
+ tgt_ndim = sh[vector_axis]
53
+ self.vector_field = vector_field
54
+ if index_transform is None:
55
+ source_ndim = in_ndim
56
+ else:
57
+ if in_ndim != index_transform.ndims.target:
58
+ raise ValueError(
59
+ "If index_transform is given, its output dimensionality must match the vector field's shape"
60
+ )
61
+ source_ndim = index_transform.ndims.source
62
+ self.index_transform = index_transform
63
+ self.vector_axis = vector_axis % len(sh)
64
+ self._mode = "constant"
65
+ self._cval = np.nan
66
+ self._order = interpolation_order
67
+ super().__init__(NDims(source_ndim, tgt_ndim), spaces=spaces)
68
+
69
+ def _vf_slices(self) -> Iterable[ArrayT]:
70
+ slicing: list[slice | int] = [slice(None)] * (self.ndims.source + 1)
71
+ for v_idx in range(self.ndims.target):
72
+ slicing[self.vector_axis] = v_idx
73
+ yield self.vector_field[tuple(slicing)] # type: ignore
74
+
75
+ def _get_vectors_inner_dask(self, index_coords_t: ArrayT) -> ArrayT:
76
+ import dask.array as da
77
+ from dask_image.ndinterp import map_coordinates
78
+
79
+ out = []
80
+ for vf in self._vf_slices():
81
+ out.append(
82
+ map_coordinates(
83
+ vf,
84
+ index_coords_t,
85
+ order=self._order,
86
+ mode=self._mode,
87
+ cval=self._cval,
88
+ )
89
+ )
90
+ stacked = da.stack(out)
91
+ return da.transpose(stacked)
92
+
93
+ def _get_vectors_inner_scipy(self, index_coords_t: ArrayT) -> ArrayT:
94
+ from scipy.ndimage import map_coordinates
95
+
96
+ set_scipy_array_api()
97
+ xp = array_namespace(index_coords_t)
98
+ out = xp.zeros_like(
99
+ self.vector_field, shape=(self.ndims.target, xp.shape(index_coords_t)[1])
100
+ )
101
+ for idx, vf in enumerate(self._vf_slices()):
102
+ map_coordinates(
103
+ vf,
104
+ index_coords_t,
105
+ order=self._order,
106
+ mode=self._mode,
107
+ cval=self._cval,
108
+ output=out[idx, :],
109
+ )
110
+ return xp.transpose(out)
111
+
112
+ def _get_vectors(self, coords: ArrayT) -> ArrayT:
113
+ if self.index_transform is not None:
114
+ coords = self.index_transform.apply(coords)
115
+ else:
116
+ coords = self._validate_coords(coords)
117
+ xp = array_namespace(coords)
118
+ c = xp.transpose(coords)
119
+ if is_dask_array(self.vector_field):
120
+ return self._get_vectors_inner_dask(c)
121
+ else:
122
+ return self._get_vectors_inner_scipy(c)
123
+
124
+ def to_device(self, xp: ModuleType, device: str | None = None) -> Self:
125
+ coords = xp.asarray(self.vector_field, device)
126
+ return type(self)(coords, spaces=self.spaces)
127
+
128
+
129
+ class Coordinates(BaseVectorField[ArrayT]):
130
+ """Look up the output coordinates in an array.
131
+
132
+ For input coordinate `(a, b, c)` and `vector_axis=-1`,
133
+ the output coordinate is `vector_field[a, b, c, :].
134
+
135
+ Input coordinates outside the vector field return NaN.
136
+ """
137
+
138
+ def __init__(
139
+ self,
140
+ vector_field: ArrayT,
141
+ index_transform: Transform[ArrayT] | None = None,
142
+ interpolation_order: int = 3,
143
+ vector_axis: int = -1,
144
+ *,
145
+ spaces: Spaces = Spaces(None, None),
146
+ ):
147
+ """Use the input coordinates as array indices to look up output coordinates.
148
+
149
+ For input coordinate `(a, b, c)`, the output coordinate is `coordinates[a, b, c, :]`.
150
+
151
+ Input coordinates outside of the `vector_field` array return `NaN` output coordinates.
152
+
153
+ Parameters
154
+ ----------
155
+ vector_field
156
+ Array with `Di + 1` dimensions, where `Di` is the input dimensionality.
157
+ index_transform
158
+ Transform the source coordinates into an array index
159
+ interpolation_order
160
+ Order of the spline interpolation used for coordinates which are not integer array indices.
161
+ vector_axis
162
+ Which axis of the `vector_field` contains the vector values; defaults to the last (`-1`).
163
+ spaces
164
+ References for source and target spaces
165
+ """
166
+ super().__init__(
167
+ vector_field,
168
+ index_transform,
169
+ interpolation_order,
170
+ vector_axis,
171
+ spaces=spaces,
172
+ )
173
+
174
+ def apply(self, coords: ArrayT) -> ArrayT:
175
+ return self._get_vectors(coords)
176
+
177
+
178
+ class Displacements(BaseVectorField[ArrayT]):
179
+ """Look up a translation in an array and apply it to the input coordinates.
180
+
181
+ For input coordinate `(a, b, c)` and `vector_axis=-1`,
182
+ the output coordinate is `(a, b, c) + vector_field[a, b, c, :].
183
+
184
+ Input coordinates outside the vector field return NaN.
185
+ """
186
+
187
+ def __init__(
188
+ self,
189
+ vector_field: ArrayT,
190
+ index_transform: Transform[ArrayT] | None = None,
191
+ interpolation_order: int = 3,
192
+ vector_axis: int = -1,
193
+ *,
194
+ spaces: Spaces = Spaces(None, None),
195
+ ):
196
+ """
197
+ Parameters
198
+ ----------
199
+ vector_field
200
+ Array with `Di + 1` dimensions, where `Di` is the input dimensionality.
201
+ index_transform
202
+ Transformation from source coordinate to array indices.
203
+ interpolation_order
204
+ Order of the spline interpolation used for coordinates which are not integer array indices.
205
+ vector_axis
206
+ Which axis of the `vector_field` contains the vector values; defaults to the last (`-1`).
207
+ spaces
208
+ References for source and target spaces
209
+
210
+ Raises
211
+ ------
212
+ ValueError
213
+ If the index transform and vector field would change the coordinates' dimensionality,
214
+ or the index transform's dimensionality does not match the vector field's.
215
+ """
216
+ super().__init__(
217
+ vector_field,
218
+ index_transform,
219
+ interpolation_order,
220
+ vector_axis,
221
+ spaces=spaces,
222
+ )
223
+ if self.ndims.source != self.ndims.target:
224
+ raise ValueError("Displacements cannot change dimensionality")
225
+
226
+ def apply(self, coords: ArrayT) -> ArrayT:
227
+ coords = self._validate_coords(coords)
228
+ vecs = self._get_vectors(coords)
229
+ return coords + vecs # type:ignore
@@ -1,12 +1,22 @@
1
1
  """Utilities used elsewhere in the package."""
2
2
 
3
- from array_api_compat import array_namespace
4
- from numpy.typing import ArrayLike
3
+ import os
4
+ from types import ModuleType
5
+ import warnings
6
+ import logging
7
+
8
+ from array_api_compat import (
9
+ array_namespace,
10
+ is_dask_namespace,
11
+ is_jax_namespace,
12
+ )
5
13
  import numpy as np
6
14
 
7
15
  from .types import SpaceRef, ArrayT
8
16
  from .constants import UNSPECIFIED_SPACE_NAME
9
17
 
18
+ logger = logging.getLogger(__name__)
19
+
10
20
 
11
21
  def none_eq[T](a: T | None, b: T | None) -> bool:
12
22
  """Check whether either is None or both are equal.
@@ -161,9 +171,56 @@ def to_single_ndim(ndim: None | int = None, ndims: None | set[int] = None) -> in
161
171
  raise ValueError(f"dimensionality conflict: {ndim} not in {ndims}")
162
172
 
163
173
 
164
- def as_floats(arr: ArrayLike):
165
- """Get array-like as a numpy array, casting to float if integral."""
166
- arr = np.asarray(arr)
167
- if not np.issubdtype(arr.dtype, np.floating):
168
- arr = arr.astype(np.float64)
169
- return arr
174
+ def as_floats(arr, *, namespace: ModuleType | None = None, device: str | None = None):
175
+ """Get array-like as an array of floats.
176
+
177
+ Convert to a particular array namespace if given;
178
+ default to keeping the same namespace one exists,
179
+ or numpy otherwise.
180
+
181
+ Cast to a float if integral.
182
+ """
183
+ try:
184
+ orig_namespace = array_namespace(arr)
185
+ except TypeError:
186
+ orig_namespace = None
187
+
188
+ if namespace is None:
189
+ namespace = orig_namespace or np
190
+
191
+ kwargs = dict()
192
+ if is_dask_namespace(namespace):
193
+ if isinstance(device, str) and device != "cpu":
194
+ logger.warning(f"Ignoring unsupported dask device: {device}")
195
+ elif device is not None:
196
+ kwargs["device"] = device
197
+
198
+ arr = namespace.asarray(arr, **kwargs) # type: ignore
199
+
200
+ # dask does not have isdtype
201
+ isdtype = getattr(namespace, "isdtype", np.isdtype)
202
+
203
+ if not isdtype(arr.dtype, "real floating"): # type:ignore
204
+ if is_jax_namespace(namespace):
205
+ dt = "float32"
206
+ else:
207
+ dt = "float64"
208
+ # N.B. a dask array wrapping over jax arrays will warn here
209
+ arr = arr.astype(dt, **kwargs) # type:ignore
210
+
211
+ return arr # type:ignore
212
+
213
+
214
+ def set_scipy_array_api() -> bool:
215
+ curr = os.environ.get("SCIPY_ARRAY_API")
216
+ match curr:
217
+ case None:
218
+ os.environ["SCIPY_ARRAY_API"] = "1"
219
+ return True
220
+ case "1":
221
+ return True
222
+ case _:
223
+ warnings.warn(
224
+ "SCIPY_ARRAY_API environment set but not '1'; certain transforms may not work with certain array types"
225
+ )
226
+ return False
@@ -1,130 +0,0 @@
1
- import numpy as np
2
- from typing import Callable
3
- from shapely.geometry import (
4
- LinearRing,
5
- LineString,
6
- MultiLineString,
7
- MultiPoint,
8
- MultiPolygon,
9
- Point,
10
- Polygon,
11
- GeometryCollection,
12
- )
13
- from shapely.geometry.base import BaseGeometry, BaseMultipartGeometry
14
- from shapely.coords import CoordinateSequence
15
-
16
- from ..base import Transform, ArrayT
17
- from .base import BaseAdapter
18
-
19
-
20
- def as_numpy(coords: CoordinateSequence) -> np.ndarray:
21
- return np.asarray(coords)
22
-
23
-
24
- class GeometryAdapter(BaseAdapter[BaseGeometry, ArrayT]):
25
- """Transform shapely geometries.
26
-
27
- As well as the generic `apply()`,
28
- there are `apply_*()` methods for transforming different geometry subclasses.
29
-
30
- N.B. some transforms may create invalid topologies
31
- (incorrect winding, self-intersections etc.)
32
- """
33
-
34
- def __init__(
35
- self,
36
- array_fn: Callable[[CoordinateSequence], ArrayT] = as_numpy, # type:ignore
37
- ) -> None:
38
- self.array_fn = array_fn
39
-
40
- def apply_point(self, transform: Transform, point: Point) -> Point:
41
- return Point(*transform.apply(self.array_fn(point.coords))[0])
42
-
43
- def apply_linestring(
44
- self, transform: Transform, linestring: LineString
45
- ) -> LineString:
46
- return LineString(transform.apply(self.array_fn(linestring.coords)))
47
-
48
- def apply_linear_ring(
49
- self, transform: Transform, linear_ring: LinearRing
50
- ) -> LinearRing:
51
- return LinearRing(transform.apply(self.array_fn(linear_ring.coords)))
52
-
53
- def apply_polygon(self, transform: Transform, polygon: Polygon) -> Polygon:
54
- return Polygon(
55
- self.apply_linear_ring(transform, polygon.exterior),
56
- [self.apply_linear_ring(transform, i) for i in polygon.interiors],
57
- )
58
-
59
- def apply_multipoint(self, transform: Transform, obj: MultiPoint) -> MultiPoint:
60
- return MultiPoint([self.apply_point(transform, o) for o in obj.geoms])
61
-
62
- def apply_multilinestring(
63
- self, transform: Transform, obj: MultiLineString
64
- ) -> MultiLineString:
65
- return MultiLineString([self.apply_linestring(transform, o) for o in obj.geoms])
66
-
67
- def apply_multipolygon(
68
- self, transform: Transform, obj: MultiPolygon
69
- ) -> MultiPolygon:
70
- return MultiPolygon([self.apply_polygon(transform, o) for o in obj.geoms])
71
-
72
- def apply_multipart(
73
- self, transform: Transform, obj: BaseMultipartGeometry
74
- ) -> BaseMultipartGeometry:
75
- """Apply the transform to any shapely multipart geometry."""
76
- if isinstance(obj, MultiPoint):
77
- return self.apply_multipoint(transform, obj)
78
- elif isinstance(obj, MultiLineString):
79
- return self.apply_multilinestring(transform, obj)
80
- elif isinstance(obj, MultiPolygon):
81
- return self.apply_multipolygon(transform, obj)
82
- elif isinstance(obj, GeometryCollection):
83
- return self.apply_collection(transform, obj)
84
- else:
85
- raise ValueError(f"Unknown multipart geometry type {type(obj)}")
86
-
87
- def apply_collection(
88
- self, transform: Transform, obj: GeometryCollection
89
- ) -> GeometryCollection:
90
- return GeometryCollection([self.apply(transform, o) for o in obj.geoms])
91
-
92
- def apply(
93
- self,
94
- transform: Transform,
95
- obj: BaseGeometry,
96
- ) -> BaseGeometry:
97
- """Transform the shapely geometry.
98
-
99
- See the other `apply_*` methods if you already know what type of geometry
100
- you're working with; this may be a bit faster.
101
-
102
- Parameters
103
- ----------
104
- transform
105
- The transformation to apply.
106
- obj
107
- Some shapely geometry in 2 or 3D
108
-
109
- Returns
110
- -------
111
- BaseGeometry
112
- An object of the same type as the input.
113
-
114
- Raises
115
- ------
116
- ValueError
117
- If the geometry type is not supported.
118
- """
119
- if isinstance(obj, BaseMultipartGeometry):
120
- return self.apply_multipart(transform, obj)
121
- elif isinstance(obj, Point):
122
- return self.apply_point(transform, obj)
123
- elif isinstance(obj, LineString):
124
- return self.apply_linestring(transform, obj)
125
- elif isinstance(obj, LinearRing):
126
- return self.apply_linear_ring(transform, obj)
127
- elif isinstance(obj, Polygon):
128
- return self.apply_polygon(transform, obj)
129
- else:
130
- raise ValueError(f"Unknown geometry type {type(obj)}")