Selective Numpy array manipulation (dependent on value)

Question

Let's say I have the following Numpy array:

array([[3, 5, 0], [7, 0, 2]])

I now want to add 2 where the value is not 0. What would be the fastest way to do this? I have to manipulate quite large multidimensional arrays?

Answer 1

It seems to me that:

a[a!=0] += 2

should work.

(for the limited case of testing for non-zero-ness), you might be able to speed things up with (you'd need to timeit to see):

mask = a.astype(bool)
a[mask] += 2

---

Of course, you can save yourself the mask calculation if you can reuse the same mask at different spots (which is a pretty restrictive constraint):

mask = a != 0
a[mask] += 2
#some more code ...
a[mask] *= 3
#more code ...

---

Of course, If this is enough of a bottleneck, you can always write a little C/Fortran extension to do this for you (using Cython or f2py respectively). This would avoid the overhead of mask creation.

Answer 2

It depends on the size of your array but you could consider using numexpr for this:

# from your exemmple
>>> a = array([[3, 5, 0], [7, 0, 2]])
>>> %timeit a[a!=0] += 2
100000 loops, best of 3: 18.6 us per loop
>>> timeit numexpr.evaluate("a + 2 * (a != 0)")
10000 loops, best of 3: 42.6 us per loop

But with a bigger array:

# make a big array with 10% of zeros :
a = np.random.rand(10000)
a[a<0.1] = 0
# same test:
>>> timeit a[a!=0] += 2
1000 loops, best of 3: 364 us per loop
>>> timeit numexpr.evaluate("a + 2 * (a != 0)")
10000 loops, best of 3: 119 us per loop

And this is with a single core. Numepxr makes use of all core available if possible.