Numpy:x和y数组的笛卡尔积指向2D点的单个数组

我有两个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数组不是最有效的形式。

 >>> 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个数组的笛卡尔乘积的一般解决方案。

一个规范的cartesian_product (几乎):

对于这个问题,有不同的属性有很多方法。 有的比别人快,有的更通用。 经过大量的测试和调整,我发现以下计算n维cartesian_product函数比许多输入的函数要快。 对于高维输入,其他版本的Python和numpy ,以及其他硬件(见下文),一个替代方案cartesian_product_transpose有时可能更快。

尽管在性能上存在这些差异,但大多数证据表明, cartesian_productcartesian_product_transpose如下所定义的应该是在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) 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 

第一个函数以不寻常的方式使用ix_ ; 而记录的ix_使用是将索引生成 一个数组中,恰好相同形状的数组可以用于广播分配。 非常感谢mgilson ,他鼓励我以这种方式尝试使用ix_ ,并且对unutbu提供了一些非常有用的反馈,包括使用numpy.result_type的建议。

测试替代品

以下是一系列测试,显示了这些功能相对于多种选择所提供的性能提升。 这里显示的所有测试是在运行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://stackoverflow.com/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

目前接受的答案使用tilerepeat播放两个阵列在一起。 但是meshgrid功能几乎是一样的。 以下是tile的输出,并在传递给转置之前repeat

 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的输出meshgriddstack然后重新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 vs. 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维产品, tilerepeat将无济于事,因为它们没有清晰的高维度类似物。 所以值得研究特定功能的行为。

大部分相关的测试都出现在这个答案的开头,但是这里是一些在早期版本的Python和numpy进行比较的测试。

在另一个答案中定义的cartesian函数用于对较大的输入表现相当好。 (与上面的cartesian_product_recursive函数相同。)为了比较cartesiandstack_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慢。

你可以在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]] 

我也对这个感兴趣,并做了一些小小的比较,可能比@Senslele的答案要清楚些。

对于两个阵列(经典案例):

在这里输入图像描述

对于四个阵列:

在这里输入图像描述

(请注意,数组的长度在这里只有几十个条目。)

在这两种情况下,@ 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://stackoverflow.com/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: 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 ) 

更一般地说,如果你有两个2d numpy数组a和b,并且你想把a的每一行连接到b的每一行(一个行的笛卡尔乘积,类似于数据库中的连接),你可以使用这个方法:

 import numpy def join_2d(a, b): assert a.dtype == b.dtype a_part = numpy.tile(a, (len(b), 1)) b_part = numpy.repeat(b, len(a), axis=0) return numpy.hstack((a_part, b_part)) 

senderle和Nico Schlomer有很好的答案。 我只是添加一个脚注。

截至2017年10月,numpy现在有一个通用的np.stack函数,它需要一个axis参数。 使用它,我们可以使用“dstack和meshgrid”技术来创建一个“广义笛卡尔积”:

 import numpy as np def cartesian_product(*arrays): ndim = len(arrays) return np.stack(np.meshgrid(*arrays), axis=-1).reshape(-1, ndim) 

请注意axis=-1参数。 这是结果中的最后一个(最内侧的)轴。 这相当于使用axis=ndim

另外一个意见是,由于笛卡尔产品很快就炸毁了,除非由于某种原因我们需要在内存中实现这​​个数组,否则如果产品非常大,我们可能需要利用itertools并且使用这些值。