numpyのpackbitで困ったのでメモ

numpyで2進数を10進数に変換する関数としてnp.packbitsがありますが、ふと使おうとしたところ、少し思った挙動とは違っていました。

そもそもやりたかったことは、「$N$ 個の $r$ bit データを$N$個の10進数表現にする」ということです。

想定外の挙動というのは以下の２つです。

1. 8bitより大きい場合には、8bit区切りで計算する
2. 8bitより小さい場合には、二進数を左シフトで0埋めをしてから計算する

想定していたのは

1. 8bitよりも大きくても特に区切らない
2. 8bitよりも小さい場合には二進数を右シフトで0埋めをしてから計算する
   という感じでした。

わかりにくいので例で示すと

import numpy as np
# 入力が8bitの場合　(思ったとおりの挙動)
A = np.array([[1, 0, 0, 0, 0, 0, 0, 0]])
print(np.packbits(A)) # ⇒ [128]
# 入力が9bitの場合 ([256]になってほしかった)
B = np.array([[1, 0, 0, 0, 0, 0, 0, 0, 0]])
print(np.packbits(B)) # ⇒ [128, 0]
# 入力が3bitの場合 ([4]になってほしかった)
C = np.array([[1, 0, 0]])
print(np.packbits(C)) # ⇒ [128]

という感じです。

これをどうにかするために、場合分けをしてごちゃごちゃやっていたのですが、結局以下の方法でいいじゃんとなりました。つまり二進数の定義どおりに計算するという当たり前の方法です。

# reverse: X[..., -1]を第1bitとする場合はTrue
#          X[..., 0] を第1bitとする場合はFalse
def mypackbits(X, reverse=True):
    p = np.power(2, np.arange(X.shape[-1]))
    if reverse:
        p = p[::-1]
    return np.dot(X, p)

計算時間はnp.packbits(X)よりもmypackbits(X)のほうが短かったです。

r = 8
N = 1000
X = np.random.randint(low=0, high=2, size=[N, r])
%timeit np.packbits(X) # 58.4 µs ± 174 ns per loop (mean ± std. dev. of 7 runs, 10000 loops each)
%timeit mypackbits(X) #  17.1 µs ± 278 ns per loop (mean ± std. dev. of 7 runs, 100000 loops each)

ニッチですが誰かの参考になればいいなーと思います。