0

If I have two vectors and a sparse mask with mostly 0's and a few 1's, is there an efficient way to compute the masked dot product of a and b? My solution is very inefficient because it computes a bunch of zeros that don't contribute to the result:

a = np.random.normal(size=1000)
b = np.random.normal(size=1000)
m = np.random.choice([0,1], size=1000, p=[0.99, 0.01])

prod = np.sum(a*b*m)

(in my actual problem I have n-dim arrays rather than 1-dim, but the question is the same)

Ziofil
  • 1,815
  • 1
  • 20
  • 30
  • 1
    Have you looked at this? https://stackoverflow.com/questions/54294577/fast-inner-product-of-two-2-d-masked-arrays-in-numpy – sagar1025 Nov 13 '20 at 02:22
  • 1
    Trying to avoid those zeros might be more expensive than using them. Be sure to time test suggestions. – hpaulj Nov 13 '20 at 02:41

1 Answers1

1
mask = np.where(m > 0)
np.sum(a[mask] * b[mask])

== UPDATE ==

As noted by commenters, writing the second line as a[mask].dot(b[mask]) is slightly faster. Presumably Numpy doesn't need to build the intermediate array.

Frank Yellin
  • 9,127
  • 1
  • 12
  • 22
  • shouldn't you use `np.sum` instead of `sum` – Kenan Nov 13 '20 at 02:47
  • @Kenan For one-dimensional arrays, they both work, but yeah, you're right. – Frank Yellin Nov 13 '20 at 03:39
  • 2
    `a[mask].dot(b[mask]` is better. Or `np.einsum('i,i,i', a,b,m)`. – hpaulj Nov 13 '20 at 04:48
  • 1
    I agree that `a[mask].dot(b[mask])` is a good alternative, but I was trying to stay as close to poster's original code. I'm not so sure about `np.einsum`. How is that in any way different from the `np.sum(a*b*m)` that poster was trying to avoid? The whole point is that we only needed to do about 1% of the multiplies. – Frank Yellin Nov 13 '20 at 06:26
  • Also `mask` can just be `m.astype(bool)` – Daniel F Nov 13 '20 at 06:45
  • @DanielF. Yeah. It's not completely obvious to me which is faster or clearer. – Frank Yellin Nov 13 '20 at 07:06
  • 1
    Actually, after some testing with 100k values, `mask = np.where(m > 0)` is the fastest method I can find. The summation is fastest with @hpaulj 's `a[mask].dot(b[mask])` though – Daniel F Nov 13 '20 at 09:01
  • [I knew I'd done this before](https://stackoverflow.com/questions/45764955/shuffling-non-zero-elements-of-each-row-in-an-array-python-numpy/45812861#45812861) - `mask = np.nonzero(m.astype(bool))` is a smidge faster, although not nearly as readable. – Daniel F Nov 13 '20 at 09:08
  • `einsum` is much slower than things like tensordot or dot though – Ziofil Nov 13 '20 at 12:57