我正在尝试在numpy中索引多维数组(4维)。阵列的形状为(125,125,125,3)。我有3个单独的2D索引数组列表。列表的大小分别为(N,4),(M,4)和(1,4)。3个单独的列表表示我要索引的4D数组中的行,列和深度值。例如,考虑以下内容:
ix = [[0,1,2,3],
[3,4,5,6]]
iy = [[2,3,4,5],
[5,6,7,8]]
iz = [[1,2,3,4]]
weights.shape = (125,125,125,3)
我想指数weights
与行,列和深度指数阵列中每一个可能的组合ix
,iy
和iz
。例如,如果我采取的第一个行中的每个索引的矩阵,这意味着我想选择的行[0,1,2,3]
,列[2,3,4,5]
,和深度值[1,2,3,4]
中weights
。我一直想选择的第4维的所有元素weights
。这意味着我实际上是在选择(4,4,4,3)
的一部分weights
。
现在,我已经通过使用以下代码对循环进行索引来实现此目的
w = np.empty(shape=(X,Y,Z,4,4,4,weights.ndim-1))
for i in range(X):
for j in range(Y):
w_ij = np.ix_(ix[i,:], iy[j,:], iz[0,:])
w[i,j,0,:,:,:,:] = weights[w_ij[0], w_ij[1], w_ij[2], :]
我的最终目标是尽快构建w
形状为(N,M,1,4,4,4,4,3)的多维数组。这部分代码将要运行多次,因此,如果有使用内置numpy函数的矢量化方法,那将是理想的选择。
如果有任何澄清的问题,请告诉我。这是我第一次问堆栈溢出问题,因此,如果有任何不清楚或令人困惑的地方,我深表歉意!
您可以将索引与广播配合使用来实现此目的。
import numpy as np
weights = np.random.rand(125, 125, 125, 3)
ix = np.array([[0,1,2,3], [3,4,5,6]])
iy = np.array([[2,3,4,5], [5,6,7,8]])
iz = np.array([[1,2,3,4]])
X = len(ix)
Y = len(iy)
Z = len(iz)
def compute1(weights):
w = np.empty(shape=(X, Y, Z, 4, 4, 4, weights.ndim-1))
for i in range(X):
for j in range(Y):
w_ij = np.ix_(ix[i,:], iy[j,:], iz[0,:])
w[i,j,0,:,:,:,:] = weights[w_ij[0], w_ij[1], w_ij[2], :]
return w
def compute2(weights):
return weights[ix[:, None, None, :, None, None], iy[None, :, None, None, :, None], iz[None, None, :, None, None, :]]
print(np.allclose(compute1(weights), compute2(weights)))
给True
。
基准测试-
%timeit compute1(weights)
%timeit compute2(weights)
给-
36.7 µs ± 897 ns per loop (mean ± std. dev. of 7 runs, 10000 loops each)
6.28 µs ± 62.8 ns per loop (mean ± std. dev. of 7 runs, 100000 loops each)
如您所见,对于这种大小的数据,广播解决方案的速度提高了约6倍。
本文收集自互联网,转载请注明来源。
如有侵权,请联系[email protected] 删除。
我来说两句