使用 numpy 数组中的索引执行计算和比较

performing calculations and comparisons using indices in numpy arrays

我有一个 3D numpy 体素数组,即每个点的索引对应于它在 3D 中的位置 space。

我希望执行许多计算,这些计算涉及了解这些点是否满足各种几何条件。这意味着通过乘以基数将索引转换为向量,然后计算点积和叉积、范数等。我正在寻找一种相当快的方法来做到这一点,因为到目前为止我的努力似乎很慢。

如果我有一个任意基础 a, b, c:

basis = np.array([[a1, a2, a3],
                  [b1, b2, b3],
                  [c1, c2, c3]])

其中 a1、a2、a3 是 a 的 x、y 和 z 分量,b 和 [=34 也是如此=]c。 我可以通过以下方式计算每个体素的笛卡尔坐标 p = (x, y, z):

for i in range(vox.shape[0]):
    for j in range(vox.shape[1]):
        for k in range(vox.shape[2]):
            p = np.dot(basis, np.array([i, j, k]))

其中 "vox" 是体素的 3D 数组。例如,如果我希望计算这些向量中的每一个与单个其他(笛卡尔)向量(例如 q = np.array([qx, qy, qz]))的点积,并在结果满足给定条件(此处大于 0.0)时存储该索引,我可以做这样的事情(在与上面相同的循环中):

            if np.dot(p, q) > 0.0:
                desired_vox_indices.append([i, j, k])

问题是这很慢。我可以用更 pythonic 的方式或使用更多的 numpy 工具来做到这一点吗?我意识到在这个阶段我什至没有访问 vox 数组的值。

编辑:根据 Divakar 的回答尝试叉积

# Get q_x (matrix multiplication version of cross product)
q_x = np.array([[0.0, -q[2], q[1]],
                [q[2], 0.0, -q[0]],
                [-q[1], q[0], 0.0]])

# transpose (as per cross product definition) and matrix multiply with basis
u = np.matmul(q_x.T, basis)

# Get open range arrays
m,n,r = vox.shape
I,J,K = np.ogrid[:m,:n,:r]

# writing everything explicitly, since I am unsure how ogrid objects behave
pxq = np.empty(3)
pxq[0] = u[0,0]*I + u[0,1]*J + u[0,2]*K
pxq[1] = u[1,0]*I + u[1,1]*J + u[1,2]*K
pxq[2] = u[2,0]*I + u[2,1]*J + u[2,2]*K

可能等同于:

pxq = np.dot(u, np.array([I, J, K]))

但我不确定...

您可以像这样以矢量化方式执行此操作:

# Array of indices for your voxel data
ind = np.indices(vox.shape).reshape(3, -1)
# Multiply the basis times each coordinate
p = basis @ ind
# Dot product of each result with vector
d = q @ p
# Select coordinates where criteria is met
desired_vox_indices = ind[:, d > 0.0].T

我们将使用用于缩放的范围数组来构建求和,而不是一次生成所有索引。我们将按照三个嵌套循环对应的三个步骤来完成。这个想法与 中探索的想法非常相似。这将是内存高效的,因此也有望提高性能。然后,我们将把总和与阈值 0.0 进行比较,并使用 np.argwhere 得到相应的指数。因此,我们会有一个解决方案,就像这样 -

# Get q scaled version
s = q.dot(basis)

# Get open range arrays and scale and sum-reduce s array
m,n,r = vox.shape
I,J,K = np.ogrid[:m,:n,:r]
sums = s[0]*I + s[1]*J + s[2]*K

# Finally compare the sums against threshold amd get corresponding indices
out = np.argwhere(sums > 0.0)

大型 vox 阵列上的计时 -

# Setup
In [371]: np.random.seed(0)
     ...: basis = np.random.randn(3,3)
     ...: vox = np.random.randn(100,100,100)
     ...: q = np.random.randn(3)

# Original soln
In [372]: %%timeit
     ...: desired_vox_indices = []
     ...: for i in range(vox.shape[0]):
     ...:     for j in range(vox.shape[1]):
     ...:         for k in range(vox.shape[2]):
     ...:             p = np.dot(basis, np.array([i, j, k]))            
     ...:             if np.dot(p, q) > 0.0:
     ...:                 desired_vox_indices.append([i, j, k])
1 loop, best of 3: 2.13 s per loop

# @jdehesa's soln
In [373]: %%timeit
     ...: ind = np.indices(vox.shape).reshape(3, -1)
     ...: p = basis.dot(ind)
     ...: d = q.dot(p)
     ...: desired_vox_indices = ind[:, d > 0.0].T
10 loops, best of 3: 35.9 ms per loop

# From this post
In [374]: %%timeit
     ...: s = q.dot(basis)
     ...: m,n,r = vox.shape
     ...: I,J,K = np.ogrid[:m,:n,:r]
     ...: sums = s[0]*I + s[1]*J + s[2]*K
     ...: out = np.argwhere(sums > 0.0)
100 loops, best of 3: 7.56 ms per loop