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找到一种方法，但我可能会遗漏一些东西。如果有另一种方法可以做到这一点，我会很高兴听到，特别是如果它不涉及循环。我在这里制作了大约2500万像素，所以循环需要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一起广播

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

1 个答案: