У меня есть 4-D тензор, triangular_distances, с размерами (3,3,2,7), который состоит из семи 3-D тензоров, каждый из которых имеет размерность (3,3,2), которые являются (3 , 3) матрицы двумерных векторов. Я нахожу норму вдоль третьего измерения normed_tri_distances, норма двухмерных векторов, в результате чего получается трехмерный тензор, имеющий размерность (3,3,7). Затем я использую np.argmin(axis=-1), который возвращает матрицу индексов 3x3, min_tri_distance_idx, для которой из семи матриц была самая низкая норма в каждой ячейке. Я хочу затем извлечь (3,3,2) тензор из исходного 4-D тензора triangular_distances, который содержит 2-D векторы, которые имели минимальные нормы. Я могу сделать это с помощью al oop, но с помощью sh, если это возможно, векторизовать операцию.
import numpy as np
mm=3
nn=3
mmm = (np.arange(mm) + (1 / mm)) / mm
nnn = (np.arange(nn) + (5 / nn)) / nn
xx, yy = np.meshgrid(mmm, nnn)
xy = np.stack([xx, yy], axis=-1)
xyt = xy.transpose((1,0,2))
Sdist = [(0,0), (-0.5, np.sqrt(3) / 2), (-0.5, -np.sqrt(3) / 2),
(0.5, np.sqrt(3) / 2), (0.5, -np.sqrt(3) / 2),
(-1, 0), (1, 0)]
distance_tensor = xy - xyt
triangular_distances = np.stack([distmat + Sdist[i] for i in range(len(Sdist))],axis=-1)
normed_tri_distances = np.linalg.norm(triangular_distances, axis=2)
min_tri_distance = np.min(normed_tri_distances, axis=-1)
min_tri_distance_idx = np.argmin(normed_tri_distances, axis=-1)
Эта последняя строка завершается неудачей.
min_distances = triangular_distances[min_tri_distance_idx]
Решение с помощью l oop :
min_distances = []
for ind, val in np.ndenumerate(min_tri_distance_idx):
dist = triangular_distances[ind[0],ind[1],:,val]
min_distances.append(dist)
min_distances_reshaped = np.array(min_distances).reshape((mm,nn,2))
sanity_check = np.linalg.norm(min_distances_reshaped, axis=2)==final_tri_distance