pythonnumpyperformanceoptimizationnumpy-ndarray

Is there a Numpy method or function to split an array of uint64 into two arrays of uint32


Say I have an array as follows:

arr = np.asarray([1, 2, 3, 4294967296, 100], dtype=np.uint64)

I now want two arrays, one array with the lower 32 bits of every element, and one with the upper 32 bits of every element, preferably by using views and minimizing copies, to get something like this:

upper = np.array([0, 0, 0, 1, 0],   dtype=np.uint32)
lower = np.array([1, 2, 3, 0, 100], dtype=np.uint32)

I tried the following:

lower = arr.view() & 0xFFFFFFFF
upper = np.bitwise_right_shift(arr.view(), 32)

But this results in a copy for the upper bits due to the bitshift, and both arrays are still of type uint64. Are there further optimizations I can try or am I out of luck and need to eat up the extra copies?


Solution

  • You can use structured arrays to split uint64 into two uint32 views without copying:

    # Create a structured view of the array (assuming little-endian system)
    view = arr.view(dtype=np.dtype([('lower', np.uint32), ('upper', np.uint32)]))
    # Extract views
    lower = view['lower']
    upper = view['upper']
    

    This creates memory views not copies, and preserves the uint32 dtypes.

    Alternative using views:

    This alternative also creates views without copying data, but my benchmarks show it can be significantly faster than the structured dtype approach.

    # View the uint64 array as uint32 (each uint64 becomes two uint32)
    arr_u32 = arr.view(np.uint32)
    # Extract the lower and upper 32 bits
    # For little-endian systems, index 0 is lower bits and index 1 is upper bits
    lower, upper = arr_u32[..., 0], arr_u32[..., 1]