1

I would like to create an index in the form of 

[:, :, :, 0, :, :, :, :]

where the position of 0 is determined by a variable, say axis to slice a NumPy array. Obviously there are two special cases that are easy to treat:

  • axis = 0 would be equivalent to [0, ...]
  • axis = -1 would be equivalent to [..., 0]

But I wonder how this can be done for any axis value?

MikeL
  • 2,369
  • 2
  • 24
  • 38
  • 2
    Would [`nump.take()`](https://docs.scipy.org/doc/numpy/reference/generated/numpy.take.html) accomplish what you're trying to do? – AMC Feb 25 '20 at 15:23
  • Indeed! Thanks! I was not aware of that routine. – MikeL Feb 25 '20 at 15:25

1 Answers1

1

You can create a tuple and use slice(None) in place of ::

def custom_index(arr, position, index):
    idx = [slice(None)] * len(arr.shape)
    idx[position] = index
    return arr[tuple(idx)]

Quick test:

mat = np.random.random((5, 3))
assert np.all(mat[2, :] == custom_index(mat, 0, 2))  # mat[(2, slice(None))]
assert np.all(mat[:, 2] == custom_index(mat, 1, 2))  # mat[(slice(None), 2))]

EDIT: as pointed out in the comment, the proper way is np.take

BlackBear
  • 22,411
  • 10
  • 48
  • 86
  • The proper way is np.take as suggested above. But that would only work if one were to slice the array as a rvalue (i.e. to read from an array). If I want to use the slice for 'writing' to the array I should use your `custom_index`? Or perhaps there is a NumPy function for that purpose too? – MikeL Feb 25 '20 at 16:02
  • @kMaster I do not know if there is a numpy function for that, but you can use this method for assignment, `arr[tuple(idx)]=x` works – BlackBear Feb 25 '20 at 16:28