What is the fastest way to get matrix < matrix results in numpy?
Suppose I have a matrix M_1
dimension (M, A) and a matrix M_2
dimension (M, B). The result of M_1 < M_2
should be a matrix of dimensions (M, B, A), where each row in M1
is being compared with each element of the corresponding row of M_2
and giving a bool vector (or 1,0 vector) for each comparison.
For example, if I have a matrix
M1 = [[1,2,3]
[3,4,5]]
M2 = [[1,2],
[3,4]]
result should be [[[False, False, False],
[True, False, False]],
[[False, False, False],
[True, False, False]]]
Currently, I’m using a for loop and it gets very slow when I have to repeat this multiple times (which takes months). Hopefully there is a vectorized way to do this. If not, what else can I do?
I’m checking that the M_1
is (500, 3000000) and the M_2
is (500, 500) and repeated about 10000 times.
Solution
For NumPy arrays, extend dims with None/np.newaxis
so that the first axis aligns and the second axis expands so that one element stylish. Finally, make a comparison using > broadcasting-for vectorized solutions-
M1[:,None,:] < M2[:,:,None]
sample run-
In [19]: M1
Out[19]:
array([[1, 2, 3],
[3, 4, 5]])
In [20]: M2
Out[20]:
array([[1, 2],
[3, 4]])
In [21]: M1[:,None,:] < M2[:,:,None]
Out[21]:
array([[[False, False, False],
[ True, False, False]],
[[False, False, False],
[ True, False, False]]])
For a list as input, use > numpy.expand_dims Then compare –
In [42]: M1 = [[1,2,3],
...: [3,4,5]]
...:
...: M2 = [[1,2],
...: [3,4]]
In [43]: np.expand_dims(M1, axis=1) < np.expand_dims(M2, axis=2)
Out[43]:
array([[[False, False, False],
[ True, False, False]],
[[False, False, False],
[ True, False, False]]])
Further elevated
Use leverage to further boost multi-core
with numexpr
module for big data –
In [44]: import numexpr as ne
In [52]: M1 = np.random.randint(0,9,(500, 30000))
In [53]: M2 = np.random.randint(0,9,(500, 500))
In [55]: %timeit M1[:,None,:] < M2[:,:,None]
1 loop, best of 3: 3.32 s per loop
In [56]: %timeit ne.evaluate('M1e<M2e',{'M1e':M1[:,None,:],'M2e':M2[:,:,None]})
1 loop, best of 3: 1.53 s per loop