numpy fancy indexing axis order

Viewed 82

I have a problem using numpy fancy indexing which I somehow can't get my head around. I know, that I can get an array of submatrices of rows like this:

A = np.array([[1,2,3],[4,5,6],[7,8,9]])
B = A[np.array([[0,1],[1,2]])] 

This gives:

array([[[1, 2, 3],
    [4, 5, 6]],

   [[4, 5, 6],
    [7, 8, 9]]])

a threedimensional numpy array containing matrices comprising the first,second row and second,third row of A, respectively. What I want is now basically the same operation for the cols of A which should give

array([[[1, 2],
      [4, 5],
      [7, 8]],

      [[2, 3],
      [5, 6],
      [8, 9]]])

But

B = A[:,np.array([[0,1],[1,2]])] 

does not work (probably because of the order of the index evaluations). It gives

array([[[1, 2],
    [2, 3]],

   [[4, 5],
    [5, 6]],

   [[7, 8],
    [8, 9]]])

How can I accomplish this in the best way? Should I work with transposed matrices?

3 Answers

You get a (3,2,2) array:

In [417]: B
Out[417]: 
array([[[1, 2],
        [2, 3]],

       [[4, 5],
        [5, 6]],

       [[7, 8],
        [8, 9]]])

The 3 is from the first axis of A. The (2,2) from B.

Swap the first 2 axes:

In [418]: B.transpose(1,0,2)
Out[418]: 
array([[[1, 2],
        [4, 5],
        [7, 8]],

       [[2, 3],
        [5, 6],
        [8, 9]]])

A (2,3,2) array

try this:

A = np.array([[1,2,3],[4,5,6],[7,8,9]])
B = A[np.array([[0,1,2],[0,1,2]])]
C = [list(), list()]
for i in range(2):
    for j in  range(3):
        C[i].append(list(B[i][j][:2]) if i==0 else list(B[i][j][1:3])) 
        
C = np.array(C)
C

output:

array([[[1, 2],
        [4, 5],
        [7, 8]],

       [[2, 3],
        [5, 6],
        [8, 9]]])

One way could be to create B from A.T and then swapaxes:

import numpy as np

A = np.array([[1,2,3],[4,5,6],[7,8,9]])

B = A.T[np.array([[0,1],[1,2]])] 

C = B.swapaxes(-2,-1)

To check intermediate step and result:

B
array([[[1, 4, 7],
        [2, 5, 8]],

       [[2, 5, 8],
        [3, 6, 9]]])

C 
array([[[1, 2],
        [4, 5],
        [7, 8]],

       [[2, 3],
        [5, 6],
        [8, 9]]])
Related