gpt4 book ai didi

python - Numba 中的笛卡尔积

转载 作者:行者123 更新时间:2023-12-03 22:40:56 27 4
gpt4 key购买 nike

我的代码使用如下列表的笛卡尔积:

import itertools

cartesian_product = itertools.product(list('ABCDEF'), repeat=n)
n可以是 0 到 4 之间的任何值。
numba当前不支持 itertools.product .到目前为止,我一直无法想出一个可行的替代方案。欢迎任何建议!

最佳答案

这是解决您的问题的一种方法。有两个 numba 函数:一个用于常规笛卡尔积(数组列表),另一个用于简单数组的重复参数。这是对一些代码的改编 from another one of my SO answers .

@nb.njit(nb.int32[:,:](nb.int32[:]))
def cproduct_idx(sizes: np.ndarray):
"""Generates ids tuples for a cartesian product"""
assert len(sizes) >= 2
tuples_count = np.prod(sizes)
tuples = np.zeros((tuples_count, len(sizes)), dtype=np.int32)
tuple_idx = 0
# stores the current combination
current_tuple = np.zeros(len(sizes))
while tuple_idx < tuples_count:
tuples[tuple_idx] = current_tuple
current_tuple[0] += 1
# using a condition here instead of including this in the inner loop
# to gain a bit of speed: this is going to be tested each iteration,
# and starting a loop to have it end right away is a bit silly
if current_tuple[0] == sizes[0]:
# the reset to 0 and subsequent increment amount to carrying
# the number to the higher "power"
current_tuple[0] = 0
current_tuple[1] += 1
for i in range(1, len(sizes) - 1):
if current_tuple[i] == sizes[i]:
# same as before, but in a loop, since this is going
# to get run less often
current_tuple[i + 1] += 1
current_tuple[i] = 0
else:
break
tuple_idx += 1
return tuples

@nb.njit
def cartesian_product(*arrays):
sizes = [len(a) for a in arrays]
sizes = np.asarray(sizes, dtype=np.int32)
tuples_count = np.prod(sizes)
array_ids = cproduct_idx(sizes)
tuples = np.zeros((tuples_count, len(sizes)))
for i in range(len(arrays)):
tuples[:, i] = arrays[i][array_ids[:, i]]
return tuples

@nb.njit
def cartesian_product_repeat(array, repeat):
sizes = [len(array) for _ in range(repeat)]
sizes = np.asarray(sizes, dtype=np.int32)
tuples_count = np.prod(sizes)
array_ids = cproduct_idx(sizes)
tuples = np.zeros((tuples_count, len(sizes)))
for i in range(repeat):
tuples[:, i] = array[array_ids[:, i]]
return tuples
这是最后一个函数执行的示例:
>>> cartesian_product_repeat(np.arange(2), 3)

array([[0., 0., 0.],
[1., 0., 0.],
[0., 1., 0.],
[1., 1., 0.],
[0., 0., 1.],
[1., 0., 1.],
[0., 1., 1.],
[1., 1., 1.]])

关于python - Numba 中的笛卡尔积,我们在Stack Overflow上找到一个类似的问题: https://stackoverflow.com/questions/57128975/

27 4 0
Copyright 2021 - 2024 cfsdn All Rights Reserved 蜀ICP备2022000587号
广告合作:1813099741@qq.com 6ren.com