Fix Python – Using numpy to build an array of all combinations of two arrays

Question

Asked By – Rafael S. Calsaverini

I’m trying to run over the parameters space of a 6 parameter function to study its numerical behavior before trying to do anything complex with it, so I’m searching for an efficient way to do this.

My function takes float values given in a 6-dim numpy array as input. What I tried to do initially was this:

First, I created a function that takes 2 arrays and generate an array with all combinations of values from the two arrays:

from numpy import *
def comb(a,b):
    c = []
    for i in a:
        for j in b:
            c.append(r_[i,j])
    return c

Then, I used reduce() to apply that to m copies of the same array:

def combs(a,m):
    return reduce(comb,[a]*m)

Finally, I evaluate my function like this:

values = combs(np.arange(0,1,0.1),6)
for val in values:
    print F(val)

This works but it’s way too slow. I know the space of parameters is huge, but this shouldn’t be so slow. I have only sampled 106 (a million) points in this example and it took more than 15 seconds just to create the array values.

Do you know any more efficient way of doing this with numpy?

I can modify the way the function F takes it’s arguments if it’s necessary.

Now we will see solution for issue: Using numpy to build an array of all combinations of two arrays


Answer

In newer version of numpy (>1.8.x), numpy.meshgrid() provides a much faster implementation:

@pv’s solution

In [113]:

%timeit cartesian(([1, 2, 3], [4, 5], [6, 7]))
10000 loops, best of 3: 135 µs per loop
In [114]:

cartesian(([1, 2, 3], [4, 5], [6, 7]))

Out[114]:
array([[1, 4, 6],
       [1, 4, 7],
       [1, 5, 6],
       [1, 5, 7],
       [2, 4, 6],
       [2, 4, 7],
       [2, 5, 6],
       [2, 5, 7],
       [3, 4, 6],
       [3, 4, 7],
       [3, 5, 6],
       [3, 5, 7]])

numpy.meshgrid() use to be 2D only, now it is capable of ND. In this case, 3D:

In [115]:

%timeit np.array(np.meshgrid([1, 2, 3], [4, 5], [6, 7])).T.reshape(-1,3)
10000 loops, best of 3: 74.1 µs per loop
In [116]:

np.array(np.meshgrid([1, 2, 3], [4, 5], [6, 7])).T.reshape(-1,3)

Out[116]:
array([[1, 4, 6],
       [1, 5, 6],
       [2, 4, 6],
       [2, 5, 6],
       [3, 4, 6],
       [3, 5, 6],
       [1, 4, 7],
       [1, 5, 7],
       [2, 4, 7],
       [2, 5, 7],
       [3, 4, 7],
       [3, 5, 7]])

Note that the order of the final resultant is slightly different.

This question is answered By – CT Zhu

This answer is collected from stackoverflow and reviewed by FixPython community admins, is licensed under cc by-sa 2.5 , cc by-sa 3.0 and cc by-sa 4.0