追溯 argmin/argmax 在布尔掩码 NumPy 数组上的原始位置 - Python

Trace back original position of argmin/argmax on boolean masked NumPy array - Python

上下文

由于使用 numpy.ma 模块进行屏蔽比直接布尔屏蔽要慢得多,因此我必须使用后者进行 argmin/argmax 计算。

一点比较:

import numpy as np

# Masked Array
arr1 = np.ma.masked_array([12,4124,124,15,15], mask=[0,1,1,0,1])

# Boolean masking
arr2 = np.array([12,4124,124,15,15])
mask = np.array([0,1,1,0,1], dtype=np.bool)

%timeit arr1.argmin()
# 16.1 µs ± 4.88 µs per loop (mean ± std. dev. of 7 runs, 100000 loops each)

%timeit arr2[mask].argmin()
# 946 ns ± 55.3 ns per loop (mean ± std. dev. of 7 runs, 1000000 loops each)

无论如何,使用 argmin/argmax returns 数组中第一次出现的索引。在布尔掩码的情况下,这意味着 arr2[mask] 中的索引而不是 arr2 中的索引。还有我的问题:我在掩码数组上计算它时需要未掩码数组中的索引


问题

如何获得未屏蔽 arr2argmin/argmax-索引,即使我将它应用于布尔屏蔽版本 [=18] =]?

我的解决方案是使用查找逻辑,其中我有第二个数组存储正确的索引。

假设我们有一个随机值数组,我们用布尔值掩码并想要应用 argmin/argmax,这看起来像:

random_array = np.random.randint(10, size=100)
random_mask  = np.random.randint(2, size=100, dtype=np.bool)

# Returns index of fist occurrence of minimum value within the masked array
random_array[random_mask].argmin()

现在我们必须创建一个查找-table 包含未屏蔽 random_array:

的索引
lookup = np.arange(len(random_array), dtype=np.int))

如果我们现在以与屏蔽 random_array 相同的方式屏蔽 lookup,我们将检索原始索引:

# Returns the index within the unmasked array
result = lookup[random_mask][random_array[random_mask].argmin()]

这是一个主要基于 masking,特别是 - mask-the-mask 并且应该是内存高效的,并且希望性能也很好,特别是在处理大型数组时 -

def reset_first_n_True(mask, n):
    # Resets (fills with False) first n True places in mask

    # Count of True in original mask array
    c = np.count_nonzero(mask)

    # Setup second mask that is to be assigned into original mask on its
    # own True positions with the idea of setting first argmin_in_masked_ar
    # True values to False
    second_mask = np.ones(c, dtype=bool)
    second_mask[:n] = False
    mask[mask] = second_mask
    return

# Use reduction function on masked data array 
idx = np.argmin(random_array[random_mask])
reset_first_n_True(random_mask, idx)
out = random_mask.argmax()

要在屏蔽数据数组上获取 argmax 并将其追溯到原始位置,只有第一步会更改为包括:

idx = np.argmax(random_array[random_mask])

因此,可以使用任何归约操作并以这种方式追溯到它们的原始位置。


如果您正在寻找紧凑的解决方案,请使用 nonzero() -

idx = np.flatnonzero(random_mask)
out = idx[random_array[random_mask].argmin()]
# Or idx[random_array[idx].argmin()]