Vectorized Run-Length Encoding

rle.py

import numpy as np

def encode(v):
    changes = np.flatnonzero(np.diff(v)) + 1
    values = np.r_[v[0], v[changes]]
    runlengths = np.diff(np.r_[0, changes, len(v)])
    return runlengths, values

def decode(runlengths, values):
    return np.repeat(values, runlengths)

def calc_message_length(runlengths, values=None):
    return 8*runlengths.nbytes \
            + (8*values.nbytes if values is not None else len(runlengths))

test_.py

import numpy as np
import rle


def test_run_length():
    M = np.array([[4, 0, 0, 1, 0, 0],
                  [0, 0, 2, 2, 2, 4]])
    runlengths, values = rle.encode(M.flatten())
    M_decoded = rle.decode(runlengths, values).reshape(M.shape)

    np.testing.assert_equal(values, [4, 0, 1, 0, 2, 4])
    np.testing.assert_equal(runlengths, [1, 2, 1, 4, 3, 1])
    np.testing.assert_equal(M_decoded, M)
    assert runlengths.sum() == 12

def test_calc_message_length():
    M = np.array([[4, 0, 0, 1, 0, 0],
                  [0, 0, 2, 2, 2, 4]])
    v = np.array([0, 0, 1, 0, 0, 0, 1])

    runlengths, values = rle.encode(M.flatten())
    
    assert rle.calc_message_length(
            runlengths.astype('uint8'), values.astype('uint16')) == (8 + 16)*6

    assert rle.calc_message_length(
            runlengths.astype('int32'), values.astype('int64')) == (32 + 64)*6

    runlengths = rle.encode_binary(v)
    assert rle.calc_message_length(runlengths.astype('int16')) == 16*4 + 4