Search code examples
pythonnumpymatrixnumpy-ndarraynumpy-slicing

Indexing ndarray with unknown number of dimensions with range dynamically


I have data array with unknown shape and array bounds of bounds for slicing data. This code is for 3D data, but is there any way of generalizing this to N-dim?

for b in bounds:
    l0, u0 = b[0]
    l1, u1 = b[1]
    l2, u2 = b[2]
    a = data[l0:u0, l1:u1, l2:u2]
    print(a)

Tried using range python object as index, did not work.

Examples for data:

data2D = np.arange(2*3).reshape((2, 3))
data3D = np.arange(2*3*4).reshape((2, 3, 4))

Corresponding bounds:

bounds2D = np.array([[[0, 2], [0, 2]], [[0, 2], [1, 3]]])
bounds3D = np.array(
        [
            [[0, 2], [0, 2], [0, 2]],
            [[0, 2], [0, 2], [2, 4]],
            [[0, 2], [1, 3], [0, 2]],
            [[0, 2], [1, 3], [2, 4]],
        ],
    )

Solution

  • You can use the slice function to create a single slice from each element in bounds. Then collect these slices into a single tuple and use it to correctly recover the wanted items of the array. You can adapt your code as follows:

    import numpy as np
    
    
    # The dimension of the slices is equal to the 
    # one specified by the bounds provided
    def create_slices(bounds):
        slices = list()
    
        # Take a single item of the bounds and create corresponding slices
        for b in bounds:
            # Slices are collected inside a single tuple
            slices.append(tuple([slice(l, u) for l, u in b]))
        return slices
       
    # 4D example data   
    data4D = np.arange(2*3*4*5).reshape((2,3,4,5))
    
    # Bounds array for 4D data
    bounds4D = np.array(
            [
                [[0, 2], [0, 2], [0, 2], [0, 2]],
                [[0, 2], [0, 2], [0, 2], [2, 4]],
                [[0, 2], [1, 3], [2, 4], [0, 2]],
                [[0, 2], [1, 3], [2, 4], [2, 4]],
            ],
       )
    
    
    slices = create_slices(bounds4D)
    
    # Each element of slices is a single slice that can be used on 
    # the corresponding data array
    for single_slice in slices:
        a = data4D[single_slice]
        print("Slice", a)