在特定子阵列处评估数组

Foo*_*Bar 6 python numpy

我提前警告:此刻我可能完全糊涂了.我讲一个关于我实际尝试实现的简短故事,因为这可能会让事情变得清晰起来.说我有f(a,b,c,d,e),我想找到arg max (d,e) f(a,b,c,d,e).考虑一个离散网格(的简单的例子)Ff:

F = np.tile(np.arange(0,10,0.1)[newaxis,newaxis,:,newaxis,newaxis], [10, 10, 1, 10, 10])
maxE = F.max(axis=-1)
argmaxD = maxE.argmax(axis=-1)
maxD = F.max(axis=-2)
argmaxE = maxD.argmax(axis=-1)
Run Code Online (Sandbox Code Playgroud)

这就是我通常如何解决离散化版本的情况.但现在假设相反,我想解决arg max d f(a,b,c,d,e=X):与其选择最佳e每隔输入,e是一个固定的和给定的(大小AxBxCxD,在这个例子是的10x10x100x10).我有麻烦解决这个问题.

我天真的做法是

X = np.tile(np.arange(0,10)[newaxis,newaxis,:,newaxis], [10,10,1,10])
maxX = F[X]
argmaxD = maxX.argmax(axis=-1)
Run Code Online (Sandbox Code Playgroud)

然而,崩溃我的IDE的大量内存意味着F[X]显然不是我想要的.

表现是关键.

小智 2

我相信你可以这样做,但也许有更好的方法..

n = 10
F = np.tile(np.arange(0,n,0.1)[None,None,:,None,None], [n, n, 1, n, n])
X = np.tile(np.arange(0,n)[None,None,:,None], [n, n, 1, n])

a,b,c,d = np.ogrid[:n,:n,:n,:n]
argmaxD = F[a,b,c,d,X].argmax(axis=-1)
Run Code Online (Sandbox Code Playgroud)

正如我们在评论中讨论的那样,上面X并没有占据整个空间。如果您想选择e所有a, b, ,cd可以这样做,例如:

X = np.tile(np.arange(0,n,0.1).astype(int)[None,None,:,None], [n, n, 1, n])
a,b,c,d = np.ogrid[:n,:n,:100,:n]
argmaxD = F[a,b,c,d,X].argmax(axis=-1)
Run Code Online (Sandbox Code Playgroud)

tile另外,请注意,您可以使用广播来代替。但然后F[a,b,c,d,X]有一个单一的维度,所以你应该提供类似的东西axis=3

X = np.arange(0,n,0.1).astype(int)[None,None,:,None]
a,b,c,d = np.ogrid[:n,:n,:100,:n]
argmaxD = F[a,b,c,d,X].argmax(axis=3)
Run Code Online (Sandbox Code Playgroud)