Python中快速获取二进制列的位运算方法

Question

在 Python 中，有没有一种高效的方法来获取位数组中第 n 个位置上的布尔值数组？

1. 创建一个值为 0 或 1 的 NumPy 数组：

import numpy as np

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

2. 使用 np.packbits 压缩大小：

pack_array = np.packbits(array, axis=1)

3. 期望结果 - 一个函数，可以从位数组中获取第 n 列的所有值。例如，如果我想要第二列，我希望获得与调用 array[:,1] 相同的结果:

array([0, 1, 0])

我尝试了以下函数的numba版本，虽然结果正确但速度非常慢：

import numpy as np
from numba import njit

@njit(nopython=True, fastmath=True)
def getVector(packed, j):
    n = packed.shape[0]
    res = np.zeros(n, dtype=np.int32)
    for i in range(n):
        res[i] = bool(packed[i, j//8] & (128>>(j%8)))
    return res

如何进行测试？

import numpy as np
import time
from numba import njit

array = np.random.choice(a=[False, True], size=(100000000,15))

pack_array = np.packbits(array, axis=1)

start = time.time()
array[:,10]
print('np array')
print(time.time()-start)

@njit(nopython=True, fastmath=True)
def getVector(packed, j):
    n = packed.shape[0]
    res = np.zeros(n, dtype=np.int32)
    for i in range(n):
        res[i] = bool(packed[i, j//8] & (128>>(j%8)))
    return res

# To initialize
getVector(pack_array, 10)

start = time.time()
getVector(pack_array, 10)
print('getVector')
print(time.time()-start)

它返回：

np array
0.00010132789611816406
getVector
0.15648770332336426

- Lukáš Sýkora

---

你可以在循环外计算j//8和128>>(j%8)，创建一个np.empty(使用dtype=np.bool?)作为res。但这些只是微小的优化，可能已经被编译器完成了。 - Michael Szczesny

除了取模运算符外，我无法想象这种实现在计算上会非常缓慢且不受内存限制。 - Michael Szczesny

numpy的方法是O(1)，你不能将其作为基准。它只返回一个调整过步幅的视图，没有任何计算。计时结果应该由print调用主导。 - Michael Szczesny

@MichaelSzczesny 是否可以使用步幅在字节内进行索引？从阅读numpy源代码来看，packbits只是对数组的循环。 - Nick ODell

令人惊讶的是，LLVM 在循环内部没有对明显的常量进行优化（在我的机器上）。将它们移到循环外部后，速度提高了约3.5倍。 - Michael Szczesny

Answer 1

除了一些微小的优化，我认为这里没有太多可以优化的地方。你的代码中还有一些小错误：

• @njit(nopython=True)重复了相同的事情（njit中的n已经代表了nopython模式）。应该使用@njit或@jit(nopython=True)
• fastMath用于在进行浮点数运算时“走捷径”，由于我们只使用整数和布尔值，因此它可以安全地删除，因为在这里不起作用。

我的更新后的代码（在我的机器上看到了微薄的40%性能提升）：

import numba as nb
import numpy as np

np.random.seed(0)
array = np.random.choice(a=[False, True], size=(10000000,15))

pack_array = np.packbits(array, axis=1)

@nb.njit(locals={'res': nb.boolean[:]})
def getVector(packed, j):
    n = packed.shape[0]
    res = np.zeros(n, dtype=nb.boolean)
    byte = j//8
    bit = 128>>(j%8)
    for i in range(n):
        res[i] = bool(packed[i, byte] & bit)
    return res

getVector(pack_array, 10)

在您的回答中，“res”是一个由32位整数组成的列表，通过将numba（而不是numpy）布尔数据类型提供给np.zeros（），我们可以将其转换为更高效的布尔值。这就是大部分性能改进的来源。在我的机器上，将j_mod和j_flr放在循环外面没有明显的影响。但对于评论者@Michael Szczesny有影响，所以它也可能对您有帮助。

我不建议使用strides，正如@Nick ODell建议的那样，因为如果使用不当，它们可能非常危险（请参阅numpy文档）。

编辑：我已经根据Michael的建议进行了一些小的修改。(谢谢)

- Rafnus

---

Colab Notebook 是我用于基准测试的优化实现。不同硬件将产生不同的结果。 - Michael Szczesny

纯numpy的(pack_array[:, j // 8] & 128>>(j%8)).astype(bool)比原始实现要快约2倍。 - Michael Szczesny

我猜function2应该被称为getVector，是吗？ - Jérôme Richard

@Jérôme Richard，你是对的，我现在已经修复了。 - Rafnus