Numpy：向量化 np.argwhere

Question

我在 numpy 中有以下数据结构：

import numpy as np

a = np.random.rand(267, 173) # dense img matrix
b = np.random.rand(199) # array of probability samples

我的目标是获取b 中的每个条目i，找到a 中所有值为<= i 的值的x,y 坐标/索引位置，然后随机选择该子集中的值之一：

from random import randint

for i in b:
  l = np.argwhere(a <= i) # list of img coordinates where pixel <= i
  sample = l[randint(0, len(l)-1)] # random selection from `l`

这“有效”，但我想对采样操作进行矢量化（即将for 循环替换为apply_along_axis 或类似的）。有谁知道如何做到这一点？任何建议将不胜感激！

Answer 1

你不能精确地矢量化np.argmax，因为你每次都有一个随机的子集大小。不过，您可以做的是通过排序显着加快计算速度。对图像进行一次排序将创建一个分配，而在每一步屏蔽图像将为提取的元素创建一个临时数组和。使用排序后的图像，您只需应用 np.searchsorted 即可获取尺寸：

a_sorted = np.sort(a.ravel())
indices = np.searchsorted(a_sorted, b, side='right')

您仍然需要一个循环来进行采样，但您可以执行类似的操作

samples = np.array([a_sorted[np.random.randint(i)] for i in indices])

使用这个系统获取 x-y 坐标而不是样本值有点复杂。您可以使用np.unravel_index 获取索引，但首先您必须将a_sorted 的参考框架转换为a.ravel()。如果您使用np.argsort 而不是np.sort 进行排序，您可以获得原始数组中的索引。幸运的是，np.searchsorted 使用 sorter 参数支持这种精确的场景：

a_ind = np.argsort(a, axis=None)
indices = np.searchsorted(a.ravel(), b, side='right', sorter=a_ind)
r, c = np.unravel_index(a_ind[[np.random.randint(i) for i in indices]], a.shape)

r 和c 的大小与b 相同，并且对应于基于b 的每个选择的a 中的行和列索引。索引转换取决于数组中的步长，因此我们假设您使用的是 C 顺序，因为默认情况下 90% 的数组都会这样做。

复杂性

假设b 的大小为M，a 的大小为N。

您当前的算法通过a 的每个元素对b 的每个元素进行线性搜索。在每次迭代中，它为匹配的元素分配一个掩码（平均为N/2），然后分配一个相同大小的缓冲区来保存被掩码的选项。也就是说时间复杂度在O(M * N)的量级，空间复杂度是一样的。

我的算法首先对a 进行排序，即O(N log N)。然后它搜索M 插入点，即O(M log N)。最后，它选择M 样本。它分配的空间是图像的一个排序副本和两个大小为M 的数组。因此它的时间复杂度为O((M + N) log N)，空间复杂度为O(M + N)。

Answer 2

这是一种替代方法 argsorting b，然后相应地使用 np.digitize 和 this post 对 a 进行分类：

import numpy as np
from scipy import sparse
from timeit import timeit
import math

def h_digitize(a,bs,right=False):
    mx,mn = a.max(),a.min()
    asz = mx-mn
    bsz = bs[-1]-bs[0]
    nbins=int(bs.size*math.sqrt(bs.size)*asz/bsz)
    bbs = np.concatenate([[0],((nbins-1)*(bs-mn)/asz).astype(int).clip(0,nbins),[nbins]])
    bins = np.repeat(np.arange(bs.size+1), np.diff(bbs))
    bbs = bbs[:bbs.searchsorted(nbins)]
    bins[bbs] = -1
    aidx = bins[((nbins-1)*(a-mn)/asz).astype(int)]
    ambig = aidx == -1
    aa = a[ambig]
    if aa.size:
        aidx[ambig] = np.digitize(aa,bs,right)
    return aidx

def f_pp():
    bo = b.argsort()
    bs = b[bo]
    aidx = h_digitize(a,bs,right=True).ravel()
    aux = sparse.csr_matrix((aidx,aidx,np.arange(aidx.size+1)),
                            (aidx.size,b.size+1)).tocsc()
    ridx = np.empty(b.size,int)
    ridx[bo] = aux.indices[np.fromiter(map(np.random.randint,aux.indptr[1:-1].tolist()),int,b.size)]
    return np.unravel_index(ridx,a.shape)

def f_mp():
    a_ind = np.argsort(a, axis=None)
    indices = np.searchsorted(a.ravel(), b, sorter=a_ind, side='right')
    return np.unravel_index(a_ind[[np.random.randint(i) for i in indices]], a.shape)


a = np.random.rand(267, 173) # dense img matrix
b = np.random.rand(199) # array of probability samples

# round to test wether equality is handled correctly
a = np.round(a,3)
b = np.round(b,3)

print('pp',timeit(f_pp, number=1000),'ms')
print('mp',timeit(f_mp, number=1000),'ms')

# sanity checks

S = np.max([a[f_pp()] for _ in range(1000)],axis=0)
T = np.max([a[f_mp()] for _ in range(1000)],axis=0)
print(f"inequality satisfied: pp {(S<=b).all()} mp {(T<=b).all()}")
print(f"largest smalles distance to boundary: pp {(b-S).max()} mp {(b-T).max()}")
print(f"equality done right: pp {not (b-S).all()} mp {not (b-T).all()}")

使用经过调整的digitize 我会快一点，但这可能会因问题大小而异。此外，@MadPhysicist 的解决方案也不那么复杂。使用标准的digitize，我们差不多。

pp 2.620121960993856 ms                                                                                                                                                                                                                                                        
mp 3.301037881989032 ms                                                                                                                                                                                                                                                        
inequality satisfied: pp True mp True
largest smalles distance to boundary: pp 0.0040000000000000036 mp 0.006000000000000005
equality done right: pp True mp True

Answer 3

对@MadPhysicist 的算法略有改进，使其更加矢量化：

%%timeit
a_ind = np.argsort(a, axis=None)
indices = np.searchsorted(a.ravel(), b, sorter=a_ind)
r, c = np.unravel_index(a_ind[[np.random.randint(i) for i in indices]], a.shape)
100 loops, best of 3: 6.32 ms per loop

%%timeit
a_ind = np.argsort(a, axis=None)
indices = np.searchsorted(a.ravel(), b, sorter=a_ind)
r, c = np.unravel_index(a_ind[(np.random.rand(indices.size) * indices).astype(int)], a.shape)
100 loops, best of 3: 4.16 ms per loop

@PaulPanzer 的解决方案仍然统治着这个领域，虽然我不确定它在缓存什么：

%timeit f_pp()
The slowest run took 14.79 times longer than the fastest. This could mean that an intermediate result is being cached.
100 loops, best of 3: 1.88 ms per loop