我正在尝试生成复杂空间中多项式根收敛的颜色映射。为此，我创建了一个点网格并将牛顿法应用于这些点，以便找到它们各自收敛到哪个复数根。这给了我一个二维复数数组，其中的元素表示它们在一定公差范围内收敛的点。我希望能够将该矩阵中的数字与元素颜色映射相匹配。
我已经通过遍历数组并逐个元素地计算颜色来完成此操作，但它非常慢，而且似乎可以从矢量化中获益。到目前为止，这是我的代码:
def colorvec(rootsmatrix, known_roots):
dim = len(known_roots)
dist = ndarray((dim, nx, ny))
for i in range(len(known_roots)):
dist[i] = abs(rootsmatrix-known_roots[i])
这将创建一个 3d 数组，其中包含每个点的计算根与每个实际根的距离。它看起来像这样，除了有 75 000 000 个元素。
[ [ [6e-15 7e-15 0.5]
[1.5 5e-15 0.5] #submatrix 1
[0.75 0.98 0.78] ]
[ [1.5 0.75 0.5]
[8e-15 5e-15 0.8] #submatrix 2
[0.75 0.98 0.78] ]
[ [1.25 0.5 5e-15]
[0.5 0.64 4e-15] #submatrix 3
[5e-15 4e-15 7e-15] ]
我想获取 dist，并为每个第二和第三维参数返回第一维参数(即 1、2 或 3)，dist 是最小值。那将是我的颜色映射。例如，比较 3 个子矩阵中的每一个的元素 (0,0) 将得出颜色 (0,0) = 0。类似地，颜色 (1,1) = 0 和颜色 (2,2) = 2。我希望能够对整个颜色矩阵执行此操作。
我一直没能找到使用 numpy.argmin 执行此操作的方法，但我可能遗漏了一些东西。如果有另一种方法可以做到这一点，我会很高兴听到，特别是如果它不涉及循环。我在这里制作约 25MP 的图像，因此循环需要整整 25 分钟来分配颜色。
提前感谢您的建议!
最佳答案
您可以将 axis 参数传递给 argmin。您想要沿第一个轴(您称之为“子矩阵”)最小化，即 axis=0:
dist.argmin(0)
dist = array([[[ 6.00e-15, 7.00e-15, 5.00e-01],
[ 1.50e+00, 5.00e-15, 5.00e-01],
[ 7.50e-01, 9.80e-01, 7.80e-01]],
[[ 1.50e+00, 7.50e-01, 5.00e-01],
[ 8.00e-15, 5.00e-15, 8.00e-01],
[ 7.50e-01, 9.80e-01, 7.80e-01]],
[[ 1.25e+00, 5.00e-01, 5.00e-15],
[ 5.00e-01, 6.40e-01, 4.00e-15],
[ 5.00e-15, 4.00e-15, 7.00e-15]]])
dist.argmin(0)
#array([[0, 0, 2],
# [1, 0, 2],
# [2, 2, 2]])
这当然会给你 0, 1, 2 作为返回，如果你想要 1, 2, 3 如前所述，使用:
dist.argmin(0) + 1
#array([[1, 1, 3],
# [2, 1, 3],
# [3, 3, 3]])
最后，如果您真的想要最小值本身(而不是它来自哪个“子矩阵”)，您可以只使用 dist.min(0):
dist.min(0)
#array([[ 6.00e-15, 7.00e-15, 5.00e-15],
# [ 8.00e-15, 5.00e-15, 4.00e-15],
# [ 5.00e-15, 4.00e-15, 7.00e-15]])
如果您想使用dist 矩阵中的最小位置从另一个矩阵中提取一个值，这有点棘手，但您可以使用
minloc = dist.argmin(0)
other[dist.argmin(0), np.arange(dist.shape[1])[:, None], np.arange(dist.shape[2])]
请注意，如果 other=dist 这会给出与调用 dist.min(0) 相同的输出:
dist[dist.argmin(0), np.arange(dist.shape[1])[:, None], np.arange(dist.shape[2])]
#array([[ 6.00e-15, 7.00e-15, 5.00e-15],
# [ 8.00e-15, 5.00e-15, 4.00e-15],
# [ 5.00e-15, 4.00e-15, 7.00e-15]])
或者如果 other 只是说它是哪个子矩阵，你会得到同样的结果:
other = np.ones((3,3,3))*np.arange(1,4).reshape(3,1,1)
other
#array([[[ 1., 1., 1.],
# [ 1., 1., 1.],
# [ 1., 1., 1.]],
# [[ 2., 2., 2.],
# [ 2., 2., 2.],
# [ 2., 2., 2.]],
# [[ 3., 3., 3.],
# [ 3., 3., 3.],
# [ 3., 3., 3.]]])
other[dist.argmin(0), np.arange(dist.shape[1])[:, None], np.arange(dist.shape[2])]
#array([[ 1., 1., 3.],
# [ 2., 1., 3.],
# [ 3., 3., 3.]])
作为一个不相关的说明，您可以在没有该循环的情况下重写 colorvec，假设 rootsmatrix.shape 是 (nx, ny) 和 known_roots.shape 是 (dim,)
def colorvec(rootsmatrix, known_roots):
dist = abs(rootsmatrix - known_roots[:, None, None])
其中 known_roots[:, None, None] 与 known_roots.reshape(len(known_roots), 1, 1) 相同，并使其通过 根矩阵
关于python - 在三个二维子矩阵中按元素查找最小值，我们在Stack Overflow上找到一个类似的问题： https://stackoverflow.com/questions/19961088/