Выборка по любой оси в массиве numpy

Question

У меня есть следующий фрагмент кода:

if axis == 0:
    return m[idx, :]
elif axis == 1:
    return m[:, idx]
else:
    raise('Argument axis currently accepts only values 0 and 1.')

Как обобщить это поведение для любой формы и любого значения axis?

Answer 1

Вы должны использовать функцию оси вращения в пакете numpy. Пример кода выглядит следующим образом:

import numpy as np
aa = np.arange(30).reshape([5, 6])
print(numpy.rollaxis(a, 0)[1])

Он распечатывает второй элемент на первой оси.

Answer 2

Я нашел ответ:

return np.take(m, idx, axis=axis)