bit-packing utility

src/kyupy/logic.py

@@ -91,6 +91,29 @@ def bit_in(a, pos):
     return a[pos >> 3] & _bit_in_lut[pos & 7]
 
 
+def unpackbits(a):
+    """Returns an uint8-array containing the unpacked bits of an array `a`.
+    Preserves the shape of `a` and adds a new last axis with the bits of each item.
+    Bits are returned in 'little'-order, i.e., a[...,0] is the least significant bit.
+    """
+    return np.unpackbits(a.view(np.uint8), bitorder='little').reshape(*a.shape, 8*a.itemsize)
+
+
+def packbits(a, dtype=np.uint8):
+    """Packs the values of a boolean-valued array `a` along its last axis into bits.
+    Returns the bits as an array of given type and the shape of `a` with the last axis removed.
+    The last axis of `a` is truncated or padded according to the bit-width of the given data type.
+    Signed integer datatypes are padded with the most significant bit, all others are padded with `0`.
+    """
+    dtype = np.dtype(dtype)
+    bits = 8 * dtype.itemsize
+    a = a[...,:bits]
+    if a.shape[-1] < bits:
+        p = [(0,0)]*(len(a.shape)-1) + [(0, bits-a.shape[-1])]
+        a = np.pad(a, p, 'edge') if dtype.name[0] == 'i' else np.pad(a, p, 'constant', constant_values=0)
+    return np.packbits(a, bitorder='little').view(dtype).reshape(a.shape[:-1])
+
+
 class MVArray:
     """An n-dimensional array of m-valued logic values.