Fastest way for boolean matrix computations
Simply use np.einsum to get all the counts -
np.einsum('ij,ik,i->jk',M,M.astype(int),N.ravel())
Feel free to play around with optimize flag with np.einsum. Also, feel free to play around with different dtypes conversion.
To leverage GPU, we can use tensorflow package that also supports einsum.
Faster alternatives with np.dot :
(M&N).T.dot(M.astype(int))
(M&N).T.dot(M.astype(np.float32))
Timings -
In [110]: np.random.seed(0)
...: M = np.random.rand(500,300)>0.5
...: N = np.random.rand(500,1)>0.5
In [111]: %timeit np.einsum('ij,ik,i->jk',M,M.astype(int),N.ravel())
...: %timeit (M&N).T.dot(M.astype(int))
...: %timeit (M&N).T.dot(M.astype(np.float32))
227 ms ± 191 µs per loop (mean ± std. dev. of 7 runs, 1 loop each)
66.8 ms ± 198 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)
3.26 ms ± 753 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)
And take it a bit further with float32 conversions for both of the boolean arrays -
In [122]: %%timeit
...: p1 = (M&N).astype(np.float32)
...: p2 = M.astype(np.float32)
...: out = p1.T.dot(p2)
2.7 ms ± 34.3 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)
EDIT: To fix the code below to the fit the corrected question, just a couple of minor changes are required in compute:
def compute(m, n):
m = np.asarray(m)
n = np.asarray(n)
# Apply mask N in advance
m2 = m & n
# Pack booleans into uint8 for more efficient bitwise operations
# Also transpose for better caching (maybe?)
mb = np.packbits(m2.T, axis=1)
# Table with number of ones in each uint8
num_bits = (np.arange(256)[:, np.newaxis] & (1 << np.arange(8))).astype(bool).sum(1)
# Allocate output array
out = np.zeros((m2.shape[1], m2.shape[1]), np.int32)
# Do the counting with Numba
_compute_nb(mb, num_bits, out)
# Make output symmetric
out = out + out.T
# Add values in diagonal
out[np.diag_indices_from(out)] = m2.sum(0)
# Scale by number of ones in n
return out
I would do this with Numba, using a few tricks. First, you can do only half of the column-wise operations, since the other half is repeated. Second, you can pack the boolean values into bytes so with each & you are operating over eight values instead of one. Third, you can use multiprocessing to parallelize it. In total, you could do it like this:
import numpy as np
import numba as nb
def compute(m, n):
m = np.asarray(m)
n = np.asarray(n)
# Pack booleans into uint8 for more efficient bitwise operations
# Also transpose for better caching (maybe?)
mb = np.packbits(m.T, axis=1)
# Table with number of ones in each uint8
num_bits = (np.arange(256)[:, np.newaxis] & (1 << np.arange(8))).astype(bool).sum(1)
# Allocate output array
out = np.zeros((m.shape[1], m.shape[1]), np.int32)
# Do the counting with Numba
_compute_nb(mb, num_bits, out)
# Make output symmetric
out = out + out.T
# Add values in diagonal
out[np.diag_indices_from(out)] = m.sum(0)
# Scale by number of ones in n
out *= n.sum()
return out
@nb.njit(parallel=True)
def _compute_nb(mb, num_bits, out):
# Go through each pair of columns without repetitions
for i in nb.prange(mb.shape[0] - 1):
for j in nb.prange(1, mb.shape[0]):
# Count common bits
v = 0
for k in range(mb.shape[1]):
v += num_bits[mb[i, k] & mb[j, k]]
out[i, j] = v
# Test
m = np.array([[ True, True, False, True],
[False, True, True, True],
[False, False, False, False],
[False, True, False, False],
[ True, True, False, False]])
n = np.array([[ True],
[False],
[ True],
[ True],
[ True]])
out = compute(m, n)
print(out)
# [[ 8 8 0 4]
# [ 8 16 4 8]
# [ 0 4 4 4]
# [ 4 8 4 8]]
As a quick comparison, here is a small benchmark against the original loop and NumPy-only methods (I am pretty sure the proposals by Divakar are the best you can get from NumPy):
import numpy as np
# Original loop
def compute_loop(m, n):
out = np.zeros((m.shape[1], m.shape[1]), np.int32)
for i in range(m.shape[1]):
for j in range(m.shape[1]):
result = m[:, i] & m[:, j]
out[i, j] = np.sum(result & n)
return out
# Divakar methods
def compute2(m, n):
return np.einsum('ij,ik,lm->jk', m, m.astype(int), n)
def compute3(m, n):
return np.einsum('ij,ik->jk',m, m.astype(int)) * n.sum()
def compute4(m, n):
return np.tensordot(m, m.astype(int),axes=((0,0))) * n.sum()
def compute5(m, n):
return m.T.dot(m.astype(int))*n.sum()
# Make random data
np.random.seed(0)
m = np.random.rand(1000, 100) > .5
n = np.random.rand(1000, 1) > .5
print(compute(m, n).shape)
# (100, 100)
%timeit compute(m, n)
# 768 µs ± 17.5 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)
%timeit compute_loop(m, n)
# 11 s ± 1.23 s per loop (mean ± std. dev. of 7 runs, 1 loop each)
%timeit compute2(m, n)
# 7.65 s ± 1.06 s per loop (mean ± std. dev. of 7 runs, 1 loop each)
%timeit compute3(m, n)
# 23.5 ms ± 1.53 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)
%timeit compute4(m, n)
# 8.96 ms ± 194 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)
%timeit compute5(m, n)
# 8.35 ms ± 266 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)