numpy np.all 轴参数的解决方法；与 numba 的兼容性答案

【问题标题】：workaround for numpy np.all axis argument; compatibility with numbanumpy np.all 轴参数的解决方法；与 numba 的兼容性
【发布时间】：2020-08-01 21:48:00
【问题描述】：

我有一个函数，给定一个 xy 坐标的 numpy 数组，它会过滤位于边 L 的盒子内的那些

import numpy as np
from numba import njit

np.random.seed(65238758)

L = 10
N = 1000
xy = np.random.uniform(0, 50, (N, 2))
box = np.array([
    [0,0],  # lower-left
    [L,L]  # upper-right
]) 

def sinjit(xy, box):
    mask = np.all(np.logical_and(xy >= box[0], xy <= box[1]), axis=1)
    return xy[mask]

如果我运行这个函数，它会返回正确的结果：

sinjit(xy, box)

Output: array([[5.53200522, 7.86890708],
       [4.60188554, 9.15249881],
       [9.072563  , 5.6874726 ],
       [4.48976127, 8.73258166],
       ...
       [6.29683131, 5.34225758],
       [2.68057087, 5.09835442],
       [5.98608603, 4.87845464],
       [2.42049857, 6.34739079],
       [4.28586677, 5.79125413]])

但是，由于我想通过使用 numba 在循环中加速此任务，因此 np.all 函数中的“axis”参数存在兼容性问题（它未在 nopython 模式下实现）。所以，我的问题是，是否有可能以任何方式避免这种争论？有什么解决办法吗？

【问题讨论】：

但是如果你的代码已经被矢量化了，你为什么还要在这里使用 numba？
因为这只是一次迭代，我想把它放在一个 for 循环中，以构建许多具有不同域范围的单元格。并且循环在纯 Python 中效率很低
我明白了，您将遍历许多 xy 数组？
我想遍历很多盒子。主要思想是我有一个边缘 L 的主要 2D 盒子，里面有很多粒子，所以我想将域划分为更小的单元
除此之外，所有的代码都进入一个带有“jitclass”的编译类，那么“numba”的使用是强制性的

标签： python numpy numba

【解决方案1】：

我真的，真的，真的希望 numba 支持可选的关键字参数。在它出现之前，我几乎忽略了它。但是，这里可能会有一些黑客攻击。

您需要特别注意二维以外或长度可能为零的任何东西。

import numpy as np
from numba import njit

@njit(cache=True)
def np_all_axis0(x):
    """Numba compatible version of np.all(x, axis=0)."""
    out = np.ones(x.shape[1], dtype=np.bool8)
    for i in range(x.shape[0]):
        out = np.logical_and(out, x[i, :])
    return out

@njit(cache=True)
def np_all_axis1(x):
    """Numba compatible version of np.all(x, axis=1)."""
    out = np.ones(x.shape[0], dtype=np.bool8)
    for i in range(x.shape[1]):
        out = np.logical_and(out, x[:, i])
    return out

@njit(cache=True)
def np_any_axis0(x):
    """Numba compatible version of np.any(x, axis=0)."""
    out = np.zeros(x.shape[1], dtype=np.bool8)
    for i in range(x.shape[0]):
        out = np.logical_or(out, x[i, :])
    return out

@njit(cache=True)
def np_any_axis1(x):
    """Numba compatible version of np.any(x, axis=1)."""
    out = np.zeros(x.shape[0], dtype=np.bool8)
    for i in range(x.shape[1]):
        out = np.logical_or(out, x[:, i])
    return out

if __name__ == '__main__':
    x = np.array([[1, 1, 0, 0], [1, 0, 1, 0]], dtype=bool)
    np.testing.assert_array_equal(np.all(x, axis=0), np_all_axis0(x))
    np.testing.assert_array_equal(np.all(x, axis=1), np_all_axis1(x))
    np.testing.assert_array_equal(np.any(x, axis=0), np_any_axis0(x))
    np.testing.assert_array_equal(np.any(x, axis=1), np_any_axis1(x))

我不确定它的性能如何，但如果你真的需要在更高级别的 jit'ed 函数中调用该函数，那么这会让你这样做。

【讨论】：

【解决方案2】：

numba 不支持 numpy.all() 的可选标志： https://numba.pydata.org/numba-doc/dev/reference/numpysupported.html

如果你坚持使用 numba，唯一的办法就是以另一种方式编码。

【讨论】：