Cartesian product of x and y array points into single array of 2D points

У меня есть два массива numpy, которые определяют оси x и y сетки. Например:

x = numpy.array([1,2,3])
y = numpy.array([4,5])

Я хотел бы сгенерировать декартово произведение этих массивов для генерации:

array([[1,4],[2,4],[3,4],[1,5],[2,5],[3,5]])

Таким образом, это не очень неэффективно, поскольку мне нужно делать это много раз в цикле. Я предполагаю, что преобразование их в список Python и использование itertools.product и обратно в массив numpy - не самая эффективная форма.

Канонический cartesian_product (почти)

Существует множество подходов к решению этой проблемы с разными свойствами. Некоторые из них быстрее других, а некоторые более универсальны. После долгих тестов и настроек я обнаружил, что следующая функция, которая вычисляет n-мерное cartesian_product значение, работает быстрее большинства других для многих входных данных. О паре подходов, которые немного сложнее, но во многих случаях даже немного быстрее, смотрите В ответе Пола Панцера .

Учитывая этот ответ, это уже не самая быстрая реализация декартова произведения в numpy, о которой я знаю. Тем не менее, я думаю, что его простота по-прежнему будет делать его полезным ориентиром для будущих улучшений:

def cartesian_product(*arrays):
    la = len(arrays)
    dtype = numpy.result_type(*arrays)
    arr = numpy.empty([len(a) for a in arrays] + [la], dtype=dtype)
    for i, a in enumerate(numpy.ix_(*arrays)):
        arr[...,i] = a
    return arr.reshape(-1, la)

Стоит упомянуть, что эта функция использует ix_ необычным образом; в то время как документированное использование ix_ заключается в генерации индексов в массиве, просто так получилось, что массивы одинаковой формы могут использоваться для широковещательного присвоения. Большое спасибо mgilson, который вдохновил меня попробовать использовать ix_ этот способ, и unutbu, который предоставил несколько чрезвычайно полезных отзывов по этому ответу, включая предложение использовать numpy.result_type.

Известные альтернативы

Иногда быстрее записывать смежные блоки памяти в порядке Fortran. Это основа этой альтернативы, cartesian_product_transpose которая оказалась быстрее на некотором оборудовании, чем cartesian_product (см. Ниже). Однако ответ Пола Панцера, использующий тот же принцип, еще быстрее. Тем не менее, я включаю это сюда для заинтересованных читателей:

def cartesian_product_transpose(*arrays):
    broadcastable = numpy.ix_(*arrays)
    broadcasted = numpy.broadcast_arrays(*broadcastable)
    rows, cols = numpy.prod(broadcasted[0].shape), len(broadcasted)
    dtype = numpy.result_type(*arrays)

    out = numpy.empty(rows * cols, dtype=dtype)
    start, end = 0, rows
    for a in broadcasted:
        out[start:end] = a.reshape(-1)
        start, end = end, end + rows
    return out.reshape(cols, rows).T

После того, как я понял подход Panzer, я написал новую версию, которая почти такая же быстрая, как у него, и почти такая же простая, как cartesian_product:

def cartesian_product_simple_transpose(arrays):
    la = len(arrays)
    dtype = numpy.result_type(*arrays)
    arr = numpy.empty([la] + [len(a) for a in arrays], dtype=dtype)
    for i, a in enumerate(numpy.ix_(*arrays)):
        arr[i, ...] = a
    return arr.reshape(la, -1).T

Похоже, что это требует постоянных затрат времени, что делает его работу медленнее, чем Panzer для небольших входных данных. Но для больших входных данных во всех тестах, которые я запускал, он работает так же хорошо, как и его самая быстрая реализация (cartesian_product_transpose_pp).

В следующих разделах я включаю некоторые тесты других альтернатив. Сейчас они несколько устарели, но вместо того, чтобы дублировать усилия, я решил оставить их здесь из исторического интереса. Для получения актуальных тестов смотрите Ответ Panzer, а также ответы Нико Шлемера.

Тесты с альтернативами

Вот набор тестов, которые показывают повышение производительности, которое обеспечивают некоторые из этих функций по сравнению с рядом альтернатив. Все показанные здесь тесты были выполнены на четырехъядерном компьютере под управлением Mac OS 10.12.5, Python 3.6.1 и numpy 1.12.1. Известно, что аппаратные и программные вариации приводят к разным результатам, поэтому YMMV. Запустите эти тесты сами, чтобы убедиться!

Определения:

import numpy
import itertools
from functools import reduce

### Two-dimensional products ###

def repeat_product(x, y):
    return numpy.transpose([numpy.tile(x, len(y)), 
                            numpy.repeat(y, len(x))])

def dstack_product(x, y):
    return numpy.dstack(numpy.meshgrid(x, y)).reshape(-1, 2)

### Generalized N-dimensional products ###

def cartesian_product(*arrays):
    la = len(arrays)
    dtype = numpy.result_type(*arrays)
    arr = numpy.empty([len(a) for a in arrays] + [la], dtype=dtype)
    for i, a in enumerate(numpy.ix_(*arrays)):
        arr[...,i] = a
    return arr.reshape(-1, la)

def cartesian_product_transpose(*arrays):
    broadcastable = numpy.ix_(*arrays)
    broadcasted = numpy.broadcast_arrays(*broadcastable)
    rows, cols = numpy.prod(broadcasted[0].shape), len(broadcasted)
    dtype = numpy.result_type(*arrays)

    out = numpy.empty(rows * cols, dtype=dtype)
    start, end = 0, rows
    for a in broadcasted:
        out[start:end] = a.reshape(-1)
        start, end = end, end + rows
    return out.reshape(cols, rows).T

# from https://pythonly.ru/a/1235363/577088

def cartesian_product_recursive(*arrays, out=None):
    arrays = [numpy.asarray(x) for x in arrays]
    dtype = arrays[0].dtype

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

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

def cartesian_product_itertools(*arrays):
    return numpy.array(list(itertools.product(*arrays)))

### Test code ###

name_func = [('repeat_product',                                                 
              repeat_product),                                                  
             ('dstack_product',                                                 
              dstack_product),                                                  
             ('cartesian_product',                                              
              cartesian_product),                                               
             ('cartesian_product_transpose',                                    
              cartesian_product_transpose),                                     
             ('cartesian_product_recursive',                           
              cartesian_product_recursive),                            
             ('cartesian_product_itertools',                                    
              cartesian_product_itertools)]

def test(in_arrays, test_funcs):
    global func
    global arrays
    arrays = in_arrays
    for name, func in test_funcs:
        print('{}:'.format(name))
        %timeit func(*arrays)

def test_all(*in_arrays):
    test(in_arrays, name_func)

# `cartesian_product_recursive` throws an 
# unexpected error when used on more than
# two input arrays, so for now I've removed
# it from these tests.

def test_cartesian(*in_arrays):
    test(in_arrays, name_func[2:4] + name_func[-1:])

x10 = [numpy.arange(10)]
x50 = [numpy.arange(50)]
x100 = [numpy.arange(100)]
x500 = [numpy.arange(500)]
x1000 = [numpy.arange(1000)]

Результаты тестирования:

In [2]: test_all(*(x100 * 2))
repeat_product:
67.5 µs ± 633 ns per loop (mean ± std. dev. of 7 runs, 10000 loops each)
dstack_product:
67.7 µs ± 1.09 µs per loop (mean ± std. dev. of 7 runs, 10000 loops each)
cartesian_product:
33.4 µs ± 558 ns per loop (mean ± std. dev. of 7 runs, 10000 loops each)
cartesian_product_transpose:
67.7 µs ± 932 ns per loop (mean ± std. dev. of 7 runs, 10000 loops each)
cartesian_product_recursive:
215 µs ± 6.01 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)
cartesian_product_itertools:
3.65 ms ± 38.7 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

In [3]: test_all(*(x500 * 2))
repeat_product:
1.31 ms ± 9.28 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)
dstack_product:
1.27 ms ± 7.5 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)
cartesian_product:
375 µs ± 4.5 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)
cartesian_product_transpose:
488 µs ± 8.88 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)
cartesian_product_recursive:
2.21 ms ± 38.4 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)
cartesian_product_itertools:
105 ms ± 1.17 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)

In [4]: test_all(*(x1000 * 2))
repeat_product:
10.2 ms ± 132 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)
dstack_product:
12 ms ± 120 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)
cartesian_product:
4.75 ms ± 57.1 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)
cartesian_product_transpose:
7.76 ms ± 52.7 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)
cartesian_product_recursive:
13 ms ± 209 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)
cartesian_product_itertools:
422 ms ± 7.77 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

Во всех случаях, cartesian_product как определено в начале этого ответа, быстрее всего.

Для тех функций, которые принимают произвольное количество входных массивов, также стоит проверить производительность, когда len(arrays) > 2. (Пока я не смогу определить, почему cartesian_product_recursive выдает ошибку в этом случае, я удалил ее из этих тестов.)

In [5]: test_cartesian(*(x100 * 3))
cartesian_product:
8.8 ms ± 138 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)
cartesian_product_transpose:
7.87 ms ± 91.5 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)
cartesian_product_itertools:
518 ms ± 5.5 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

In [6]: test_cartesian(*(x50 * 4))
cartesian_product:
169 ms ± 5.1 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)
cartesian_product_transpose:
184 ms ± 4.32 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)
cartesian_product_itertools:
3.69 s ± 73.5 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

In [7]: test_cartesian(*(x10 * 6))
cartesian_product:
26.5 ms ± 449 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)
cartesian_product_transpose:
16 ms ± 133 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)
cartesian_product_itertools:
728 ms ± 16 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

In [8]: test_cartesian(*(x10 * 7))
cartesian_product:
650 ms ± 8.14 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
cartesian_product_transpose:
518 ms ± 7.09 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
cartesian_product_itertools:
8.13 s ± 122 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

Как показывают эти тесты, cartesian_product остается конкурентоспособным до тех пор, пока количество входных массивов не превысит (примерно) четырех. После этого у cartesian_product_transpose действительно есть небольшое преимущество.

Стоит повторить, что пользователи с другим оборудованием и операционными системами могут видеть другие результаты. Например, unutbu сообщает, что видит следующие результаты для этих тестов с использованием Ubuntu 14.04, Python 3.4.3 и numpy 1.14.0.dev0 + b7050a9:

>>> %timeit cartesian_product_transpose(x500, y500) 
1000 loops, best of 3: 682 µs per loop
>>> %timeit cartesian_product(x500, y500)
1000 loops, best of 3: 1.55 ms per loop

Ниже я расскажу несколько подробностей о более ранних тестах, которые я проводил в этом направлении. Относительная производительность этих подходов со временем менялась для разного оборудования и разных версий Python и numpy. Хотя это не сразу полезно для людей, использующих актуальные версии numpy, это иллюстрирует, как все изменилось со времени выхода первой версии этого ответа.

Простая альтернатива: meshgrid + dstack

Принятый в настоящее время ответ использует tile и repeat для объединения двух массивов. Но meshgrid функция делает практически то же самое. Вот выходные данные tile и repeat перед передачей в transpose:

In [1]: import numpy
In [2]: x = numpy.array([1,2,3])
   ...: y = numpy.array([4,5])
   ...: 

In [3]: [numpy.tile(x, len(y)), numpy.repeat(y, len(x))]
Out[3]: [array([1, 2, 3, 1, 2, 3]), array([4, 4, 4, 5, 5, 5])]

И вот результат meshgrid:

In [4]: numpy.meshgrid(x, y)
Out[4]: 
[array([[1, 2, 3],
        [1, 2, 3]]), array([[4, 4, 4],
        [5, 5, 5]])]

Как вы можете видеть, это почти идентично. Нам нужно только изменить результат, чтобы получить точно такой же результат.

In [5]: xt, xr = numpy.meshgrid(x, y)
   ...: [xt.ravel(), xr.ravel()]
Out[5]: [array([1, 2, 3, 1, 2, 3]), array([4, 4, 4, 5, 5, 5])]

Однако вместо изменения формы на этом этапе мы могли бы передать выходные данные из meshgrid в dstack и изменить форму впоследствии, что сэкономит некоторую работу:

In [6]: numpy.dstack(numpy.meshgrid(x, y)).reshape(-1, 2)
Out[6]: 
array([[1, 4],
       [2, 4],
       [3, 4],
       [1, 5],
       [2, 5],
       [3, 5]])

Вопреки утверждению в этом комментарии, я не видел никаких доказательств того, что разные входные данные будут выдавать выходные данные разной формы, и, как показано выше, они делают очень похожие вещи, поэтому было бы довольно странно, если бы они это делали. Пожалуйста, дайте мне знать, если найдете контрпример.

Тестирование meshgrid + dstack против repeat + transpose

Относительная производительность этих двух подходов со временем менялась. В более ранней версии Python (2.7) результат с использованием meshgrid + dstack был заметно быстрее при небольших входных данных. (Обратите внимание, что эти тесты взяты из старой версии этого ответа.) Определения:

>>> def repeat_product(x, y):
...     return numpy.transpose([numpy.tile(x, len(y)), 
                                numpy.repeat(y, len(x))])
...
>>> def dstack_product(x, y):
...     return numpy.dstack(numpy.meshgrid(x, y)).reshape(-1, 2)
...     

При вводе данных среднего размера я заметил значительное ускорение. Но я повторил эти тесты с более свежими версиями Python (3.6.1) и numpy (1.12.1) на более новой машине. Сейчас эти два подхода практически идентичны.

Старый тест

>>> x, y = numpy.arange(500), numpy.arange(500)
>>> %timeit repeat_product(x, y)
10 loops, best of 3: 62 ms per loop
>>> %timeit dstack_product(x, y)
100 loops, best of 3: 12.2 ms per loop

Новый тест

In [7]: x, y = numpy.arange(500), numpy.arange(500)
In [8]: %timeit repeat_product(x, y)
1.32 ms ± 24.7 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)
In [9]: %timeit dstack_product(x, y)
1.26 ms ± 8.47 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)

Как всегда, YMMV, но это говорит о том, что в последних версиях Python и numpy они взаимозаменяемы.

Обобщенные функции произведения

В общем, мы могли бы ожидать, что использование встроенных функций будет быстрее для небольших входных данных, в то время как для больших входных данных специальная функция может быть быстрее. Кроме того, для обобщенного n-мерного произведения tile и repeat не помогут, потому что у них нет четких многомерных аналогов. Поэтому стоит изучить поведение специально созданных функций.

Большинство соответствующих тестов приведены в начале этого ответа, но вот несколько тестов, выполненных на более ранних версиях Python и numpy для сравнения.

cartesian Функция, определенная в другом ответе, довольно хорошо работала при больших входных данных. (Это то же самое, что и функция, вызванная cartesian_product_recursive выше.) Для сравнения cartesian с dstack_prodct мы используем только два измерения.

И здесь старый тест показал существенную разницу, в то время как новый тест не показывает почти никакой.

Старый тест

>>> x, y = numpy.arange(1000), numpy.arange(1000)
>>> %timeit cartesian([x, y])
10 loops, best of 3: 25.4 ms per loop
>>> %timeit dstack_product(x, y)
10 loops, best of 3: 66.6 ms per loop

Новый тест

In [10]: x, y = numpy.arange(1000), numpy.arange(1000)
In [11]: %timeit cartesian([x, y])
12.1 ms ± 199 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)
In [12]: %timeit dstack_product(x, y)
12.7 ms ± 334 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

Как и раньше, dstack_product все еще превосходит cartesian в меньших масштабах.

Новый тест (избыточный старый тест не показан)

In [13]: x, y = numpy.arange(100), numpy.arange(100)
In [14]: %timeit cartesian([x, y])
215 µs ± 4.75 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)
In [15]: %timeit dstack_product(x, y)
65.7 µs ± 1.15 µs per loop (mean ± std. dev. of 7 runs, 10000 loops each)

Эти различия, я думаю, интересны и их стоит записать; но, в конце концов, они академичны. Как показали тесты в начале этого ответа, все эти версии почти всегда медленнее, чем cartesian_product, определенный в самом начале этого ответа, который сам по себе немного медленнее, чем самые быстрые реализации среди ответов на этот вопрос.

>>> numpy.transpose([numpy.tile(x, len(y)), numpy.repeat(y, len(x))])
array([[1, 4],
       [2, 4],
       [3, 4],
       [1, 5],
       [2, 5],
       [3, 5]])

Смотрите Использование numpy для построения массива из всех комбинаций двух массивов, чтобы узнать общее решение для вычисления декартова произведения N массивов.

Вы можете просто выполнить обычное понимание списка на python

x = numpy.array([1,2,3])
y = numpy.array([4,5])
[[x0, y0] for x0 in x for y0 in y]

что должно дать вам

[[1, 4], [1, 5], [2, 4], [2, 5], [3, 4], [3, 5]]

Меня это тоже заинтересовало, и я провел небольшое сравнение производительности, возможно, несколько более четкое, чем в ответе @senderle.

Для двух массивов (классический случай):

Для четырех массивов:

(Обратите внимание, что длина массивов здесь составляет всего несколько десятков записей.)

Код для воспроизведения графиков:

from functools import reduce
import itertools
import numpy
import perfplot


def dstack_product(arrays):
    return numpy.dstack(numpy.meshgrid(*arrays, indexing="ij")).reshape(-1, len(arrays))


# Generalized N-dimensional products
def cartesian_product(arrays):
    la = len(arrays)
    dtype = numpy.find_common_type([a.dtype for a in arrays], [])
    arr = numpy.empty([len(a) for a in arrays] + [la], dtype=dtype)
    for i, a in enumerate(numpy.ix_(*arrays)):
        arr[..., i] = a
    return arr.reshape(-1, la)


def cartesian_product_transpose(arrays):
    broadcastable = numpy.ix_(*arrays)
    broadcasted = numpy.broadcast_arrays(*broadcastable)
    rows, cols = reduce(numpy.multiply, broadcasted[0].shape), len(broadcasted)
    dtype = numpy.find_common_type([a.dtype for a in arrays], [])

    out = numpy.empty(rows * cols, dtype=dtype)
    start, end = 0, rows
    for a in broadcasted:
        out[start:end] = a.reshape(-1)
        start, end = end, end + rows
    return out.reshape(cols, rows).T


# from https://pythonly.ru/a/1235363/577088
def cartesian_product_recursive(arrays, out=None):
    arrays = [numpy.asarray(x) for x in arrays]
    dtype = arrays[0].dtype

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

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


def cartesian_product_itertools(arrays):
    return numpy.array(list(itertools.product(*arrays)))


perfplot.show(
    setup=lambda n: 2 * (numpy.arange(n, dtype=float),),
    n_range=[2 ** k for k in range(13)],
    # setup=lambda n: 4 * (numpy.arange(n, dtype=float),),
    # n_range=[2 ** k for k in range(6)],
    kernels=[
        dstack_product,
        cartesian_product,
        cartesian_product_transpose,
        cartesian_product_recursive,
        cartesian_product_itertools,
    ],
    logx=True,
    logy=True,
    xlabel="len(a), len(b)",
    equality_check=None,
)