How do I convert numpy array to the form ((value0, row0, column0), (value1, row0, column1)...)?

Question

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.

Answer 1

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]]]])

Benchmarking

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.

Answer 2

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)]