根据 2D 索引配置文件分配 3D 数据值

Question

我有一个 3D numpy 数组：

data0 = np.random.rand(30, 50, 50)

我有一个二维表面：

surf = np.random.rand(50, 50) * 30
surf = surf.astype(int)

现在我想沿着表面轮廓将“0”分配给 data0。我知道循环可以实现这个：

for xx in range(50):
    for yy in range(50):
        data0[0:surf[xx, yy], xx, yy] = 0

Data0 是一个 3D 体积，大小为 30 * 50 * 50。surf 是一个 2D 表面轮廓，大小为 50 * 50。我想做的是从顶部到表面填充“0”（轴=0) 卷

这里，'for'循环很慢，data0很大时效率很低。有人可以建议如何根据冲浪配置文件有效地分配值吗？

Answer 1

如果您想使用 numpy，您可以创建一个遮罩，其中 z-index 值低于您的 surf 值设置为 True，然后填充这些单元格带零：

import numpy as np

np.random.seed(123)
x, y, z = 4, 5, 3
data0 = np.random.rand(z, x, y)

surf = np.random.rand(x, y) * z
surf = surf.astype(int)

#your attempt
#we copy the data just for the comparison
data_loop = data0.copy()
for xx in range(x):
    for yy in range(y):
        data_loop[0:surf[xx, yy], xx, yy] = 0

#again, we copy the data just for the comparison
data_np = data0.copy()
#masking the cells according to your index comparison criteria
mask = np.broadcast_to(np.arange(data0.shape[0])[:,None, None], data0.shape) < surf[None, :]
#set masked values to zero
data_np[mask] = 0

#check for equivalence of resulting arrays
print((data_np==data_loop).all())

我确信有更好的、更加数字化的方法来生成索引号掩码。事实上，这个版本不一定更快。这取决于阵列的形状。对于 x=500、y=200 和 z=3000，您的循环需要 1.42 秒，而我的 numpy 方法需要 1.94 秒。对于相同的数组大小但形状为 x=5000、y=2000 和 z=30，您的循环方法需要 7.06 秒，而 numpy 方法需要 1.95 秒。

根据 2D 索引配置文件分配 3D 数据值

Assign 3D data value based on a 2D index profile

python

numpy

numpy-ndarray