Compressing integers to 12-bit packed format with zarr and numcodecs

Question

I am working with 12-bit integer image data (i.e. data aquired from a camera that uses a 12-bit ADC) which I ultimately store in a Zarr array. Currently I store the images as 16-bit integers, which means I am wasting 30% extra memory. I would like to avoid this extra memory use by storing the images using a "12-bit packed format", i.e. by storing two pixel values in 3 bytes. See the description of MONO12_PACKED for a complete description of this format.

The most convenient way to do this would be if zarr/numcodecs had a compressor for 16-bit integer data which converted it to this 12-bit packed format. Then, assuming my image data consists of nt time points and each image is of size ny x nx, I would run something like

import zarr

img = load() # load image of size nt x ny x nx

z = zarr("test.zarr", "w")
z.array("images", img, chunks=(1, ny, nx), dtype="uint16", compressor=mono12_packed)

However, currently numcodecs does not appear to implement such a compressor. A quick look at their packbits compressor suggests it would be easy to implement a 12-bit packing compressor. However, the downside of writing this myself is I would then need to distribute that code together with my zarr file for someone else to be able to read it.

After some experimenting, creating a mono12 packed codec is doable (see below).

I am working with datasets which are ~90GB when stored as uint16. When saving these to hard disk I am IO bound, and using compression speeds up the saving.

I find that the below implementation of mono12 compresses the data by the expected factor of 4:3 and runs ~20% faster. I find that zlib compresses ~2:1 and runs ~25% faster and bz2 compresses ~3:1 and runs ~2x slower.

However, none of this directly addresses my question, which is: is there an existing implementation of the mono12 compressor which is compatible with zarr/numcodecs? Or if not, is this kind of compressor general enough to be included in numcodecs (since it only works on a very specific subclass of 16-bit integers)?

import numpy as np
import zarr
from numcodecs.abc import Codec
from numcodecs.compat import ensure_ndarray, ndarray_copy
from numcodecs.registry import register_codec

def pack12(data):
    """
    Convert data from uint16 integers to 12-bit packed format

    :param data: even sized array of uint16
    :return bytes:
    """

    # most significant 8 bits of first integer
    first_bit = (data[::2] >> 4).astype(np.uint8)
    # least significant 4 bits of first integer and most significant 4 bits of second
    second_bit = ((np.bitwise_and(15, data[::2]) << 4) + (data[1::2] >> 8)).astype(np.uint8)
    # least significant 8 bits of second integer
    third_bit = np.bitwise_and(255, data[1::2]).astype(np.uint8)

    return np.stack((first_bit, second_bit, third_bit), axis=1).ravel()

def unpack12(data: np.ndarray) -> np.ndarray:
    """
    Convert data from 12-bit packed integers to uint16 integers

    :param data: an array of uint8 integers
    :return img: an array of uint16 integers
    """

    # convert from 12L packed (two 12-bit pixels packed into three bytes)
    a_uint8 = data[::3].astype(np.uint16)
    b_uint8 = data[1::3].astype(np.uint16)
    c_uint8 = data[2::3].astype(np.uint16)

    # middle byte contains least-significant bits of first integer
    # and most-significant bits of second integer
    first_int = (a_uint8 << 4) + (b_uint8 >> 4)
    second_int = (np.bitwise_and(15, b_uint8) << 8) + c_uint8

    img = np.stack((first_int, second_int), axis=1).ravel()

    return img

class mono12(Codec):
    codec_id = 'mono12'

    def init__(self):
        pass

    def encode(self, buf):
        # normalise input
        arr = ensure_ndarray(buf)

        # flatten to simplify implementation
        arr = arr.reshape(-1, order='A')

        # determine size of packed data
        n = arr.size
        n_bytes_packed = (n // 2) * 3
        n_ints_leftover = n % 2
        if n_ints_leftover > 0:
            n_bytes_packed += 3

        # setup output
        enc = np.empty(n_bytes_packed + 1, dtype='u1')

        # store how many bits were padded
        if n_ints_leftover > 0:
            n_ints_padded = 1
            # apply encoding
            enc[1:] = pack12(np.concatenate((arr, np.array([0], dtype=np.uint16))))
        else:
            n_ints_padded = 0
            # apply encoding
            enc[1:] = pack12(arr)
        enc[0] = n_ints_padded

        return enc

    def decode(self, buf, out=None):
        # normalise input
        enc = ensure_ndarray(buf).view('u1')

        # flatten to simplify implementation
        enc = enc.reshape(-1, order='A')

        # find out how many integers were padded
        n_ints_padded = int(enc[0])

        # apply decoding
        dec = unpack12(enc[1:])

        # remove padded bits
        if n_ints_padded:
            dec = dec[:-n_ints_padded]

        # handle destination
        return ndarray_copy(dec, out)

# need to register codec if want to load data later
register_codec(mono12)

img = load() # load image of size nt x ny x nx

z = zarr("test.zarr", "w")
z.array("images", img, chunks=(1, ny, nx), dtype="uint16", compressor=mono12)