0

I have a function that takes two arguments. I want to run the function for every possible combination of my inputs, and store each returned value. For example:

def foo(a, b):
    return (a + b)

if __name__ == "__main__":
    a = np.array([1., 2., 3.])
    b = np.array([5., 6.])
    f1 = foo(a[0], b[0]) #6
    f2 = foo(a[0], b[1]) #7
    f3 = foo(a[1], b[0]) #etc
    f4 = foo(a[1], b[1])
    f5 = foo(a[2], b[0])
    f6 = foo(a[2], b[1])

How can I call f1 through f6 in a more elegant way, like a loop? It can't be a direct loop, because a and b have differing numbers of elements.

StatsSorceress
  • 3,019
  • 7
  • 41
  • 82

1 Answers1

3

itertools.product is the way to go in a Pythonic manner. However, if you are looking for a more raw, basic way, you'll need two for loops, which is the basic property of a cross product.

for i in a:
    for j in b:
        foo(i, j)

Edit:

Example usages of itertools.product:

for i, j in itertools.product(a, b):
    foo(i, j)

for tup in itertools.product(a, b):
    foo(*tup)

I prefer the first one because tuple's elements can be used explicitly if needed.

Rockybilly
  • 2,938
  • 1
  • 13
  • 38
  • Thanks @Rockybilly, but I think I'm missing something. itertools.product returns a tuple, and then I have to somehow feed that in to my function. I saw [this](http://stackoverflow.com/questions/1993727/expanding-tuples-into-arguments) but didn't really understand. Could you show a MWE, using itertools.product and then feeding it in to the function? – StatsSorceress Dec 09 '16 at 02:30