3

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

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

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为每个第 2 维和第 3 维参数返回第 1 维参数（即 1、2 或 3），这dist是最小值。那将是我的颜色映射。例如，比较 3 个子矩阵中每一个的元素 (0,0) 将得出 color(0,0) = 0。类似地，color(1,1) = 0 和 color (2,2) = 2。希望能够对整个颜色矩阵执行此操作。

我一直无法找到使用 的方法来做到这一点numpy.argmin，但我可能会遗漏一些东西。如果有另一种方法可以做到这一点，我很高兴听到，特别是如果它不涉及循环。我在这里制作约 25MP 的图像，因此循环需要 25 分钟来分配颜色。

提前感谢您的建议！

4

1 回答 1

3

您可以将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.shapeis(nx, ny)和known_roots.shapeis (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

于 2013-11-13T18:41:30.667 回答