How can I iterate over an array of numpy arrays row by row efficiently?

Question

I have a large list of lists that I'm converting to a numpy array, and then iterating through, something like this:

a = [ [0,1,2], [3,4,5], [6,7,8] ]
np_a = np.array(a)

for i in np_a:
    print i # prints [0,1,2], then [3,4,5], then [6,7,8] as expected
    # do a bunch of stuff to each row

Of course this works, but since I'm dealing with a large body of data, this can take up to a minute to run. I've been looking into various ways to speed up this code, with the number one recommendation being to vectorize the operations I'm performing within the loop. However, the operations are not trivial to vectorize (think a bunch of dot products and transforms and such on each row...). One method I've found is to use the np.nditer method, like this:

a = [ [0,1,2], [3,4,5], [6,7,8] ]
np_a = np.array(a)

for i in np.nditer(np_a):
    print i # prints 0, then 1, then 2, etc...
    # do a bunch of stuff to each rowon

But this seems to flatten the array into a 1-D list. I've looked through the documentation pretty extensively, and haven't found what I'm looking for here: https://docs.scipy.org/doc/numpy-1.13.0/reference/arrays.nditer.html#iterating-as-a-specific-data-type

Does anyone have any specific recommendations on how to make this work as I expect with np.nditer, or some other implementation that's more efficient than the standard 'for r in a' method?

@coldspeed Curious what you mean here? Do you mean that vectorizing is the way to go? How would you go about it? — stumbling.fool, Dec 25 '17 at 19:27
@Divakar it seems that the actual operations I'm performing within the loop are relatively quick, it's just the volume of operations. I guess what I'm wondering is if there is a way to loop faster with numpy or something similar (people have said that np.nditer is significantly faster for this kind of thing) — stumbling.fool, Dec 25 '17 at 19:29
`nditer` is not faster. Usually the iteration mechanism isn't the slow thing; t's calling your function many times. True vectorizing means rewriting your function so that it is called once - processing all rows at once. It may still iterate in compiled code, but what you want to avoid is thousands of Python function calls. — hpaulj, Dec 25 '17 at 19:35
Read that `nditer` tutorial all the way to the end. It may, on occasion, help with complex broadcasting, but its real value is as a stepping stone to writing your code in `cython` (or other direct `C` code). — hpaulj, Dec 25 '17 at 19:38
@hpaulj thank you, that's helpful. I'm looking into vectorizing so that all of these operations are performed at once, but it's not so straightforward. But that does answer my question (instead of just saying that I'm doing it wrong, e.g. non-constructive criticism), since I won't pursue improving the iteration mechanism. — stumbling.fool, Dec 25 '17 at 19:40
To do a depth limited iteration (i.e. just rows) with `nditer` you have use it as `np.ndindex` does. — hpaulj, Dec 25 '17 at 19:44
Searching my answers for `frompyfunc` turned up some previous tests: https://stackoverflow.com/questions/45999814/how-to-apply-function-which-returns-vector-to-each-numpy-array-element-and-get, https://stackoverflow.com/questions/44239498/how-to-apply-a-generic-function-over-numpy-rows, https://stackoverflow.com/questions/43018527/numpy-iterate-over-outer-dimension-of-numpy-array-using-nditer, https://stackoverflow.com/questions/25096600/shallow-iteration-with-nditer (see sidebar for titles). — hpaulj, Dec 25 '17 at 20:46

How can I iterate over an array of numpy arrays row by row efficiently?

0 Answers0