Numar subarray для произвольного числа измерений

Question

В Numpy предположим, что у вас есть Nd-массив A, вы можете нарезать последнее измерение, выполнив A[...,0], или вы можете нарезать первое измерение, выполнив A[0]. Я хочу обобщить эту операцию для всей размерности (не только первой или последней), и я хочу это для Nd-массива произвольного N, так что если A3 представляет собой 3d-массив, а A4 - это 4d-массив, func(A3, dim = 1, slice = 0) дайте мне A3[ : , 0 , :] и func(A4, dim = 1, slice = 0) дайте мне A4[ : , 0 , : , : ].

Я искал это в течение некоторого времени и, наконец, понял, как это сделать, не делая ужасного взлома (например, меняя размер, пока интерес не окажется на последнем месте). Таким образом, код я отправляю здесь работает для моей потребности, но

1) Я всегда искал советы, чтобы лучше себя

2) Как я уже сказал, у меня есть взгляд на это в течение некоторого времени и никогда нашел что-нибудь, чтобы оно могло быть полезным для других.

def fancy_subarray(farray, fdim, fslice): 
    # Return the farray slice at position fslice in dimension fdim 

    ndim = farray.ndim 

    # Handle negative dimension and slice indexing 
    if fdim < 0: 
     fdim += ndim 
    if fslice < 0: 
     fslice += v.shape[fdim] 

    # Initilize slicing tuple 
    obj =() 

    for i in range(ndim): 
     if i == fdim: 
      # Only element fslice in that dimension 
      obj += (slice(fslice, fslice+1, 1),) 
     else: 
      # "Full" dimension 
      obj += (slice(None,None,1),) 

    return farray[obj].copy() 

Так эта маленькая функция просто строит нарезки кортежа конкатенации slice(None,None,1) в положении измерения, мы не хотим, чтобы нарезать и slice(fslice, fslice+1, 1) в измерении интереса. Это, чем возвращает подмассиву. Он обрабатывает отрицательные индексы.

Есть небольшая разница в этом и прямое индексирование: если A3 - 3x4x5, A3[:,0,:] будет 3x5, тогда как fancy_subarray(A3, fdim = 1, fslice = 0) будет 3x1x5. Кроме того, функция обрабатывает связанные измерения и индексы «естественно». Если fdim >= farray.ndim функция просто возвращает полный массив, потому что условие if внутри цикла for никогда не является истинным, а если fslice >= farray.shape[fdim], то возвращаемый субарабрий будет иметь размер 0 в размерном fdim.

Это может быть легко распространено не только на выбор одного элемента в одном измерении, конечно.

Спасибо!

Answer 1

Я думаю, что вы чрезмерны, особенно при обработке fslice. Если вы просто сделали:

def fancy_subarray(farray, fdim, fslice): 
    fdim += farray.ndim if fdim < 0 else 0 
    index = ((slice(None),) * fdim + (fslice,) + 
      (slice(None),) * (farray.ndim - fdim - 1)) 
    return farray[index] 

Тогда не только вы делаете ваш код более компактным, но та же функция может принимать отдельные индексы, ломтики или даже списки индексов:

>>> fancy_subarray(a, 1, 2) 
array([[10, 11, 12, 13, 14], 
     [30, 31, 32, 33, 34], 
     [50, 51, 52, 53, 54]]) 
>>> fancy_subarray(a, 1, slice(2,3)) 
array([[[10, 11, 12, 13, 14]], 

     [[30, 31, 32, 33, 34]], 

     [[50, 51, 52, 53, 54]]]) 
>>> fancy_subarray(a, 1, [2, 3]) 
array([[[10, 11, 12, 13, 14], 
     [15, 16, 17, 18, 19]], 

     [[30, 31, 32, 33, 34], 
     [35, 36, 37, 38, 39]], 

     [[50, 51, 52, 53, 54], 
     [55, 56, 57, 58, 59]]])