How to rewrite this nested for loop in python?-CodePudding

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?

CodePudding user response：

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.

CodePudding user response：

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)

CodePudding user response：

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