Finding minimum value element-wise in three 2d submatrices

Question

I'm trying to produce a color mapping of the convergence of a polynomial's roots in complex space. In order to do this, I have created a grid of points and applied Newton's method to those points, in order to find to which complex root they each converge. This gives me a 2d array of complex numbers, the elements of which denote the point to which they converge, within some tolerance. I want to be able to match the numbers in that matrix to an element-wise color mapping.

I have done this by iterating over the array and computing colors element-by-element, but it is very slow, and seems it would benefit from vectorizing. Here's my code so far:

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])

This creates a 3d array with the distances of each point's computed root to each of the actual roots. It looks something like this, except with 75 000 000 elements.

    [ [ [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] ]

I want to take dist, and return the 1st dimension argument (i.e., 1, 2, or 3) for each 2nd- and 3rd-dimension argument, for which dist is minimum. That will be my color mapping. For example, comparing the element (0,0) of each of the 3 submatrices would yield that color(0,0) = 0. Similarly, color(1,1) = 0 and color (2,2) = 2. I want to be able to do this for the entire color matrix.

I haven't been able to find a way to do this using numpy.argmin, but I could be missing something. If there's another way to do this, I'd be happy to hear, especially if it doesn't involve loops. I'm making ~25MP images here, so looping takes fully 25 minutes to assign colors.

Thanks in advance for your advice!

Answer 1

You can pass an axis argument to argmin. You want to minimize along the first axis (what you're calling 'submatrices'), which is 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]])

This of course gives you 0, 1, 2 as the returns, if you want 1, 2, 3 as stated, use:

dist.argmin(0) + 1
#array([[1, 1, 3],
#       [2, 1, 3],
#       [3, 3, 3]])

Finally, if you actually want the minimum value itself (instead of which 'submatrix' it comes from), you can just use 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]])

If you want to use the minimum location from the dist matrix to pull a value out of another matrix, it's a little tricky, but you can use

minloc = dist.argmin(0)
other[dist.argmin(0), np.arange(dist.shape[1])[:, None], np.arange(dist.shape[2])]

Note that if other=dist this gives the same output as just calling 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]])

or if other just says which submatrix it is, you get the same thing back:

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.]])

---

As an unrelated note, you can rewrite colorvec without that loop, assuming rootsmatrix.shape is (nx, ny) and known_roots.shape is (dim,)

def colorvec(rootsmatrix, known_roots):
    dist = abs(rootsmatrix - known_roots[:, None, None])

where known_roots[:, None, None] is the same as known_roots.reshape(len(known_roots), 1, 1) and causes it to broadcast with rootsmatrix