Im looking for an elegant way to extract the values of an single axis of a numpy array by an index. For example:
x = np.arange(16).reshape((4,4))
a = x[0]
b = x[:, 0]
Is what i usually do, however i am looking for something like:
a = get( x, axis=0, index=0)
b = get( x, axis=1, index=0)
Is there some fancy function to do this maybe?
You can use np.rollaxis to move the axis you're interested in to the front, then just index into it as normal:
def get(x, axis=0, index=0):
return np.rollaxis(x, axis, 0)[index]
x = np.arange(27).reshape(3, 3, 3)
assert np.all(get(x, 1, 2) == x[:, 2, :])
As Joe correctly pointed out, this returns a view onto x. In order to force a copy to be made, you could use the .copy() method:
cpy = get(x, 1, 2).copy()
rollaxis in this manner is a common idiom in numpy. np.rollaxis makes a view rather than a copy, so this technique is just as memory-efficient and fast as writing the indexing expression directly. It's useful in a lot of cases where you want to operate on an arbitrary axis of a N-dimensional array. Roll and then operate on the first axis. Because rolling creates a view, if modifications are made to the "rolled" version, they're made in the original as well.You could use the __getitem__ magic method directly to get the same functionally, but with easier support for dynamic arguments. Here is an example that uses the itertools package:
def get(matrix, axis, index):
return a.__getitem__(tuple(chain(repeat(slice(None), axis), (index,))))
This creates a tuple with a slice object that represents the colon in a[:] repeated axis times with index at the end. I think the tuple generation could be cleaned up, but I can't think of a cleaner way at the moment.
An example usage would be this:
a = np.arange(9).reshape(3, 3) # [[0 1 2], [3 4 5], [6 7 8]]
get(a, axis=0, index=0) # [0 1 2]
get(a, axis=1, index=0) # [0 3 6]
get(a, axis=0, index=1) # [3 4 5]
get(a, axis=1, index=1) # [1 4 7]
