An efficient numpy implementation of GLSL packHalf2x16 and unpackHalf2x16

packHalf2x16.py

"""
An efficient numpy implementation of GLSL packHalf2x16 and unpackHalf2x16.
Author: Elie Michel <elie.michel@exppad.com>
License: MIT - 2023
"""

def packHalf2x16(x, y):
    """
    Pack two floats into a single uint32, like GLSL's packHalf2x16 function
    (but vectorized for numpy arrays).
    """
    assert(x.shape == y.shape)
    assert(len(x.shape) == 1)
    
    # Allocate output data
    n = x.shape[0]
    uint32_data = np.ndarray((n,), dtype=np.uint32)

    # Create view
    f16_data = np.lib.stride_tricks.as_strided(
        uint32_data.view(dtype=np.float16),
        shape=(2, n),
        strides=(1*2, 2*2),
        writeable=True,
    )

    # Convert from whatever type x and y use
    f16_data[0] = x
    f16_data[1] = y

    return uint32_data

def unpackHalf2x16(uint32_data):
    """Inverse of packHalf2x16"""
    n = uint32_data.shape[0]
    return np.lib.stride_tricks.as_strided(
        uint32_data.view(dtype=np.float16),
        shape=(2, n),
        strides=(1*2, 2*2),
    )

def test():
    a = np.array([3.14, 42.0, 1.23])
    b = np.array([9.81, 2.34, 27.0])
    packed = packHalf2x16(a, b)
    print(packed)
    unpacked = unpackHalf2x16(packed)
    print(unpacked)
    assert(np.isclose(unpacked[0], a, atol=1e-02).all())
    assert(np.isclose(unpacked[1], b, atol=1e-02).all())