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.

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