CODEWITHSUNDEEP
pythonarraysnumpy
generic_user
generic_user

Reputation: 3562

Good way to slice a numpy array based on the shape of another array

Take two arrays of arbitrary shape, but where each of the dimensions of the second is less than or equal to the dimensions of the first. For example:

np.random.seed(8675309)
a = np.random.choice(10, 3**3).reshape(3,3,3)
b = np.zeros(2**3).reshape(2,2,2)

What I want is the following:

c = a[:b.shape[0], :b.shape[1], :b.shape[2]]

but for an array b with arbitrary shape, potentially with fewer dimensions. How could I do this programmatically? Such that 

def reference_slicer(a, b):
    ???
    return c

reference_slicer(a,b) == c

Upvotes: 1

Views: 501

Answers (1)

ExplodingGayFish
ExplodingGayFish

Reputation: 2897

You mean something like this?

def reference_slicer(a, b):
    index = [slice(0, dim) for dim in b.shape]
    for i in range(len(b.shape), len(a.shape)):
        index.append(slice(0,a.shape[i]))
    index = tuple(index)
    return a[index]

#array([[[ True,  True],
#        [ True,  True]],
#       [[ True,  True],
#        [ True,  True]]])

Upvotes: 2

Related Questions