Make Python 2D masked array similar to MATLAB's

user_5714890 0
Viewed 77

I'm currently working in translating a MATLAB script to python. There is a step on the MATLAB code where I need to select entries from a matrix (2D array) based on a boolean matrix with the same dimensions. When I tried to write the equivalent code in Python I noticed that the elements in the resulting array were in a different order compared to the MATLAB. More precisely, MATLAB seems to select the elements in column-wise order, and Python in row-wise order. Is there a way to make Python output the array in MATLAB order?

Tiny example:

MATLAB:

a = [1, 2, 3; 4, 5, 6; 7, 8, 9];
b = [false, false, true; false, false, false; true, false, false];
a(b) % outputs [7;3] or [a(3,1); a(1,3)]

PYTHON:

import numpy as np
a = np.array([[1, 2, 3], [4, 5, 6], [7, 8, 9]]))
b= np.array([[False, False, True], [False, False, False], [True, False, False]]))
a[b] # outputs array([3,7]) or array([a[0,2], a[2,0])
matlab numpy arrays python
3 Answers

Expanding on the answer of Mikhail Genkin, one could use

a.T[b.T]

(a.T[b] gives the same result in this particular case since b is a symmetric matrix). Another option is

a.flatten('F')[b.flatten('F')]
user_151877280

How about transposing the array a in Python first?

 a.T[b]
user_43722370

If we flatten the MATLAB matrices:

>> a(:)
ans =

   1
   4
   7
   2
   5
   8
   3
   6
   9

>> b(:)
ans =

  0
  0
  1
  0
  0
  0
  1
  0
  0

the selection becomes more obvious.

>> a(b)
ans =

   7
   3

a.T reorders the numpy array, so does the use of order 'F' (which is what matlab uses)

In [112]: a=np.arange(1,10).reshape(3,3,order='F')
In [113]: a
Out[113]: 
array([[1, 4, 7],
       [2, 5, 8],
       [3, 6, 9]])
In [114]: a[b]
Out[114]: array([7, 3])

numpy default is 'C' order; the last dimension changes fast. When converting between the two we need to keep that order difference in mind.

Or look at the order of elements when the transpose is raveled:

In [126]: a = np.array([[1, 2, 3], [4, 5, 6], [7, 8, 9]])
In [127]: a
Out[127]: 
array([[1, 2, 3],
       [4, 5, 6],
       [7, 8, 9]])
In [128]: a.T.ravel()
Out[128]: array([1, 4, 7, 2, 5, 8, 3, 6, 9])
user_9019250
Related