使用Numba优化NumPy数组二进制值到浮点数的映射

import numpy as np
import timeit

# 生成测试数据
random_bit = np.random.randint(0, 2, size=(10000), dtype=np.uint64)

# 方法1: 直接类型转换与算术运算
def np_cast(arr):
    return 1.0 - 2.0 * np.float64(arr)

# 方法2: 产品运算（NumPy自动类型提升）
def product(arr):
    return 1.0 - 2.0 * arr

# 方法3: 数组索引
np_one_minus_one = np.array([1.0, -1.0]).astype(np.float64)
def _array(arr):
    return np_one_minus_one[arr]

# 方法4: 显式astype转换
one = np.float64(1)
minus_two = np.float64(-2)
def astype_explicit(arr):
    return one + minus_two * arr.astype(np.float64)

print("--- NumPy 方法性能基准测试 ---")
# 使用 %timeit 模拟，实际运行请在Jupyter或IPython中执行
# %timeit np_cast(random_bit)
# %timeit product(random_bit)
# %timeit _array(random_bit)
# %timeit astype_explicit(random_bit)

# 模拟的性能结果（基于原问题答案的 %timeit 数据，仅供参考，实际运行可能略有差异）
# np_cast: ~6.58 µs
# product: ~7.58 µs
# _array:  ~11 µs
# astype_explicit: ~7.32 µs

import numba as nb

@nb.vectorize(['float64(uint64)']) # 指定输入输出类型
def numba_if(x):
    # 根据条件返回不同值
    return -1.0 if x else 1.0

@nb.vectorize(['float64(uint64)']) # 指定输入输出类型
def numba_product(x):
    # 使用数学公式
    return 1.0 - 2.0 * x

@nb.njit
def numba_if_loop(arr):
    # 确保输入是一维数组
    assert arr.ndim == 1
    # 预分配结果数组
    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):
    # 确保输入是一维数组
    assert arr.ndim == 1
    # 预分配结果数组
    result = np.empty_like(arr, dtype=np.float64)
    # 遍历数组元素进行转换
    for i in range(arr.size):
        result[i] = 1.0 - 2.0 * arr[i]
    return result

# 确保所有函数返回相同的结果，以验证正确性
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 优化方法性能基准测试 ---")
# 使用 %timeit 模拟，实际运行请在Jupyter或IPython中执行
# %timeit numba_if(random_bit)
# %timeit numba_product(random_bit)
# %timeit numba_if_loop(random_bit)
# %timeit numba_product_loop(random_bit)

# 模拟的性能结果（基于原问题答案的 %timeit 数据，仅供参考）
# numba_if: ~1.89 µs
# numba_product: ~2.07 µs
# numba_if_loop: ~1.6 µs
# numba_product_loop: ~1.78 µs