python - Numpy:查找两个不同数组中值的索引条件(来自 R)

Question

我有一个由 3D ndarray X 表示的体积，其值介于 0 和 255 之间，并且我有另一个 3D ndarray Y，即第一个数组的任意掩码，值为 0 或 1。

我想找到 50 个体素的随机样本的索引，这些样本的 X(“图像”)大于零，Y(“掩模”)等于 1。

我的经验是使用 R，其中以下内容可行:

idx <- sample(which(X>0 & Y==1), 50)

也许 R 的优点是我可以线性索引 3D 数组，因为例如，仅在 numpy 中使用单个索引就可以得到一个 2D 矩阵。

我猜它可能涉及numpy.random.choice，但似乎我不能有条件地使用它，更不用说以两个不同的数组为条件了。我应该使用另一种方法吗？

Answer 1

这是一种方法 -

N = 50 # number of samples needed (50 for your actual case)

# Get mask based on conditionals
mask = (X>0) & (Y==1)

# Get corresponding linear indices (easier to random sample in next step)
idx = np.flatnonzero(mask)

# Get random sample
rand_idx = np.random.choice(idx, N)

# Format into three columnar output (each col for each dim/axis)
out = np.c_[np.unravel_index(rand_idx, X.shape)]

如果您需要随机样本而不需要替换，请使用带有可选参数replace=False的np.random.choice()。

示例运行 -

In [34]: np.random.seed(0)
    ...: X = np.random.randint(0,4,(2,3,4))
    ...: Y = np.random.randint(0,2,(2,3,4))

In [35]: N = 5 # number of samples needed (50 for your actual case)
    ...: mask = (X>0) & (Y==1)
    ...: idx = np.flatnonzero(mask)
    ...: rand_idx = np.random.choice(idx, N)
    ...: out = np.c_[np.unravel_index(rand_idx, X.shape)]

In [37]: mask
Out[37]: 
array([[[False,  True,  True, False],
        [ True, False,  True, False],
        [ True, False,  True,  True]],

       [[False,  True,  True, False],
        [False, False, False,  True],
        [ True,  True,  True,  True]]], dtype=bool)

In [38]: out
Out[38]: 
array([[1, 0, 1],
       [0, 0, 1],
       [0, 0, 2],
       [1, 1, 3],
       [1, 1, 3]])

将输出 out 与 mask 中 True 值的位置相关联，以进行快速验证。

<小时/>

如果您不想展平以获得线性索引并直接获取每个暗淡/轴的索引，我们可以这样做 -

i0,i1,i2 = np.where(mask)
rand_idx = np.random.choice(len(i0), N)
out = np.c_[i0,i1,i2][rand_idx]

为了提高性能，首先索引，然后在最后一步与 np.c_ 连接 -

out = np.c_[i0[rand_idx], i1[rand_idx], i2[rand_idx]]