基于参考 n-dim 数组对 n-dim 数组进行操作的最有效方法

Question

我有两个相同形状的 numpy 数组：dat_ara 和 ref_ara。

我想对 dat_ara 的 axis = -1 执行操作 op_func，但是我只想对每个数组中选定的值切片进行操作，该切片是指定的当参考数组 ref_ara.

超过阈值 thres 时

为了说明，在数组只是 2-dim 的简单情况下，我有：

thres = 4

op_func = np.average

ref_ara = array([[1, 2, 1, 4, 3, 5, 1, 5, 2, 5],
                 [1, 2, 2, 1, 1, 1, 2, 7, 5, 8],
                 [2, 3, 2, 5, 1, 6, 5, 2, 7, 3]]) 

dat_ara = array([[1, 0, 0, 1, 1, 1, 1, 0, 1, 1],
                 [1, 1, 1, 1, 1, 1, 1, 0, 1, 0],
                 [1, 0, 1, 1, 1, 1, 0, 1, 1, 1]])

我们看到 thres 在 ref_ara 的 axis=0 中第 1、2 和 3 数组的第 5、7 和 3 个索引中被破坏。因此我想要的结果是

out_ara = array([op_func(array([1, 0, 0, 1, 1, 1]), 
                 op_func(array([1, 1, 1, 1, 1, 1, 1, 0]),
                 op_func(array([1, 0, 1, 1])])

这道题很难，因为它需要参考ref_ara。如果不是这种情况，我可以简单地使用 numpy.apply_along_axis.

我试过扩展两个数组的维度以将它们关联起来进行计算，即：

assos_ara = np.append(np.expand_dims(dat_ara, axis=-1), np.expand_dims(ref_ara, axis=-1), axis=-1)

但同样，numpy.apply_along_axis 要求输入函数只能在 1-dim 数组上运行，因此我仍然无法使用该函数。

我知道的唯一其他方法是通过数组索引明智地迭代，但是，由于数组具有两个数组的不断变化的维度，这是一个棘手的问题，而且，它在计算上效率不高。

我很想使用矢量化函数来帮助这个过程。最有效的方法是什么？

Answer 1

这是屏蔽数组的一个很好的用例，因为它们允许您对部分数据执行正常的 numpy 操作。

让我们假设每一行都包含至少一个大于阈值的值。您可以将断点的索引计算为

breaks = np.argmax(ref_ara > thres, axis=-1)   # 5, 7, 3

然后您可以使用我之前链接的 to the 创建一个遮罩。掩码通常是处理 numpy 中形状不规则数据的最佳方式。

mask = np.arange(ref_ara.shape[-1]) <= breaks.reshape(*breaks.shape, 1)

在这里，我们不需要对 arange 做任何花哨的事情，因为它沿着最后一个维度。如果不是这种情况，您可能希望将 1 插入到范围所在的中断形状中，并用 1 填充范围形状的尾部。

现在掩码数组和 ufunc 解决方案略有不同。 masked array版本比较通用，所以排在第一位：

data = np.ma.array(data_ara, mask=~mask)

掩码数组从正常布尔索引的方式向后解释掩码，因此我们反转掩码。或者，您可以使用 > 而不是 <= 来计算掩码。计算现在很简单：

out_ara = np.ma.average(data, axis=-1).data

一个不太通用的替代方法是将您的操作分解为 ufunc，并使用它们提供的掩码。这对于 np.average, which is just np.sum and np.divide 来说很容易，但对于更复杂的操作来说可能会更难。

从 numpy 1.17.0 开始，np.sum 有一个 where 关键字：

out_ara = np.sum(dat_ara, where=mask, axis=-1) / breaks

基于参考 n-dim 数组对 n-dim 数组进行操作的最有效方法

Most efficient way to operate on a n-dim array based on a reference n-dim array

python

arrays

numpy

vectorization