上的可变轴

Question

欲索引通过多维阵列这样numpy的索引：

a = range(12).reshape(3, 2, 2) 
def fun(axis, state): 
    # if axis=0 
    return a[state, :, :] 
    # if axis=1 it should return a[:, state, :] 

样本输出：

fun(0, 1) 
array([[4, 5],         
     [6, 7]]) 

fun(1, 1) 
array([[2, 3], 
     [6, 7], 
     [10, 11]]) 

总之我希望接受轴作为参数。

我想不出有办法做到这一点。任何可能的解

Answer 1

可以采取与指定轴线的阵列的视图移到前面使用numpy.rollaxis：

def fun(a, axis, state): 
    return numpy.rollaxis(a, axis)[state] 

演示：

>>> a = numpy.arange(12).reshape([3, 2, 2]) 
>>> def fun(a, axis, state): 
...  return numpy.rollaxis(a, axis)[state] 
... 
>>> fun(a, 0, 1) 
array([[4, 5], 
     [6, 7]]) 
>>> fun(a, 1, 1) 
array([[ 2, 3], 
     [ 6, 7], 
     [10, 11]]) 

numpy.rollaxis还支持移动轴到另一个位置时，虽然它解释为什么这是奇怪的论据。