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