Python: Index an array using the colon operator in an arbitrary dimension

后端 未结 2 930
暖寄归人
暖寄归人 2021-02-20 02:17

I have a numpy nd array. A simplified version of my task is to take a vector from along each axis. To illustrate:

import numpy
x = numpy.array(range(24)).reshape         


        
相关标签:
2条回答
  • 2021-02-20 02:47

    You could compose an string with the code selecting the dimension you want and use eval to execute that code string.

    An start is:

    n = 2
    sel = "0,"*(n-1) + ":"
    eval('x[' + sel + ']')
    

    To get exactly what you want, thinks are a little bit more complicated (but not so much):

    ind = 2
    n = 3
    sel = "".join([ ("0" if i != ind else ":") + ("," if i < n-1 else "") for i in xrange(n)])
    eval('x[' + sel + ']')
    

    It is the same strategy that is used for Dynamic SQL.

    0 讨论(0)
  • 2021-02-20 02:51

    As suggested from numpy's documentation about indexing you can use the slice built-in function and tuple concatenation to create variable indexes.

    In fact the : in the subscript is simply the literal notation for a slice literal.

    In particular : is equivalent to slice(None) (which, itself, is equivalent to slice(None, None, None) where the arguments are start, stop and step).

    For example:

    a[(0,) * N + (slice(None),)]
    

    is equivalent to:

    a[0, 0, ..., 0, :]   # with N zeros
    

    The : notation for slices can only be used directly inside a subscript. For example this fails:

    In [10]: a[(0,0,:)]
      File "<ipython-input-10-f41b33bd742f>", line 1
        a[(0,0,:)]
               ^
    SyntaxError: invalid syntax
    

    To allow extracting a slice from an array of arbitrary dimensions you can write a simple function such as:

    def make_index(num_dimension, slice_pos):
        zeros = [0] * num_dimension
        zeros[slice_pos] = slice(None)
        return tuple(zeros)
    

    And use it as in:

    In [3]: a = np.array(range(24)).reshape((2, 3, 4))
    
    In [4]: a[make_index(3, 2)]
    Out[4]: array([0, 1, 2, 3])
    
    In [5]: a[make_index(3, 1)]
    Out[5]: array([0, 4, 8])
    
    In [6]: a[make_index(3, 0)]
    Out[6]: array([ 0, 12])
    

    You can generalize make_index to do any kind of things. The important thing to remember is that it should, in the end, return a tuple containing either integers or slices.

    0 讨论(0)
提交回复
热议问题