使用 numpy.frompyfunc 将广播添加到带参数的 python 函数

Use numpy.frompyfunc to add broadcasting to a python function with argument

从像 db 这样的数组(大约 (1e6, 300))和一个 mask = [1, 0, 1] 向量中,我将目标定义为第一列中的 1。

我想创建一个 out 向量,其中包含 db 中对应行与 masktarget==1 相匹配的向量,其他地方都为零。

db = np.array([       # out for mask = [1, 0, 1]
# target,  vector     #
  [1,      1, 0, 1],  # 1
  [0,      1, 1, 1],  # 0 (fit to mask but target == 0)
  [0,      0, 1, 0],  # 0
  [1,      1, 0, 1],  # 1
  [0,      1, 1, 0],  # 0
  [1,      0, 0, 0],  # 0
  ])

我定义了一个 vline 函数,它使用 np.array_equal(mask, mask & vector)mask 应用于每个数组行,以检查向量 101 和 111 是否符合掩码,然后仅保留索引target == 1.

out初始化为array([0, 0, 0, 0, 0, 0])

out = [0, 0, 0, 0, 0, 0]

vline函数定义为:

def vline(idx, mask):
    line = db[idx]
    target, vector = line[0], line[1:]
    if np.array_equal(mask, mask & vector):
        if target == 1:
            out[idx] = 1

通过在 for 循环中逐行应用此函数,我得到了正确的结果:

def check_mask(db, out, mask=[1, 0, 1]):
    # idx_db to iterate over db lines without enumerate
    for idx in np.arange(db.shape[0]):
        vline(idx, mask=mask)
    return out

assert check_mask(db, out, [1, 0, 1]) == [1, 0, 0, 1, 0, 0] # it works !

现在我想通过创建 ufunc:

来矢量化 vline
ufunc_vline = np.frompyfunc(vline, 2, 1)
out = [0, 0, 0, 0, 0, 0]
ufunc_vline(db, [1, 0, 1])
print out

但是 ufunc 抱怨广播具有这些形状的输入:

In [217]:     ufunc_vline(db, [1, 0, 1])
---------------------------------------------------------------------------
ValueError                                Traceback (most recent call last)
<ipython-input-217-9008ebeb6aa1> in <module>()
----> 1 ufunc_vline(db, [1, 0, 1])
ValueError: operands could not be broadcast together with shapes (6,4) (3,)
In [218]:

正在将 vline 转换为相同形状的 numpy ufunc fundamentally doesn't make sense, since ufuncs are always applied to numpy arrays in an elementwise fashion. Because of this, the input arguments must either have the same shape, or must be broadcastable。您将两个形状不兼容的数组传递给 ufunc_vline 函数(db.shape == (6, 4)mask.shape == (3,)),因此您看到的是 ValueError

ufunc_vline 还有一些其他问题:

  • np.frompyfunc(vline, 2, 1) 指定 vline 应该 return 单个输出参数,而 vline 实际上 return 什么都没有(但修改out到位).

  • 您将 db 作为第一个参数传递给 ufunc_vline,而 vline 期望第一个参数是 idx,它用于作为 db.

  • 行的索引

此外,请记住,使用 np.frompyfunc 从 Python 函数创建 ufunc 不会比标准 Python for 循环产生任何明显的性能优势。要看到任何重大改进,您可能需要使用 low-level 语言(例如 C)编写 ufunc 代码(请参阅文档中的 this example)。


话虽如此,您的 vline 函数可以使用标准布尔数组运算轻松向量化:

def vline_vectorized(db, mask): 
    return db[:, 0] & np.all((mask & db[:, 1:]) == mask, axis=1)

例如:

db = np.array([       # out for mask = [1, 0, 1]
# target,  vector     #
  [1,      1, 0, 1],  # 1
  [0,      1, 1, 1],  # 0 (fit to mask but target == 0)
  [0,      0, 1, 0],  # 0
  [1,      1, 0, 1],  # 1
  [0,      1, 1, 0],  # 0
  [1,      0, 0, 0],  # 0
  ])

mask = np.array([1, 0, 1])

print(repr(vline_vectorized(db, mask)))
# array([1, 0, 0, 1, 0, 0])