如何使用 Numba 并行化代码，其中数组必须由每个循环中产生的不同形状的数组填充

Question

我写了一个可以正常工作的代码。我使用 numba njit 来让它更快。我很想知道这段代码是否可以并行化以尽可能快。很明显，ints 有不同的数字，因此内部循环将在每个主循环中产生不同长度的列表。由于这个问题，零数组不能像往常一样使用（在循环中填充）并且应该是 object 类型；但是，object 类型在这方面没有用。我想知道 在 python 中使用 numba 并行此代码的最有效适用方法是什么，并知道是否有这样做的限制。可重现的例子如下：

n = 10
a = np.random.rand(n, 2)
a_min = a - 0.7           # 0.7 is just an arbitrary float to create a lower limits
a_max = a + 0.7
ints = np.random.randint(17, 21, n, dtype=np.int32)

@nb.njit("float64[:, ::1], float64[:, ::1], int32[::1]")
def rand_(a_min, a_max, ints):
    np.random.seed(85)
    main_ = []                       # np.zeros((len(ints),), dtype=object)
    for i, j in enumerate(ints):
        min_x, min_y = a_min[i]
        max_x, max_y = a_max[i]
        temp_ = []                   # np.zeros((j, 2), dtype=nb.float64)
        for m in range(j):
            rand_x = np.random.uniform(min_x, max_x, 1)
            rand_y = np.random.uniform(min_y, max_y, 1)
            temp_.append([rand_x, rand_y])        

        main_.append(temp_)          # temp_ shape: (j, 2); j will be different in each main loop
    
    return main_

自从使用 np.random.uniform 以来，并行化是否会影响结果的可重复性？

过程说明：
a_min是一个数组（形状（n*2）），每行包含x、y的下限值，以及a_max的上限值。在每个循环中，我们的内部循环将为每个 x、y 产生一个介于下限和上限之间的随机浮点数，并将这两个浮点数存储为一个新组合 ((x, y))在列表或数组中。根据 ints 中的整数值（例如 9），这些组合中的一些将在每个主循环中创建，并将被分组并存储到另一个列表或数组中。

Answer 1

您的代码生成数组列表的列表，其中每个数组包含一个浮点值。

相反，生成一个二维数组列表是合理的。此外，您可以使用对 np.random.uniform():

的单个调用来创建数组中的每一列

def rand_2(a_min, a_max, ints):
    np.random.seed(85)
    main_ = []
    for i, length in enumerate(ints):
        min_x, min_y = a_min[i]
        max_x, max_y = a_max[i]
        temp_ = np.empty((length, 2))
        temp_[:, 0] = np.random.uniform(min_x, max_x, length)
        temp_[:, 1] = np.random.uniform(min_y, max_y, length)
        main_.append(temp_)
    return main_

此代码在没有使用 Numba jit 的情况下，使用 n=100000 的速度是原始代码的两倍。使用 jit，它比原来的快 20 倍左右。由于随机生成器的使用不同，结果和原来的不一样，但是在不同的调用之间会保持一致。

然后您可以用 Numba 类型的列表替换 Python 列表（Numba“反射列表”）：

Array = nb.types.Array(dtype=nb.float64, ndim=2, layout="C")

@nb.njit
def rand_4(a_min, a_max, ints):
    np.random.seed(85)
    main_ = nb.typed.List.empty_list(Array)     # Typed List
    for i, length in enumerate(ints):
        min_x, min_y = a_min[i]
        max_x, max_y = a_max[i]
        temp_ = np.empty((length, 2))
        temp_[:, 0] = np.random.uniform(min_x, max_x, length)
        temp_[:, 1] = np.random.uniform(min_y, max_y, length)
        main_.append(temp_)
    return main_

现在比原来快 50 倍左右。

循环可以并行化，但是列表的追加方法不再可用，因为it's not thread-safe。所以列表必须是 pre-allocated。只要它是类型列表，就不能 pre-allocated as [None]*n:

dummy_array = np.array([[0.]])

@nb.njit(parallel=True)
def rand_5(a_min, a_max, ints):
    np.random.seed(85)
    n = len(ints)
    main_ = nb.typed.List([dummy_array] * n)
    for i in nb.prange(n):                     # Parallel loop
        length = ints[i]
        min_x, min_y = a_min[i]
        max_x, max_y = a_max[i]
        temp_ = np.empty((length, 2))
        temp_[:, 0] = np.random.uniform(min_x, max_x, length)
        temp_[:, 1] = np.random.uniform(min_y, max_y, length)
        main_[i] = temp_                       # Avoiding append
    return main_

当 ints 非常大时，这只比以前的版本快一点。此外，不同调用之间的结果永远不会一致。

如何使用 Numba 并行化代码，其中数组必须由每个循环中产生的不同形状的数组填充

How could parallelize a code with Numba, where an array must be filled by different shaped arrays produced in each loop

python

optimization

parallel-processing

numba

jit