Question

例如，我有一个复杂的形状数组（3,3,3），我希望找到沿着轴2的数组大小最大时元素的角度。也就是说，

a = np.random.random((3, 3, 3)) + 1j * np.random.random((3, 3, 3))

amp3d = np.abs(a)
amp2d = np.max(amp3d, axis=2)

angle2d =  np.zeros_like(amp2d)

for i in range(a.shape[0]):
    for j in range(a.shape[1]):
        for k in range(a.shape[2]):
            if amp3d[i, j, k] == amp2d[i,j]:
                angle2d[i, j] = np.angle(a[i, j, k])

当然我可以循环，但我想有一种非常好的方式来完成这项任务

在Matlab中，它看起来与np.angle(a[np.argmax(a, axis=2)])

类似

谢谢！

Answer 1

通过展开amp2d的维度以匹配amp3d的维度，并将其与a相乘，为相关值创建mask/boolean array并选择最大值axis=2。

a = np.random.random((3, 3, 3)) + 1j * np.random.random((3, 3, 3))

amp3d = np.abs(a)
amp2d = np.max(amp3d, axis=2)

mask = amp3d == amp2d[...,np.newaxis]
angle2d = np.angle((a*mask).max(axis=2))

查看numpy documentation on indexing以获取np.newaxis的参考内容，或查看this stackoverflow question。

如何根据条件获得子阵列？

1 个答案: