Using numpy to build an array of all combinations of two arrays

Question

Hi,

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

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

1) 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

The I used reduce to apply that to m copies of the same array:

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

And then 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 waaaay too slow. I know the space of parameters is huge, but this shouldn't be so slow. I have only sampled (10)^6 = a million points in this example and it took more then 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 take it's arguments if it's necessary.

Answer 1

itertools.combinations is in general the fastest way to get combinations from a Python container (if you do in fact want combinations, i.e., arrangements WITHOUT repetitions and independent of order; that's not what your code appears to be doing, but I can't tell whether that's because your code is buggy or because you're using the wrong terminology).

If you want something different than combinations perhaps other iterators in itertools, product or permutations, might serve you better. For example, it looks like your code is roughly the same as:

for val in itertools.product(np.arange(0, 1, 0.1), repeat=6):
    print F(val)

All of these iterators yield tuples, not lists or numpy arrays, so if your F is picky about getting specifically a numpy array you'll have to accept the extra overhead of constructing or clearing and re-filling one at each step.

Answer 2

Here's a pure-numpy implementation. It's ca. 5× faster than using itertools.

import numpy as np

def cartesian(arrays, out=None):
    """
    Generate a cartesian product of input arrays.

    Parameters
    ----------
    arrays : list of array-like
        1-D arrays to form the cartesian product of.
    out : ndarray
        Array to place the cartesian product in.

    Returns
    -------
    out : ndarray
        2-D array of shape (M, len(arrays)) containing cartesian products
        formed of input arrays.

    Examples
    --------
    >>> cartesian(([1, 2, 3], [4, 5], [6, 7]))
    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]])

    """

    arrays = [np.asarray(x) for x in arrays]
    dtype = arrays[0].dtype

    n = np.prod([x.size for x in arrays])
    if out is None:
        out = np.zeros([n, len(arrays)], dtype=dtype)

    m = n / arrays[0].size
    out[:,0] = np.repeat(arrays[0], m)
    if arrays[1:]:
        cartesian(arrays[1:], out=out[0:m,1:])
        for j in xrange(1, arrays[0].size):
            out[j*m:(j+1)*m,1:] = out[0:m,1:]
    return out