在 Python 中使用掩码选择子向量的最佳方法答案

【问题标题】：Best way of using masks to select sub-vectors in Python在 Python 中使用掩码选择子向量的最佳方法
【发布时间】：2017-08-25 01:29:07
【问题描述】：

我们想要选择一维数组v_data 的一些元素。我们需要做的处理需要遍历v_data的子向量。

现在我使用像Gosper's hack 这样的字节逻辑来创建一个整数n_mask，它的二进制表示对应于我想要的v_data 的索引。 n_mask可以通过一个方法转换成二进制向量：

def num2bv(num, n_len):
    """Convert a number to a binary vector of some length"""
    return [bool((2**ii & num)//(2**ii)) for ii in reversed(range(0, n_len))]

设置bv_mask = num2bv(n_mask, len(v_data))，运行v_data[bv_using]可以恢复子向量

这是一个不好的方法吗？我特别担心：

在实践中使用二进制向量进行索引会很慢
num2bv 在实践中会很慢
能否将此技术用于任何长度的向量取决于 Python 的任意精度整数，这可能很慢或不可移植

这些担忧是否合理？

【问题讨论】：

我认为num2bv 确实会很慢。所以你可以考虑用np.unpackbits替换它。但是，这可能不适用于任意精度整数:(
itertools 有一个功能可以做到这一点，它被称为compress。虽然它不使用位，但它采用任何“真/假”值序列（您应该能够通过生成器生成）。
n_len 很大（如 1000000），num2bv 非常缓慢。这是使用n_len 进行的列表理解缩放。相比之下，np.ix_ 只是np.nonzero，找到所有 True 的索引，并在编译后的代码中进行迭代。 v_data[idx] 生成的索引列表也比较快。
vdata[bv_mask] 隐式执行 vdata[np.nonzero(bv_mask)] （可以在 timeits 中看到。如果 vdata 是一个列表，itertools.compress 很有用，但如果它已经是一个 numpy 数组，则不需要。数组索引更快。

标签： python performance numpy indexing mask

【解决方案1】：

num2bv 的数组版本是：

def foo(num, n_len):
    c = np.int64(2)**np.arange(n_len)[::-1]
    return (np.bitwise_and(c,num)//c).astype(bool)

In [713]: N=16
In [714]: foo(2**(N//2)-1,N)
Out[714]: 
array([False, False, False, False, False, False, False, False,  True,
        True,  True,  True,  True,  True,  True,  True], dtype=bool)
In [715]: np.array(num2bv(2**(N//2)-1,N))
Out[715]: 
array([False, False, False, False, False, False, False, False,  True,
        True,  True,  True,  True,  True,  True,  True], dtype=bool)

In [716]: N=32
In [717]: timeit np.array(num2bv(2**(N//2)-1,N))
10000 loops, best of 3: 39.7 µs per loop
In [718]: timeit foo(2**(N//2)-1,N)
The slowest run took 4.72 times longer than the fastest. This could mean that an intermediate result is being cached.
100000 loops, best of 3: 15.7 µs per loop

但是对于较大的N，数组版本开始遇到 2**i 的整数表示问题。

对于较大的N，num2bv 时间支配索引，例如N=1024

In [734]: timeit num2bv(2**(N//2)-1,N)
100 loops, best of 3: 4.34 ms per loop
In [735]: timeit np.nonzero(num2bv(2**(N//2)-1,N))
100 loops, best of 3: 4.4 ms per loop
In [736]: A=np.ones(N)
In [738]: timeit A[num2bv(2**(N//2)-1,N)]
100 loops, best of 3: 4.39 ms per loop

np.binary_repr 返回一个字符串表示。它使用 Python bin。创建掩码的一种方法是用list 拆分它，然后让np.array 将其转换为布尔数组

In [846]: np.binary_repr(100,16)
Out[846]: '0000000001100100'
In [848]: np.array(list(np.binary_repr(100,16)),bool)
Out[848]: 
array([ True,  True,  True,  True,  True,  True,  True,  True,  True,
        True,  True,  True,  True,  True,  True,  True], dtype=bool)

它没有to_numpy_mask_2 快，但仍然比num2bv 有很大改进。并且不受我之前foo的大小限制。

In [842]: timeit np.array(num2bv(100,120))
10000 loops, best of 3: 166 µs per loop
In [843]: timeit to_numpy_mask_2(100,120)
The slowest run took 5.07 times longer than the fastest. This could mean that an intermediate result is being cached.
100000 loops, best of 3: 7.6 µs per loop
In [850]: timeit np.array(list(np.binary_repr(100,120)),bool)
100000 loops, best of 3: 18 µs per loop

【讨论】：

【解决方案2】：

将.to_bytes 整数方法与np.unpackbits 结合起来相当快。一个 1,000,000 位的字在不到一毫秒的时间内被转换：

import numpy as np
import random
from timeit import timeit

# a slower approach using builtin 'bin' function
def to_numpy_mask_1(n, bits = 120):
    return (np.frombuffer(bin(n + 2**bits)[-bits:].encode('utf8'),
                          dtype=np.uint8)-48).view(bool)

# the real thing based on '.to_bytes'
def to_numpy_mask_2(n, bits = 120):
    return np.unpackbits(np.frombuffer(n.to_bytes((bits-1)//8 + 1, 'big'),
                                       dtype=np.uint8)).view(bool)[-bits:]

# check

base = 2**np.arange(120)[::-1].astype(object)
n = random.randint(0, 2**120)
print(n, base[to_numpy_mask_1(n)].sum(), base[to_numpy_mask_2(n)].sum())

# benchmark

# translation only, no indexing
N = 10**6
n = random.randint(0, 2**N)
print('{:8.6g} secs'.format(timeit(lambda: to_numpy_mask_1(n, bits = N),
                                   number=10)/10))
print('{:8.6g} secs'.format(timeit(lambda: to_numpy_mask_2(n, bits = N),
                                   number=10)/10))

# including indexing
data = np.random.randn(N)
print('{:8.6g} secs'.format(timeit(lambda: data[to_numpy_mask_1(n, bits = N)],
                                   number=10)/10))
print('{:8.6g} secs'.format(timeit(lambda: data[to_numpy_mask_2(n, bits = N)],
                                   number=10)/10))

样本输出：

# 303734588154968662776606530859339928 303734588154968662776606530859339928 303734588154968662776606530859339928
# 0.00622677 secs
# 0.00051558 secs
# 0.0121338 secs
# 0.00697929 secs

【讨论】：

【解决方案3】：

根据向量的长度，使用itertools.combinations 可能会快得多：

In [2] v = np.array(range(15))

In [3]: %time x = [v[num2bv(i,15)] for i in range(2**15)]
CPU times: user 498 ms, sys: 5.83 ms, total: 504 ms
Wall time: 506 ms

In [4]: %time x = [c for i in range(15) for c in combinations(v,i)]
CPU times: user 5.67 ms, sys: 1.53 ms, total: 7.2 ms
Wall time: 6.91 ms

【讨论】：