Combining slicing and broadcasted indexing for multi-dimensional numpy arrays

Question

I have a ND numpy array (let say for instance 3x3x3) from wich I'd like to extract a sub-array, combining slices and index arrays. For instance:

import numpy as np  
A = np.arange(3*3*3).reshape((3,3,3))
i0, i1, i2 = ([0,1], [0,1,2], [0,2])
ind1 = j0, j1, j2 = np.ix_(i0, i1, i2)
ind2 = (j0, slice(None), j2)
B1 = A[ind1]
B2 = A[ind2]

I would expect that B1 == B2, but actually, the shapes are different

>>> B1.shape
(2, 3, 2)
>>> B2.shape
(2, 1, 2, 3)
>>> B1
array([[[ 0,  2],
        [ 3,  5],
        [ 6,  8]],

       [[ 9, 11],
        [12, 14],
        [15, 17]]])
>>> B2
array([[[[ 0,  3,  6],
         [ 2,  5,  8]]],

       [[[ 9, 12, 15],
         [11, 14, 17]]]])

Someone understands why? Any idea of how I could get 'B1' by manipulating only 'A' and 'ind2' objects? The goal is that it would work for any nD arrays, and that I would not have to look for the shape of dimensions I want to keep entirely (hope I'm clear enough:)). Thanks!!
---EDIT---
To be clearer, I would like to have a function 'fun' such that

A[fun(ind2)] == B1

Answer 1

The indexing subspaces of ind1 are (2,),(3,),(2,), and the resulting B is (2,3,2) . This is a simple case of advanced indexing.

ind2 is a case of (advanced) partial indexing. There are 2 indexed arrays, and 1 slice. The advanced indexing documentation states:

If the indexing subspaces are separated (by slice objects), then the broadcasted indexing space is first, followed by the sliced subspace of x.

In this case advanced indexing constructs a (2,2) array (from the 1st and 3rd indexes), and appends the slice dimension at the end, resulting in a (2,2,3) array.

I explain the reasoning in more detail in https://stackoverflow.com/a/27097133/901925

A way to fix a tuple like ind2 , is to expand each slice into an array. I recently saw this done in np.insert .

np.arange(*ind2[1].indices(3))

expands : to [0,1,2] . But the replacement has to have the right shape.

ind=list(ind2)
ind[1]=np.arange(*ind2[1].indices(3)).reshape(1,-1,1)
A[ind]

I'm leaving off the details of determining which term is a slice, its dimension, and the relevant reshape. The goal is to reproduce i1 .

If indices were generated by something other than ix_ , reshaping this slice could be more difficult. For example

A[np.array([0,1])[None,:,None],:,np.array([0,2])[None,None,:]] # (1,2,2,3)
A[np.array([0,1])[None,:,None],np.array([0,1,2])[:,None,None],np.array([0,2])[None,None,:]]
# (3,2,2)

The expanded slice has to be compatible with the other arrays under broadcasting.

Swapping axes after indexing is another option. The logic, though, might be more complex. But in some cases transposing might actually be simpler:

A[np.array([0,1])[:,None],:,np.array([0,2])[None,:]].transpose(2,0,1)
# (3,2,2)
A[np.array([0,1])[:,None],:,np.array([0,2])[None,:]].transpose(0,2,1)
# (2, 3, 2)

Answer 2

This is the closer I can get to your specs, I haven't been able to devise a solution that can compute the correct indices without knowing A (or, more precisely, its shape...).

import numpy as np  

def index(A, s):
    ind = []
    groups = s.split(';')
    for i, group in enumerate(groups):
        if group == ":":
            ind.append(range(A.shape[i]))
        else:
            ind.append([int(n) for n in group.split(',')])
    return np.ix_(*ind)

A = np.arange(3*3*3).reshape((3,3,3))

ind2 = index(A,"0,1;:;0,2")
print A[ind2]

A shorter version

def index2(A,s):return np.ix_(*[range(A.shape[i])if g==":"else[int(n)for n in g.split(',')]for i,g in enumerate(s.split(';'))])

ind3 = index2(A,"0,1;:;0,2")
print A[ind3]

Answer 3

In restricted indexing cases like this using ix_ , it is possible to do the indexing in successive steps.

A[ind1]

is the same as

A[i1][:,i2][:,:,i3]

and since i2 is the full range,

A[i1][...,i3]

If you only have ind2 available

A[ind2[0].flatten()][[ind2[2].flatten()]

In more general contexts you have to know how j0,j1,j2 broadcast with each other, but when they are generated by ix_ , the relationship is simple.

I can imagine circumstances in which it would be convenient to assign A1 = A[i1] , followed by a variety of actions involving A1 , including, but not limited to A1[...,i3] . You have to be aware of when A1 is a view, and when it is a copy.

Another indexing tool is take :

A.take(i0,axis=0).take(i2,axis=2)