时间戳

首页 > 解决方案 > 一次生成 numpy 数组元素的组合，并且在运行中

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

问题描述

我知道这是一个非常微不足道的问题，但我无法找到合适的解决方案。我有两个从两个不同的元组列表创建的 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)->()')

我想要达到的目标：

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

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

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

我愿意使用，numpy甚至dask愿意numba

标签： pythonarraysnumpyvectorization

解决方案

您可以像这样处理索引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])

推荐阅读