如何将3D NumPy位数组沿第三轴打包为uint16/uint32类型？

好问题！numpy.packbits确实只支持将位数据打包成uint8类型，但我们完全可以手动实现类似逻辑，把3D数组沿指定轴打包成uint16或uint32。核心思路是利用位运算把多个二进制位组合成目标宽度的整数，下面给你具体的实现方案：

我们可以写一个可复用的函数，通过调整参数来支持不同的目标类型（uint16/uint32）：

import numpy as np

def packbits_custom(arr, dtype=np.uint16, axis=-1):
    # 确定目标类型对应的位数
    dtype = np.dtype(dtype)
    bits_per_element = dtype.itemsize * 8
    
    # 确保输入是0/1的整数数组（布尔数组会自动转换）
    arr = arr.astype(np.uint8)
    
    # 补零使目标轴长度为位数的整数倍（可选，根据需求调整）
    pad_length = (bits_per_element - arr.shape[axis] % bits_per_element) % bits_per_element
    if pad_length > 0:
        pad_tuple = [(0, 0)] * axis + [(0, pad_length)] + [(0, 0)] * (arr.ndim - axis - 1)
        arr = np.pad(arr, pad_tuple, mode='constant')
    
    # 重塑数组，将目标轴拆分为「组数 × 每组位数」的结构
    new_shape = list(arr.shape)
    new_shape[axis] = new_shape[axis] // bits_per_element
    new_shape.insert(axis + 1, bits_per_element)
    arr_reshaped = arr.reshape(new_shape)
    
    # 创建位移数组：从最高位到最低位（如uint16对应15,14,...,0）
    shifts = np.arange(bits_per_element - 1, -1, -1, dtype=dtype)
    # 广播位移数组以匹配输入数组的维度
    shifts = np.expand_dims(shifts, tuple(range(arr_reshaped.ndim - 1)))
    
    # 通过左移求和完成打包
    packed = (arr_reshaped << shifts).sum(axis=axis + 1, dtype=dtype)
    return packed

使用示例 #

打包成uint16 #

# 创建一个测试用的3D布尔数组（形状：2行×3列×20位）
test_arr = np.random.randint(0, 2, size=(2, 3, 20), dtype=bool)

# 沿第三轴打包成uint16
packed_uint16 = packbits_custom(test_arr, dtype=np.uint16, axis=2)
print(f"原数组形状：{test_arr.shape}")
print(f"打包后数组形状：{packed_uint16.shape}")  # 输出 (2, 3, 2)，因为20位补到32位，分成2个uint16元素

# 验证打包正确性（解包对比）
unpacked = np.unpackbits(packed_uint16.view(np.uint8), axis=2)[:, :, :20].astype(bool)
print(f"打包解包后与原数组是否一致：{np.array_equal(test_arr, unpacked)}")

打包成uint32 #

只需要修改dtype参数即可：

packed_uint32 = packbits_custom(test_arr, dtype=np.uint32, axis=2)
print(f"打包成uint32后的形状：{packed_uint32.shape}")  # 输出 (2,3,1)，20位补到32位，1个uint32元素

• 输入要求：输入数组的元素必须是0或1（布尔数组会自动转换为0/1），否则位运算结果会出错。
• 补零逻辑：如果不需要补零，可删除函数中的补零代码，但需确保目标轴的长度是目标类型位数的整数倍（如打包成uint16时，第三轴长度必须是16的倍数）。
• 位序调整：默认是将数组中靠前的位作为整数的最高位，如果需要将靠前的位作为最低位，只需把位移数组改为shifts = np.arange(0, bits_per_element, dtype=dtype)即可。

内容的提问来源于stack exchange，提问作者ZZM