python: making array index generation more efficient/elegant

Question

I am trying to get some array indices with python. At the moment, the code looks very cumbersome and I was wondering if I am doing it inefficiently or in unpythonic style. So, I have an n-dimensional array and I am trying to generate some indexes as follows. Here is an isolated code sample that demonstrates what I am trying to do. I am demonstrating this in a simple 2D code segment but the array can be arbitrary dimension.

import numpy as np

a = np.random.rand(5, 5)
shape = a.shape

for i in range(len(shape)):
    shape_temp = np.zeros(len(shape), dtype=np.int)
    shape_temp[i] = 1
    p = np.meshgrid(*[np.arange (shape_temp[l], shape[l]) for l in range(len(shape))])

    # Do something with p

I was wondering if there was a more elegant and hopefully efficient way to generate these indices? In particular the use of this shape_temp variable looks ugly especially how it is created every time in the loop.

It would be clearer if you didn't use the same index variable in 2 places. But, can you describe in words, or with an example, what you want to accomplish. Without running the code I can't tell what is going on. — hpaulj
– hpaulj, Commented Aug 16, 2014 at 2:18
Sorry the indexing thing was an oversight. I have fixed it. One should be able to run this code as is. — Luca
– Luca, Commented Aug 16, 2014 at 2:31
Sorry the indexing thing was an oversight. I have fixed it. Basically it generates an array of indexes where it runs from some arbitrary position, I have set it to 1 through shape_temp variable for an axes to the end and it basically does this for an n-dimensional array with dynamic dimension — Luca
– Luca, Commented Aug 16, 2014 at 2:39
So every time the outer loop runs $p$ will be a 2D array of array indexes. You should be able to run this code in ipython. — Luca
– Luca, Commented Aug 16, 2014 at 2:46

hpaulj · Accepted Answer · 2014-08-16 03:18:23Z

1

shape_temp = np.zeros_like(shape)

You could avoid shape_temp with an expression like:

[np.arange(1 if l==i else 0, e) for l,e in enumerate(shape)]

Whether it is prettier or more efficient is debatable

Another snippet

temp = np.eye(len(shape))
[np.arange(j,k) for j,k in zip(temp[i,:],shape)]

An alternative to meshgrid is mgrid or ogrid (though I have to change in indexing to get match). The expressions are more compact because they take slices rather than ranges. Internally those functions use arange.

meshgrid(*[np.arange(j,k) for j,k in zip(temp[i,:],shape)],indexing='ij')
np.mgrid[[slice(j,k) for j,k in zip(temp[i,:],shape)]]

meshgrid(*[np.arange(j,k) for j,k in zip(temp[i,:],shape)],sparse=True,indexing='ij')
np.ogrid[[slice(j,k) for j,k in zip(temp[i,:],shape)]]

edited Aug 16, 2014 at 3:18

answered Aug 16, 2014 at 2:54

hpaulj

233k14 gold badges260 silver badges392 bronze badges

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

Collectives™ on Stack Overflow

python: making array index generation more efficient/elegant

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