Python: Fastest way of packing a 2d array of binary values into UINT64 array

Question

I have a 2D UINT8 numpy array of size (149797, 64). Each of the elements are either 0 or 1. I want to pack these binary values in each row into a UINT64 value so that i get a UINT64 array of shape 149797 as a result. I tried the following code using numpy bitpack function.

test = np.random.randint(0, 2, (149797, 64),dtype=np.uint8)
col_pack=np.packbits(test.reshape(-1, 8, 8)[:, ::-1]).view(np.uint64)

The packbits function takes about 10 ms to execute. A simple reshaping of this array itself seems to take around 7 ms.I also tried iterating over 2d numpy array using shifting operations to achieve the same result; but there was no speed improvement.

Finally i also want to compile it using numba for CPU.

@njit
def shifting(bitlist):
    x=np.zeros(149797,dtype=np.uint64)  #54
    rows,cols=bitlist.shape
    for i in range(0,rows):             #56
      out=0
      for bit in range(0,cols):
         out = (out << 1) | bitlist[i][bit] # If i comment out bitlist, time=190 microsec
      x[i]=np.uint64(out)  # Reduces time to microseconds if line is commented in njit
    return x

It takes about 6 ms using njit.

Here is the parallel njit version

@njit(parallel=True)
def shifting(bitlist): 
    rows,cols=149797,64
    out=0
    z=np.zeros(rows,dtype=np.uint64)
    for i in prange(rows):
      for bit in range(cols):
         z[i] = (z[i] * 2) + bitlist[i,bit] # Time becomes 100 micro if i use 'out' instead of 'z[i] array'

    return z

It's slightly better wit 3.24ms execution time(google colab dual core 2.2Ghz) Currently, the python solution with swapbytes(Paul's) method seems to be the best one i.e 1.74 ms.

How can we further speed up this conversion? Is there scope for using any vectorization(or parallelization), bitarrays etc, for achieving speedup?

Ref: numpy packbits pack to uint16 array

On a 12 core machine(Intel(R) Xeon(R) CPU E5-1650 v2 @ 3.50GHz),

Pauls method: 1595.0 microseconds (it does not use multicore, i suppose)

Numba code: 146.0 microseconds (aforementioned parallel-numba)

i.e around 10x speedup !!!

Answer 1

You can get a sizeable speedup by using byteswap instead of reshaping etc.:

test = np.random.randint(0, 2, (149797, 64),dtype=np.uint8)

np.packbits(test.reshape(-1, 8, 8)[:, ::-1]).view(np.uint64)
# array([ 1079982015491401631,   246233595099746297, 16216705265283876830,
#        ...,  1943876987915462704, 14189483758685514703,
       12753669247696755125], dtype=uint64)
np.packbits(test).view(np.uint64).byteswap()
# array([ 1079982015491401631,   246233595099746297, 16216705265283876830,
#        ...,  1943876987915462704, 14189483758685514703,
       12753669247696755125], dtype=uint64)

timeit(lambda:np.packbits(test.reshape(-1, 8, 8)[:, ::-1]).view(np.uint64),number=100)
# 1.1054180909413844

timeit(lambda:np.packbits(test).view(np.uint64).byteswap(),number=100)
# 0.18370431219227612

Answer 2

A bit Numba solution (version 0.46/Windows).

Code

import numpy as np
import numba as nb

#with memory allocation
@nb.njit(parallel=True)
def shifting(bitlist):
    assert bitlist.shape[1]==64
    x=np.empty(bitlist.shape[0],dtype=np.uint64)

    for i in nb.prange(bitlist.shape[0]):
        out=np.uint64(0)
        for bit in range(bitlist.shape[1]):
            out = (out << 1) | bitlist[i,bit] 
        x[i]=out
    return x

#without memory allocation
@nb.njit(parallel=True)
def shifting_2(bitlist,x):
    assert bitlist.shape[1]==64

    for i in nb.prange(bitlist.shape[0]):
        out=np.uint64(0)
        for bit in range(bitlist.shape[1]):
            out = (out << 1) | bitlist[i,bit] 
        x[i]=out
    return x

Timings

test = np.random.randint(0, 2, (149797, 64),dtype=np.uint8)

#If you call this function multiple times, only allocating memory 
#once may be enough
x=np.empty(test.shape[0],dtype=np.uint64)

#Warmup first call takes significantly longer
res=shifting(test)
res=shifting_2(test,x)

%timeit res=shifting(test)
#976 µs ± 41.6 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)
%timeit res=shifting_2(test,x)
#764 µs ± 63 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)
%timeit np.packbits(test).view(np.uint64).byteswap()
#8.07 ms ± 52.5 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)
%timeit np.packbits(test.reshape(-1, 8, 8)[:, ::-1]).view(np.uint64)
#17.9 ms ± 91 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)