torch-geopooling 1.3.0__cp310-cp310-manylinux_2_28_x86_64.whl

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.
@@ -0,0 +1,37 @@
1
+ # SPDX-License-Identifier: GPL-3.0-or-later
2
+ # SPDX-FileCopyrightText: 2025 Yakau Bubnou
3
+ # SPDX-FileType: SOURCE
4
+
5
+ import pytest
6
+ import torch
7
+ from torch import nn
8
+ from torch.optim import SGD
9
+
10
+ from torch_geopooling.nn.embedding import Embedding2d
11
+
12
+
13
+ def test_embedding2d_optimize() -> None:
14
+ embedding = Embedding2d(
15
+ (2, 2, 1),
16
+ padding=(0, 0),
17
+ exterior=(-180.0, -90.0, 360.0, 180.0),
18
+ )
19
+
20
+ x_true = torch.tensor(
21
+ [[90.0, 45.0], [90.0, -45.0], [-90.0, -45.0], [-90.0, 45.0]], dtype=torch.float64
22
+ )
23
+ y_true = torch.tensor([[10.0], [20.0], [30.0], [40.0]], dtype=torch.float64)
24
+
25
+ optim = SGD(embedding.parameters(), lr=0.1)
26
+ loss_fn = nn.L1Loss()
27
+
28
+ for i in range(10000):
29
+ optim.zero_grad()
30
+
31
+ y_pred = embedding(x_true)
32
+ loss = loss_fn(y_pred[:, 0, 0, :], y_true)
33
+ loss.backward()
34
+
35
+ optim.step()
36
+
37
+ assert pytest.approx(0.0, abs=1e-1) == loss.detach().item()
@@ -0,0 +1,385 @@
1
+ # SPDX-License-Identifier: GPL-3.0-or-later
2
+ # SPDX-FileCopyrightText: 2025 Yakau Bubnou
3
+ # SPDX-FileType: SOURCE
4
+
5
+ from typing import Optional, Union
6
+
7
+ import torch
8
+ from shapely.geometry import Polygon
9
+ from torch import Tensor, nn
10
+
11
+ from torch_geopooling import functional as F
12
+ from torch_geopooling.tiling import Exterior, ExteriorTuple, regular_tiling
13
+
14
+ __all__ = [
15
+ "AdaptiveAvgQuadPool2d",
16
+ "AdaptiveQuadPool2d",
17
+ "AdaptiveMaxQuadPool2d",
18
+ "AvgQuadPool2d",
19
+ "MaxQuadPool2d",
20
+ "QuadPool2d",
21
+ ]
22
+
23
+
24
+ _Exterior = Union[Exterior, ExteriorTuple]
25
+
26
+
27
+ _exterior_doc = """
28
+ Note:
29
+ Input coordinates must be within a specified exterior geometry (including boundaries).
30
+ For input coordinates outsize of the specified exterior, module throws an exception.
31
+ """
32
+
33
+
34
+ _terminal_group_doc = """
35
+ Note:
36
+ A **terminal group** refers to a collection of terminal nodes within the quadtree that
37
+ share the same parent tile.
38
+ """
39
+
40
+
41
+ class _AdaptiveQuadPool(nn.Module):
42
+ __doc__ = f"""
43
+ Args:
44
+ feature_dim: Size of each feature vector.
45
+ exterior: Geometrical boundary of the learning space in (X, Y, W, H) format.
46
+ max_terminal_nodes: Optional maximum number of terminal nodes in a quadtree. Once a
47
+ maximum is reached, internal nodes are no longer sub-divided and tree stops growing.
48
+ max_depth: Maximum depth of the quadtree. Default: 17.
49
+ capacity: Maximum number of inputs, after which a quadtree's node is subdivided and
50
+ depth of the tree grows. Default: 1.
51
+ precision: Optional rounding of the input coordinates. Default: 7.
52
+
53
+ Shape:
54
+ - Input: :math:`(*, 2)`, where 2 comprises longitude and latitude coordinates.
55
+ - Output: :math:`(*, H)`, where * is the input shape and :math:`H = \\text{{feature_dim}}`.
56
+
57
+ {_exterior_doc}
58
+ {_terminal_group_doc}
59
+ """
60
+
61
+ def __init__(
62
+ self,
63
+ feature_dim: int,
64
+ exterior: _Exterior,
65
+ max_terminal_nodes: Optional[int] = None,
66
+ max_depth: int = 17,
67
+ capacity: int = 1,
68
+ precision: Optional[int] = 7,
69
+ ) -> None:
70
+ super().__init__()
71
+ self.feature_dim = feature_dim
72
+ self.exterior = tuple(map(float, exterior))
73
+ self.max_terminal_nodes = max_terminal_nodes
74
+ self.max_depth = max_depth
75
+ self.capacity = capacity
76
+ self.precision = precision
77
+
78
+ self.initialize_parameters()
79
+
80
+ def initialize_parameters(self) -> None:
81
+ # The weight for adaptive operation should be sparse, since training operation
82
+ # results in a dynamic change of the underlying quadtree.
83
+ weight_size = (
84
+ self.max_depth + 1,
85
+ 1 << self.max_depth,
86
+ 1 << self.max_depth,
87
+ self.feature_dim,
88
+ )
89
+ self.weight = nn.Parameter(torch.sparse_coo_tensor(size=weight_size, dtype=torch.float64))
90
+
91
+ @property
92
+ def tiles(self) -> torch.Tensor:
93
+ """Return tiles of the quadtree."""
94
+ return self.weight.coalesce().detach().indices().t()[:, :-1]
95
+
96
+ def extra_repr(self) -> str:
97
+ return (
98
+ "{feature_dim}, "
99
+ "exterior={exterior}, capacity={capacity}, max_depth={max_depth}, "
100
+ "precision={precision}".format(**self.__dict__)
101
+ )
102
+
103
+
104
+ class AdaptiveQuadPool2d(_AdaptiveQuadPool):
105
+ __doc__ = f"""Adaptive lookup index over quadtree decomposition of input 2D coordinates.
106
+
107
+ This module constructs an internal lookup quadtree to organize closely situated 2D points.
108
+ Each terminal node in the resulting quadtree is paired with a weight. Thus, when providing
109
+ an input coordinate, the module retrieves the corresponding terminal node and returns its
110
+ associated weight.
111
+
112
+ {_AdaptiveQuadPool.__doc__}
113
+
114
+ Examples:
115
+
116
+ >>> # Feature vectors of size 4 over a 2d space.
117
+ >>> pool = nn.AdaptiveQuadPool2d(4, (-10, -5, 20, 10))
118
+ >>> # Grow tree up to 4-th level and sub-divides a node after 8 coordinates in a quad.
119
+ >>> pool = nn.AdaptiveQuadPool2d(4, (-10, -5, 20, 10), max_depth=4, capacity=8)
120
+ >>> # Create 2D coordinates and query associated weights.
121
+ >>> input = torch.rand((1024, 2), dtype=torch.float64) * 10 - 5
122
+ >>> output = pool(input)
123
+ """
124
+
125
+ def forward(self, input: Tensor) -> Tensor:
126
+ result = F.adaptive_quad_pool2d(
127
+ self.weight,
128
+ input,
129
+ self.exterior,
130
+ training=self.training,
131
+ max_terminal_nodes=self.max_terminal_nodes,
132
+ max_depth=self.max_depth,
133
+ capacity=self.capacity,
134
+ precision=self.precision,
135
+ )
136
+ if self.training:
137
+ self.weight.data = result.weight
138
+ return result.values
139
+
140
+
141
+ class AdaptiveMaxQuadPool2d(_AdaptiveQuadPool):
142
+ __doc__ = f"""Adaptive maximum pooling over quadtree decomposition of input 2D coordinates.
143
+
144
+ This module constructs an internal lookup quadtree to organize closely situated 2D points.
145
+ Each terminal node in the resulting quadtree is paired with a weight. Thus, when providing
146
+ an input coordinate, the module retrieves a **terminal group** of nodes and calculates the
147
+ maximum value for each ``feature_dim``.
148
+
149
+ {_AdaptiveQuadPool.__doc__}
150
+
151
+ Examples:
152
+
153
+ >>> pool = nn.AdaptiveMaxQuadPool2d(3, (-10, -5, 20, 10), max_depth=5)
154
+ >>> # Create 2D coordinates and feature vector associated with them.
155
+ >>> input = torch.rand((2048, 2), dtype=torch.float64) * 10 - 5
156
+ >>> output = pool(input)
157
+ """
158
+
159
+ def forward(self, input: Tensor) -> Tensor:
160
+ result = F.adaptive_max_quad_pool2d(
161
+ self.weight,
162
+ input,
163
+ self.exterior,
164
+ training=self.training,
165
+ max_terminal_nodes=self.max_terminal_nodes,
166
+ max_depth=self.max_depth,
167
+ capacity=self.capacity,
168
+ precision=self.precision,
169
+ )
170
+ if self.training:
171
+ self.weight.data = result.weight
172
+ return result.values
173
+
174
+
175
+ class AdaptiveAvgQuadPool2d(_AdaptiveQuadPool):
176
+ __doc__ = f"""Adaptive average pooling over quadtree decomposition of input 2D coordinates.
177
+
178
+ This module constructs an internal lookup quadtree to organize closely situated 2D points.
179
+ Each terminal node in the resulting quadtree is paired with a weight. Thus, when providing
180
+ an input coordinate, the module retrieves a **terminal group** of nodes and calculates an
181
+ average value for each ``feature_dim``.
182
+
183
+ {_AdaptiveQuadPool.__doc__}
184
+
185
+ Examples:
186
+
187
+ >>> # Create pool with 7 features.
188
+ >>> pool = nn.AdaptiveAvgQuadPool2d(7, (0, 0, 1, 1), max_depth=12)
189
+ >>> input = torch.rand((2048, 2), dtype=torch.float64)
190
+ >>> output = pool(input)
191
+ """
192
+
193
+ def forward(self, input: Tensor) -> Tensor:
194
+ result = F.adaptive_avg_quad_pool2d(
195
+ self.weight,
196
+ input,
197
+ self.exterior,
198
+ training=self.training,
199
+ max_terminal_nodes=self.max_terminal_nodes,
200
+ max_depth=self.max_depth,
201
+ capacity=self.capacity,
202
+ precision=self.precision,
203
+ )
204
+ if self.training:
205
+ self.weight.data = result.weight
206
+ return result.values
207
+
208
+
209
+ class _QuadPool(nn.Module):
210
+ __doc__ = f"""
211
+ Args:
212
+ feature_dim: Size of each feature vector.
213
+ polygon: Polygon that resembles boundary for the terminal nodes of a quadtree.
214
+ exterior: Geometrical boundary of the learning space in (X, Y, W, H) format.
215
+ max_terminal_nodes: Optional maximum number of terminal nodes in a quadtree. Once a
216
+ maximum is reached, internal nodes are no longer sub-divided and tree stops growing.
217
+ max_depth: Maximum depth of the quadtree. Default: 17.
218
+ precision: Optional rounding of the input coordinates. Default: 7.
219
+
220
+ Shape:
221
+ - Input: :math:`(*, 2)`, where 2 comprises longitude and latitude coordinates.
222
+ - Output: :math:`(*, H)`, where * is the input shape and :math:`H = \\text{{feature_dim}}`.
223
+
224
+ {_exterior_doc}
225
+ {_terminal_group_doc}
226
+
227
+ Note:
228
+ All terminal nodes that have an intersection with the specified polygon boundary are
229
+ included into the quadtree.
230
+ """
231
+
232
+ def __init__(
233
+ self,
234
+ feature_dim: int,
235
+ polygon: Polygon,
236
+ exterior: _Exterior,
237
+ max_terminal_nodes: Optional[int] = None,
238
+ max_depth: int = 17,
239
+ precision: Optional[int] = 7,
240
+ ) -> None:
241
+ super().__init__()
242
+ self.feature_dim = feature_dim
243
+ self.polygon = polygon
244
+ self.exterior = tuple(map(float, exterior))
245
+ self.max_terminal_nodes = max_terminal_nodes
246
+ self.max_depth = max_depth
247
+ self.precision = precision
248
+
249
+ # Generate regular tiling for the provided polygon and build from those
250
+ # tiles a quadtree from terminal nodes all way up to the root node.
251
+ tiles_iter = regular_tiling(
252
+ polygon, Exterior.from_tuple(exterior), z=max_depth, internal=True
253
+ )
254
+ tiles = torch.tensor(list(tiles_iter), dtype=torch.int64)
255
+
256
+ self.register_buffer("tiles", tiles)
257
+ self.tiles: Tensor
258
+
259
+ self.initialize_parameters()
260
+ self.reset_parameters()
261
+
262
+ def initialize_parameters(self) -> None:
263
+ weight_size = [self.tiles.size(0), self.feature_dim]
264
+ self.weight = nn.Parameter(torch.empty(weight_size, dtype=torch.float64))
265
+
266
+ def reset_parameters(self) -> None:
267
+ nn.init.uniform_(self.weight)
268
+
269
+ def extra_repr(self) -> str:
270
+ return (
271
+ "{feature_dim}, exterior={exterior}, max_depth={max_depth}, "
272
+ "precision={precision}".format(**self.__dict__)
273
+ )
274
+
275
+
276
+ class QuadPool2d(_QuadPool):
277
+ __doc__ = f"""Lookup index over quadtree decomposition of input 2D coordinates.
278
+
279
+ This module constructs an internal lookup tree to organize closely situated 2D points using
280
+ a specified polygon and exterior, where polygon is treated as a *boundary* of terminal
281
+ nodes of a quadtree.
282
+
283
+ Each terminal node in the resulting quadtree is paired with a weight. Thus, when providing
284
+ an input coordinate, the module retrieves the corresponding terminal node and returns its
285
+ associated weight.
286
+
287
+ {_QuadPool.__doc__}
288
+
289
+ Examples:
290
+
291
+ >>> from shapely.geometry import Polygon
292
+ >>> # Create a pool for squared exterior 100x100 and use only a portion of that
293
+ >>> # exterior isolated by a square 10x10.
294
+ >>> poly = Polygon([(0, 0), (10, 0), (10, 10), (0, 10)])
295
+ >>> pool = nn.QuadPool2d(5, poly, exterior=(0, 0, 100, 100))
296
+ >>> input = torch.rand((2048, 2), dtype=torch.float64)
297
+ >>> output = pool(input)
298
+ """
299
+
300
+ def forward(self, input: Tensor) -> Tensor:
301
+ result = F.quad_pool2d(
302
+ self.tiles,
303
+ self.weight,
304
+ input,
305
+ self.exterior,
306
+ # This is not a mistake, since we already know the shape of the
307
+ # quadtree, there is no need to learn it.
308
+ training=False,
309
+ max_terminal_nodes=self.max_terminal_nodes,
310
+ max_depth=self.max_depth,
311
+ precision=self.precision,
312
+ )
313
+ return result.values
314
+
315
+
316
+ class MaxQuadPool2d(_QuadPool):
317
+ __doc__ = f"""Maximum pooling over quadtree decomposition of input 2D coordinates.
318
+
319
+ This module constructs an internal lookup tree to organize closely situated 2D points using
320
+ a specified polygon and exterior, where polygon is treated as a *boundary* of terminal nodes
321
+ of a quadtree.
322
+
323
+ Each terminal node in the resulting quadtree is paired with a weight. Thus, when providing
324
+ an input coordinate, the module retrieves a **terminal group** of nodes and calculates the
325
+ maximum value for each ``feature_dim``.
326
+
327
+ {_QuadPool.__doc__}
328
+
329
+ Examples:
330
+
331
+ >>> from shapely.geometry import Polygon
332
+ >>> poly = Polygon([(0, 0), (10, 0), (10, 10), (0, 10)])
333
+ >>> pool = nn.MaxQuadPool2d(3, poly, exterior=(0, 0, 100, 100))
334
+ >>> input = torch.rand((2048, 2), dtype=torch.float64)
335
+ >>> output = pool(input)
336
+ """
337
+
338
+ def forward(self, input: Tensor) -> Tensor:
339
+ result = F.max_quad_pool2d(
340
+ self.tiles,
341
+ self.weight,
342
+ input,
343
+ self.exterior,
344
+ training=False,
345
+ max_terminal_nodes=self.max_terminal_nodes,
346
+ max_depth=self.max_depth,
347
+ precision=self.precision,
348
+ )
349
+ return result.values
350
+
351
+
352
+ class AvgQuadPool2d(_QuadPool):
353
+ __doc__ = f"""Average pooling over quadtree decomposition of input 2D coordinates.
354
+
355
+ This module constructs an internal lookup tree to organize closely situated 2D points using
356
+ a specified polygon and exterior, where polygon is treated as a *boundary* of terminal
357
+ nodes of a quadtree.
358
+
359
+ Each terminal node in the resulting quadtree is paired with a weight. Thus, when providing
360
+ an input coordinate, the module retrieves a **terminal group** of nodes and calculates an
361
+ average value for each ``feature_dim``.
362
+
363
+ {_QuadPool.__doc__}
364
+
365
+ Examples:
366
+
367
+ >>> from shapely.geometry import Polygon
368
+ >>> poly = Polygon([(0, 0), (10, 0), (10, 10), (0, 10)])
369
+ >>> pool = nn.AvgQuadPool2d(4, poly, exterior=(0, 0, 100, 100))
370
+ >>> input = torch.rand((2048, 2), dtype=torch.float64)
371
+ >>> output = pool(input)
372
+ """
373
+
374
+ def forward(self, input: Tensor) -> Tensor:
375
+ result = F.avg_quad_pool2d(
376
+ self.tiles,
377
+ self.weight,
378
+ input,
379
+ self.exterior,
380
+ training=False,
381
+ max_terminal_nodes=self.max_terminal_nodes,
382
+ max_depth=self.max_depth,
383
+ precision=self.precision,
384
+ )
385
+ return result.values
@@ -0,0 +1,147 @@
1
+ # SPDX-License-Identifier: GPL-3.0-or-later
2
+ # SPDX-FileCopyrightText: 2025 Yakau Bubnou
3
+ # SPDX-FileType: SOURCE
4
+
5
+ from typing import Type
6
+
7
+ import pytest
8
+ import torch
9
+ from shapely.geometry import Polygon
10
+ from torch import nn
11
+ from torch.optim import SGD
12
+ from torch.nn import L1Loss
13
+
14
+ from torch_geopooling.nn.pooling import (
15
+ AdaptiveAvgQuadPool2d,
16
+ AdaptiveMaxQuadPool2d,
17
+ AdaptiveQuadPool2d,
18
+ AvgQuadPool2d,
19
+ MaxQuadPool2d,
20
+ QuadPool2d,
21
+ )
22
+
23
+
24
+ @pytest.mark.parametrize(
25
+ "module_class",
26
+ [
27
+ AdaptiveQuadPool2d,
28
+ AdaptiveMaxQuadPool2d,
29
+ AdaptiveAvgQuadPool2d,
30
+ ],
31
+ ids=["id", "max", "avg"],
32
+ )
33
+ def test_adaptive_quad_pool2d_gradient(module_class: Type[nn.Module]) -> None:
34
+ pool = module_class(5, (-180, -90, 360, 180))
35
+
36
+ input = torch.rand((100, 2), dtype=torch.float64) * 90
37
+ y = pool(input)
38
+
39
+ assert pool.weight.grad is None
40
+
41
+ loss_fn = L1Loss()
42
+ loss = loss_fn(y, torch.ones_like(y))
43
+ loss.backward()
44
+
45
+ assert pool.weight.grad is not None
46
+ assert pool.weight.grad.sum().item() == pytest.approx(-1)
47
+
48
+
49
+ def test_adaptive_quad_pool2d_optimize() -> None:
50
+ pool = AdaptiveQuadPool2d(1, (-180, -90, 360, 180), max_depth=1)
51
+
52
+ # Input coordinates are simply centers of the level-1 quads.
53
+ x_true = torch.tensor(
54
+ [[90.0, 45.0], [90.0, -45.0], [-90.0, -45.0], [-90.0, 45.0]], dtype=torch.float64
55
+ )
56
+ y_true = torch.tensor([[10.0], [20.0], [30.0], [40.0]], dtype=torch.float64)
57
+ y_tile = [[1, 1, 1], [1, 1, 0], [1, 0, 0], [1, 0, 1]]
58
+
59
+ optim = SGD(pool.parameters(), lr=0.01)
60
+ loss_fn = nn.L1Loss()
61
+
62
+ for i in range(20000):
63
+ optim.zero_grad()
64
+
65
+ y_pred = pool(x_true)
66
+ loss = loss_fn(y_pred, y_true)
67
+ loss.backward()
68
+
69
+ optim.step()
70
+
71
+ # Ensure that model converged with a small loss.
72
+ assert pytest.approx(0.0, abs=1e-1) == loss.detach().item()
73
+
74
+ # Ensure that weights that pooling operation learned are the same as in the
75
+ # target matrix (y_true).
76
+ weight = pool.weight.to_dense()
77
+
78
+ for i, tile in enumerate(y_tile):
79
+ z, x, y = tile
80
+ expect_weight = y_true[i].item()
81
+ actual_weight = weight[z, x, y].detach().item()
82
+
83
+ assert pytest.approx(expect_weight, abs=1e-1) == actual_weight, f"tile {tile} is wrong"
84
+
85
+
86
+ @pytest.mark.parametrize(
87
+ "module_class",
88
+ [
89
+ QuadPool2d,
90
+ MaxQuadPool2d,
91
+ AvgQuadPool2d,
92
+ ],
93
+ ids=["id", "max", "avg"],
94
+ )
95
+ def test_quad_pool2d_gradient(module_class: Type[nn.Module]) -> None:
96
+ poly = Polygon([(0.0, 0.0), (1.0, 0.0), (1.0, 1.1), (0.0, 1.0)])
97
+ exterior = (0.0, 0.0, 1.0, 1.0)
98
+
99
+ pool = module_class(4, poly, exterior, max_depth=5)
100
+ assert pool.weight.size() == torch.Size([pool.tiles.size(0), 4])
101
+
102
+ input = torch.rand((100, 2), dtype=torch.float64)
103
+ y = pool(input)
104
+
105
+ assert pool.weight.grad is None
106
+
107
+ loss_fn = L1Loss()
108
+ loss = loss_fn(y, torch.ones_like(y))
109
+ loss.backward()
110
+
111
+ assert pool.weight.grad is not None
112
+ assert pool.weight.grad.sum().item() == pytest.approx(-1)
113
+
114
+
115
+ def test_quad_pool2d_optimize() -> None:
116
+ poly = Polygon([(-180, -90), (-180, 90), (180, 90), (180, -90)])
117
+ pool = QuadPool2d(1, poly, (-180, -90, 360, 180), max_depth=1)
118
+
119
+ x_true = torch.tensor(
120
+ [[90.0, 45.0], [90.0, -45.0], [-90.0, -45.0], [-90.0, 45.0]], dtype=torch.float64
121
+ )
122
+ y_true = torch.tensor([[10.0], [20.0], [30.0], [40.0]], dtype=torch.float64)
123
+ y_tile = [(1, 1, 1), (1, 1, 0), (1, 0, 0), (1, 0, 1)]
124
+
125
+ optim = SGD(pool.parameters(), lr=0.01)
126
+ loss_fn = nn.L1Loss()
127
+
128
+ for i in range(20000):
129
+ optim.zero_grad()
130
+
131
+ y_pred = pool(x_true)
132
+ loss = loss_fn(y_pred, y_true)
133
+ loss.backward()
134
+
135
+ optim.step()
136
+
137
+ # Ensure that model converged with a small loss.
138
+ assert pytest.approx(0.0, abs=1e-1) == loss.detach().item()
139
+
140
+ actual_tiles = {}
141
+ for i in range(pool.tiles.size(0)):
142
+ tile = tuple(pool.tiles[i].detach().tolist())
143
+ actual_tiles[tile] = pool.weight[i, 0].detach().item()
144
+
145
+ for tile, expect_weight in zip(y_tile, y_true[:, 0].tolist()):
146
+ actual_weight = actual_tiles[tile]
147
+ assert pytest.approx(expect_weight, abs=1e-1) == actual_weight, f"tile {tile} is wrong"
File without changes
@@ -0,0 +1,18 @@
1
+ # SPDX-License-Identifier: GPL-3.0-or-later
2
+ # SPDX-FileCopyrightText: 2025 Yakau Bubnou
3
+ # SPDX-FileType: SOURCE
4
+
5
+ from typing import NamedTuple
6
+
7
+ from torch import Tensor
8
+
9
+
10
+ class quad_pool2d(NamedTuple):
11
+ tiles: Tensor
12
+ weight: Tensor
13
+ values: Tensor
14
+
15
+
16
+ class adaptive_quad_pool2d(NamedTuple):
17
+ weight: Tensor
18
+ values: Tensor
@@ -0,0 +1,101 @@
1
+ # SPDX-License-Identifier: GPL-3.0-or-later
2
+ # SPDX-FileCopyrightText: 2025 Yakau Bubnou
3
+ # SPDX-FileType: SOURCE
4
+
5
+ from __future__ import annotations
6
+
7
+ from collections import deque
8
+ from itertools import product
9
+ from typing import Iterator, NamedTuple, Tuple
10
+
11
+ from shapely.geometry import Polygon
12
+
13
+ __all__ = ["Exterior", "ExteriorTuple", "Tile", "regular_tiling"]
14
+
15
+
16
+ class Tile(NamedTuple):
17
+ z: int
18
+ x: int
19
+ y: int
20
+
21
+ @classmethod
22
+ def root(cls) -> Tile:
23
+ return cls(0, 0, 0)
24
+
25
+ def child(self, x: int, y: int) -> Tile:
26
+ return Tile(self.z + 1, self.x * 2 + x, self.y * 2 + y)
27
+
28
+ def children(self) -> Iterator[Tile]:
29
+ for x, y in product(range(2), range(2)):
30
+ yield self.child(x, y)
31
+
32
+
33
+ ExteriorTuple = Tuple[float, float, float, float]
34
+
35
+
36
+ class Exterior(NamedTuple):
37
+ xmin: float
38
+ ymin: float
39
+ width: float
40
+ height: float
41
+
42
+ @classmethod
43
+ def from_tuple(cls, exterior_tuple: ExteriorTuple) -> Exterior:
44
+ return cls(*exterior_tuple)
45
+
46
+ @property
47
+ def xmax(self) -> float:
48
+ return self.xmin + self.width
49
+
50
+ @property
51
+ def ymax(self) -> float:
52
+ return self.ymin + self.height
53
+
54
+ def slice(self, tile: Tile) -> Exterior:
55
+ w = self.width / (1 << tile.z)
56
+ h = self.height / (1 << tile.z)
57
+ return Exterior(self.xmin + tile.x * w, self.ymin + tile.y * h, w, h)
58
+
59
+ def as_polygon(self) -> Polygon:
60
+ return Polygon(
61
+ [
62
+ (self.xmin, self.ymin),
63
+ (self.xmax, self.ymin),
64
+ (self.xmax, self.ymax),
65
+ (self.xmin, self.ymax),
66
+ ]
67
+ )
68
+
69
+
70
+ def regular_tiling(
71
+ polygon: Polygon, exterior: Exterior, z: int, internal: bool = False
72
+ ) -> Iterator[Tile]:
73
+ """Returns a regular quad-tiling (tiles of the same size).
74
+
75
+ Method returns all tiles of level (z) that have a common intersection with a specified
76
+ polygon.
77
+
78
+ Args:
79
+ polygon: A polygon to cover with tiles.
80
+ exterior: Exterior (bounding box) of the quadtree. For example, for geospatial
81
+ coordinates, this will be `(-180.0, -90.0, 360.0, 180.0)`.
82
+ z: Zoom level of the tiles.
83
+ internal: When `True`, returns internal tiles (nodes) of the quadtree up to a root
84
+ tile (0,0,0).
85
+
86
+ Returns:
87
+ Iterator of tiles.
88
+ """
89
+ queue = deque([Tile.root()])
90
+
91
+ while len(queue) > 0:
92
+ tile = queue.pop()
93
+
94
+ tile_poly = exterior.slice(tile).as_polygon()
95
+ if not tile_poly.intersects(polygon):
96
+ continue
97
+
98
+ if internal or tile.z >= z:
99
+ yield tile
100
+ if tile.z < z:
101
+ queue.extend(tile.children())