How to convert numpy array to form such as ((value0, row0, column0), (value1, row0, column1)...)? e.g. if the input is
a = np.array([[10, 15],
[20, 25]])
The output should be:
((10, 0, 0), (15, 0, 1), (20, 1, 0), (25, 1, 1)). Moreover, this function should be able to be applied to constants, 1D, 3D, 4D, etc. arrays, e.g. for 4D the output should be (value0, dim1, dim2, dim3, dim4). I can create such a function only for a specific dimensionality, e.g. only for 2D.
One idea is to flatten a and pair with meshgrid or product of the indices:
a = np.array([[10, 15, 12],
[20, 25, 13]])
[(u,)+ v for u,v in zip(a.ravel(), product(*[np.arange(x) for x in a.shape]) )]
Output:
[(10, 0, 0), (15, 0, 1), (12, 0, 2), (20, 1, 0), (25, 1, 1), (13, 1, 2)]
Here's a vectorized one for n-dim arrays based on np.indices -
def stack_indices(a):
i = np.moveaxis(np.indices(a.shape),0,-1)
return np.concatenate((a[...,None],i),axis=-1)
Sample runs -
1D :
In [70]: a = np.array([4,5,2])
In [71]: stack_indices(a)
Out[71]:
array([[4, 0],
[5, 1],
[2, 2]])
2D :
In [62]: a
Out[62]:
array([[10, 15, 17],
[20, 25, 30]])
In [63]: stack_indices(a)
Out[63]:
array([[[10, 0, 0],
[15, 0, 1],
[17, 0, 2]],
[[20, 1, 0],
[25, 1, 1],
[30, 1, 2]]])
3D :
In [68]: a
Out[68]:
array([[[51, 67],
[45, 21],
[56, 92]],
[[10, 24],
[63, 22],
[52, 94]]])
In [69]: stack_indices(a)
Out[69]:
array([[[[51, 0, 0, 0],
[67, 0, 0, 1]],
[[45, 0, 1, 0],
[21, 0, 1, 1]],
[[56, 0, 2, 0],
[92, 0, 2, 1]]],
[[[10, 1, 0, 0],
[24, 1, 0, 1]],
[[63, 1, 1, 0],
[22, 1, 1, 1]],
[[52, 1, 2, 0],
[94, 1, 2, 1]]]])
Other posted approaches :
from itertools import product
# @Quang Hoang's soln
def prod(a):
return [(u,)+ v for u,v in zip(a.ravel(), product(*[np.arange(x) for x in a.shape]) )]
# @Andy L.'s soln
def ravel(a):
return [z for z in zip(a.ravel(), *np.unravel_index(np.arange(a.size), a.shape))]
Using benchit package (few benchmarking tools packaged together; disclaimer: I am its author) to benchmark proposed solutions.
import benchit
funcs = [stack_indices, prod, ravel]
in_ = [np.random.randint(10,100,([5]*n)) for n in range(1,10)]
t = benchit.timings(funcs, in_, indexby='shape')
t.rank()
t.plot(logx=False, save='timings.png', rot=90)
To gain better performance on higher-dim arrays, I suppose one could look into array-assignment.
use zip with unravel_index
arr = [z for z in zip(a.ravel(), *np.unravel_index(np.arange(a.size), a.shape))]
A more concise code is:
arr = list(zip(a.ravel(), *np.unravel_index(np.arange(a.size),a.shape)))
Out[582]: [(10, 0, 0), (15, 0, 1), (12, 0, 2), (20, 1, 0), (25, 1, 1), (13, 1, 2
)]