Numpy Subarray за произвольное количество измерений

Question

В Numpy, допустим, у вас есть ND-Array 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,:] будет 3х5, пока fancy_subarray(A3, fdim = 1, fslice = 0) будет 3x1x5. Также функциональная ручка из граничного измерения и индексы «естественно». Если fdim >= farray.ndim Функция просто возвращает полный массив, потому что условие if внутри цикла никогда не верно, и если fslice >= farray.shape[fdim] Возвращенный Subarray будет иметь размер 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]]])