Can this for loop be written in a simpler way?

import itertools
import numpy as np

def f(a, b, c):
  print(a b c)

a = np.arange(12).reshape(3, 4)

for y, x in itertools.product(range(a.shape[0]-1), range(a.shape[1]-1)):
  f(a[y, x], a[y, x 1], a[y 1, x])

The other options I tried, look more convoluted, e.g.:

it = np.nditer(a[:-1, :-1], flags=['multi_index'])

for e in it:
  y, x = it.multi_index
  f(a[y, x], a[y, x 1], a[y 1, x])

CodePudding user response：

Posting it as an answer, and sorry if this is too obvious, but isn't this simply

for y in range(a.shape[0]-1):
     for x in range(a.shape[1]-1):
         f(a[y, x], a[y, x 1], a[y 1, x])

CodePudding user response：

If I use you method I got:

expected = [5, 8, 11, 17, 20, 23]

but you can vectorize the computation by generating an array containing the data in a more suitable way:

a_stacked = np.stack([a[:-1, :-1], a[:-1, 1:], a[1:, :-1]], axis=0)

From there multiple solutions:

• If you already know the function will be the sum:

>>> a_stacked.sum(axis=0)
array([[ 5,  8, 11],
       [17, 20, 23]])

• If you know that your function is already vectorized:

>>> f(*a_stacked)
array([[ 5,  8, 11],
       [17, 20, 23]])

• If your function does not vectorize, you can use np.vectorize for convenience (no performance improvement):

>>> np.vectorize(f)(*a_stacked)
array([[ 5,  8, 11],
       [17, 20, 23]])

Obviously you can flatten the array next.