Явная нарезка по определенному измерению

Question

У меня есть 3D-тензор x (например, 4x4x100).Я хочу получить подмножество этого, явно выбирая элементы в последнем измерении.Это было бы легко, если бы я выбирал одни и те же элементы в последнем измерении (например, x[:,:,30:50], но я хочу нацеливать разные элементы в этом измерении, используя 2D-тензор indices, который задает idx для третьего измерения. Есть ли простой способсделать это в numpy?

Более простой 2D-пример:

x = [[1,2,3,4,5,6],[10,20,30,40,50,60]]
indices = [1,3]

Допустим, я хочу перехватить два элемента в третьем измерении x, начиная с точек, указанных indices.Итак, мой желаемый результат:

[[2,3],[40,50]]

Обновление: я думаю, что я мог бы использовать комбинацию take() и ravel_multi_index(), но некоторые из платформ, которые вдохновлены Numpy (например, PyTorch), не кажутсяиметь ravel_multi_index, поэтому я ищу альтернативные решения

Answer 1

Итерация по idx и сбор фрагментов - неплохой вариант, если число «строк» ​​не слишком велико (а размер размеров относительно велик).

In [55]: x = np.array([[1,2,3,4,5,6],[10,20,30,40,50,60]])                      
In [56]: idx = [1,3]                                                            
In [57]: np.array([x[j,i:i+2] for j,i in enumerate(idx)])                       
Out[57]: 
array([[ 2,  3],
       [40, 50]])

Соединение ломтиков, как это работает, только если они все одного размера.

Альтернативой является сбор индексов в массив и выполнение одной индексации.

Например, с похожей итерацией:

idxs = np.array([np.arange(i,i+2) for i in idx])

Но вещание может быть лучше:

In [58]: idxs = np.array(idx)[:,None]+np.arange(2)                              
In [59]: idxs                                                                   
Out[59]: 
array([[1, 2],
       [3, 4]])
In [60]: x[np.arange(2)[:,None], idxs]                                          
Out[60]: 
array([[ 2,  3],
       [40, 50]])

---

ravel_multi_index не сложно скопировать (если вам не нужно вырезать и т. Д.):

In [65]: np.ravel_multi_index((np.arange(2)[:,None],idxs),x.shape)              
Out[65]: 
array([[ 1,  2],
       [ 9, 10]])
In [66]: x.flat[_]                                                              
Out[66]: 
array([[ 2,  3],
       [40, 50]])
In [67]: np.arange(2)[:,None]*x.shape[1]+idxs                                   
Out[67]: 
array([[ 1,  2],
       [ 9, 10]])

Answer 2

Я придумал ниже, используя трансляцию:

x = np.array([[1,2,3,4,5,6,7,8,9,10],[10,20,30,40,50,60,70,80,90,100]]) 
i = np.array([1,5])
N = 2 # number of elements I want to extract along each dimension. Starting points specified in i

r = np.arange(x.shape[-1])
r = np.broadcast_to(r, x.shape)

ii = i[:, np.newaxis]
ii = np.broadcast_to(ii, x.shape)

mask = np.logical_and(r-ii>=0, r-ii<=N) 

output = x[mask].reshape(2,3)

Это выглядит хорошо?

Answer 3

Вот как это сделать в numpy, теперь подсказка о факеле.

Следующий фрагмент выбирает отрезок длины n вдоль третьего измерения, начиная с точек idx, в зависимости от двух других измерений:

# example
a = np.arange(60).reshape(2, 3, 10)
idx = [(1,2,3),(4,3,2)]
n = 4

# build auxiliary 4D array where the last two dimensions represent
# a sliding n-window of the original last dimension
j,k,l = a.shape
s,t,u = a.strides
aux = np.lib.stride_tricks.as_strided(a, (j,k,l-n+1,n), (s,t,u,u))

# pick desired offsets from sliding windows
aux[(*np.ogrid[:j, :k], idx)]
# array([[[ 1,  2,  3,  4],
#         [12, 13, 14, 15],
#         [23, 24, 25, 26]],

#        [[34, 35, 36, 37],
#         [43, 44, 45, 46],
#         [52, 53, 54, 55]]])

Answer 4

вдоль оси 3D:

x = [x[:,i].narrow(2,index,2) for i,index in enumerate(indices)]
x = torch.stack(x,dim=1)

перечисляя, вы получаете индекс оси и индекс, с которого вы хотите начать нарезку по одной.

узкий дает нулевую копию length длинного среза из начального индекса start вдоль определенной оси

Вы сказали, что хотели:

dim = 2
start = index
length = 2

тогда вам просто нужно сложить эти тензоры обратно в одно 3D.

Это наименее трудоемкая вещь, которую я могу придумать для pytorch.

EDIT

, если вам просто нужны разные индексы вдоль разных осей, а indices - это 2D-тензор, который вы можете сделать:

x = [x[:,i,index] for i,index in enumerate(indices)]
x = torch.stack(x,dim=1)

Вы действительно должны были привести правильный рабочий пример, из-за чего он излишне запутался.