2

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?

Improve this question
3
  • 2
    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 Commented Oct 27, 2021 at 7:55
  • Using-numpy-to-build-an-array-of-all-combinations-of-two-arrays on SO has various methods -- the term is cartesian product. (Was itertools.product at all faster in your case ?) Commented Mar 29, 2022 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 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. Commented Jun 25, 2022 at 2:39

4 Answers 4

Reset to default
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)
Improve this answer
Sign up to request clarification or add additional context in comments.

Comments

1

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

0 0
0 1
0 2
0 3
1 0
1 1
1 2
1 3
2 0
2 1
2 2
2 3

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.

Improve this answer

Comments

1

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
Improve this answer

Comments

0

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

Improve this answer

Comments

Start asking to get answers

Find the answer to your question by asking.

Ask question

Explore related questions

See similar questions with these tags.