For the 6.438 project part I

supplementary.py

"""Translated to python by Eric Wieser. Use with caution - untested"""
import numpy as np

def bin2gray(x):
    return x ^ (x >> 1)

def gray2bin(x):
    out = x
    mask = x >> 1
    while np.any(mask != 0):
        out ^= mask
        mask >>= 1
    return out


# equivalent to numpy.unpackbits?
def vals_1to8(s, height, width, bits=8):
    Sb = np.zeros(height, width)

    for i in range(bits):
        Sb |= s[i::bits].reshape(height, width) << i

    return Sb

# equivalent to numpy.packbits?
def vals_8to1(Sb, bits=8):
    """converts 8bit (grayscale) images to 1bit data vectors"""

    height, width = Sb.shape
    res = np.zeros(height*width*bits)
    inp = Sb.ravel()
    for i in range(bits):
        res[i::bits] = (inp >> i) & 1
    return res

def msgs_8to1(M_in):
    """
    converts 8bit messages on an image to 1bit messages on a single vector
    M_in: height x width x M
    M_out: (height*width*bits) x 2
    """
    N = np.prod(M_in.shape[:-1])
    M = M_in.shape[-1]
    M_in = M_in.reshape(N, M)

    bits = np.ceil(np.log2(M))
    M_out = zeros(N*bits, 2)

    for i in range(bits):
        index_mask = (np.arange(M) >> i) & 1 # i-th bit of index
        M_out[i::bits,1] = np.sum(M_in[:,index_mask == 0], axis=1) # sumprob of 0's on i-th bit of index
        M_out[i::bits,2] = np.sum(M_in[:,index_mask == 1], axis=1) # sumprob of 1's on i-th bit of index

    return M_out

def msgs_1to8(M_in, height, width):
    """ converts 1bit messages on a single vector to 8bit messages on an image
    % M_in: (heigh*width*bits) x 2
    % M_out: height x width x M
    """

    temp = M_in.reshape([-1, height, width, 2])
    bits = temp.shape[0]
    M = 1 << bits
    M_out = np.ones(height, width, M)

    for i in range(bits):
        index_mask = (np.arange(M) >> i) & 1  # i-th bit of index
        M_out[...,index_mask == 0] *= temp[i,:,:,0]
        M_out[...,index_mask == 1] *= temp[i,:,:,1]

    return M_out

# utility combined functions follow

def msgs_8to1_gray(M_in):
    M_in = np.copy(M_in)
    M_in[...,bin2gray(np.arange(M_in.shape[-1]))] = M_in
    return msgs_8to1(M_in)


def msgs_1to8_gray(M_in, height, width):
    M_out = msgs_1to8(M_in, height, width)
    M_out[...,gray2bin(np.arange(M_out.shape[-1]))] = M_out
    if width == 1: M_out = np.squeeze(M_out)
    return M_out

def vals_8to1_gray(Sb, bits=8):
    return vals_8to1(bin2gray(Sb), bits)

def vals_1to8_gray(Sb, height, width, bits=8):
    return gray2bin(vals_8to1(Sb, height, width, bits))