/
_regular_grid.py
72 lines (66 loc) · 2.42 KB
/
_regular_grid.py
1
2
3
4
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
import numpy as np
def regular_grid(ar_shape, n_points):
"""Find `n_points` regularly spaced along `ar_shape`.
The returned points (as slices) should be as close to cubically-spaced as
possible. Essentially, the points are spaced by the Nth root of the input
array size, where N is the number of dimensions. However, if an array
dimension cannot fit a full step size, it is "discarded", and the
computation is done for only the remaining dimensions.
Parameters
----------
ar_shape : array-like of ints
The shape of the space embedding the grid. ``len(ar_shape)`` is the
number of dimensions.
n_points : int
The (approximate) number of points to embed in the space.
Returns
-------
slices : list of slice objects
A slice along each dimension of `ar_shape`, such that the intersection
of all the slices give the coordinates of regularly spaced points.
Examples
--------
>>> ar = np.zeros((20, 40))
>>> g = regular_grid(ar.shape, 8)
>>> g
[slice(5, None, 10), slice(5, None, 10)]
>>> ar[g] = 1
>>> ar.sum()
8.0
>>> ar = np.zeros((20, 40))
>>> g = regular_grid(ar.shape, 32)
>>> g
[slice(2, None, 5), slice(2, None, 5)]
>>> ar[g] = 1
>>> ar.sum()
32.0
>>> ar = np.zeros((3, 20, 40))
>>> g = regular_grid(ar.shape, 8)
>>> g
[slice(1, None, 3), slice(5, None, 10), slice(5, None, 10)]
>>> ar[g] = 1
>>> ar.sum()
8.0
"""
ar_shape = np.asanyarray(ar_shape)
ndim = len(ar_shape)
unsort_dim_idxs = np.argsort(np.argsort(ar_shape))
sorted_dims = np.sort(ar_shape)
space_size = float(np.prod(ar_shape))
if space_size <= n_points:
return [slice(None)] * ndim
stepsizes = (space_size / n_points) ** (1.0 / ndim) * np.ones(ndim)
if (sorted_dims < stepsizes).any():
for dim in range(ndim):
stepsizes[dim] = sorted_dims[dim]
space_size = float(np.prod(sorted_dims[dim + 1:]))
stepsizes[dim + 1:] = ((space_size / n_points) **
(1.0 / (ndim - dim - 1)))
if (sorted_dims >= stepsizes).all():
break
starts = (stepsizes // 2).astype(int)
stepsizes = np.round(stepsizes).astype(int)
slices = [slice(start, None, step) for
start, step in zip(starts, stepsizes)]
slices = [slices[i] for i in unsort_dim_idxs]
return slices