Extract sub-array from 2D array using logical indexing - python

Question

I am trying to extract a sub-array using logical indexes as,

a = np.array([[1, 2, 3, 4], [5, 6, 7, 8], [9, 10, 11, 12], [13, 14, 15, 16]])
a
Out[45]: 
array([[ 1,  2,  3,  4],
       [ 5,  6,  7,  8],
       [ 9, 10, 11, 12],
       [13, 14, 15, 16]])
b = np.array([False, True, False, True])
a[b, b]
Out[49]: array([ 6, 16])

python evaluates the logical indexes in b per element of a. However in matlab you can do something like

>> a = [1 2 3 4; 5 6 7 8; 9 10 11 12; 13 14 15 16]

a =

     1     2     3     4
     5     6     7     8
     9    10    11    12
    13    14    15    16

>> b = [2 4]

b =

     2     4

>> a(b, b)

ans =

     6     8
    14    16

how can I achieve the same result in python without doing,

c = a[:, b]
c[b,:]
Out[51]: 
array([[ 6,  8],
       [14, 16]])

Answer 1

Numpy supports logical indexing, though it is a little different than what you are familiar in MATLAB. To get the results you want you can do the following:

a[b][:,b]  # first brackets isolates the rows, second brackets isolate the columns
Out[27]: 
array([[ 6,  8],
       [14, 16]])

The more "numpy" method will be understood after you will understand what happend in your case. b = np.array([False, True, False, True]) is similar to b=np.array([1,3]) and will be easier for me to explain. When writing a[[1,3],[1,3]] what happens is that numpy crates a (2,1) shape array, and places a[1,1] in the [0] location and a[3,3] in the second location. To create an output of shape (2,2), the indexing must have the same dimensionality. Therefore, the following will get your result:

a[[[1,1],[3,3]],[[1,3],[1,3]]]
Out[28]: 
array([[ 6,  8],
       [14, 16]])

Explanation :

The indexing arrays are:

temp_rows = np.array([[1,1],
                      [3,3]])
temp_cols = np.array([[1,3],
                      [1,3])

both arrays have dimensions of (2,2) and therefore, numpy will create an output of shape (2,2). Then, it places a[1,1] in location [0,0], a[1,3] in [0,1], a[3,1] in location [1,0] and a[3,3] in location [1,1]. This can be expanded to any shape but for your purposes, you wanted a shape of (2,2)

After figuring this out, you can make things even simpler by utilizing the fact you if you insert a (2,1) array in the 1st dimension and a (1,2) array in the 2nd dimension, numpy will perform the broadcasting, similar to the MATLAB operation. This means that by using:

temp_rows = np.array([[1],[3]])
temp_cols = np.array([1,3])

you can do:

a[[[1],[3]], [1,3])
Out[29]: 
array([[ 6,  8],
       [14, 16]])

Answer 2

You could usenp.ix_ here.

a[np.ix_(b, b)]

# array([[ 6,  8],
#        [14, 16]])

---

Output returned by np.ix_

>>> np.ix_(b, b)
(array([[1],
        [3]]),
 array([[1, 3]]))

Answer 3

You could make use of a outer product of the b vector. The new dimesion you can obtain from the number of True values using a sum.

import numpy as np
a = np.array([[1, 2, 3, 4], [5, 6, 7, 8], [9, 10, 11, 12], [13, 14, 15, 16]])
b = np.array([False, True, False, True])
#
M = np.outer(b, b)
new_dim = b.sum()
new_shape = (new_dim, new_dim)

selection = a[M].reshape(new_shape)

The result looks like

[[ 6  8]
 [14 16]]