CODEWITHSUNDEEP
pythonnumpy
Kashif
Kashif

Reputation: 3327

How do I make a function that slices an array depending on parameters

So let's say I have a 2x2x2x2x2 numpy array G. I want to create a function that slices depending on the parameters a and b (where a and b are indices).

For example, I want the function to return G[0,:,0,:,:] if a=0 and b=2. Is this possible?

Upvotes: 3

Views: 107

Answers (2)

Divakar
Divakar

Reputation: 221584

I would think @unutbu's slice based method to be the go-to approach for being short and low on memory usage. Alternatively, I would like to propose an approach based on transposing and reshaping, like so - 

# Get axes IDs for remaining axes
o_axes = np.setdiff1d(np.arange(G.ndim),axes)

# Transpose multi-dimensional array such that input axes are brough at th front
sliced_arr = G.transpose(np.concatenate((axes,o_axes)))

# Finally reshape to merge axes into one axis & slice to get first index from it
out = sliced_arr.reshape(np.append(-1,np.array(G.shape)[o_axes]))[0]

Verify output -

In [23]: G = np.random.randint(0,9,(5,2,4,3,6,4,2))
    ...: axes = [0,2,5]
    ...: out_direct = G[0,:,0,:,:,0,:]
    ...: 

In [24]: o_axes = np.setdiff1d(np.arange(G.ndim),axes)
    ...: sliced_arr = G.transpose(np.concatenate((axes,o_axes)))
    ...: out = sliced_arr.reshape(np.append(-1,np.array(G.shape)[o_axes]))[0]
    ...: 

In [25]: np.allclose(out_direct,out)
Out[25]: True

Upvotes: 1

unutbu
unutbu

Reputation: 879869

You could create a list of slices:

idx = [0 if i in axes else slice(None) for i in range(G.ndim)]

and then return G[idx]:

import numpy as np
np.random.seed(2015)

def getslice(G, axes):
    idx = [0 if i in axes else slice(None) for i in range(G.ndim)]
    return G[idx]

G = np.random.randint(10, size=(2,2,2,2,2,))
assert np.allclose(getslice(G, [0,2]), G[0,:,0,:,:])

Upvotes: 5

Related Questions