CODEWITHSUNDEEP
pythonmatlabnumpymultidimensional-arraymasking
Developer
Developer

Reputation: 8400

Python: multidimensional array masking

What is equivalent pythonic implementation for the following simple piece of code in Matlab.

Matlab:


B = 2D array of integers as indices [1...100]
A = 2D array of numbers: [10x10]
A[B] = 0

which works well as for example for B[i]=42 it finds the position 2 of column 5 to be set. In Python it causes an error: out of bound which is logical. However to translate the above Matlab code into Python we are looking for pythonic ways. Please also consider the problem for higher dimensions such as:


B = 2D array of integers as indices [1...3000]
C = 3D array of numbers: [10x10x30]
C[B] = 0

One way we thought about it is to reform indices array elements as i,j instead of being absolute position. That is, position 42 to divmod(42,m=10)[::-1] >>> (2,4). So we will have a nx2 >>> ii,jj vectors of indices which can be used for indexing A easily. We thought that it might be a better way, efficient also for higher dimensions in Python.

Upvotes: 2

Views: 1040

Answers (1)

wim
wim

Reputation: 363456

You can use .ravel() on the array (A) before indexing it, and then .reshape() after.

Alternatively, since you know A.shape, you can use np.unravel_index on the other array (B) before indexing.

Example 1:

>>> import numpy as np
>>> A = np.ones((5,5), dtype=int)
>>> B = [1, 3, 7, 23]
>>> A
array([[1, 1, 1, 1, 1],
       [1, 1, 1, 1, 1],
       [1, 1, 1, 1, 1],
       [1, 1, 1, 1, 1],
       [1, 1, 1, 1, 1]])
>>> A_ = A.ravel()
>>> A_[B] = 0
>>> A_.reshape(A.shape)
array([[1, 0, 1, 0, 1],
       [1, 1, 0, 1, 1],
       [1, 1, 1, 1, 1],
       [1, 1, 1, 1, 1],
       [1, 1, 1, 0, 1]])

Example 2: 

>>> b_row, b_col = np.vstack([np.unravel_index(b, A.shape) for b in B]).T
>>> A[b_row, b_col] = 0
>>> A
array([[1, 0, 1, 0, 1],
       [1, 1, 0, 1, 1],
       [1, 1, 1, 1, 1],
       [1, 1, 1, 1, 1],
       [1, 1, 1, 0, 1]])

Discovered later: you can use numpy.put

>>> import numpy as np
>>> A = np.ones((5,5), dtype=int)
>>> B = [1, 3, 7, 23]
>>> A.put(B, [0]*len(B))
>>> A
array([[1, 0, 1, 0, 1],
       [1, 1, 0, 1, 1],
       [1, 1, 1, 1, 1],
       [1, 1, 1, 1, 1],
       [1, 1, 1, 0, 1]])

Upvotes: 4

Related Questions