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.
@@ -1,11 +1,11 @@
1
1
  Metadata-Version: 2.4
2
2
  Name: gridpoints
3
- Version: 1.0.7
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.com/Neighborhood-Grid/CubeNet
8
- Project-URL: Repository, https://github.com/Neighborhood-Grid/CubeNet
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
- ## ❒ Gridpoints — Multi dimensional sort for point clouds
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 linear time.
41
- → Geometric operations (convolution, clustering, interpolation …) can then be applied **in linear time** directly on the `Pgrid` view instead of relying on complex graph or point-cloud methods
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: CPU: < 10s (numpy), GPU: < 500 ms (torch cuda)
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 linear time is not worth the overhead, because naive quadratic implementations will probably be simultaneously simpler and faster.
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 linear grid kernels. On GPU, tiling the grid and leveraging a custom `triton` kernel might be particularly efficient.
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
- ## ❒ Gridpoints — Multi dimensional sort for point clouds
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 linear time.
25
- → Geometric operations (convolution, clustering, interpolation …) can then be applied **in linear time** directly on the `Pgrid` view instead of relying on complex graph or point-cloud methods
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: CPU: < 10s (numpy), GPU: < 500 ms (torch cuda)
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 linear time is not worth the overhead, because naive quadratic implementations will probably be simultaneously simpler and faster.
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 linear grid kernels. On GPU, tiling the grid and leveraging a custom `triton` kernel might be particularly efficient.
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"
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.com/Neighborhood-Grid/CubeNet"
30
- Repository = "https://github.com/Neighborhood-Grid/CubeNet"
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"]
@@ -0,0 +1,10 @@
1
+
2
+ from .quick_demo import sample, show_result, print_zoom
3
+
4
+ __all__ = (
5
+ "sample",
6
+ "show_result",
7
+ "print_zoom",
8
+ )
9
+
10
+ __version__ = "1.0.8"
@@ -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
- __all__ = (
6
- "argsort",
7
- "sort",
8
- "invert_permutation",
9
- )
10
-
11
- __version__ = "1.0.6"
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
- from .array_libs import (
9
- get_backend,
10
- arange,
11
- inplace_sort,
12
- all_true
13
- )
14
-
15
- def sort_diag_trick(X):
16
- """
17
- we rotate the tensor along the main diagonal
18
- to place diagonals on the last axis, then
19
- perform a standard sort(axis = -1), and
20
- rotate back the diagonals
21
- """
22
- xp, _ = get_backend(X)
23
- #bigM for safety, protect
24
- #overlapping diagonals
25
- bigM = (X.max() - X.min()) + 1
26
-
27
- n, m = X.shape[-2:]
28
-
29
- #move the smallest axis to last position
30
- if m > n:
31
- X = xp.moveaxis(X, [-1, -2], [-2, -1])
32
- n, m = m, n #now n >= m
33
-
34
- i = arange(n, 'int', X)[:, None]
35
- j = arange(m, 'int', X)[None, :]
36
-
37
- di = (i - j)
38
- i_ = di % n
39
- ires = di // n
40
- groupkey = bigM * ires
41
-
42
- X[..., i_, j] = X[..., i, j] - groupkey #Rotation
43
- # ========================
44
- # here is the trick: as we rotated the tensor
45
- # sorting on the diagonal is just working on
46
- # axis -1 of the rotated view
47
- # ========================
48
- was_sorted = all_true(xp.diff(X, axis=-1) >= 0)
49
- inplace_sort(X, axis=-1)
50
- X[..., i, j] = X[..., i_, j] + groupkey
51
-
52
- return was_sorted
53
-
54
- def sort_tridiag_trick(X):
55
- """
56
- we rotate the tensor along the tridiagonal,
57
- then perform a standard sort(axis = -1),
58
- and rotate back the tridiagonals
59
- """
60
- xp, _ = get_backend(X)
61
- #bigM for safety, protect
62
- #overlapping diagonals
63
- bigM = (X.max() - X.min()) + 1
64
-
65
- #move the smallest axes to last positions
66
- abc = X.shape
67
- dst_axes = (-3, -2, -1)
68
- src_axes = tuple(sorted(dst_axes, key=lambda ax: -abc[ax]))
69
- X = xp.moveaxis(X, src_axes, dst_axes)
70
- a, b, c = X.shape[-3:] # now a >= b >= c
71
-
72
- i = arange(a, 'int', X)[:, None, None]
73
- j = arange(b, 'int', X)[None, :, None]
74
- k = arange(c, 'int', X)[None, None, :]
75
-
76
- di = i - k
77
- dj = j - k
78
- di_ = di % a
79
- dj_ = dj % b
80
- dires = di // a
81
- djres = dj // b
82
- groupkey = bigM * (dires * 4 + djres)
83
-
84
- X[..., di_, dj_, k] = X[..., i, j, k] - groupkey #Trirotation
85
- # ========================
86
- # Again, thanks to the rotation trick
87
- # we just need to .sort(X, axis = -1)
88
- # to sort the tridiagonal
89
- # ========================
90
- was_sorted = all_true(xp.diff(X, axis=-1) >= 0)
91
- inplace_sort(X, axis=-1)
92
- X[..., i, j, k] = X[..., di_, dj_, k] + groupkey
93
-
94
- return was_sorted
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