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

Question

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?

Answer 1

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