Question

假设我在batches和channels上有一组坐标numpy数组。

每个数组具有：

从索引N到{0}和lim的x-y-z坐标的0:2项
与索引3上的坐标相关联的第四个索引上的值。

如何将该坐标数组广播到“图形”数组graph中，以使coords中包含的每个唯一坐标都表示为该索引处所有坐标的总和？

即：

coords = np.random.randint(0,lim,(batches, channels,N, 4))
graph = foo(coords)
graph.shape = (batches, channels, lim, lim, lim)

在一种简单的情况下，可以执行以下操作：

def foo(coords):
    graph = np.zeros(batches, channels, lim, lim, lim)
    for b_i in range(batches):
        for c_i in range(channels):
            for n_i in range(N):
                elem = coords[b_i, c_i, n_i]
                if elem[3] > 0:
                   graph[b_i, c_i, elem[0], elem[1], elem[2]]+=elem[3]
    return graph

但是，我正在寻找一种涉及广播的解决方案，因此我可以将该技术移植到PyTorch上，在那里广播是提高速度的必需条件。

Answer 1

以下是循环的矢量化解决方案：

mask = np.argwhere(coords[:,:,:,3]>0)
graph = np.zeros((batches, channels, lim, lim, lim))
idx = coords[tuple(mask.T)]
graph[tuple(np.hstack((mask[:,0:2], idx[:,:3])).T)] += idx[:,3]

脾气暴躁：坐标数组到数组

1 个答案: