Search code examples
pythonnumpymultidimensional-arrayindexingdimensions

Accessing the j-th dimension of an n-dimensional matrix in python


Given:

  • an integer input j,d such that 0 < j < d+1

  • integer vectors -1 < a < b of dimension d.

  • d-dimensional matrix (i.e. tensor) T as a numpy array

I would like to read certain information that would variably depend on the integer j.

For example,

  c[1]>u[,1]

I would like to access

  T[(a[0]):(b[0]),...,(a[j]-1):(b[j]+1),...,(a[n-1]):(b[n-1])]

I am wondering if there is a generic way of doing this, especially in the case where d and j can be variable.

A similar question can be found here: Access n-th dimension in python.


Solution

  • Constructing an indexing tuple from slices:

    In [88]: a = [1,0,4]; b = [4,1,None]
    In [89]: idx = [slice(i,j) for i,j in zip(a,b)]
    In [90]: idx
    Out[90]: [slice(1, 4, None), slice(0, 1, None), slice(4, None, None)]
    In [91]: arr = np.arange(5*3*7).reshape(5,3,7)
    In [92]: arr[tuple(idx)]
    Out[92]: 
    array([[[25, 26, 27]],
    
           [[46, 47, 48]],
    
           [[67, 68, 69]]])
    In [93]: _.shape
    Out[93]: (3, 1, 3)