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

Question

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

我想创建一个 out 向量，其中包含 db 中的相应行与 mask 和 target= =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]:

Answer 1

将 vline 转换为 numpy ufunc从根本上讲没有意义，因为 ufunc 总是以元素方式应用于 numpy 数组。因此，输入参数必须具有相同的形状，或者必须是 broadcastable。到相同的形状。您将两个形状不兼容的数组传递给 ufunc_vline 函数(db.shape == (6, 4) 和 mask.shape == (3,) )，因此您看到的是 ValueError。

ufunc_vline 还有一些其他问题:

• np.frompyfunc(vline, 2, 1) 指定 vline 应该返回单个输出参数，而 vline 实际上什么都不返回(但就地修改 out)。

• 您正在将 db 作为第一个参数传递给 ufunc_vline，而 vline 期望第一个参数是 idx ，用作 db 行的索引。

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

---

话虽如此，您的 vline 函数可以使用标准 bool 数组操作轻松矢量化:

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])