Indexing ndarray with unknown number of dimensions with range dynamically-CodePudding

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]],
        ],
    )

CodePudding user response：

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)