gridpoints 1.0.7__tar.gz → 1.1.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.
- {gridpoints-1.0.7 → gridpoints-1.1.0}/PKG-INFO +13 -9
- {gridpoints-1.0.7 → gridpoints-1.1.0}/README.md +10 -6
- {gridpoints-1.0.7 → gridpoints-1.1.0}/pyproject.toml +3 -3
- gridpoints-1.1.0/src/griddemo/__init__.py +10 -0
- {gridpoints-1.0.7 → gridpoints-1.1.0}/src/gridpoints/__init__.py +13 -11
- gridpoints-1.1.0/src/gridpoints/array_libs.py +183 -0
- {gridpoints-1.0.7 → gridpoints-1.1.0}/src/gridpoints/diagonal_tricks.py +106 -94
- gridpoints-1.1.0/src/gridpoints/kdtree.py +176 -0
- gridpoints-1.1.0/src/gridpoints/mlg_sort.py +289 -0
- gridpoints-1.1.0/src/gridpoints/padding.py +48 -0
- {gridpoints-1.0.7 → gridpoints-1.1.0}/src/gridpoints.egg-info/PKG-INFO +13 -9
- {gridpoints-1.0.7 → gridpoints-1.1.0}/src/gridpoints.egg-info/SOURCES.txt +1 -0
- gridpoints-1.0.7/src/griddemo/__init__.py +0 -10
- gridpoints-1.0.7/src/gridpoints/array_libs.py +0 -175
- gridpoints-1.0.7/src/gridpoints/kdtree.py +0 -203
- gridpoints-1.0.7/src/gridpoints/mlg_sort.py +0 -280
- {gridpoints-1.0.7 → gridpoints-1.1.0}/LICENSE +0 -0
- {gridpoints-1.0.7 → gridpoints-1.1.0}/setup.cfg +0 -0
- {gridpoints-1.0.7 → gridpoints-1.1.0}/src/griddemo/quick_demo.py +0 -0
- {gridpoints-1.0.7 → gridpoints-1.1.0}/src/gridpoints.egg-info/dependency_links.txt +0 -0
- {gridpoints-1.0.7 → gridpoints-1.1.0}/src/gridpoints.egg-info/requires.txt +0 -0
- {gridpoints-1.0.7 → gridpoints-1.1.0}/src/gridpoints.egg-info/top_level.txt +0 -0
|
@@ -1,11 +1,11 @@
|
|
|
1
1
|
Metadata-Version: 2.4
|
|
2
2
|
Name: gridpoints
|
|
3
|
-
Version: 1.0
|
|
3
|
+
Version: 1.1.0
|
|
4
4
|
Summary: fast multidimensional sorting tool for point clouds using the monotonic lagrangian grid
|
|
5
5
|
Author: Armand de Cacqueray
|
|
6
6
|
License: MIT
|
|
7
|
-
Project-URL: Homepage, https://github.
|
|
8
|
-
Project-URL: Repository, https://github.com/Neighborhood-Grid/
|
|
7
|
+
Project-URL: Homepage, https://neighborhood-grid.github.io/gridpoints/
|
|
8
|
+
Project-URL: Repository, https://github.com/Neighborhood-Grid/gridpoints
|
|
9
9
|
Keywords: sorting,point-clouds,multi-indexing,grid,neighborhood-grid,monotonic-lagrangian-grid,squarenet,nearest-neighbor
|
|
10
10
|
Requires-Python: >=3.9
|
|
11
11
|
Description-Content-Type: text/markdown
|
|
@@ -14,7 +14,8 @@ Provides-Extra: demo
|
|
|
14
14
|
Requires-Dist: squarenet; extra == "demo"
|
|
15
15
|
Dynamic: license-file
|
|
16
16
|
|
|
17
|
-
## ❒
|
|
17
|
+
## ❒ CubeNet: Multi dimensional sort for point clouds
|
|
18
|
+
# with the gridpoints python package
|
|
18
19
|
|
|
19
20
|
`Gridpoints` maps raw indexes of unstructured **point clouds** to structured grids multi-index through a **bijective transformation**:
|
|
20
21
|
One point, one cell,
|
|
@@ -37,8 +38,8 @@ z[i, j, k+1] >= z[i, j, k]
|
|
|
37
38
|
```text
|
|
38
39
|
neighborhood(Pgrid[i, j, k]) = {Pgrid[i±di, j±dj, k±dk] | (di, dj, dk) ≤ R},
|
|
39
40
|
````
|
|
40
|
-
where `R` is a radius cutoff to determine, allowing local operations in
|
|
41
|
-
→ Geometric operations (convolution, clustering, interpolation …) can then be applied **in
|
|
41
|
+
where `R` is a radius cutoff to determine, allowing local operations in O(N) time.
|
|
42
|
+
→ Geometric operations (convolution, clustering, interpolation …) can then be applied **in O(N) time** directly on the `Pgrid` view instead of relying on complex graph or point-cloud methods
|
|
42
43
|
|
|
43
44
|
`P` can be a NumPy, PyTorch or CuPy array of any dimension (N, D). To allow natural padding when the grid has more cells than they are points in the cloud, NaNs and Infs are supported in a consistent manner:
|
|
44
45
|
- nans → random position
|
|
@@ -46,12 +47,15 @@ where `R` is a radius cutoff to determine, allowing local operations in linear t
|
|
|
46
47
|
|
|
47
48
|
This allow to gridsort prime or variable number of points, as long as one is ready to deal with void/special grid cells.
|
|
48
49
|
|
|
49
|
-
Expected runtime for sorting 1 million points
|
|
50
|
+
Expected runtime for sorting 1 million points, tested on google colab GPU T4 for 2D/3D point clouds:
|
|
51
|
+
- Numpy: 10s
|
|
52
|
+
- Torch (gpu): 500 ms
|
|
53
|
+
- Cupy (gpu): 1 s (after compilation e.g. not on the first cold call, which take 10 s)
|
|
50
54
|
|
|
51
55
|
`When not to use gridpoints ?`
|
|
52
56
|
- high dimension: the package is implemented to support arbitrary dimension, but sweetspot is really 2D/3D. dimensions 4-6 might still be reasonable depending on the task, but anything above 8D is generally too high dimensional for gridpoints.
|
|
53
57
|
- weird geometries. Supported datasets goes beyond smooth convex distributions: map of Indonesia, a sponge, a donuts, an elephant, an eggshell (by specifying a suitable gridshape e.g. (128,128,2) to `gridpoints.sort()`). You can see various exemples in the [plot](https://github.com/Neighborhood-Grid/CubeNet/blob/main/plots) folder But with some limits. Bad fits: a spider web, a wind turbine, same eggshell with a naive 3d sort (gridshape = (32,32,32)) ... the issue is not that gridpoints cannot sort these distributions, but that it will produce a poor representation of the geometry.
|
|
54
|
-
- small point clouds: beyond a few hundred points, local operations in
|
|
58
|
+
- small point clouds: beyond a few hundred points, local operations in O(N) time is not worth the overhead, because naive quadratic implementations will probably be simultaneously simpler and faster.
|
|
55
59
|
|
|
56
60
|
---
|
|
57
61
|
|
|
@@ -109,4 +113,4 @@ The typical use-case of `Gridpoints` is to allow fast local operations on arbitr
|
|
|
109
113
|
```text
|
|
110
114
|
output(i, j, k) = f( Pgrid[i±di, j±dj, k±dk] | di, dj, dk in local window )
|
|
111
115
|
```
|
|
112
|
-
To go beyond standard (slow) python loops, this kind of kernel computation can be accelerated with native grid convolution operations of standard libraries whenever possible, or with `pystencils` or `taichi` compilers for complex/non
|
|
116
|
+
To go beyond standard (slow) python loops, this kind of kernel computation can be accelerated with native grid convolution operations of standard libraries whenever possible, or with `pystencils` or `taichi` compilers for complex/non O(N) grid kernels. On GPU, tiling the grid and leveraging a custom `triton` kernel might be particularly efficient.
|
|
@@ -1,4 +1,5 @@
|
|
|
1
|
-
## ❒
|
|
1
|
+
## ❒ CubeNet: Multi dimensional sort for point clouds
|
|
2
|
+
# with the gridpoints python package
|
|
2
3
|
|
|
3
4
|
`Gridpoints` maps raw indexes of unstructured **point clouds** to structured grids multi-index through a **bijective transformation**:
|
|
4
5
|
One point, one cell,
|
|
@@ -21,8 +22,8 @@ z[i, j, k+1] >= z[i, j, k]
|
|
|
21
22
|
```text
|
|
22
23
|
neighborhood(Pgrid[i, j, k]) = {Pgrid[i±di, j±dj, k±dk] | (di, dj, dk) ≤ R},
|
|
23
24
|
````
|
|
24
|
-
where `R` is a radius cutoff to determine, allowing local operations in
|
|
25
|
-
→ Geometric operations (convolution, clustering, interpolation …) can then be applied **in
|
|
25
|
+
where `R` is a radius cutoff to determine, allowing local operations in O(N) time.
|
|
26
|
+
→ Geometric operations (convolution, clustering, interpolation …) can then be applied **in O(N) time** directly on the `Pgrid` view instead of relying on complex graph or point-cloud methods
|
|
26
27
|
|
|
27
28
|
`P` can be a NumPy, PyTorch or CuPy array of any dimension (N, D). To allow natural padding when the grid has more cells than they are points in the cloud, NaNs and Infs are supported in a consistent manner:
|
|
28
29
|
- nans → random position
|
|
@@ -30,12 +31,15 @@ where `R` is a radius cutoff to determine, allowing local operations in linear t
|
|
|
30
31
|
|
|
31
32
|
This allow to gridsort prime or variable number of points, as long as one is ready to deal with void/special grid cells.
|
|
32
33
|
|
|
33
|
-
Expected runtime for sorting 1 million points
|
|
34
|
+
Expected runtime for sorting 1 million points, tested on google colab GPU T4 for 2D/3D point clouds:
|
|
35
|
+
- Numpy: 10s
|
|
36
|
+
- Torch (gpu): 500 ms
|
|
37
|
+
- Cupy (gpu): 1 s (after compilation e.g. not on the first cold call, which take 10 s)
|
|
34
38
|
|
|
35
39
|
`When not to use gridpoints ?`
|
|
36
40
|
- high dimension: the package is implemented to support arbitrary dimension, but sweetspot is really 2D/3D. dimensions 4-6 might still be reasonable depending on the task, but anything above 8D is generally too high dimensional for gridpoints.
|
|
37
41
|
- weird geometries. Supported datasets goes beyond smooth convex distributions: map of Indonesia, a sponge, a donuts, an elephant, an eggshell (by specifying a suitable gridshape e.g. (128,128,2) to `gridpoints.sort()`). You can see various exemples in the [plot](https://github.com/Neighborhood-Grid/CubeNet/blob/main/plots) folder But with some limits. Bad fits: a spider web, a wind turbine, same eggshell with a naive 3d sort (gridshape = (32,32,32)) ... the issue is not that gridpoints cannot sort these distributions, but that it will produce a poor representation of the geometry.
|
|
38
|
-
- small point clouds: beyond a few hundred points, local operations in
|
|
42
|
+
- small point clouds: beyond a few hundred points, local operations in O(N) time is not worth the overhead, because naive quadratic implementations will probably be simultaneously simpler and faster.
|
|
39
43
|
|
|
40
44
|
---
|
|
41
45
|
|
|
@@ -93,4 +97,4 @@ The typical use-case of `Gridpoints` is to allow fast local operations on arbitr
|
|
|
93
97
|
```text
|
|
94
98
|
output(i, j, k) = f( Pgrid[i±di, j±dj, k±dk] | di, dj, dk in local window )
|
|
95
99
|
```
|
|
96
|
-
To go beyond standard (slow) python loops, this kind of kernel computation can be accelerated with native grid convolution operations of standard libraries whenever possible, or with `pystencils` or `taichi` compilers for complex/non
|
|
100
|
+
To go beyond standard (slow) python loops, this kind of kernel computation can be accelerated with native grid convolution operations of standard libraries whenever possible, or with `pystencils` or `taichi` compilers for complex/non O(N) grid kernels. On GPU, tiling the grid and leveraging a custom `triton` kernel might be particularly efficient.
|
|
@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
|
|
|
4
4
|
|
|
5
5
|
[project]
|
|
6
6
|
name = "gridpoints"
|
|
7
|
-
version = "1.0
|
|
7
|
+
version = "1.1.0"
|
|
8
8
|
description = "fast multidimensional sorting tool for point clouds using the monotonic lagrangian grid "
|
|
9
9
|
readme = "README.md"
|
|
10
10
|
requires-python = ">=3.9"
|
|
@@ -26,8 +26,8 @@ demo = [
|
|
|
26
26
|
|
|
27
27
|
|
|
28
28
|
[project.urls]
|
|
29
|
-
Homepage = "https://github.
|
|
30
|
-
Repository = "https://github.com/Neighborhood-Grid/
|
|
29
|
+
Homepage = "https://neighborhood-grid.github.io/gridpoints/"
|
|
30
|
+
Repository = "https://github.com/Neighborhood-Grid/gridpoints"
|
|
31
31
|
|
|
32
32
|
[tool.setuptools.packages.find]
|
|
33
33
|
where = ["src"]
|
|
@@ -1,11 +1,13 @@
|
|
|
1
|
-
"""Monotonic Lagrangian sorting for NumPy, CuPy, and PyTorch arrays."""
|
|
2
|
-
|
|
3
|
-
from .mlg_sort import argsort, sort, invert_permutation
|
|
4
|
-
|
|
5
|
-
|
|
6
|
-
|
|
7
|
-
"
|
|
8
|
-
"
|
|
9
|
-
|
|
10
|
-
|
|
11
|
-
|
|
1
|
+
"""Monotonic Lagrangian sorting for NumPy, CuPy, and PyTorch arrays."""
|
|
2
|
+
|
|
3
|
+
from .mlg_sort import argsort, sort, invert_permutation
|
|
4
|
+
from .padding import pad
|
|
5
|
+
|
|
6
|
+
__all__ = (
|
|
7
|
+
"argsort",
|
|
8
|
+
"sort",
|
|
9
|
+
"invert_permutation",
|
|
10
|
+
"pad",
|
|
11
|
+
)
|
|
12
|
+
|
|
13
|
+
__version__ = "1.1.0"
|
|
@@ -0,0 +1,183 @@
|
|
|
1
|
+
"""
|
|
2
|
+
Helpers:
|
|
3
|
+
Unified array syntax to support numpy, cupy and pytorch arrays
|
|
4
|
+
|
|
5
|
+
Here we just define basic primitives that we will need later
|
|
6
|
+
like arange, copy, shuffle etc.
|
|
7
|
+
"""
|
|
8
|
+
|
|
9
|
+
|
|
10
|
+
def precision(N):
|
|
11
|
+
"""Names of the (float, int) dtypes suited to point clouds of N points.
|
|
12
|
+
|
|
13
|
+
- float32 resolves the gaps (~1/N) between neighbouring points only up to
|
|
14
|
+
a few 10**4 points, float64 is used beyond.
|
|
15
|
+
- the diagonal sorts store values up to 6N: int32 is safe up to N ~ 3.5e8.
|
|
16
|
+
"""
|
|
17
|
+
return ('float32' if N <= 2**16 else 'float64'), ('int32' if 6 * N < 2**31 else 'int64')
|
|
18
|
+
|
|
19
|
+
|
|
20
|
+
def get_dtype(dtp, xp):
|
|
21
|
+
if dtp == 'bool':
|
|
22
|
+
return getattr(xp, 'bool_', None) or xp.bool
|
|
23
|
+
return getattr(xp, dtp) if isinstance(dtp, str) else dtp
|
|
24
|
+
|
|
25
|
+
|
|
26
|
+
def get_backend(X):
|
|
27
|
+
module = type(X).__module__.split('.')[0]
|
|
28
|
+
if module == 'torch':
|
|
29
|
+
import torch
|
|
30
|
+
return torch, 'torch'
|
|
31
|
+
elif module == 'cupy':
|
|
32
|
+
import cupy
|
|
33
|
+
return cupy, 'cupy'
|
|
34
|
+
else:
|
|
35
|
+
import numpy
|
|
36
|
+
return numpy, 'numpy'
|
|
37
|
+
|
|
38
|
+
def astype(X, dtype_str, reference_array=None):
|
|
39
|
+
xp, backend_name = get_backend(X)
|
|
40
|
+
dtype = get_dtype(dtype_str, xp)
|
|
41
|
+
if backend_name == 'torch':
|
|
42
|
+
device = reference_array.device if reference_array is not None else X.device
|
|
43
|
+
return X.to(dtype=dtype, device=device)
|
|
44
|
+
return X.astype(dtype, copy=False)
|
|
45
|
+
|
|
46
|
+
|
|
47
|
+
def inplace_sort(X, axis):
|
|
48
|
+
xp, backend_name = get_backend(X)
|
|
49
|
+
if backend_name == 'torch':
|
|
50
|
+
X.copy_(X.sort(dim=axis)[0])
|
|
51
|
+
elif backend_name == "cupy":
|
|
52
|
+
X[...] = xp.sort(xp.ascontiguousarray(X), axis=axis)
|
|
53
|
+
else:
|
|
54
|
+
X.sort(axis=axis)
|
|
55
|
+
|
|
56
|
+
|
|
57
|
+
def sort(X, axis):
|
|
58
|
+
xp, backend_name = get_backend(X)
|
|
59
|
+
return X.sort(dim=axis)[0] if backend_name == 'torch' else xp.sort(X, axis=axis)
|
|
60
|
+
|
|
61
|
+
|
|
62
|
+
def _alloc(fn, shape, dtype_str, reference_array):
|
|
63
|
+
xp, backend_name = get_backend(reference_array)
|
|
64
|
+
dtype = get_dtype(dtype_str, xp)
|
|
65
|
+
if backend_name == 'torch':
|
|
66
|
+
return getattr(xp, fn)(shape, dtype=dtype, device=reference_array.device)
|
|
67
|
+
return getattr(xp, fn)(shape, dtype=dtype)
|
|
68
|
+
|
|
69
|
+
def empty(shape, dtype_str, reference_array):
|
|
70
|
+
return _alloc('empty', shape, dtype_str, reference_array)
|
|
71
|
+
|
|
72
|
+
def zeros(shape, dtype_str, reference_array):
|
|
73
|
+
return _alloc('zeros', shape, dtype_str, reference_array)
|
|
74
|
+
|
|
75
|
+
def ones(shape, dtype_str, reference_array):
|
|
76
|
+
return _alloc('ones', shape, dtype_str, reference_array)
|
|
77
|
+
|
|
78
|
+
|
|
79
|
+
def arange(N, dtype_str, reference_array):
|
|
80
|
+
xp, backend_name = get_backend(reference_array)
|
|
81
|
+
dtype = get_dtype(dtype_str, xp)
|
|
82
|
+
if backend_name == 'torch':
|
|
83
|
+
return xp.arange(N, dtype=dtype, device=reference_array.device)
|
|
84
|
+
return xp.arange(N, dtype=dtype)
|
|
85
|
+
|
|
86
|
+
def random(shape, dtype_str, reference_array):
|
|
87
|
+
"""Uniform random numbers in [0, 1)."""
|
|
88
|
+
xp, backend_name = get_backend(reference_array)
|
|
89
|
+
dtype = get_dtype(dtype_str, xp)
|
|
90
|
+
if backend_name == 'torch':
|
|
91
|
+
return xp.rand(shape, dtype=dtype, device=reference_array.device)
|
|
92
|
+
elif backend_name == 'cupy':
|
|
93
|
+
return xp.random.rand(*shape).astype(dtype)
|
|
94
|
+
return xp.random.default_rng(42).random(shape, dtype=dtype)
|
|
95
|
+
|
|
96
|
+
def copy(arr):
|
|
97
|
+
_, backend_name = get_backend(arr)
|
|
98
|
+
if backend_name == 'torch':
|
|
99
|
+
return arr.clone()
|
|
100
|
+
return arr.copy()
|
|
101
|
+
|
|
102
|
+
def amax(X, axis):
|
|
103
|
+
return X.amax(axis) if get_backend(X)[1] == 'torch' else X.max(axis)
|
|
104
|
+
|
|
105
|
+
def amin(X, axis):
|
|
106
|
+
return X.amin(axis) if get_backend(X)[1] == 'torch' else X.min(axis)
|
|
107
|
+
|
|
108
|
+
def all_true(tensor):
|
|
109
|
+
return bool(tensor.all())
|
|
110
|
+
|
|
111
|
+
def any_true(tensor):
|
|
112
|
+
return bool(tensor.any())
|
|
113
|
+
|
|
114
|
+
|
|
115
|
+
def nan_infs_to_num_and_tiebreaker(arr):
|
|
116
|
+
"""Make a batch of clouds (B, N, D) finite, and break ties with a tiny noise.
|
|
117
|
+
|
|
118
|
+
+inf / -inf are put just beyond the finite points of their cloud and axis
|
|
119
|
+
(they end up on the border of the grid), NaN are replaced by the value of a
|
|
120
|
+
random finite point of their cloud and axis (they end up anywhere).
|
|
121
|
+
"""
|
|
122
|
+
xp, _ = get_backend(arr)
|
|
123
|
+
out = arr
|
|
124
|
+
finite = xp.isfinite(arr)
|
|
125
|
+
if not all_true(finite):
|
|
126
|
+
hi = amax(xp.where(finite, arr, -xp.inf), 1)[:, None, :]
|
|
127
|
+
lo = amin(xp.where(finite, arr, xp.inf), 1)[:, None, :]
|
|
128
|
+
out = xp.where(xp.isposinf(arr), hi + 1, arr)
|
|
129
|
+
out = xp.where(xp.isneginf(arr), lo - 1, out)
|
|
130
|
+
nan = xp.isnan(arr)
|
|
131
|
+
if any_true(nan):
|
|
132
|
+
finite_first = sort(xp.where(finite, arr, xp.inf), 1)
|
|
133
|
+
n_finite = finite.sum(1)[:, None, :]
|
|
134
|
+
pick = astype(random(arr.shape, arr.dtype, arr) * n_finite, 'int64')
|
|
135
|
+
out = xp.where(nan, take_along_axis(finite_first, pick, 1), out)
|
|
136
|
+
eps = 1e-5 if '32' in str(arr.dtype) else 1e-10
|
|
137
|
+
return out + eps * random(out.shape, arr.dtype, out)
|
|
138
|
+
|
|
139
|
+
def flip(X, axis=None):
|
|
140
|
+
xp, _ = get_backend(X)
|
|
141
|
+
if axis is None or axis == () or axis == []:
|
|
142
|
+
return X
|
|
143
|
+
|
|
144
|
+
if isinstance(axis, int):
|
|
145
|
+
axis = (axis,)
|
|
146
|
+
else:
|
|
147
|
+
axis = tuple(axis)
|
|
148
|
+
|
|
149
|
+
if hasattr(X, "flip") and callable(X.flip):
|
|
150
|
+
return X.flip(axis)
|
|
151
|
+
|
|
152
|
+
return xp.flip(X, axis=axis)
|
|
153
|
+
|
|
154
|
+
def argsort(X):
|
|
155
|
+
xp, backend_name = get_backend(X)
|
|
156
|
+
return xp.argsort(X, dim=-1) if backend_name == 'torch' else xp.argsort(X, axis=-1)
|
|
157
|
+
|
|
158
|
+
def argpartition(X, mid):
|
|
159
|
+
"""Indices that put the `mid` smallest values of the last axis first."""
|
|
160
|
+
xp, backend_name = get_backend(X)
|
|
161
|
+
if backend_name == 'torch':
|
|
162
|
+
return xp.argsort(X, dim=-1)
|
|
163
|
+
else:
|
|
164
|
+
return xp.argpartition(X, mid - 1, axis=-1)
|
|
165
|
+
|
|
166
|
+
def take_along_axis(arr, idx, axis):
|
|
167
|
+
xp, backend_name = get_backend(arr)
|
|
168
|
+
if backend_name == 'torch':
|
|
169
|
+
return xp.take_along_dim(arr, idx.long(), dim=axis)
|
|
170
|
+
else:
|
|
171
|
+
return xp.take_along_axis(arr, idx, axis=axis)
|
|
172
|
+
|
|
173
|
+
def concatenate(arrays, axis = 0):
|
|
174
|
+
xp, name = get_backend(arrays[0])
|
|
175
|
+
if name == "torch":
|
|
176
|
+
return xp.cat(list(arrays), dim=axis)
|
|
177
|
+
return xp.concatenate(arrays, axis=axis)
|
|
178
|
+
|
|
179
|
+
def transpose(X, axes):
|
|
180
|
+
_, name = get_backend(X)
|
|
181
|
+
if name == "torch":
|
|
182
|
+
return X.permute(*axes)
|
|
183
|
+
return X.transpose(axes)
|
|
@@ -1,94 +1,106 @@
|
|
|
1
|
-
"""
|
|
2
|
-
Helpers:
|
|
3
|
-
Diagonal trick is based on a tensor rotation that allow
|
|
4
|
-
to sort arrays efficiently along arbitrary diagonal fibers.
|
|
5
|
-
|
|
6
|
-
|
|
7
|
-
|
|
8
|
-
|
|
9
|
-
|
|
10
|
-
|
|
11
|
-
|
|
12
|
-
|
|
13
|
-
|
|
14
|
-
|
|
15
|
-
|
|
16
|
-
|
|
17
|
-
|
|
18
|
-
|
|
19
|
-
|
|
20
|
-
|
|
21
|
-
|
|
22
|
-
|
|
23
|
-
|
|
24
|
-
|
|
25
|
-
|
|
26
|
-
|
|
27
|
-
|
|
28
|
-
|
|
29
|
-
|
|
30
|
-
|
|
31
|
-
|
|
32
|
-
|
|
33
|
-
|
|
34
|
-
|
|
35
|
-
|
|
36
|
-
|
|
37
|
-
|
|
38
|
-
|
|
39
|
-
|
|
40
|
-
|
|
41
|
-
|
|
42
|
-
|
|
43
|
-
|
|
44
|
-
|
|
45
|
-
|
|
46
|
-
|
|
47
|
-
|
|
48
|
-
|
|
49
|
-
|
|
50
|
-
|
|
51
|
-
|
|
52
|
-
|
|
53
|
-
|
|
54
|
-
|
|
55
|
-
|
|
56
|
-
|
|
57
|
-
|
|
58
|
-
|
|
59
|
-
|
|
60
|
-
|
|
61
|
-
|
|
62
|
-
|
|
63
|
-
|
|
64
|
-
|
|
65
|
-
|
|
66
|
-
|
|
67
|
-
|
|
68
|
-
|
|
69
|
-
|
|
70
|
-
|
|
71
|
-
|
|
72
|
-
|
|
73
|
-
|
|
74
|
-
|
|
75
|
-
|
|
76
|
-
|
|
77
|
-
|
|
78
|
-
|
|
79
|
-
|
|
80
|
-
|
|
81
|
-
|
|
82
|
-
|
|
83
|
-
|
|
84
|
-
|
|
85
|
-
|
|
86
|
-
|
|
87
|
-
|
|
88
|
-
|
|
89
|
-
|
|
90
|
-
|
|
91
|
-
|
|
92
|
-
|
|
93
|
-
|
|
94
|
-
|
|
1
|
+
"""
|
|
2
|
+
Helpers:
|
|
3
|
+
Diagonal trick is based on a tensor rotation that allow
|
|
4
|
+
to sort arrays efficiently along arbitrary diagonal fibers.
|
|
5
|
+
|
|
6
|
+
All functions work on a batch of grids X of shape (B, ...), and tell for each
|
|
7
|
+
grid of the batch whether it was already sorted.
|
|
8
|
+
"""
|
|
9
|
+
|
|
10
|
+
import math
|
|
11
|
+
|
|
12
|
+
from .array_libs import (
|
|
13
|
+
get_backend,
|
|
14
|
+
arange,
|
|
15
|
+
inplace_sort,
|
|
16
|
+
)
|
|
17
|
+
|
|
18
|
+
|
|
19
|
+
def sorted_flags(X, axis):
|
|
20
|
+
"""Bool array of shape (B,): is each grid of the batch X sorted along `axis`?"""
|
|
21
|
+
axis %= X.ndim
|
|
22
|
+
sl = lambda s: (slice(None),) * axis + (s,)
|
|
23
|
+
ok = X[sl(slice(1, None))] >= X[sl(slice(None, -1))]
|
|
24
|
+
return ok.reshape(len(X), math.prod(ok.shape[1:])).all(1)
|
|
25
|
+
|
|
26
|
+
|
|
27
|
+
def sort_diag_trick(X):
|
|
28
|
+
"""
|
|
29
|
+
we rotate the tensor along the main diagonal
|
|
30
|
+
to place diagonals on the last axis, then
|
|
31
|
+
perform a standard sort(axis = -1), and
|
|
32
|
+
rotate back the diagonals
|
|
33
|
+
"""
|
|
34
|
+
xp, _ = get_backend(X)
|
|
35
|
+
#bigM for safety, protect
|
|
36
|
+
#overlapping diagonals
|
|
37
|
+
bigM = (X.max() - X.min()) + 1
|
|
38
|
+
|
|
39
|
+
n, m = X.shape[-2:]
|
|
40
|
+
|
|
41
|
+
#move the smallest axis to last position
|
|
42
|
+
if m > n:
|
|
43
|
+
X = xp.moveaxis(X, [-1, -2], [-2, -1])
|
|
44
|
+
n, m = m, n #now n >= m
|
|
45
|
+
|
|
46
|
+
i = arange(n, X.dtype, X)[:, None]
|
|
47
|
+
j = arange(m, X.dtype, X)[None, :]
|
|
48
|
+
|
|
49
|
+
di = (i - j)
|
|
50
|
+
i_ = di % n
|
|
51
|
+
ires = di // n
|
|
52
|
+
groupkey = bigM * ires
|
|
53
|
+
|
|
54
|
+
X[..., i_, j] = X[..., i, j] - groupkey #Rotation
|
|
55
|
+
# ========================
|
|
56
|
+
# here is the trick: as we rotated the tensor
|
|
57
|
+
# sorting on the diagonal is just working on
|
|
58
|
+
# axis -1 of the rotated view
|
|
59
|
+
# ========================
|
|
60
|
+
was_sorted = sorted_flags(X, -1)
|
|
61
|
+
inplace_sort(X, axis=-1)
|
|
62
|
+
X[..., i, j] = X[..., i_, j] + groupkey
|
|
63
|
+
|
|
64
|
+
return was_sorted
|
|
65
|
+
|
|
66
|
+
def sort_tridiag_trick(X):
|
|
67
|
+
"""
|
|
68
|
+
we rotate the tensor along the tridiagonal,
|
|
69
|
+
then perform a standard sort(axis = -1),
|
|
70
|
+
and rotate back the tridiagonals
|
|
71
|
+
"""
|
|
72
|
+
xp, _ = get_backend(X)
|
|
73
|
+
#bigM for safety, protect
|
|
74
|
+
#overlapping diagonals
|
|
75
|
+
bigM = (X.max() - X.min()) + 1
|
|
76
|
+
|
|
77
|
+
#move the smallest axes to last positions
|
|
78
|
+
abc = X.shape
|
|
79
|
+
dst_axes = (-3, -2, -1)
|
|
80
|
+
src_axes = tuple(sorted(dst_axes, key=lambda ax: -abc[ax]))
|
|
81
|
+
X = xp.moveaxis(X, src_axes, dst_axes)
|
|
82
|
+
a, b, c = X.shape[-3:] # now a >= b >= c
|
|
83
|
+
|
|
84
|
+
i = arange(a, X.dtype, X)[:, None, None]
|
|
85
|
+
j = arange(b, X.dtype, X)[None, :, None]
|
|
86
|
+
k = arange(c, X.dtype, X)[None, None, :]
|
|
87
|
+
|
|
88
|
+
di = i - k
|
|
89
|
+
dj = j - k
|
|
90
|
+
di_ = di % a
|
|
91
|
+
dj_ = dj % b
|
|
92
|
+
dires = di // a
|
|
93
|
+
djres = dj // b
|
|
94
|
+
groupkey = bigM * (dires * 4 + djres)
|
|
95
|
+
|
|
96
|
+
X[..., di_, dj_, k] = X[..., i, j, k] - groupkey #Trirotation
|
|
97
|
+
# ========================
|
|
98
|
+
# Again, thanks to the rotation trick
|
|
99
|
+
# we just need to .sort(X, axis = -1)
|
|
100
|
+
# to sort the tridiagonal
|
|
101
|
+
# ========================
|
|
102
|
+
was_sorted = sorted_flags(X, -1)
|
|
103
|
+
inplace_sort(X, axis=-1)
|
|
104
|
+
X[..., i, j, k] = X[..., di_, dj_, k] + groupkey
|
|
105
|
+
|
|
106
|
+
return was_sorted
|