使用Numba高效转换NumPy二进制数组到浮点数

import numpy as np
import timeit

# 模拟一个包含0和1的uint64数组
random_bit = np.random.randint(0, 2, size=(1000000), dtype=np.uint64) # 增加数组大小以更明显地体现性能差异

def np_cast(arr):
    """直接使用浮点数乘法和减法进行转换"""
    return 1.0 - 2.0 * np.float64(arr)

def product(arr):
    """直接使用整数数组进行算术运算，NumPy会自动处理类型提升"""
    return 1.0 - 2.0 * arr

def _array(arr):
    """使用预定义的映射数组作为索引"""
    np_one_minus_one = np.array([1.0, -1.0], dtype=np.float64)
    return np_one_minus_one[arr]

def astype_method(arr):
    """先转换为float64再进行算术运算"""
    one = np.float64(1)
    minus_two = np.float64(-2)
    return one + minus_two * arr.astype(np.float64)

# 性能基准测试（以微秒为单位，表示每次操作的平均时间）
print("--- 纯NumPy方法性能基准测试 ---")
print(f"np_cast: {timeit.timeit(lambda: np_cast(random_bit), number=1000):.6f} seconds")
print(f"product: {timeit.timeit(lambda: product(random_bit), number=1000):.6f} seconds")
print(f"_array: {timeit.timeit(lambda: _array(random_bit), number=1000):.6f} seconds")
print(f"astype_method: {timeit.timeit(lambda: astype_method(random_bit), number=1000):.6f} seconds")

import numba as nb

@nb.vectorize(['float64(uint64)']) # 明确指定输入输出类型，有助于Numba优化
def numba_if(val):
    """使用条件判断进行映射：0 -> 1.0, 1 -> -1.0"""
    return -1.0 if val else 1.0

@nb.vectorize(['float64(uint64)'])
def numba_product(val):
    """使用算术运算进行映射：1.0 - 2.0 * val"""
    return 1.0 - 2.0 * val

@nb.njit
def numba_if_loop(arr):
    """使用njit和显式循环进行条件判断映射"""
    assert arr.ndim == 1, "Input array must be 1-dimensional"
    result = np.empty_like(arr, dtype=np.float64)
    for i in range(arr.size):
        result[i] = -1.0 if arr[i] else 1.0
    return result

@nb.njit
def numba_product_loop(arr):
    """使用njit和显式循环进行算术运算映射"""
    assert arr.ndim == 1, "Input array must be 1-dimensional"
    result = np.empty_like(arr, dtype=np.float64)
    for i in range(arr.size):
        result[i] = 1.0 - 2.0 * arr[i]
    return result

# 重新定义random_bit，确保所有测试使用相同大小和内容的数据
random_bit = np.random.randint(0, 2, size=(1000000), dtype=np.uint64)

# 确保所有方法结果一致性
assert np.array_equal(np_cast(random_bit), numba_if(random_bit))
assert np.array_equal(np_cast(random_bit), numba_product(random_bit))
assert np.array_equal(np_cast(random_bit), numba_if_loop(random_bit))
assert np.array_equal(np_cast(random_bit), numba_product_loop(random_bit))

print("\n--- Numba优化方法性能基准测试 ---")
# Numba函数首次调用会进行编译，因此第一次运行可能较慢，后续调用加速明显
# timeit会自动多次运行并取平均，因此编译开销会被分摊
print(f"numba_if: {timeit.timeit(lambda: numba_if(random_bit), number=1000):.6f} seconds")
print(f"numba_product: {timeit.timeit(lambda: numba_product(random_bit), number=1000):.6f} seconds")
print(f"numba_if_loop: {timeit.timeit(lambda: numba_if_loop(random_bit), number=1000):.6f} seconds")
print(f"numba_product_loop: {timeit.timeit(lambda: numba_product_loop(random_bit), number=1000):.6f} seconds")