How to rewrite this nested for loop in python?

Question

I have a nested for loop, and I'm doing some operations on it, however it's take about 20 minutes to run. Hoping to reduce the wall time. This is just a simple example to make a reproducible one, but how could I re-write this nested for loop to make it more efficient?

I tried using the zip function in python but it doesn't print out the order of i,j the same way as the nested loop, which is necessary for my calculation.

This question is similar to this stackoverflow one, but I'm struggling to reproduce the answer: Automatically nested for loops in python

array = np.random.rand(3,4,10)

x_vals = np.array([22,150,68,41]) #size of 4 (same as size as 2nd axis of array above)

new_arr = np.zeros((array.shape))

for i in range(3): 
     for j in range(4): 
            print(i,j)
            new_arr[i,:,j] = array[i,:,j]*x_vals
            
0 0
0 1
0 2
0 3
1 0
1 1
1 2
1 3
2 0
2 1
2 2
2 3

I tried:

for i,j in zip(range(3),range(4)):
    print(i,j) # output is i,j is not the same as it is for the nested loop above

0 0
1 1
2 2

I was thinking maybe the function enumerate would work as well, but I also get an error:

for idx,i in enumerate(range(3),range(4)):
    print(idx,i)

---------------------------------------------------------------------------
TypeError                                 Traceback (most recent call last)
/tmp/ipykernel_2101160/1601005143.py in <module>
----> 1 for idx,i in enumerate(range(3),range(4)):
      2     print(idx,i)

TypeError: 'range' object cannot be interpreted as an integer

Any ideas on how to either vectorize or speed up the nested loop?

itertools is your friend, `for i, j in itertools.product(range(3), range(4))` although I am not sure it will really speed up your loop. Since you will still have to iterate over all the combinations — Chris Doyle, Oct 27 '21 at 07:55
[Using-numpy-to-build-an-array-of-all-combinations-of-two-arrays](https://stackoverflow.com/questions/1208118/using-numpy-to-build-an-array-of-all-combinations-of-two-arrays) on SO has various methods -- the term is [cartesian product](https://stackoverflow.com/questions/tagged/cartesian-product). (Was `itertools.product` at all faster in your case ?) — denis, Mar 29 '22 at 15:37
The tool you're looking for is `itertools.product` in order to do that kind of iteration - this is a *very* common duplicate question. `zip` is wrong because it iterates in parallel, solving a different problem - see https://stackoverflow.com/questions/1663807/. However, there is no real performance benefit to be had here. For the given code example, I don't see why not just slice in all three dimensions, instead of slicing in the middle dimension and manually iterating over the others. — Karl Knechtel, Jun 25 '22 at 02:39

score 2 · Accepted Answer · answered Oct 27 '21 at 07:57

2

You can use itertools.product, equivalent to nested for-loops in a generator expression. https://docs.python.org/3/library/itertools.html#itertools.product

import itertools

for i, j in itertools.product(range(3), range(4)):
    print(i, j)

answered Oct 27 '21 at 07:57

4d61726b

427
5
26

score 1 · Answer 2 · answered Oct 27 '21 at 07:57

You might use itertools.product as follows to avoid nesting fors

import itertools
for i,j in itertools.product(range(3),range(4)):
    print(i,j)

output

However it would rather insignificant influence on time required. Your mention

order of i,j(...)which is necessary for my calculation.

suggest that you need effect of previous computation i.e. have certain rules regarding order which could not be broken. You need to consider how you can change without breaking rules and if it would provide required speed boost.

Dani Mesejo · Answer 3 · 2021-10-27T08:03:42.503

One approach is to use np.mgrid to generate the indices and take advantage of numpy indexing (see here) to avoid any for-loops:

import numpy as np

array = np.random.rand(3, 4, 10)
x_vals = np.array([22, 150, 68, 41])  # size of 4 (same as size as 2nd axis of array above)

new_arr = np.zeros(array.shape)
rows, cols = map(np.ndarray.flatten, np.mgrid[0:3, 0:4])
new_arr[rows, :, cols] = array[rows, :, cols] * x_vals

score 0 · Answer 4 · answered Oct 27 '21 at 08:27

As long as the printing is optional, the following will do the trick in vectorized fashion:

array = np.random.rand(3,4,10)

x_vals = np.array([22,150,68,41])

new_arr = array*x_vals.reshape((1,4,1))

This is also assuming you actually want j to run all the way up to 10, and not stop at 4 as in your example

How to rewrite this nested for loop in python?

4 Answers4