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

