一次生成 numpy 数组元素的组合，并且在运行中答案

【问题标题】：Generate combinations of numpy array elements in one go, and on the fly一次生成 numpy 数组元素的组合，并且在运行中
【发布时间】：2021-02-06 10:18:27
【问题描述】：

我知道这是一个非常琐碎的问题，但我无法找到合适的解决方案。我有两个从两个不同的元组列表创建的 numpy 数组。

    a = np.array([(1, 2), (3, 4), (5, 6)....])
    b = np.array([(100, 200), (300, 400), (500, 600)....])

我有一个函数，它将两个元组作为输入并返回单个值作为输出

    def myfunc(x, y):
        .....
        return val

    vectorized_myfunc = np.vectorize(myfunc, signature='(n),(n)->()')

我想要达到的目标：

使用b 中的每个元素生成a 中的所有元素组合。喜欢((1, 2), (100, 200)), ((1, 2), (300, 400))..
将生成的组合传递给vectorized_myfunc。
如果我想动态生成这些组合的块，然后执行 1) 和 2)，该怎么办？

PS：我知道我可以使用itertools.product(..)，但它很慢。

PPS：我尝试使用 np.stack(np.meshgrid(a, b), -1).reshape(-1, 2, 3)，然后尝试将其传递给 vectorized_myfunc，但无法正常工作

我愿意为它使用numpy、dask 甚至numba

【问题讨论】：

有些行为不符合您的要求。当您创建 a 和 b 时，它们失去了不可变的属性。所以直接用元组创建一个元组列表是不可能的，因为它们不再是元组了。
那是一个错字。已更正
您的问题仍然包括带有元组的 np.arrays。在你告诉你想要实现什么的那一行，你想要的结果是一个元组的元组。

标签： python arrays numpy vectorization

【解决方案1】：

您可以像这样处理a 和b 的索引组合：

def product(arrays):
    # the same as itertools.product(*arrays)
    return np.stack(np.meshgrid(*arrays), axis=-1).reshape(-1, len(arrays))

a = np.array([(1, 2), (3, 4), (5, 6)])
b = np.array([(100, 200), (300, 400), (500, 600)])
idx = product([np.arange(len(a)), np.arange(len(b))])

检查：

>>> idx
array([[0, 0],
       [1, 0],
       [2, 0],
       [0, 1],
       [1, 1],
       [2, 1],
       [0, 2],
       [1, 2],
       [2, 2]])

我对签名没有经验，但以下示例对我有用：

vectorized_myfunc = np.vectorize(lambda x,y: x+y)
>>> vectorized_myfunc(a[idx[:, 0]], b[idx[:, 1]])
array([[101, 202],
      [103, 204],
      [105, 206],
      [301, 402],
      [303, 404],
      [305, 406],
      [501, 602],
      [503, 604],
      [505, 606]])

更新：我发现了另一个适用于 signature 的示例：

vectorized_myfunc = np.vectorize(lambda x,y: np.sum(np.array([x,y])), signature='(j),(j)->()')
>>> vectorized_myfunc(a[idx[:, 0]], b[idx[:, 1]])
array([ 303,  307,  311,  703,  707,  711, 1103, 1107, 1111])

【讨论】：