Поиск минимального значения по элементам в трех подматрицах 2d

Question

Я пытаюсь создать цветное отображение сходимости корней многочлена в сложном пространстве. Чтобы сделать это, я создал сетку точек и применил метод Ньютона к этим точкам, чтобы найти, к какому сложному корню они сходятся. Это дает мне 2d массив комплексных чисел, элементы которого обозначают точку, с которой они сходятся, в пределах некоторого допуска. Я хочу иметь возможность сопоставлять числа в этой матрице с цветовым отображением по элементам.

Я сделал это, итерируя по массиву и вычисляя цвета поэтапно, но он очень медленный и кажется, что это принесет пользу от векторизации. Вот мой код до сих пор:

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-й размер аргумента (то есть 1, 2, или 3) для каждого 2nd- и 3-размерности аргумента, для которых dist минимальна. Это будет мое цветовое отображение. Например, сравнение элемента (0,0) каждой из 3 подматриц даст такой цвет (0,0) = 0. Аналогично, цвет (1,1) = 0 и цвет (2,2) = 2. I хотите иметь возможность сделать это для всей цветовой матрицы.

Я не смог найти способ сделать это, используя 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