如何在不使用 python 循环的情况下将元素插入 3d numpy 数组中的特定索引？

Question

我有一个形状为 (480, 640, 30) 的 3 维 NumPy 数组，其中每个 2d 数组的形状为 (480,640)，并且此 2d 数组上的每个元素都是沿深度有效的 FIFO 缓冲区的一部分尺寸（因此实际上存在 480x640 个大小为 30 的缓冲区）。对于接收到的每个新的 2d 帧，每个缓冲区都会根据布尔掩码 2d 数组进行更新（如果掩码数组中的相应索引为 False，则使用此新帧中的值更新相应的缓冲区，如果为 True，则不要t)。我正在使用另一个 2d 数组来跟踪将新元素从新帧添加到每个缓冲区的位置，如果将元素添加到缓冲区，则该特定索引的计数器会增加（使用时会采用缓冲区的中位数，所以元素的顺序无关紧要） 我已经能够以这种方式使用循环在 O(n^2) 时间内实现这一点：

buffer = np.full((30,480,640), -1, dtype=int) #initialized with -1, values added range from 0 -> 255
bool_mask = np.random.choice([True, False], size=(480,640), replace=True, p=None) 
#here I am using the same mask for every loop but in reality mask would change with each input frame
counter = np.zeros((480,640), dtype=np.uint8) 
rand_frame = np.random.randint(1, 5, size=(480,640), dtype=int) #just a random array to replicate the new frame that is received

for m in range(100): #in the final implementation, frames are received continuously and buffer is updated, I've used this loop to replicate that process
    for i in range(buffer.shape[1]):
        for k in range(buffer.shape[2]):
            if not bool_mask[i][k]:
                buffer[counter[i][k]%30][i][k] = rand_frame[i][k] #modulus used to replicate a FIFO buffer insertion
                counter[i][k] += 1

此逻辑按预期工作，但在 850 毫秒时非常慢。我需要它以 1 毫秒运行，所以循环不会。我试过使用 np.where，但是每个元素的更新可能会或可能不会发生，导致计数器中的深度索引对于 2d 数组上的每个索引都不同，所以 np.where 不会这样做，加上如果无论哪种方式，我都必须添加元素，而在我的情况下，fifo 缓冲区的更新是有条件的。我一直在寻找对相应索引进行选择性更新但没有运气的替代方案。 非常感谢您的帮助，谢谢！

Answer 1

我想你想要Boolean indexing:

import numpy as np

# use smaller dimensions so problem runs in reasonable time
w = 50
h = 70
d = 7

bool_mask = np.random.choice([True, False], size=(h,w), replace=True, p=None) 
rand_frame = np.random.randint(1, 5, size=(h,w), dtype=int)

# for-loop solution
def ex1():
    counter = np.zeros((h,w), dtype=np.uint8) 
    buffer1 = np.full((d,h,w), -1, dtype=int)
    for m in range(d):
        for i in range(buffer1.shape[1]):
            for k in range(buffer1.shape[2]):
                if not bool_mask[i][k]:
                    buffer1[counter[i][k]%d][i][k] = rand_frame[i][k] #modulus used to replicate a FIFO buffer insertion
                    counter[i][k] += 1
    return buffer1, counter


# boolean indexing solution
def ex2():
    counter = np.zeros((h,w), dtype=np.uint8) 
    buffer2 = np.full((d,h,w), -1, dtype=int) #initialized with -1, values added range from 0 -> 255
    for m in range(d):
        buffer2[m%d][~bool_mask] = rand_frame[~bool_mask]
        counter[~bool_mask] += 1
    return buffer2, counter


b1, c1 = ex1()
b2, c2 = ex2()

# verify results are the same
print(f'{(b1 == b2).all() = }')
print(f'{(c1 == c2).all() = }')

在我的系统（Python 3.9.7，numpy 1.21.4）上，这给出了

(b1 == b2).all() = True
(c1 == c2).all() = True

和

In [57]: timeit ex1()
80.2 ms ± 159 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)

In [58]: timeit ex2
23.8 ns ± 0.0931 ns per loop (mean ± std. dev. of 7 runs, 10000000 loops each)