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

Question

我有一个3D numpy体素阵列，即每个点的索引对应于它在3D空间中的位置。

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

如果我有任意基础a，b，c：

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

其中a1，a2，a3是a的x，y和z分量，同样对于b和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]))

但我不确定......

Answer 1

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

# 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

Answer 2

我们将使用范围数组来构建求和以进行缩放，而不会立即生成所有索引。我们将分三步对应三个嵌套循环。这个想法与this answer to - Python vectorizing nested for loops探索的想法非常相似。这将是内存效率，因此希望性能也很好。然后，我们将总结与阈值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