Numpy subarray for arbitrary number of dimension

Question

In Numpy, let's say you have a Nd-array A, you can slice the last dimension by doing A[...,0] or you can slice the first dimension by doing A[0]. I want to generalize this operation for all dimension (not just first or last) and i want this for Nd-array of arbitrary N so that if A3 is a 3d-array and A4 is a 4d-array, func(A3, dim = 1, slice = 0) give me A3[ : , 0 , :] and func(A4, dim = 1, slice = 0) give me A4[ : , 0 , : , : ].

I've look for this for some time and finally figured out how to do it without doing horrible hack (like swapping dimension until the one of interest is in last place). So the code i'm posting here works for my need but

1) I'm always looking for advice to better myself

2) As I said, I've look for this for some time and never found anything so it could be usefull for others.

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()

So this little function just builds a slicing tuple by concatenating slice(None,None,1) in the position of the dimension we don't want to slice and slice(fslice, fslice+1, 1) in the dimension of interest. It than returns a subarray. It handles negative indexes.

There's a slight difference with this and straight up indexing: if A3 is 3x4x5, A3[:,0,:] is gonna be 3x5 while fancy_subarray(A3, fdim = 1, fslice = 0) is gonna be 3x1x5. Also the function handle out of bound dimension and indexes "naturally". If fdim >= farray.ndim the function just returns the full array because the if condition inside the for loop is never True and if fslice >= farray.shape[fdim] the returned subarray will have a size of 0 in dimension fdim.

This can easily be extended to more than just selecting one element in one dimension of course.

Thanks!

Answer 1

I think you are overcomplicating things, especially when handling fslice. If you simply did:

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]

Then not only do you make your code more compact, but the same function can take single indices, slices or even lists of indices:

>>> 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]]])