Combining slicing and broadcasted indexing for multi-dimensional numpy arrays

Question

I have a ND numpy array (let say for instance 3x3x3) from wich I\'d like to extract a sub-array, combining slices and index arrays. For instance:

import numpy as         


        

Answer 1

This is the closer I can get to your specs, I haven't been able to devise a solution that can compute the correct indices without knowing A (or, more precisely, its shape...).

import numpy as np  

def index(A, s):
    ind = []
    groups = s.split(';')
    for i, group in enumerate(groups):
        if group == ":":
            ind.append(range(A.shape[i]))
        else:
            ind.append([int(n) for n in group.split(',')])
    return np.ix_(*ind)

A = np.arange(3*3*3).reshape((3,3,3))

ind2 = index(A,"0,1;:;0,2")
print A[ind2]

A shorter version

def index2(A,s):return np.ix_(*[range(A.shape[i])if g==":"else[int(n)for n in g.split(',')]for i,g in enumerate(s.split(';'))])

ind3 = index2(A,"0,1;:;0,2")
print A[ind3]