在三个二维子矩阵中逐元素查找最小值

Question

我正在尝试生成复数空间中多项式根收敛的颜色映射。为了做到这一点，我创建了一个点网格并将牛顿法应用于这些点，以便找到它们各自收敛到哪个复根。这给了我一个二维复数数组，其中的元素表示它们收敛的点，在一定的公差范围内。我希望能够将该矩阵中的数字与元素颜色映射相匹配。

我通过遍历数组并逐个元素计算颜色来做到这一点，但它非常慢，而且似乎它会从矢量化中受益。到目前为止，这是我的代码：

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 分钟来分配颜色。

提前感谢您的建议！

Answer 1

您可以将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) 相同，并使其与rootsmatrix 一起广播