如何分别用Numpy（无Numba）和Numba并行处理数组计算逻辑

给定需求：接收含两个数值的数组[a, b]，返回[a², a*b]，以下针对函数内独立的两行计算逻辑，提供两种并行实现方法：

1. 不使用Numba的并行实现 #

可以利用Python标准库的线程/进程池实现并行，因为两个计算逻辑完全独立，适合拆分任务并行执行。

线程池实现（轻量计算首选） #

import numpy as np
from concurrent.futures import ThreadPoolExecutor

def func_parallel_non_numba(array_1):
    a, b = array_1
    
    # 定义独立计算任务
    def calc_square():
        return a ** 2
    def calc_product():
        return a * b
    
    # 并行执行任务
    with ThreadPoolExecutor(max_workers=2) as executor:
        future_square = executor.submit(calc_square)
        future_product = executor.submit(calc_product)
        result = np.array([future_square.result(), future_product.result()])
    return result

# 测试
array_1 = np.array([3., 4.])
print(array_1)
print(func_parallel_non_numba(array_1))

进程池实现（重计算场景可选） #

import numpy as np
import multiprocessing as mp

def calc_square(a):
    return a ** 2

def calc_product(a, b):
    return a * b

def func_parallel_non_numba_mp(array_1):
    a, b = array_1
    with mp.Pool(processes=2) as pool:
        res_square = pool.apply_async(calc_square, (a,))
        res_product = pool.apply_async(calc_product, (a, b))
        result = np.array([res_square.get(), res_product.get()])
    return result

# 测试
array_1 = np.array([3., 4.])
print(func_parallel_non_numba_mp(array_1))

2. 使用Numba的并行实现 #

Numba通过JIT编译结合并行指令，可实现低开销的并行计算，适合数值计算场景。

基于prange的循环并行 #

import numpy as np
from numba import njit, prange

@njit(parallel=True)
def func_parallel_numba(array_1):
    array_2 = np.zeros_like(array_1)
    # 用prange标记可并行的循环迭代
    for i in prange(2):
        if i == 0:
            array_2[i] = array_1[0] ** 2
        else:
            array_2[i] = array_1[0] * array_1[1]
    return array_2

# 测试
array_1 = np.array([3., 4.])
print(array_1)
print(func_parallel_numba(array_1))

基于vectorize的向量化并行（适合批量处理） #

import numpy as np
from numba import vectorize

@vectorize(['float64(float64, float64)'], target='parallel')
def func_vectorize_numba(a, b):
    return a**2, a*b

# 测试单个输入
array_1 = np.array([3., 4.])
print(func_vectorize_numba(array_1[0], array_1[1]))
# 批量测试多组输入
batch_a = np.array([3., 5., 7.])
batch_b = np.array([4., 6., 8.])
print(func_vectorize_numba(batch_a, batch_b))

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