Getting a slice of a numpy ndarray (for arbitary dimensions)

Question

I have a Numpy array of arbitrary dimensions, and an index vector containing one number for each dimension. I would like to get the slice of the array corresponding to the set of indices less than the value in the index array for all dimensions, e.g.

A = np.array([[1, 2, 3, 4],
              [5, 6, 7, 8],
              [9,10,11,12]])
index = [2,3]

result = [[1,2,3],
          [5,6,7]]

The intuitive syntax for this would be something like A[:index], but this doesn't work for obvious reasons.

If the dimension of the array were fixed, I could write A[:index[0],:index[1],...:index[n]]; is there some kind of list comprehension I could use, like A[:i for i in index]?

Answer 1

You can slice multiple dimensions in one go:

result = A[:2,:3]

that slices dimension one up to the index 2 and dimension two up to the index 3.

If you have arbitary dimensions you can also create a tuple of slices:

slicer = tuple(slice(0, i, 1) for i in index)
result = A[slicer]

A slice defines the start(0), stop(the index you specified) and step(1) - basically like a range but useable for indexing. And the i-th entry of the tuple slices the i-th dimension of your array.

If you only specify stop-indices you can use the shorthand:

slicer = tuple(slice(i) for i in index)

I would recommend the first option if you know the number of dimensions and the last one if you don't.