Numpy индексация по диапазону массивов

Question

Скажите, у меня есть массив myarr такой, что myarr.shape = (2,64,64,2). Теперь, если я определю myarr2 = myarr[[0,1,0,0,1],...], тогда верно следующее:

myarr2.shape #(5,64,64,2)
myarr2[0,...] == myarr[0,...] # = True
myarr2[1,...] == myarr[1,...] # = True
myarr2[2,...] == myarr[0,...] # = True
...

Можно ли это обобщить, чтобы срезы были массивами? То есть есть ли способ заставить следующий гипотетический код работать?

myarr2 = myarr[...,[20,30,40]:[30,40,50],[15,25,35]:[25,35,45],..]
myarr2[0,] == myarr[...,20:30,15:25,...] # = True
myarr2[1,] == myarr[...,30:40,25:35,...] # = True
myarr2[2,] == myarr[...,40:50,35:45,...] # = True

Answer 1

Давайте немного упростим задачу; сначала удалив два внешних измерения, которые не влияют на проблему индексации ядра; и уменьшив размер, чтобы мы могли видеть и понимать результаты.

Настройка

In [540]: arr = np.arange(7*7).reshape(7,7)                                                      
In [541]: arr                                                                                    
Out[541]: 
array([[ 0,  1,  2,  3,  4,  5,  6],
       [ 7,  8,  9, 10, 11, 12, 13],
       [14, 15, 16, 17, 18, 19, 20],
       [21, 22, 23, 24, 25, 26, 27],
       [28, 29, 30, 31, 32, 33, 34],
       [35, 36, 37, 38, 39, 40, 41],
       [42, 43, 44, 45, 46, 47, 48]])
In [542]: idx =np.array([[0,2,4,6],[1,3,5,7]])                                                   

Теперь простой итерационный подход:

In [543]: alist = [] 
     ...: for i in range(idx.shape[1]-1): 
     ...:     j,k = idx[:,i] 
     ...:     sub = arr[j:j+2, k:k+2] 
     ...:     alist.append(sub) 
     ...:                                                                                        
In [544]: np.array(alist)                                                                        
Out[544]: 
array([[[ 1,  2],
        [ 8,  9]],

       [[17, 18],
        [24, 25]],

       [[33, 34],
        [40, 41]]])
In [545]: _.shape                                                                                
Out[545]: (3, 2, 2)

Я упростил итерацию from:

 ...: for i in range(idx.shape[1]-1): 
 ...:     sub = arr[idx[0,i]:idx[0,i+1],idx[1,i]:idx[1,i+1]] 
 ...:     alist.append(sub) 

, чтобы подчеркнуть тот факт, что мы генерируем диапазоны согласованного размера, и сделать следующее преобразование более очевидным.

Итак, я начну с массива (7,7) и создайте 3 (2,2) среза.

Как я продемонстрировал в Нарезая разные диапазоны для каждого индекса многомерного numpy массива , мы можем использовать linspace для расширения набор фрагментов или диапазонов.

In [567]: ranges = np.linspace(idx[:,:3],idx[:,:3]+1,2).astype(int)                              
In [568]: ranges                                                                                 
Out[568]: 
array([[[0, 2, 4],
        [1, 3, 5]],

       [[1, 3, 5],
        [2, 4, 6]]])

Таким образом, ranges[0] расширяется на idx[0] фрагментов и т. д. c. Но если я просто индексирую их, я получаю «диагональные» значения из Out[554]:

In [569]: arr[ranges[0], ranges[1]]                                                              
Out[569]: 
array([[ 1, 17, 33],
       [ 9, 25, 41]])

, чтобы получить блоки, мне нужно добавить измерение к первым индексам:

In [570]: arr[ranges[0,:,None], ranges[1]]                                                       
Out[570]: 
array([[[ 1, 17, 33],
        [ 2, 18, 34]],

       [[ 8, 24, 40],
        [ 9, 25, 41]]])

это те же значения, что и в Out[554], но их необходимо транспонировать:

In [571]: _.transpose(2,0,1)                                                                     
Out[571]: 
array([[[ 1,  2],
        [ 8,  9]],

       [[17, 18],
        [24, 25]],

       [[33, 34],
        [40, 41]]])

Код немного неуклюж и нуждается в обобщении, но дает общее представление о том, как можно заменить один индекс на итеративный, обеспечить срезы достаточно регулярно. Для этого небольшого примера это, вероятно, не быстрее, но, вероятно, это произойдет, когда размер проблемы станет больше.

Answer 2

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

idx_list = [[[20,30,40],[30,40,50]],[[15,25,35]:[25,35,45]]]  # assuming 2D cutouts
idx_array = np.array([k for i in idx_list for j in i for k in j]) # unpack
idx_array = idx_array .reshape(-1,2).T  # reshape
myarray2 = np.array([myarray[a:b,c:d] for a,b,c,d in i2])  # cut and combine