-1

I use to concatenate my numpy arrays of arbitrary shape to make my code cleaner, however, it seems pretty hard to me to iterate over it in a pythonesque way.

Lets consider a 4 dimension array x (thus len(x.shape) = 4), and that the index I want to iterate is 2, the naive solution that I usually use is something like

y = np.array([my_operation(x[:, :, i, :])
              for i in range(x.shape[2])])

I'm looking for something more readable, because it is annoying to have so many ":" and any changes in the dimensions of x would require rewrite a part of my code. Something magic like

y = np.array([my_operation(z) for z in magic_function(x, 2)])

Is there a numpy method that allows me to iterate over any arbitrary dimension of an array ?

Alexandre
  • 14
  • 2
  • 1
    Iteration on any axis is sub-optimal. The fast compiled `ufuncs` provide a `axis` or `axes` parameter so you don't have to explicitly iterate. There are some 'apply...' functions, but they are convenience tools, not performance enhancers. In order to achieve much generality the Python code tends to be quite long. So your own 'magic_function', tuned to your own subset of possible iterations is the best choice. – hpaulj Jun 05 '19 at 16:19

3 Answers3

0

One possible solution is to use a dict().

What you can do is:

x = dict()
x['param1'] = [1, 1, 1, 1]
x['param2'] = [2, 2, 2, 2]

print(x['param1']) 
# > [1, 1, 1, 1]
spassador
  • 393
  • 1
  • 12
0

I don't know any standard way to do that. Anyway, your trick is good. We can elaborate it a little, and we get an implementation for the "magic function" you are looking for:

def magic_function(x, n):
    slices = [slice(w) for w in x.shape]
    for i in range(x.shape[n]):
        slices[n] = i
        z = x[tuple(slices)]
        yield z
luca.vercelli
  • 898
  • 7
  • 24
0

You can temporarily move the desired axis to the front and then iterate over the array. Afterwards you move the axis back:

x = np.moveaxis(x, 2, 0)
x = np.array([my_operation(sub_x) for sub_x in x])
x = np.moveaxis(x, 0, 2)
Nils Werner
  • 34,832
  • 7
  • 76
  • 98