Build a Numpy array where datas match their own coordinates

Question

I want to write a function with a tuple s (of length n) as an argument, which should return an array of shape s concatenated with (n,) (thus adding an extra dimension to the array), for which the data indexed by any tuple of length n (thus forgetting the last dimension) returns the tuple itself.

Here is an example of what it should do:

>>> a=f((2,3,4))
>>> a[1,1,1]
ndarray([1, 1, 1])
>>> a.shape
(2,3,4,3)

I managed to do it with the following code (if I am not wrong), but I am pretty sure it can be achieved in a more simple way:

a = np.transpose(
        np.meshgrid(
            *(np.arange(0, x) for x in s)),
        axes = (2,1) + tuple(range(3, n+1)) + (0,))

(I hope I correctly copied my code since I had to simplify variables here.)

Divakar · Accepted Answer · 2016-12-04 17:26:17Z

1

You are looking for np.indices to generate those ranged-indices. To bring to the same format as the one proposed in the code listed in the question, we need to push back the first axis to the end.

So, we would have an implementation for s of generic length, like so -

s = (2,3,4)  # Input
out = np.indices(s).transpose(range(1,len(s)+1) + [0])

Alternatively, this pushing back of dimension could be achieved with np.rollaxis as well -

out = np.rollaxis(np.indices(s),0,len(s)+1)

edited Dec 4, 2016 at 17:26

answered Dec 4, 2016 at 17:15

Divakar

222k19 gold badges273 silver badges374 bronze badges

Sign up to request clarification or add additional context in comments.

Collectives™ on Stack Overflow

Build a Numpy array where datas match their own coordinates

1 Answer 1

Comments

Your Answer

Hot Network Questions

Collectives™ on Stack Overflow

1 Answer 1

Comments

Your Answer

Sign up or log in

Post as a guest

Related