0

Can this for loop be written in a simpler way?

import itertools
import numpy as np

def f(a, b, c): # placeholder for a complex function
  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])
Paul Jurczak
  • 7,008
  • 3
  • 47
  • 72

2 Answers2

1

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])
nenad.popovic
  • 146
  • 5
1

If I use your 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.

paime
  • 2,901
  • 1
  • 6
  • 17
  • Function `f` is just a placeholder for something much more complex. I was looking for the loop options. – Paul Jurczak Jul 11 '22 at 15:34
  • So last option should do the job (for loop under the hood), but you might want to check if you can't vectorize `f` so option 2 can be used, which will be much faster. – paime Jul 11 '22 at 16:05