slicing array numpy lungo una dimensione arbitraria

Question

dire che ho una matrice numpy (40,20,30) e che ho una funzione che dopo un po 'di lavoro restituirà metà della matrice di input lungo un asse di input selezionato. C'è un modo automatico per farlo? Vorrei evitare un codice così brutta:

def my_function(array,axis=0): 

    ... 

    if axis == 0: 
     return array[:array.shape[0]/2,:,:] --> (20,20,30) array 
    elif axis = 1: 
     return array[:,:array.shape[1]/2,:] --> (40,10,30) array 
    elif axis = 2: 
     return array[:,:,:array.shape[2]/2] --> (40,20,15) array 

grazie per il vostro aiuto

Eric

Answer 1

Penso che si possa utilizzare np.split per questo [docs], e semplicemente prendere il primo o il secondo elemento è tornato, dipende da quale vuoi Ad esempio:

>>> a = np.random.random((40,20,30)) 
>>> np.split(a, 2, axis=0)[0].shape 
(20, 20, 30) 
>>> np.split(a, 2, axis=1)[0].shape 
(40, 10, 30) 
>>> np.split(a, 2, axis=2)[0].shape 
(40, 20, 15) 
>>> (np.split(a, 2, axis=0)[0] == a[:a.shape[0]/2, :,:]).all() 
True 

Answer 2

grazie per il vostro aiuto, DSM. Userò il tuo approccio.

Nel frattempo, ho trovato un (sporco?) Incidere:

>>> a = np.random.random((40,20,30)) 
>>> s = [slice(None),]*a.ndim 
>>> s[axis] = slice(f,l,s) 
>>> a1 = a[s] 

Forse un po 'più generale di np.split ma molto meno elegante!

Answer 3

numpy.rollaxis è un buon strumento per questo:

def my_func(array, axis=0): 
    array = np.rollaxis(array, axis) 
    out = array[:array.shape[0] // 2] 
    # Do stuff with array and out knowing that the axis of interest is now 0 
    ... 

    # If you need to restore the order of the axes 
    if axis == -1: 
     axis = out.shape[0] - 1 
    out = np.rollaxis(out, 0, axis + 1)