CODEWITHSUNDEEP
pythonnumpy
Man
Man

Reputation: 51

numpy mask for 2d array with all values in 1d array

I want to convert a 2d matrix of dates to boolean matrix based on dates in a 1d matrix. i.e., 

[[20030102, 20030102, 20070102],
 [20040102, 20040102, 20040102].,
 [20050102, 20050102, 20050102]]

should become

[[True, True, False],
 [False, False, False].,
 [True, True, True]]

if I provide a 1d array [20010203, 20030102, 20030501, 20050102, 20060101]

Upvotes: 2

Views: 1209

Answers (2)

wwii
wwii

Reputation: 23743

a = np.array([[20030102, 20030102, 20070102],
              [20040102, 20040102, 20040102],
              [20050102, 20050102, 20050102]])

b = np.array([20010203, 20030102, 20030501, 20050102, 20060101])

>>> a.shape
(3, 3)
>>> b.shape
(5,)
>>>

For the comparison, you need to broadcast b onto a by adding an axis to a. - this compares each element of a with each element of b

>>> mask = a[...,None] == b
>>> mask.shape
(3, 3, 5)
>>>

Then use np.any() to see if there are any matches

>>> np.any(mask, axis = 2, keepdims = False)
array([[ True,  True, False],
       [False, False, False],
       [ True,  True,  True]], dtype=bool)

timeit.Timer comparison with in1d:

>>> 
>>> t = Timer("np.any(a[...,None] == b, axis = 2)","from __main__ import np, a, b")
>>> t.timeit(10000)
0.13268041338812964
>>> t = Timer("np.in1d(a.ravel(), b).reshape(a.shape)","from __main__ import np, a, b")
>>> t.timeit(10000)
0.26060646913566643
>>>

Upvotes: 1

ode2k
ode2k

Reputation: 2723

import numpy as np

dateValues = np.array(
    [[20030102, 20030102, 20030102],
     [20040102, 20040102, 20040102],
     [20050102, 20050102, 20050102]])

requestedDates = [20010203, 20030102, 20030501, 20050102, 20060101]

ix = np.in1d(dateValues.ravel(), requestedDates).reshape(dateValues.shape)

print(ix)

Returns:

[[ True  True  True]
 [False False False]
 [ True  True  True]]

Refer to numpy.in1d for more information (documentation): http://docs.scipy.org/doc/numpy/reference/generated/numpy.in1d.html

Upvotes: 4

Related Questions